[WIP] Hessenberg: Use fori_loop() instead of python for-loop

[paolot-gc]

In the Hessenberg transformation, this PR replaces the python for-loop that scan the columns of the matrix with a jax.lax.fori_loop().
Avoids code bloating associated with the loop unrolling of the python for-loop.
In its first version, this PR does not use the "smart indexing" that avoids some matrix processing.
Currently, the code crashes for sizes larger than 128x128.

[paolot-gc]

The crash was due to memory not being aligned properly. I fixed the issue by slapping a MIN_ALIGN on all vertex Vectors (as suggested by @balancap). I compiled and it ran correctly with size 736.
Unfortunately, it appears to take twice as many cycles as the unrolled version, but it was expected anyway, because I removed for the time being the smart indexing.

On the plus side, the compile time is just 55s (vs. 15 minutes) and the memory usage is much smaller than for the unrolled-loop version.

[paolot-gc]

With the use of the gather_p primitive (commit) suggested by @balancap, the cost of the slicing a column of the R matrix and sdiag_full matrix has reduced considerably.
For size 1472, the overall number of cycles for the computation is 37.6e6 (vs. 46.6e6 with direct indexing).
The two charts below illustrate the execution trace for the fori_loop() (implemented by Jax with a while() loop) before (30,676 cycles) and after (24,507).

[paolot-gc]

With the last commit, the mapping to tiles has been changed to allow more than one row per tile.
The memory chart below is for size 1472*3 = 4416, which is probably the largest we can handle.

One iteration takes ~164k cycles, whereas the entire computation ~726e6 cycles