lax.cond is much slower than calling true_fn itself when cond is always true.

[hanzhi713]

Reproducing example:

import os
import timeit
import jax
import jax.sharding
from jax.experimental.compute_on import compute_on
import jax.numpy as jnp

sharding = jax.sharding.SingleDeviceSharding(jax.devices()[0])
p_sharding = jax.sharding.SingleDeviceSharding(jax.devices()[0], memory_kind="pinned_host")


@compute_on("device_host")
@jax.jit
def host_fn(is_valid_step, gradient, opt_state):
    def true_fn(gradient, opt_state):
        opt_state = opt_state + jnp.sin(gradient)
        delta = opt_state * gradient
        return delta, opt_state
    
    def false_fn(gradient, opt_state):
        return gradient, opt_state
    
    # 4.7s per it
    return true_fn(gradient, opt_state)
    # 11.22s per it. Much slower!
    return jax.lax.cond(is_valid_step, true_fn, false_fn, gradient, opt_state)

def test_fn(gradient, opt_state):
    gradient = jnp.cos(gradient)
    return host_fn(jnp.abs(jnp.sum(gradient)) > 0, gradient, opt_state)


x = jnp.arange(0, 1024*1024*100, dtype=jnp.float32)
y = jnp.arange(0, 1024*1024*100, dtype=jnp.float32)
y = jax.device_put(y, p_sharding)

jit_fn = jax.jit(test_fn, in_shardings=(sharding, p_sharding), out_shardings=(sharding, p_sharding), donate_argnums=(0,1))
jit_fn = jit_fn.lower(x, y).compile()


def fn():
    global x, y
    x, y = jit_fn(x, y)
    jax.block_until_ready((x, y))


t = timeit.Timer(fn)
print(t.timeit(10))