Is compiling lax.scan creating constant folding?

[alonfnt]

Hi, lately I've been seeing the constant folding warning much more frequently and I was wondering what could cause it (code is run on GPU).
The error message:

2022-10-11 08:11:05.780606: E external/org_tensorflow/tensorflow/compiler/xla/service/slow_operation_alarm.cc:65] Constant folding an instruction is taking > 1s:

  dynamic-update-slice.1 (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.

Looking at following example, I am able to get the warning when I jit the outer most function. AFAIK lax.scan should be fair game to compile. However, if instead I compile only the inner function, I do not get the message but the same performance after compilation.

Code with warning message:

@jax.jit
def eden_growth_simulation(key):
    shape = (1000, 1000, 1000)
    acceptance_prob = 0.2
    lattice = jnp.zeros(shape, dtype=jnp.int8).at[(500, 500, 500)].set(1)

    max_nn = 6
    window = [3,3,3]
    strides = [1,1,1]

    def eden_step(carry, key):
        lattice, count = carry
        nn = lax.reduce_window(lattice, jnp.int8(0), jnp.add, window, strides, "SAME") - lattice
        boundary = (nn < max_nn) & (nn > 0) & (lattice == 0)
        growers = jr.uniform(key, shape) < acceptance_prob
        new_lattice = jnp.where(growers & boundary, 1, lattice)
        count += jnp.sum(growers & boundary)
        return (new_lattice, count), None

    keys = jr.split(key, 100)
    (lattice, total_updates), _ = lax.scan(eden_step, init=(lattice, 0), xs=keys)
    return lattice, total_updates

output:

2022-10-11 08:21:57.003922: E external/org_tensorflow/tensorflow/compiler/xla/service/slow_operation_alarm.cc:65] Constant folding an instruction is taking > 1s:

  dynamic-update-slice.1 (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-11 08:22:12.902382: E external/org_tensorflow/tensorflow/compiler/xla/service/slow_operation_alarm.cc:133] The operation took 16.898548854s
Constant folding an instruction is taking > 1s:

  dynamic-update-slice.1 (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.
Simulation took 26.95s to run 2014240 updates.
Simulation took 5.57s to run 2022828 updates.

code without (and same performance)

def eden_growth_simulation(key):
    shape = (1000, 1000, 1000)
    acceptance_prob = 0.2
    lattice = jnp.zeros(shape, dtype=jnp.int8).at[(500, 500, 500)].set(1)

    max_nn = 6
    window = [3,3,3]
    strides = [1,1,1]

    @jax.jit
    def eden_step(carry, key):
        lattice, count = carry
        nn = lax.reduce_window(lattice, jnp.int8(0), jnp.add, window, strides, "SAME") - lattice
        boundary = (nn < max_nn) & (nn > 0) & (lattice == 0)
        growers = jr.uniform(key, shape) < acceptance_prob
        new_lattice = jnp.where(growers & boundary, 1, lattice)
        count += jnp.sum(growers & boundary)
        return (new_lattice, count), None

    keys = jr.split(key, 100)
    (lattice, total_updates), _ = lax.scan(eden_step, init=(lattice, 0), xs=keys)
    return lattice, total_updates

output:

Simulation took 5.88s to run 2014240 updates.
Simulation took 5.64s to run 2022828 updates.

and if you need to run it

if __name__ == "__main__":
    import time

    for i in range(2):
        key = jr.PRNGKey(i)
        start_time = time.time()
        L, count = eden_growth_simulation(key)
        L = L.block_until_ready()
        end_time  = time.time()

        print(f'Simulation took {end_time-start_time:.2f}s to run {count} updates.')

Is there something I should be aware when compiling scan?

[alonfnt]

This issue seems to be related to the dynamic_update_slice part instead of the lax.scan. I've reported the issue #12789 so I can close this.