A case which would benefit from dynamic shape compilation

[oliverdutton]

I was looking at use cases popping foldcomp on GPU by combining with nerfax but compilation times kills it from being competitive due to slightly different length arrays everywhere.

In meandering towards this I have code which does run on CPU, and has the potential to be competitive with the (in my opinion) more complex C++ code if compiled.

A colab showing the exact function and variation in shapes to compile is here.

I discussed this in iree-org/iree-jax#57 where iree seems like the project soonest to make it possible. @jpienaar @mattjj - I promised to share this as a general toy problem about 2 months ago (sorry for the ping-ing)

[jpienaar]

Hey, sorry I missed this while I was out and just noticed. Thanks for the example! @rsuderman we should dump this one a bit and see if can handle it

[oliverdutton]

Amazing, unfortunately one of the two dependencies, mdtraj, seems to fail installing by pip (even with older versions) as cython has changed, and pinning it doesn't seem to be working to fix it.

I'll track the issue and re-ping when the colab actually works

[oliverdutton]

The colab example now runs happily again (the cython version issues in the dependency is fixed)

As you can see

138 residues, 17.05 seconds with compilation, 0.036 Angstrom RMSD, 1.33 ms when compiled (with some compile time eval)
360 residues, 10.54 seconds with compilation, 0.030 Angstrom RMSD, 1.05 ms when compiled (with some compile time eval)
360 residues, 11.29 seconds with compilation, 0.030 Angstrom RMSD, 1.09 ms when compiled (with some compile time eval)
333 residues, 9.14 seconds with compilation, 0.033 Angstrom RMSD, 1.40 ms when compiled (with some compile time eval)

Compilation times are high and the program is run with many different shapes. Things could be padded, but then you pay the run time costs of excess computation and the time to do the padding itself.