Repository: google/jax

Issue:
Slow gradient calculation when indexing a complex-valued object inside a loss.
I have observed that there is large increase in computation time when differentiating through a loss that involves index a complex128-valued array. Consider the following demo script.
```python
import argparse
import time

import jax.numpy as jnp
from jax import jit, grad
from jax import random
from jax.config import config


parser = argparse.ArgumentParser(description='Complex indexing and gradients in JAX')
parser.add_argument('--complex', default=True, dest='complex', action='store_true', help='Index a complex-valued object')
parser.add_argument('--no-complex', dest='complex', action='store_false')
parser.add_argument('--double', default=True, dest='double', action='store_true', help='Use double precision')
parser.add_argument('--no-double', dest='double', action='store_false')
parser.add_argument('--index', default=True, dest='index', action='store_true', help='Index the object inside the loss')
parser.add_argument('--no-index', dest='index', action='store_false')
args = parser.parse_args()

if args.double:
    config.update("jax_enable_x64", True)

def loss(fvol, idx):
    # Define a loss function that either indexes or does not index the array.
    if args.index:
        vv = fvol[idx].sum([-1])
        # Doing this is fast again in complex-128.
        # fvol_real = fvol.real[idx]
        # fvol_imag = fvol.imag[idx]
        # vv = jnp.sum(fvol_real + 1j*fvol_imag, [-1])
    else:
        vv = fvol
    return jnp.square(jnp.linalg.norm(vv))

@jit
def loss_grad(fvol, idx):
    return grad(loss)(fvol, idx)

# Create a large array.
rng = random.PRNGKey(0)
rng_real, rng_imag = random.split(rng, 2)
L = 256
fvol = random.normal(rng_real, [L**3])
if args.complex:
    fvol += 1j * random.normal(rng_imag, fvol.shape)

dtype = fvol.dtype
idx = random.randint(rng, [L*L*8], minval=0, maxval=L)

# Time how long it takes to compute the loss.
for i in range(3):
    start = time.time()
    lv = loss(fvol, idx)
    lv.block_until_ready()
    elapsed = time.time() - start
    print('object data: {} - time elapsed to evaluate loss: {:.5f}'.format(dtype, elapsed))

# Time how long it takes to compute the gradient of the loss.
for i in range(3):
    start = time.time()
    g = loss_grad(fvol, idx)
    g.block_until_ready()
    elapsed = time.time() - start
    print('object data: {} - time elapsed to evaluate gradient: {:.5f}'.format(dtype, elapsed))
```
I am running this demo script on a V100 GPU with jaxlib 0.1.57 with Python 3.6.9. When I invoke the demo as `python demo.py --complex --double --index` (meaning that I have a complex-valued array, double precision, and a loss function that involves indexing) I get the following result:
```
object data: complex128 - time elapsed to evaluate loss: 0.63968
object data: complex128 - time elapsed to evaluate loss: 0.03179
object data: complex128 - time elapsed to evaluate loss: 0.03082
object data: complex128 - time elapsed to evaluate gradient: 4.70996
object data: complex128 - time elapsed to evaluate gradient: 4.55858
object data: complex128 - time elapsed to evaluate gradient: 4.56153
```
We see that it takes significantly longer to compute the gradient than the loss itself. If I invoke the script as `python demo.py --complex --double --no-index` (complex-valued array, double precision, and a loss that does not involve indexing) then things are better; the gradient calculation is ~3x the timing of the loss calculation.
```
object data: complex128 - time elapsed to evaluate loss: 0.18214
object data: complex128 - time elapsed to evaluate loss: 0.00098
object data: complex128 - time elapsed to evaluate loss: 0.00059
object data: complex128 - time elapsed to evaluate gradient: 0.13375
object data: complex128 - time elapsed to evaluate gradient: 0.00175
object data: complex128 - time elapsed to evaluate gradient: 0.00148
```

What about in single precision? Invoking the script as `python demo.py --complex --no-double --index` yields
```
object data: complex64 - time elapsed to evaluate loss: 0.60020
object data: complex64 - time elapsed to evaluate loss: 0.00300
object data: complex64 - time elapsed to evaluate loss: 0.00164
object data: complex64 - time elapsed to evaluate gradient: 0.18140
object data: complex64 - time elapsed to evaluate gradient: 0.03122
object data: complex64 - time elapsed to evaluate gradient: 0.03108
```
The gradient calculation is around an order of magnitude slower. What about when we have an unindexed loss function using the invocation `python demo.py --complex --no-double --no-index`?
```
object data: complex64 - time elapsed to evaluate loss: 0.17156
object data: complex64 - time elapsed to evaluate loss: 0.00114
object data: complex64 - time elapsed to evaluate loss: 0.00057
object data: complex64 - time elapsed to evaluate gradient: 0.11934
object data: complex64 - time elapsed to evaluate gradient: 0.00087
object data: complex64 - time elapsed to evaluate gradient: 0.00077
```
Finally, what happens if I consider indexing a real-valued array instead of a complex one? Using the invocation `python demo.py --no-complex --double --index`
```
object data: float64 - time elapsed to evaluate loss: 0.60958
object data: float64 - time elapsed to evaluate loss: 0.00578
object data: float64 - time elapsed to evaluate loss: 0.00180
object data: float64 - time elapsed to evaluate gradient: 0.13558
object data: float64 - time elapsed to evaluate gradient: 0.00106
object data: float64 - time elapsed to evaluate gradient: 0.00082
```
and compare this to the first output I showed where the gradient calculation was around four seconds.

