Repository: google/jax

Issue:
Step length dependent compilation time of lax.scan.
This is related to #1130 and was originally brought to my attention by @cpgoodri. Here is an example repro run in google cloud colab GPU. 
```python
!pip install --upgrade -q https://storage.googleapis.com/jax-wheels/cuda$(echo $CUDA_VERSION | sed -e 's/\.//' -e 's/\..*//')/jaxlib-$(pip search jaxlib | grep -oP '[0-9\.]+' | head -n 1)-cp36-none-linux_x86_64.whl
!pip install --upgrade -q jax

import jax.numpy as np
from jax import lax
from jax import jit


def f(x, s):
  for i in range(2000):
    x = x + 1.0
  return x, np.array([0])

def g(x, steps):
  xs = lax.iota(np.int32, steps)
  return lax.scan(f, x, xs)[0]

q = jit(g, static_argnums=(1,))

for p in range(1, 14):
  x = 2 ** p
  
  time1 = time.clock()
  temp = q(np.array(2.0), x).block_until_ready()
  time2 = time.clock()
  temp = q(np.array(2.0), x).block_until_ready()
  time3 = time.clock()

  first_run_time = (time2 - time1)
  second_run_time = (time3 - time2)
  compile_time = first_run_time - second_run_time
  
  print('x = {}, compile time = {} (run times = {}, {})'.format(
      x, compile_time, first_run_time, second_run_time))
```
which gives the output,
```
x = 2, compile time = 0.9660040000001118 (run times = 0.967151000000058, 0.0011469999999462743)
x = 4, compile time = 0.761009999999942 (run times = 0.7619959999999537, 0.0009860000000116997)
x = 8, compile time = 1.0294560000000956 (run times = 1.0305190000000266, 0.0010629999999309803)
x = 16, compile time = 1.0474190000001045 (run times = 1.0485160000000633, 0.0010969999999588254)
x = 32, compile time = 0.8571110000001454 (run times = 0.8583040000000892, 0.0011929999999438223)
x = 64, compile time = 1.1593259999999646 (run times = 1.1608429999999998, 0.0015170000000352957)
x = 128, compile time = 1.331446999999912 (run times = 1.333646999999928, 0.002200000000016189)
x = 256, compile time = 1.6889869999999974 (run times = 1.6926300000000083, 0.003643000000010943)
x = 512, compile time = 2.2191810000000487 (run times = 2.225097000000005, 0.005915999999956512)
x = 1024, compile time = 3.9208149999999478 (run times = 3.9315219999999726, 0.010707000000024891)
x = 2048, compile time = 6.794198999999935 (run times = 6.814807999999971, 0.02060900000003585)
x = 4096, compile time = 13.265449999999987 (run times = 13.306156999999985, 0.040706999999997606)
x = 8192, compile time = 24.92249199999992 (run times = 25.00054499999999, 0.0780530000000681)
```



