How to prune the computations unrelated to the output?

[llan-ml]

I'm just wondering if we can ignore the computations unrelated to the output of a jitted function?
For example, the following function

import jax
import jax.numpy as jnp

def foo():
  x = jnp.ones((10, 2))
  x = x @ x.T
  x = x.sum()
  y = jnp.ones((10, 2)) + 1
  y = y @ y.T
  y = y.sum()
  return x

corresponds to the following jaxpr

{ lambda ; . let
    a:f32[10,2] = broadcast_in_dim[broadcast_dimensions=() shape=(10, 2)] 1.0
    b:f32[2,10] = transpose[permutation=(1, 0)] a
    c:f32[10,10] = dot_general[
      dimension_numbers=(((1,), (0,)), ((), ()))
      precision=None
      preferred_element_type=None
    ] a b
    d:f32[] = reduce_sum[axes=(0, 1)] c
    e:f32[10,2] = broadcast_in_dim[broadcast_dimensions=() shape=(10, 2)] 1.0
    f:f32[10,2] = add e 1.0
    g:f32[2,10] = transpose[permutation=(1, 0)] f
    h:f32[10,10] = dot_general[
      dimension_numbers=(((1,), (0,)), ((), ()))
      precision=None
      preferred_element_type=None
    ] f g
    _:f32[] = reduce_sum[axes=(0, 1)] h
  in (d,) }

The computations about the variable y are also compiled, but they do not make any effect on the output of that function.

[jakevdp]

Thanks for the question! The short answer is: use jit and unnecessary computations will be pruned by XLA as part of the compilation process. Note that because this happens at the XLA level, the dead code elimination (DCE) will not be reflected in the printed jaxpr. There's been some discussion in the past about doing DCE also at the jaxpr lavel; you can read more at #4818.

If you want to see the compiled HLO with the dead code eliminated, you can do something like this:

print(jax.jit(foo).lower().compile().as_text())

%region_0.7 (Arg_0.8: f32[], Arg_1.9: f32[]) -> f32[] {
  %Arg_0.8 = f32[] parameter(0)
  %Arg_1.9 = f32[] parameter(1)
  ROOT %add.10 = f32[] add(f32[] %Arg_0.8, f32[] %Arg_1.9), metadata={op_name="jit(foo)/jit(main)/reduce_sum[axes=(0, 1)]" source_file="<ipython-input-1-c9b7a6784b2a>" source_line=7}
}

ENTRY %main.12 () -> f32[] {
  %constant.1 = f32[] constant(1)
  %broadcast.4 = f32[10,2]{1,0} broadcast(f32[] %constant.1), dimensions={}
  %broadcast.2 = f32[2,10]{1,0} broadcast(f32[] %constant.1), dimensions={}
  %dot.6 = f32[10,10]{1,0} dot(f32[10,2]{1,0} %broadcast.4, f32[2,10]{1,0} %broadcast.2), lhs_contracting_dims={1}, rhs_contracting_dims={0}, metadata={op_name="jit(foo)/jit(main)/dot_general[dimension_numbers=(((1,), (0,)), ((), ())) precision=None preferred_element_type=None]" source_file="<ipython-input-1-c9b7a6784b2a>" source_line=6}
  %constant.5 = f32[] constant(0)
  ROOT %reduce.11 = f32[] reduce(f32[10,10]{1,0} %dot.6, f32[] %constant.5), dimensions={0,1}, to_apply=%region_0.7, metadata={op_name="jit(foo)/jit(main)/reduce_sum[axes=(0, 1)]" source_file="<ipython-input-1-c9b7a6784b2a>" source_line=7}
}

[llan-ml]

Thanks for the explanation and example!