What does jax.jit knows about arithmetic?

[mariogeiger]

Hi all,

I love how jax.jit is able to boost the performance of my code and I was wondering which kind of optimization it was able to do under the hood.

Playing with printing optimized code by jit I observed that:

• x + 0 and 1 * x are both simplified into x
• multiple calls to the same function with the same arguments are replaced by only one call
• but x + y and y + x are considered different (same for *)
• and also x - x is not replaced by 0

code:

def f(x, y):
    return x + 0.0, 1.0 * y, jnp.tanh(x), jnp.tanh(x), (x * y), (y * x), (y - y)

ENTRY (parameter: f32[], parameter: f32[]) -> (f32[], f32[], f32[], f32[], f32[], /*index=5*/f32[], f32[]) {
  tmp_0 = parameter(0)
  tmp_1 = copy(tmp_0)  # x + 0.0
  tmp_2 = parameter(1)
  tmp_3 = copy(tmp_2)  # 1.0 * y
  tmp_4 = tanh(tmp_0)
  tmp_5 = copy(tmp_4)  # only one call to tanh
  tmp_6 = multiply(tmp_0, tmp_2)
  tmp_7 = multiply(tmp_2, tmp_0)  # no copy here :cry: 
  tmp_8 = subtract(tmp_2, tmp_2)  # no zeros here :cry: 
  ROOT tmp_9 = tuple(tmp_1, tmp_3, tmp_4, tmp_5, ...(+3))
}

Are those simplifications still performed somehow at another stage?
If no, is there good reasons to not perform them?

(for instance I know that floating point rounding is affected by the operations order but the same happen with integers)

[jakevdp]

I don't know of any documentation that summarizes the algebraic simplification rules implemented by XLA, but you can see it in the source code. For example, I believe this is the relevant implementation for XLA on CPU: https://github.com/tensorflow/tensorflow/blob/master/tensorflow/compiler/xla/service/algebraic_simplifier.cc

The header file may be easier to skim through, and has some useful comments: https://github.com/tensorflow/tensorflow/blob/master/tensorflow/compiler/xla/service/algebraic_simplifier.h

For example, here we see that A - A is simplified to 0 only for integer A.

[mariogeiger]

Thanks a lot, these source files are very interesting.
I see that somehow advanced simplifications are performed.

For instance this one that I can see it in action with jax.jit 🎉

def f(A, B, C, D):
    return (A / B) / (C / D)

ENTRY {
  tmp_0 = parameter(1)
  tmp_1 = parameter(2)
  tmp_2 = parameter(0)
  tmp_3 = parameter(3)
  tmp_4 = fusion(tmp_0, tmp_1, tmp_2, tmp_3), kind=kLoop, calls=
  {
    tmp_0 = parameter(2)
    tmp_1 = parameter(3)
    tmp_2 = multiply(tmp_0, tmp_1)
    tmp_3 = parameter(0)
    tmp_4 = parameter(1)
    tmp_5 = multiply(tmp_3, tmp_4)
    ROOT tmp_6 = divide(tmp_2, tmp_5)
  }
  ROOT tmp_5 = tuple(tmp_4)
}

Since I see that advanced simplification are done, I'm even more wondering why x+y and y+x are both computed 🤷‍♂️.

Also A - A = 0 for integers exists but does not happens with jax.jit

def f(A):
    return A - A

ENTRY (parameter: s32[2]) -> (s32[2]) {
  tmp_0 = parameter(0)
  tmp_1 = subtract(tmp_0, tmp_0)
  ROOT tmp_2 = tuple(tmp_1)
}

[jakevdp]

It looks like that A - A -> 0 simplification is only about a week old, and hasn't made it into a release yet: tensorflow/tensorflow@8eb46e2

[mariogeiger]

Now it works :)

def f(A):
    return A - A

tmp_0 = parameter(0)
tmp_1 = constant(0)
tmp_2 = broadcast(tmp_1), dimensions={}
ROOT tmp_3 = tuple(tmp_2)