Early termination for any/all boolean reduction functions

[dan-zheng]

Here is the current definition of any and all in the prelude:

def scan (init:a) (body:n->a->(a&b)) : (a & n=>b) =
  swap $ withState init \s. for i.
    c = get s
    (c', y) = body i c
    s := c'
    y

def fold (init:a) (body:(n->a->a)) : a = fst $ scan init \i x. (body i x, ())
def reduce (identity:a) (combine:(a->a->a)) (xs:n=>a) : a =
  -- `combine` should be a commutative and associative, and form a
  -- commutative monoid with `identity`
  -- TODO: implement with a parallelizable monoid-parameterized writer
  fold identity (\i c. combine c xs.i)

def any (xs:n=>Bool) : Bool = reduce False (||) xs
def all (xs:n=>Bool) : Bool = reduce True  (&&) xs

In imperative languages, implementations of any and all support early termination logic:

def all(f, xs):
  for x in xs:
    if not f(x):
      return False # return False on first occurrence of f(x) == False
  return True

Can Dex support this too - either in the language or as an optimization on the IR?

[danieldjohnson]

In principle, I think this could be implemented fairly easily using a while loop instead of a for loop (and this would only work for short-circuit-able reductions, so any and all would no longer be special-cases of reduce).

I think early-termination is possibly less friendly to parallel hardware like GPUs, since it requires dynamic branching. (Although, that being said, if we really care about parallel hardware for any and all then we would eventually want to implement them using withAccum or similar instead of using withState.)

There's also a bigger question of "how early" we want to terminate. Using a while loop would avoid looping over the indices after the first False, but probably wouldn't prevent us from computing the rest of the table. If we had a version like

def lazyAny (xs:n->Bool) : Bool = ...

where we use a regular arrow instead of a table arrow, that might be more efficient, because we wouldn't even bother computing the value for indices we aren't going to check. Although maybe an IR pass could convert the table arrow into (effectively) a regular arrow if it saw that the table was never used again (I think there might be an optimization like this already, in fact).

[dan-zheng]

Interesting distinction between while- and for-loops. I wonder what the performance difference looks like in Dex, across backends.

If we had a version like [def lazyAny (xs:n->Bool) : Bool = ...] where we use a regular arrow instead of a table arrow, that might be more efficient, because we wouldn't even bother computing the value for indices we aren't going to check.

That's neat. It reminds of me call-by-name evaluation semantics (e.g. Scala's by-name parameters). Such a table-to-lambda optimization seems quite interesting, since it's not wholly obvious to me when the optimization is beneficial: materializing arrays for easy access to computed elements, vs keeping arrays represented as functions for lazy evaluation. Some use/escape analysis seems relevant.

Although maybe an IR pass could convert the table arrow into (effectively) a regular arrow if it saw that the table was never used again (I think there might be an optimization like this already, in fact).

If the table is never used again, but every element of the table is used (at least once), wouldn't it be better more efficient to materialize the table than to compute each element individually n times? Or, maybe both are roughly equivalent (modulo allocation differences and scalar-replacement-of-aggregates optimization) if the table is used exactly once.

[apaszke]

Just FYI we already sometimes promote tables to lambdas (although to table lambdas to keep the program well-typed; those are unaccessible in the surface level language) when we decide that they shouldn't be materialized. This includes e.g. arrays with broadcast dimensions (in particular all-constant arrays).

I agree with @danieldjohnson in that making any and all shortcut would be very bad for parallel hardware, and with his observation that shortcutting doesn't really make sense to make them terminate early. According to (eager) language semantics we cannot even being to compute any/all until the full table is evaluated, but by that time all conditional expressions have been evaluated too. So there might be a slight performance improvement on sequential platforms at best, with a potential strong downside on parallel platforms.

But, I do think that it would be a good idea for && and || to shortcut, since that's what people generally expect from other languages.

[dan-zheng]

Just FYI we already sometimes promote tables to lambdas (although to table lambdas to keep the program well-typed; those are unaccessible in the surface level language) when we decide that they shouldn't be materialized. This includes e.g. arrays with broadcast dimensions (in particular all-constant arrays).

Neat!

Thanks folks for the explanations!