I note in the demo script that there is a correction to the loss that will make the code fast again when run in double precision and indexing a complex-valued array. Replacing `vv = fvol[idx].sum([-1])` (line 25) with
```python
fvol_real = fvol.real[idx]
fvol_imag = fvol.imag[idx]
vv = jnp.sum(fvol_real + 1j*fvol_imag, [-1])
```
 produces the output
```
object data: complex128 - time elapsed to evaluate loss: 0.92206
object data: complex128 - time elapsed to evaluate loss: 0.03364
object data: complex128 - time elapsed to evaluate loss: 0.03140
object data: complex128 - time elapsed to evaluate gradient: 0.20444
object data: complex128 - time elapsed to evaluate gradient: 0.06708
object data: complex128 - time elapsed to evaluate gradient: 0.06507
```
which a significant improvement over the four seconds the code was taking before. This is just speculation, but in researching this issue my colleagues and I thought this bug could be at play here as well: https://github.com/google/jax/issues/4115




Discussion:
Yes, I strongly suspect this is a duplicate of #4115.

The issue is that XLA's implementation of `scatter-add`, which is the gradient of the `gather` we use to compute indexing, relies on atomic operations, but GPUs don't have atomic operations wide enough for complex128 values. In this case, XLA falls back to a slow loop of updates without any parallelism.

That said, for this *particular* scatter, which is a scatter-add, we can use the trick you have commented out, which is to perform separate scatters on the real and imaginary parts. We could do it inside the XLA translation rule easily enough.

Resolving patch:
diff --git a/jax/_src/lax/lax.py b/jax/_src/lax/lax.py
--- a/jax/_src/lax/lax.py
+++ b/jax/_src/lax/lax.py
@@ -4315,6 +4315,35 @@ def _scatter_translation_rule(c, operand, scatter_indices, updates, *,
                       _scatter_dimensions_proto(indices_shape, dimension_numbers),
                       indices_are_sorted, unique_indices)
 
+def _scatter_add_translation_rule(
+    c, operand, scatter_indices, updates, *, update_jaxpr, update_consts,
+    dimension_numbers, indices_are_sorted, unique_indices,
+    expand_complex128=False):
+  dtype = c.get_shape(operand).numpy_dtype()
+  scatter_dims = _scatter_dimensions_proto(c.get_shape(scatter_indices),
+                                           dimension_numbers)
+
+  def _make_reducer(dtype):
+    subc = xla_bridge.make_computation_builder("scatter_add_reducer")
+    shape = xc.Shape.array_shape(np.dtype(dtype), ())
+    args = [xb.parameter(subc, 0, shape), xb.parameter(subc, 1, shape)]
+    out = xops.Add(args[0], args[1])
+    return subc.build(out)
+
+  if expand_complex128 and dtype == np.complex128:
+    update_computation = _make_reducer(np.float64)
+    re = xops.Scatter(xops.Real(operand), scatter_indices, xops.Real(updates),
+                      update_computation, scatter_dims, indices_are_sorted,
+                      unique_indices)
+    im = xops.Scatter(xops.Imag(operand), scatter_indices, xops.Imag(updates),
+                      update_computation, scatter_dims, indices_are_sorted,
+                      unique_indices)
+    return xops.Complex(re, im)
+  else:
+    update_computation = _make_reducer(dtype)
+    return xops.Scatter(operand, scatter_indices, updates, update_computation,
+                        scatter_dims, indices_are_sorted, unique_indices)
+
 def _scatter_add_jvp(primals, tangents, *, update_jaxpr, update_consts,
                      dimension_numbers, indices_are_sorted, unique_indices):
   operand, scatter_indices, updates = primals
@@ -4454,12 +4483,14 @@ def _scatter_batching_rule(scatter_op, batched_args, batch_dims, *,
 
 scatter_add_p = standard_primitive(
     _scatter_shape_rule, _scatter_dtype_rule, 'scatter-add',
-    _scatter_translation_rule)
+    _scatter_add_translation_rule)
 ad.primitive_jvps[scatter_add_p] = _scatter_add_jvp
 ad.primitive_transposes[scatter_add_p] = _scatter_add_transpose_rule
 batching.primitive_batchers[scatter_add_p] = (
   partial(_scatter_batching_rule, scatter_add))
 
+xla.backend_specific_translations['gpu'][scatter_add_p] = partial(
+    _scatter_add_translation_rule, expand_complex128=True)
 
 scatter_mul_p = standard_primitive(
     _scatter_shape_rule, _scatter_dtype_rule, 'scatter-mul',


Test patch:
diff --git a/tests/lax_test.py b/tests/lax_test.py
--- a/tests/lax_test.py
+++ b/tests/lax_test.py
@@ -1934,7 +1934,7 @@ def testGatherShapeCheckingRule(self, operand_shape, start_indices_shape,
        "arg_shape": arg_shape, "dtype": dtype, "idxs": idxs,
        "update_shape": update_shape, "dnums": dnums,
        "rng_factory": rng_factory, "rng_idx_factory": rng_idx_factory}
-      for dtype in float_dtypes
+      for dtype in inexact_dtypes
       for arg_shape, idxs, update_shape, dnums in [
           ((5,), np.array([[0], [2]]), (2,), lax.ScatterDimensionNumbers(
             update_window_dims=(), inserted_window_dims=(0,),