[MFMA] Refactor dot pipeline to reduce code duplication (#400)

This PR: - simplifies data types generated by `shared->mfma dot op` layout conversions. Do not pack data types in int32 or int64 - reduce code duplication between fast/normal path - reduce code duplication between operand A and operand B Co-authored-by: Shucai Xiao <[email protected]> Co-authored-by: Lixun Zhang <[email protected]>
ROCm · Dec 13, 2023 · f2afd65 · f2afd65
1 parent 605a90c
commit f2afd65
Show file tree

Hide file tree

Showing 6 changed files with 132 additions and 308 deletions.
diff --git a/include/triton/Dialect/TritonGPU/IR/TritonGPUAttrDefs.td b/include/triton/Dialect/TritonGPU/IR/TritonGPUAttrDefs.td
@@ -855,8 +855,7 @@ section 9.7.13.4.1 for more details.
                                      int bitwidth) const;
 #ifdef USE_ROCM
     SmallVector<int64_t> getMFMAElemsPerInstr() const;
-    SmallVector<int64_t> getMFMARep(ArrayRef<int64_t> operandShape,
-                                    Type elemType) const;
+    SmallVector<int64_t> getMFMARep(ArrayRef<int64_t> operandShape) const;
 #endif
   }];
 }

diff --git a/lib/Conversion/TritonGPUToLLVM/ConvertLayoutOpToLLVM.cpp b/lib/Conversion/TritonGPUToLLVM/ConvertLayoutOpToLLVM.cpp
@@ -884,7 +884,7 @@ struct ConvertLayoutOpConversion
     return success();
   }
 
-  // shared -> mma_operand
+  // shared -> mma_operand/mfma_operand
   LogicalResult
   lowerSharedToDotOperand(triton::gpu::ConvertLayoutOp op, OpAdaptor adaptor,
                           ConversionPatternRewriter &rewriter) const {
@@ -1060,7 +1060,7 @@ struct ConvertLayoutOpConversion
   }
 
 #ifdef USE_ROCM
-  // shared -> dot_operand if the result layout is mma
+  // shared -> dot_operand if the result layout is mfma
   Value lowerSharedToDotOperandMFMA(
       triton::gpu::ConvertLayoutOp op, OpAdaptor adaptor,
       ConversionPatternRewriter &rewriter, const MfmaEncodingAttr &mfmaLayout,