Repository: google/jax

Issue:
Redundant computation in expm?
Hi all:

It looks to me there are redundant calculations in the `expm` function, but wanted to ask to see if I'm missing something (perhaps something clever is happening in compilation). 

E.g., for `float64` and `complex128`, the numerator/denominator for the Pade approximant is generated as follows:

```
if A.dtype == 'float64' or A.dtype == 'complex128':
   U3, V3 = _pade3(A)
   U5, V5 = _pade5(A)
   U7, V7 = _pade7(A)
   U9, V9 = _pade9(A)
   maxnorm = 5.371920351148152
   n_squarings = jnp.maximum(0, jnp.floor(jnp.log2(A_L1 / maxnorm)))
   A = A / 2**n_squarings
   U13, V13 = _pade13(A)
   conds=jnp.array([1.495585217958292e-002, 2.539398330063230e-001,
                    9.504178996162932e-001, 2.097847961257068e+000])
   U = jnp.select((A_L1<conds), (U3, U5, U7, U9), U13)
   V = jnp.select((A_L1<conds), (V3, V5, V7, V9), V13)
```

The algorithm being used (I believe by Al-Mohy and Higham, same as in `scipy`) prescribes choosing between 1 of 5 different Pade approximations (of various orders) based on the norm of `A`. In the above, it appears all potential Pade approximations are being computed, and only afterwards is one chosen based on the norm of `A`. Hence, the computations for the others are discarded and ultimately unnecessary.

It seems this function could benefit from changing the above logic to use nested `lax.cond` statements (or perhaps `switch`?) to avoid computing all possible versions. To motivate this: I quickly tried modifying the `expm` function to simply always do the `U13, V13` version of of the computation (the most expensive) and not compute any of the lower order versions, and observed roughly a factor of 2 speed increase in compiled execution. My speed tests weren't rigorous - I didn't test it for a variety of norms - but I believe this should speed things up in all cases (again, unless I'm missing something with lax behaviour). Edit: This is on CPU.

Thoughts?

P.s. In some of my use cases of jax, the most expensive piece is a ton of matrix exponentials, so I'm very interested in potential speed gains for this function :D.


Discussion:
This sounds right to me! Indeed using `jnp.select` is wasting compute, and `cond`/`switch` could save it. I don't remember any context around why `expm` was written this way (@zhangqiaorjc do you remember?) but it may just have been for simplicity. In that case, what you describe would be a welcome upgrade!
I guess that on GPU it _might_ be better to compute everything rather than pass control flow back to the host. But that's really an optimization problem for XLA GPU, not something we should worry about in JAX.

In general, matrix-matrix multiplication (L3 BLAS) seems like the level of computation where we should be using `cond` rather than `where`.
Does `cond` always pass control flow back to the host on GPU? That sounds plausible, though also it seems like at least in some cases XLA:GPU could avoid that. (At the very least, it could decide to optimize a `cond` into a `select` if needed!) On TPU (where control flow can stay on the device) and CPU (where there's no host/device separation) it seems profitable.

Good rule of thumb!
If there are no objections, given how small this is, and my personal interest in seeing it go through asap, I'm willing to implement this change. My plan would be to just change it to use nested `cond` statements. In a PR I can show some local CPU benchmarks.
@DanPuzzuoli Feel free to send in a PR. I don't think anyone from JAX team is currently working on this. Thanks!

Resolving patch:
diff --git a/jax/_src/scipy/linalg.py b/jax/_src/scipy/linalg.py
--- a/jax/_src/scipy/linalg.py
+++ b/jax/_src/scipy/linalg.py
@@ -289,28 +289,20 @@ def _calc_P_Q(A):
   A_L1 = np_linalg.norm(A,1)
   n_squarings = 0
   if A.dtype == 'float64' or A.dtype == 'complex128':
-   U3, V3 = _pade3(A)
-   U5, V5 = _pade5(A)
-   U7, V7 = _pade7(A)
-   U9, V9 = _pade9(A)
    maxnorm = 5.371920351148152
    n_squarings = jnp.maximum(0, jnp.floor(jnp.log2(A_L1 / maxnorm)))
    A = A / 2**n_squarings
-   U13, V13 = _pade13(A)
    conds = jnp.array([1.495585217958292e-002, 2.539398330063230e-001,
                       9.504178996162932e-001, 2.097847961257068e+000])
-   U = jnp.select((A_L1 < conds), (U3, U5, U7, U9), U13)
-   V = jnp.select((A_L1 < conds), (V3, V5, V7, V9), V13)
+   idx = jnp.digitize(A_L1, conds)
+   U, V = lax.switch(idx, [_pade3, _pade5, _pade7, _pade9, _pade13], A)
   elif A.dtype == 'float32' or A.dtype == 'complex64':
-    U3,V3 = _pade3(A)
-    U5,V5 = _pade5(A)
     maxnorm = 3.925724783138660
     n_squarings = jnp.maximum(0, jnp.floor(jnp.log2(A_L1 / maxnorm)))
     A = A / 2**n_squarings
-    U7,V7 = _pade7(A)
     conds = jnp.array([4.258730016922831e-001, 1.880152677804762e+000])
-    U = jnp.select((A_L1 < conds), (U3, U5), U7)
-    V = jnp.select((A_L1 < conds), (V3, V5), V7)
+    idx = jnp.digitize(A_L1, conds)
+    U, V = lax.switch(idx, [_pade3, _pade5, _pade7], A)
   else:
     raise TypeError("A.dtype={} is not supported.".format(A.dtype))
   P = U + V  # p_m(A) : numerator


Test patch: