Why does SVD no longer determine whether the matrix contains infinite numbers?

[mayuyuace]

b2654c0?diff=split&w=0
Is this a hardware improvement or was it determined in XLA?
If it was a change in XLA, could you tell me the corresponding commits?

[jakevdp]

Hi - thanks for the question! It looks like this change was purely on the JAX side, and not associated with any XLA commits. Pinging for more info @rmlarsen, who authored the change.

[mayuyuace]

I found that NAN is checked at the end of the svd calculation, but before that, it might have returned early here:

jax/jax/_src/lax/svd.py

Line 214 in be6e154

return _svd_tall_and_square_input(

However, the function _svd_tall_and_square_input called here did not check NAN on the input.
I think the same check on input is needed in this function

[rmlarsen]

When using JIT compilation there is no such thing as "returning early", so we decided to get of rid of the special casing to simplify the code and slightly speed up the common case.

[mayuyuace]

Thank you, I understand that although not checking for Nan values ​​increases the overhead when the input contains Nan values, in general this will reduce a certain amount of overhead.

[Zantares]

Hi, is there any exact performance or case available for this optimization? I tried jax v0.4.26 (before opt) and jax v0.4.28 (after opt) with public jax UT svd_test.py on A100, but I got strange performance as below:

Case name	0.4.26 duration	0.4.28 duration
SvdTest.testSvdvals0	2.309 s	2.267 s
SvdTest.testSvdWithRectangularInput0	29.072s	39.460 s

I looped each case 100 times to watch the performance, but there's no performance improvement after opt...

[rmlarsen]

Actually, I don't think the implementation in jax/_src/lax/linalg is used for SVD on GPU. @jakevdp can you verify?

[jakevdp]

The GPU svd lowering is found here:

jax/jax/_src/lax/linalg.py

Lines 2189 to 2192 in ea5fd29

	mlir.register_lowering(
	svd_p, partial(_svd_cpu_gpu_lowering, gpu_solver.cuda_gesvd,
	platform='cuda'),
	platform='cuda')

The changes referenced in this question (in jax/_src/lax/svd.py and jax/_src/lax/qdwh.py) are only used for the TPU lowering rule I believe.

[rmlarsen]

(Just for completeness: jax/_src/lax/eigh.py is also called by jax/_src/lax/svd.py.)

[rmlarsen]

But yeah, changes to the GPU SVD seem unrelated.

[Zantares]

@Zantares what happens if you install JAX v0.4.26 and copy the files in jax/_src/lax/linalg from v0.4.28 there? This should verify that the changes to the SVD code are to blame and not some other unrelated change in the GPU compiler stack.

Thanks for the suggestion, I will try it.

[Zantares]

Hi @rmlarsen , @jakevdp , I did below changes and retested again:

• Revert the svd opt in jax v0.4.28
• Cherry-pick svd opt to jax v0.4.26 and handle the conflict manually

It seems that the svd opt still impacted part of the GPU performance:

Case name	0.4.26 duration	0.4.28 duration	0.4.28 W/O opt duration	0.4.26 W/ opt duration
SvdTest.testSvdWithRectangularInput0	29.072 s	39.460 s	34.407 s	35.394 s

As you mentioned, jax v0.4.28 has other changes which impacted performance, but this svd opt still brought additional ~5 s regression here.

[rmlarsen]

@jakevdp thanks for figuring that out. This code was not optimized for GPU. It probably could be, but that was not the goal.

[rmlarsen]

I would recommend calling the path that gets lowered to the cuSolver SVD, which has been carefully optimized by NVIDIA.

[Zantares]

Got it, now I understand this question completely, thanks!

[Zantares]

One last question, is possible to add TPU device limitation here to fix this issue first?

[jakevdp]

There is no way to hit this code path on anything but TPU via public JAX APIs. The only reference to it is here, in the TPU-specific lowering rule: https://github.com/google/jax/blob/7d2f0a75c1c10aeae63ed1a0ced17b7aeb979989/jax/_src/lax/linalg.py#L2193-L2197

The test file runs it on GPU because it calls the private API directly. I suppose we could mark the test to be skipped unless running on a TPU backend, but that would be a net loss in test coverage, so I'm not sure it's warranted.