[TPU][Pallas]Fix example/cross_entropy.py on Pallas TPU

[yarongmu-google]

The kernels currently has 2 common issues that need support:

1. Long types are not supported in Pallas/Mosaic (XLA does support it but Helion doesn't go through XLA).
2. Directly indexing into vectors on HBM.
   Add CI workflow #2 is the bigger fix here. More about it:

The issue was that evaluating hl.load(logits_flat, [flat_indices]) maps to random reads from HBM, which TensorCores do not support. By changing the Helion code to apply the label == v_chunk_index boolean mask directly across the logits_rows which are sequentially streaming into VMEM, we eliminated the 1D sparse gather entirely.

The updated cross_entropy.py is now verified mathematically correct, fully functional, and autotunes on the TPU v7s on smaller shapes.

After this PR:

=================================================================
Benchmark Results
=================================================================
Implementation       Time (ms)    Speedup        
-----------------------------------------------------------------
helion               0.3826       1.10x          
torch                0.4208       1.00x (ref)    
=================================================================

Note: this PR depends on pytorch/pytorch#180252

[norx1991]

FYI, there is #1950 for the long type part.

[yarongmu-google]

FYI, there is #1950 for the long type part.

Thanks. Any idea why those type mapping classes were not updated? Were those actially not needed?

[norx1991]

FYI, there is #1950 for the long type part.

Thanks. Any idea why those type mapping classes were not updated? Were those actially not needed?

The idea is that if a long type is not really needed on TPU, the user can use the newly added LONG_INT_TYPE. If they are really needed, then casting will not help. Therefore, the type mapping does not need to be updated. We would reject 64 bit data type directly.

[AmesingFlank]

I'd recommend creating a new example instead of modifying the existing one. Ideally, we would make the compiler smart enough so that even for the original example, the compiler could generate masked indexings which works on TPU, so there's value in keeping that example around

[yarongmu-google]

FYI, there is #1950 for the long type part.

Thanks. Any idea why those type mapping classes were not updated? Were those actially not needed?

The idea is that if a long type is not really needed on TPU, the user can use the newly added LONG_INT_TYPE. If they are really needed, then casting will not help. Therefore, the type mapping does not need to be updated. We would reject 64 bit data type directly.

gotcha. Reverted.