Discussion:
Thanks for raising this. Can you dump the HLO? That will tell us if we're staging out something different to XLA for these different cases (i.e. it's fixable on the JAX side) or if this is a phenomenon related to XLA compile times of a fixed-size HLO graph (i.e. it would be fixed on the XLA side).
Two ways to dump the HLO: use `jax.xla_computation` or set the environment variable `XLA_FLAGS="--xla_dump_to=DIRECTORY_PATH"`. Googlers can read more at [go/jax-dump-xla-hlo](go/jax-dump-xla-hlo) thanks to @levskaya. 
Following up on this, after looking into this with the help of @levskaya and @mattjj, it seems like an XLA issue. I reported a bug to XLA:GPU.
@hawkinsp found that this is a JAX bug after all: we're hitting the "impl rule" when we're not supposed to! This could explain #1130 as well. I'm on it!
I think I see the issue, though this explanation isn't going to make much sense unless you're already deeply read into the JAX internals: when `jit` a `scan`, we actually do [partial evaluation which "unzips" the scan into two scans](https://github.com/google/jax/blob/a6ec5a916cae5b71e0232633b89e7490f77d81ef/jax/lax/lax_control_flow.py#L656), one that evaluates everything known at compile time (the "jaxpr_1 side") and one staged out to a jaxpr (and ultimately  an XLA computation) to be evaluated later (the "jaxpr_2 side"). The computation done by that first scan is essentially compile-time constant folding. Under a `jit` like this one, we actually stage all the real work out to XLA: there's no real constant folding / partial evaluation to do.  That is, the first "jaxpr_1" scan is trivial: it's just shuffling empty tuples around.

The problem here is that even when the first scan is trivial, [we're still executing it](https://github.com/google/jax/blob/a6ec5a916cae5b71e0232633b89e7490f77d81ef/jax/lax/lax_control_flow.py#L687-L688). That means we're still hitting the while_loop impl, and those are currently slow (and in any case running a number of steps that scales with the length of the scan). That's why JAX is doing work "at compile time" that scales with the length of the scanned array.

I'm not sure yet what the right fix is, but one option is just to notice when the jaxpr_1 computation is trivial and not run it in that case!
Here's a look at how the jaxpr_1 side is trivial:

```
ipdb> l
    684   ys_aval = _promote_aval_rank(length, y_aval)
    685   out_aval = core.AbstractTuple((carry_aval, ys_aval))
    686   out_pv = _put_known_pvs(sc_out, out_aval)
    687
    688   out_carry, (ys, residuals) = scan_p.bind(
--> 689       *in_consts, forward=forward, length=length, jaxpr=jaxpr_1)
    690   out_const = core.pack((out_carry, ys))
    691   residuals_tracer = trace.new_instantiated_const(core.pack(residuals))
    692   d, c, a = lifted_tracers
    693   new_tracers = (d, c, (a, residuals_tracer))
    694   eqn = core.JaxprEqn(new_tracers, None, scan_p, (), True, False,

ipdb> p jaxpr_1
{ lambda d f ;  ; a b c.
  let (e) = id a
      g = pack e f
      h = pack d g
  in h }

ipdb> in_consts[0]  # argument for the `a` parameter
JaxTuple(DeviceArray([0], dtype=int32))
ipdb> jaxpr_1.literals  # arguments for the `d` and `f` parameters
(*, JaxTuple())
```

So that's what we're running for 8192 iterations in a slow op-by-op mode while loop at compile time...

Resolving patch:
diff --git a/jax/interpreters/xla.py b/jax/interpreters/xla.py
--- a/jax/interpreters/xla.py
+++ b/jax/interpreters/xla.py
@@ -79,11 +79,17 @@ def primitive_computation(prim, *shapes, **params):
   c = xb.make_computation_builder("primitive_computation")
   platform = xb.get_backend().platform
   xla_args = map(c.ParameterWithShape, shapes)
-  try:
-    rule = backend_specific_translations[platform].get(prim) or translations[prim]
-  except KeyError:
+  if prim in backend_specific_translations[platform]:
+    rule = backend_specific_translations[platform][prim]
+    rule(c, *xla_args, **params)  # return val set as a side-effect on c
+  elif prim in translations:
+    rule = translations[prim]
+    rule(c, *xla_args, **params)  # return val set as a side-effect on c
+  elif prim in initial_style_translations:
+    rule = initial_style_translations[prim]
+    rule(c, AxisEnv(1, [], []), *xla_args, **params)  # side-effect on c
+  else:
     raise NotImplementedError("XLA translation rule for {} not found".format(prim))
-  rule(c, *xla_args, **params)  # return val set as a side-effect on c
   try:
     return c.Build()
   except RuntimeError as e:
diff --git a/jax/lax/lax_control_flow.py b/jax/lax/lax_control_flow.py
--- a/jax/lax/lax_control_flow.py
+++ b/jax/lax/lax_control_flow.py
@@ -156,16 +156,6 @@ def while_loop(cond_fun, body_fun, init_val):
   return build_tree(out_tree(), out_flat)
 
 
-def _while_loop_impl(init_val, cond_consts, body_consts, aval_out, cond_jaxpr,
-                     body_jaxpr):
-  cond_fun = partial(core.eval_jaxpr, cond_jaxpr, cond_consts, ())
-  body_fun = partial(core.eval_jaxpr, body_jaxpr, body_consts, ())
-
-  val = init_val
-  while cond_fun(val):
-    val = body_fun(val)
-  return val
-
 def _while_loop_abstract_eval(init_val, cond_consts, body_consts, aval_out,
                               cond_jaxpr, body_jaxpr):
   return _maybe_tracer_tuple_to_abstract_tuple(aval_out)
@@ -261,7 +251,7 @@ def _jaxtupletree_select(pred, on_true, on_false):
 
 
 while_p = lax.Primitive('while')
-while_p.def_impl(_while_loop_impl)
+while_p.def_impl(partial(xla.apply_primitive, while_p))
 while_p.def_abstract_eval(_while_loop_abstract_eval)
 xla.initial_style_translations[while_p] = _while_loop_translation_rule
 batching.primitive_batchers[while_p] = _while_loop_batching_rule


Test patch: