Repository: google-deepmind/dm-haiku

Issue:
Constant folding on `hk.avg_pool` with `padding='SAME'` on large arrays
Hi, lately my model has increased its compilation time due to constant folding and I discovered that the issue comes from using `padding='SAME'` in `hk.avg_pool`.

I tried to create a dummy code to showcase the issue (at the bottom), but the main point is that when `hk.avg_pool` uses `padding='VALID'` I get these output.

```
Wall time [Initialisation]: 0.112s
Wall time [Compilation]: 1.33s
Wall time [Train step]: 0.0024s
```

but when I use `padding='SAME'`
```
Wall time [Initialisation]: 0.119s
2022-10-12 16:01:47.590034: E external/org_tensorflow/tensorflow/compiler/xla/service/slow_operation_alarm.cc:65] Constant folding an instruction is taking > 1s:

  reduce-window.16 (displaying the full instruction incurs a runtime overhead. Raise your logging level to 4 or above).

This isn't necessarily a bug; constant-folding is inherently a trade-off between compilation time and speed at runtime.  XLA has some guards that attempt to keep constant folding from taking too long, but fundamentally you'll always be able to come up with an input program that takes a long time.

If you'd like to file a bug, run with envvar XLA_FLAGS=--xla_dump_to=/tmp/foo and attach the results.
2022-10-12 16:02:05.867327: E external/org_tensorflow/tensorflow/compiler/xla/service/slow_operation_alarm.cc:133] The operation took 19.277362298s
Constant folding an instruction is taking > 1s:

  reduce-window.16 (displaying the full instruction incurs a runtime overhead. Raise your logging level to 4 or above).

This isn't necessarily a bug; constant-folding is inherently a trade-off between compilation time and speed at runtime.  XLA has some guards that attempt to keep constant folding from taking too long, but fundamentally you'll always be able to come up with an input program that takes a long time.

If you'd like to file a bug, run with envvar XLA_FLAGS=--xla_dump_to=/tmp/foo and attach the results.
Wall time [Compilation]: 21.8s
Wall time [Train step]: 0.00221s
```
with `21.8s` of compilation!

The code:
```python
import haiku as hk
import jax
import contextlib
import timeit

@contextlib.contextmanager
def time_eval(task):
  start = timeit.default_timer()
  try:
    yield
  finally:
    end = timeit.default_timer()
    print(f'Wall time [{task}]: {(end - start):.3}s')

class Model(hk.Module):
    def __call__(self, x):
        x = hk.avg_pool(x, window_shape=(1,3,3,1), strides=(1,2,2,1), padding='SAME')
        x = hk.Conv2D(32, 4, 2)(x)
        x = jax.nn.relu(x)
        return x

def forward(x):
    return Model()(x)

forward = hk.without_apply_rng(hk.transform(forward))

rng = hk.PRNGSequence(jax.random.PRNGKey(42))
x = jax.random.uniform(next(rng), ([128, 512, 512, 3]))

with time_eval('Initialisation'):
    params = jax.jit(forward.init)(next(rng), x)

forward_apply = jax.jit(forward.apply)
with time_eval('Compilation'):
    logits = forward_apply(params, x).block_until_ready()

with time_eval('Train step'):
    logits = forward_apply(params, x).block_until_ready()

```


Discussion:

Resolving patch:
diff --git a/haiku/_src/pool.py b/haiku/_src/pool.py
--- a/haiku/_src/pool.py
+++ b/haiku/_src/pool.py
@@ -140,9 +140,9 @@ def avg_pool(
   else:
     # Count the number of valid entries at each input point, then use that for
     # computing average. Assumes that any two arrays of same shape will be
-    # padded the same.
-    window_counts = lax.reduce_window(jnp.ones_like(value), *reduce_window_args)
-    assert pooled.shape == window_counts.shape
+    # padded the same. Avoid broadcasting on axis where pooling is skipped. 
+    _shape = tuple(vd if wd!=1 else 1 for (vd, wd) in zip(value.shape, window_shape))
+    window_counts = lax.reduce_window(jnp.ones(_shape), *reduce_window_args)
     return pooled / window_counts
 
 


Test patch: