Transform: support sdpa to flash attention kernel conversion

[yifeizh2]

Tracking issue #147.

TODO:

• Check correctness
• Align performance
• Allow tuning for default config

[yifeizh2]

As for performance evaluation, there are two issues

• deep tiled matmul will not be in effect if the parent op of matmul is scf::forall, so we need to either support this scenario or directly invoke the deep tiled matmul in flash attention kernel
• 5x performance gap

<style> </style>

SEQ LENGTH / DTYPE	graph compiler v1 (ms)	v2, block 32, brgemm invoked (ms)	Ratio
384 / fp32	8.28744	48.996	5.912079
768 / fp32	39.1193	190.557	4.871176
1536 / fp32	177.446	752.762	4.242203
2304 / fp32	389.382	1837.45	4.718888
3072 / fp32	682.228	3273.3813	4.798075

Next steps for performance alignment is

• Compare the precise brgemm config used in both cases (v1 v.s. mlir)
• Perform more detailed performance breakdown

[ZhennanQin]

As for performance evaluation, there are two issues

• deep tiled matmul will not be in effect if the parent op of matmul is scf::forall, so we need to either support this scenario or directly invoke the deep tiled matmul in flash attention kernel
• 5x performance gap

<style> </style>

SEQ LENGTH / DTYPE graph compiler v1 (ms) v2, block 32, brgemm invoked (ms) Ratio
384 / fp32 8.28744 48.996 5.912079
768 / fp32 39.1193 190.557 4.871176
1536 / fp32 177.446 752.762 4.242203
2304 / fp32 389.382 1837.45 4.718888
3072 / fp32 682.228 3273.3813 4.798075
Next steps for performance alignment is

• Compare the precise brgemm config used in both cases (v1 v.s. mlir)
• Perform more detailed performance breakdown

Please try brgemm instead of matmul, which can provide better performance result.

[yifeizh2]

As for performance evaluation, there are two issues

• deep tiled matmul will not be in effect if the parent op of matmul is scf::forall, so we need to either support this scenario or directly invoke the deep tiled matmul in flash attention kernel
• 5x performance gap

<style> </style>

SEQ LENGTH / DTYPE graph compiler v1 (ms) v2, block 32, brgemm invoked (ms) Ratio
384 / fp32 8.28744 48.996 5.912079
768 / fp32 39.1193 190.557 4.871176
1536 / fp32 177.446 752.762 4.242203
2304 / fp32 389.382 1837.45 4.718888
3072 / fp32 682.228 3273.3813 4.798075
Next steps for performance alignment is

• Compare the precise brgemm config used in both cases (v1 v.s. mlir)
• Perform more detailed performance breakdown

Please try brgemm instead of matmul, which can provide better performance result.

I dumped the final llvm IR, and verified that the current performance is collected with brgemm invoked. Previously when brgemm was not in effect, the performance is 10x worse. I think I need to do more detailed analysis to find where the performance gap exists exactly.

[yifeizh2]

Latest performance:

<style> </style>

SEQ LENGTH / DTYPE	graph compiler v1 (ms)	v2, block 64, brgemm invoked
384 / fp32	8.28744	22.482	2.71
768 / fp32	39.1193	93.392	2.387
1536 / fp32	177.446	377.7458	2.128
2304 / fp32	389.382	810.249	2.080
3072 / fp32	682.228	1514.491	2.220

[yifeizh2]

Current observed gap from v1 are the following:

• No vectorization
• No fast transpose
• No post op fusion
• linalg.exp takes 1/3 of total execution time; needs to convert to optimized version