WMMA instructions are not supported for GEMM

[joviliast]

Here is highlevel tasks for WMMA enabling for GEMM:

• [x] Implement BlockedToWMMA in TritonAMDGPUAccelerateMatmulPass according to required layout (https://github.com/ROCmSoftwarePlatform/triton/pull/423)
• [x] Support WMMA in ConvertTritonGPUToLLVMPass:
  • [x] Implement conversion from WMMA layout
  • [x] Implement SharedToDotOperandWMMA
  • [x] Implement packMfmaOperand for ElementwiseOpToLLVM
• [x] Support WMMA in TritonGPURemoveLayoutConversionsPass
• [x] Fix bugs and pitfalls
  • [x] Fix Flash attention kernel
  • [x] Fix swizzling (wip)

This issue includes https://github.com/ROCmSoftwarePlatform/triton/issues/250.

[joviliast]

Support WMMA in ConvertTritonGPUToLLVMPass is in progress for now. Here are auxiliary PRs: https://github.com/ROCmSoftwarePlatform/triton/pull/435 https://github.com/ROCmSoftwarePlatform/triton/pull/436

working branch: https://github.com/joviliast/triton/tree/wmma-convertions

[joviliast]

WMMA is partially enabled in draft PR: https://github.com/ROCmSoftwarePlatform/triton/pull/481

[joviliast]

Performance results on working branch using WMMA instructions:

matmul-performance-tflops.txt

Unfortunately, couldn't wait for results without WMMA to compare. fused-attention-tflops.txt

[joviliast]

fused-attention mismatches for now

[joviliast]

All the works are moved to the upstream for now. Please check PRs:

https://github.com/openai/triton/pull/3112 Introduce WMMA layout attr https://github.com/openai/triton/pull/3170 Emit Indices for WMMA https://github.com/openai/triton/pull/3171 Support SharedToDotOperandWMMA https://github.com/openai/triton/pull/3199 Convert WMMA dot op to LLVM

[joviliast]

fused-attention mismatches for now

Fused attention is fixed here https://github.com/joviliast/triton/tree/wmma-convertions

[joviliast]

All required PRs were merged to the upstream: https://github.com/openai/triton Swizzling is not working for now. Fix in progress,