If the cur_eqn is program_order_p (or any HOP in the future) add an opt_barrier for every cur_inp i.e. cur_inps = [opt_barrier((token, c))[1] for c in cur_inps] where token, prev_outs_new = optimization_barrier((token, prev_outs)). This is strictly better than the approach before this change (described below).

[copybara-service]

If the cur_eqn is program_order_p (or any HOP in the future) add an opt_barrier for every cur_inp i.e. cur_inps = [opt_barrier((token, c))[1] for c in cur_inps] where token, prev_outs_new = optimization_barrier((token, prev_outs)). This is strictly better than the approach before this change (described below).

One thing to note here is that opt_barrier won't DCE tokens even if they are unused!

Why do this?

Currently we add an opt_barrier of form opt_barrier((prev_outs, all_cur_inps)). This means that in a program_order(False) scope, we block all equations until all the inputs are ready. This is sub-optimal because if there are equations depending on different inputs, they can execute when their inputs are ready.

program_order's semantics are that every equation will begin after it's previous equation has finished. This change does not violate that. In fact, it makes the scheduling better.

For example:

@program_order(True)
def f(x, y):
  x = jnp.sin(x)

  @program_order(False)
  def g(x, y):
    x = jax.lax.all_gather(x, ...)
    y = jnp.dot(y, y)
    return x, y
  return g(x, y)

Before this change, all_gather won't execute until both x and y are ready even though all_gather only needs x to be ready. Same for the dot op.

Now, since we do a per-input opt_barrier, program_order's semantics are respected i.e. g (all_gather and dot) won't execute until jnp.sin is done but all_gather and dot are free to be scheduled independently once their inputs are ready (since they are in a program_order(False) scope).