[Linalg][Transform] An Approach for Tiling and Fusing Dependent Reductions in Linalg

[charithaintc]

This is the lighthouse-based implementation of the [Linalg][RFC] An Approach for Tiling and Fusing Dependent Reductions in Linalg.

The proposed structured.fuse_dependent_reduction_op is implemented as lighthouse transform extension. Major difference in the implementation include re-implementing IR utilities that are not available in python bindings (like the scf::tileAndFuseConsumer).

Fused attention example and nanoGPT examples are modified to use the new fusion extension. PR also provides isolated tests for the transform extension inside test/transform/test_fuse_dependant_reduction_ops.py. Implementation can be extended to other intricate usecases as the needs evolve.

Preliminary performance looks encouraging. Fused attention automatic fusion can match the hand-written version. For the time being I did not remove the hand-written fused loop emission path (available via --reference-flash flag)

On PVC:

$ python examples/xegpu/fused_attention.py --check-result 
Result is correct. Proceeding to benchmark...
batch-size=2 n-head=8 n-ctx=4096 d-head=64 dt=f16 time(us): 693.02 GFLOPS: 99159.39

$ python examples/xegpu/fused_attention.py --check-result --reference-flash 
Result is correct. Proceeding to benchmark...
batch-size=2 n-head=8 n-ctx=4096 d-head=64 dt=f16 time(us): 693.92 GFLOPS: 99030.75 

Assisted by: Claude

[charithaintc]

cc @tkarna