Repository: google/jax

Issue:
slice-based indexing is slow for repeated indexing
This issue is the result of some pair-debugging with @sharadmv.

In the course of trying to re-land #12091, I found that it was leading to some test timeouts due to replacing `lax.gather` with `lax.slice`. This issue actually can be reproduced in the current release of JAX; given the current implementation, we can define two functions, one of which takes the `gather` path and one of which takes the `slice` path:
```python
import jax.numpy as jnp
import jax

def index_via_gather(x):
  return x[(0,)]

def index_via_slice(x):
  return x[0]

x = jnp.arange(100)

print(jax.make_jaxpr(index_via_gather)(x))
# { lambda ; a:i32[100]. let
#     b:i32[1] = broadcast_in_dim[broadcast_dimensions=() shape=(1,)] 0
#     c:i32[] = gather[
#       dimension_numbers=GatherDimensionNumbers(offset_dims=(), collapsed_slice_dims=(0,), start_index_map=(0,))
#       fill_value=None
#       indices_are_sorted=True
#       mode=GatherScatterMode.PROMISE_IN_BOUNDS
#       slice_sizes=(1,)
#       unique_indices=True
#     ] a b
#   in (c,) }

print(jax.make_jaxpr(index_via_slice)(x))
# { lambda ; a:i32[100]. let
#     b:i32[1] = slice[limit_indices=(1,) start_indices=(0,) strides=(1,)] a
#     c:i32[] = squeeze[dimensions=(0,)] b
#   in (c,) }
```
The latter slice behavior was added in #11867 because it can be more efficient.

Unfortunately, there is a catch: when indexing is done repeatedly on different static indices, the `slice` approach has a cache miss on each new index, making it much slower. The test that led to the rollback of #12091 exposed this issue, which we can reproduce for the current implementation like this:
```python
def f_gather(x):
  return jnp.array([x[(i,)] for i in range(len(x))])

def f_slice(x):
  return jnp.array([x[i] for i in range(len(x))])

x = jnp.arange(500)

%timeit -n 1 -r 1 f_gather(x).block_until_ready()
# 611 ms ± 0 ns per loop (mean ± std. dev. of 1 run, 1 loop each)

%timeit -n 1 -r 1 f_slice(x).block_until_ready()
# 6.78 s ± 0 ns per loop (mean ± std. dev. of 1 run, 1 loop each)
```
Repeated indexing via static slice is *very* slow comparted to repeated indexing via gather.

We could instead consider using `lax.dynamic_slice`, which has comparable performance to `gather` for this repeated operation:
```python
def f_dynamic_slice(x):
  return jnp.array([jax.lax.dynamic_index_in_dim(x, i, 0) for i in range(len(x))])

%timeit -n 1 -r 1 f_dynamic_slice(x).block_until_ready()
# 420 ms ± 0 ns per loop (mean ± std. dev. of 1 run, 1 loop each)
```

Any thoughts? cc/ @mattjj, who contributed #11867


Discussion:
Chatted offline with @mattjj; we agreed that switching from `slice` to `dynamic_slice` is probably the best option here.

Resolving patch:
diff --git a/benchmarks/api_benchmark.py b/benchmarks/api_benchmark.py
--- a/benchmarks/api_benchmark.py
+++ b/benchmarks/api_benchmark.py
@@ -663,6 +663,19 @@ def bench_slicing_compilation2(state):
   while state:
     jax.jit(lambda x: (x[:1], x[1:2], x[2:3])).lower(x).compile()
 
+@google_benchmark.register
+@google_benchmark.option.unit(google_benchmark.kMillisecond)
+def bench_repeated_static_indexing(state):
+  x = jnp.arange(500)
+  while state:
+    jax.block_until_ready([x[i] for i in range(500)])
+
+@google_benchmark.register
+@google_benchmark.option.unit(google_benchmark.kMillisecond)
+def bench_repeated_static_slicing(state):
+  x = jnp.arange(1000)
+  while state:
+    jax.block_until_ready([x[i:i + 2] for i in range(0, 1000, 2)])
 
 def pjit_simple_benchmark(state, num_devices, num_args, cpp_jit, use_aot=False):
   spec = pjit_lib.PartitionSpec('x')
diff --git a/jax/_src/numpy/lax_numpy.py b/jax/_src/numpy/lax_numpy.py
--- a/jax/_src/numpy/lax_numpy.py
+++ b/jax/_src/numpy/lax_numpy.py
@@ -3776,10 +3776,9 @@ def _rewriting_take(arr, idx, indices_are_sorted=False, unique_indices=False,
   if (arr.ndim > 0 and isinstance(idx, (int, np.integer)) and
       not isinstance(idx, (bool, np.bool_)) and isinstance(arr.shape[0], int)):
     if 0 <= idx < arr.shape[0]:
-      if _any(isinstance(d, core.Tracer) for d in arr.shape[1:]):
-        return lax.dynamic_index_in_dim(arr, idx, keepdims=False)
-      else:
-        return lax.index_in_dim(arr, idx, keepdims=False)
+      # Use dynamic rather than static index here to avoid slow repeated execution:
+      # See https://github.com/google/jax/issues/12198
+      return lax.dynamic_index_in_dim(arr, idx, keepdims=False)
   if (arr.ndim > 0 and isinstance(arr.shape[0], int) and
       isinstance(idx, slice) and
       (type(idx.start) is int or idx.start is None) and
@@ -3794,6 +3793,10 @@ def _rewriting_take(arr, idx, indices_are_sorted=False, unique_indices=False,
       if _any(isinstance(d, core.Tracer) for d in arr.shape[1:]):
         if step == 1:  # TODO(mattjj, sharadmv): handle step != 1
           return lax.dynamic_slice_in_dim(arr, start, _max(0, stop - start), 0)
+      elif step == 1:
+        # Use dynamic rather than static slice here to avoid slow repeated execution:
+        # See https://github.com/google/jax/issues/12198
+        return lax.dynamic_slice_in_dim(arr, start, _max(0, stop - start), 0)
       else:
         return lax.slice_in_dim(arr, start, stop, step)
 


Test patch: