Add support for conversion fp16 to fp32 using ExtFOp

triton-lang · May 11, 2024 · 5d76f19 · 5d76f19
1 parent 161f7a4
commit 5d76f19
Show file tree

Hide file tree

Showing 2 changed files with 12 additions and 4 deletions.
diff --git a/lib/Dialect/TritonGPU/Transforms/AccelerateMatmul.cpp b/lib/Dialect/TritonGPU/Transforms/AccelerateMatmul.cpp
@@ -339,8 +339,12 @@ class BlockedToMMA : public mlir::RewritePattern {
 
 static Value promoteOperand(OpBuilder &builder, Location loc, Value operand,
                             Type promotedType) {
-  Type tensorPromotedType = cast<RankedTensorType>(operand.getType())
-                                .cloneWith(std::nullopt, promotedType);
+  RankedTensorType tensor = cast<RankedTensorType>(operand.getType());
+  Type tensorElementType = tensor.getElementType();
+  Type tensorPromotedType = tensor.cloneWith(std::nullopt, promotedType);
+  if (tensorElementType.isF16() && promotedType.isF32()) {
+    return builder.create<arith::ExtFOp>(loc, tensorPromotedType, operand);
+  }
   return builder.create<tt::FpToFpOp>(loc, tensorPromotedType, operand);
 }
 

diff --git a/third_party/nvidia/lib/TritonNVIDIAGPUToLLVM/TritonGPUToLLVM.cpp b/third_party/nvidia/lib/TritonNVIDIAGPUToLLVM/TritonGPUToLLVM.cpp
@@ -207,8 +207,12 @@ struct ConvertTritonGPUToLLVM
 
   static Value promoteOperand(OpBuilder &builder, Location loc, Value operand,
                               Type promotedType) {
-    Type tensorPromotedType = cast<RankedTensorType>(operand.getType())
-                                  .cloneWith(std::nullopt, promotedType);
+    RankedTensorType tensor = cast<RankedTensorType>(operand.getType());
+    Type tensorElementType = tensor.getElementType();
+    Type tensorPromotedType = tensor.cloneWith(std::nullopt, promotedType);
+    if (tensorElementType.isF16() && promotedType.isF32()) {
+      return builder.create<arith::ExtFOp>(loc, tensorPromotedType, operand);
+    }
     return builder.create<triton::FpToFpOp>(loc, tensorPromotedType, operand);
   }
 };