[AMD] Introduce Scaled Upcast Ops (#8088)

knwng · web-flow · commit a25e06d275ed · 2025-09-05T12:42:05.000-07:00
This PR introduced scaled upcast ops including `ScaledUpcastFp4Op` and
`ScaledUpcastFp8Op`

This is one of a series of PRs to decompose scaled dot on AMD backend.
diff --git a/include/triton/Dialect/TritonGPU/IR/TritonGPUOps.td b/include/triton/Dialect/TritonGPU/IR/TritonGPUOps.td
@@ -419,6 +419,14 @@ def TTG_Fp4ToFpOp : TTG_Op<"fp4_to_fp", [Pure]> {
   let arguments = (ins RankedTensorOf<[I8]>:$src, I32Attr:$axis);
   let results = (outs TT_FloatTensor:$result);
 
+  let extraClassDeclaration = [{
+      static LogicalResult verifyFp4ToFp(
+        mlir::Operation *op,
+        RankedTensorType srcTy,
+        RankedTensorType resTy,
+        unsigned axis);
+  }];
+
   let assemblyFormat = [{
     $src attr-dict `:` type($src) `->` type($result)
   }];
diff --git a/lib/Dialect/TritonGPU/IR/Ops.cpp b/lib/Dialect/TritonGPU/IR/Ops.cpp
@@ -376,36 +376,44 @@ void ConvertLayoutOp::getCanonicalizationPatterns(RewritePatternSet &patterns,
 LogicalResult Fp4ToFpOp::verify() {
   auto srcTy = cast<RankedTensorType>(getSrc().getType());
   auto resTy = cast<RankedTensorType>(getResult().getType());
+  auto axis = getAxis();
+
+  auto elemType = resTy.getElementType();
+  if (!(elemType.isBF16() || elemType.isF16()))
+    return emitError() << "only bf16 or f16 is supported for now, got "
+                       << elemType;
+
+  return verifyFp4ToFp(*this, srcTy, resTy, axis);
+}
+
+LogicalResult Fp4ToFpOp::verifyFp4ToFp(mlir::Operation *op,
+                                       RankedTensorType srcTy,
+                                       RankedTensorType resTy, unsigned axis) {
   auto rank = srcTy.getRank();
 
   if (rank != resTy.getRank())
-    return emitError() << "source rank " << rank << " != result rank "
-                       << resTy.getRank();
+    return op->emitError() << "source rank " << rank << " != result rank "
+                           << resTy.getRank();
 
   auto srcShape = srcTy.getShape();
   auto resShape = resTy.getShape();
-  auto axis = getAxis();
 
   if (!(0 <= axis && axis < rank))
-    return emitError() << "axis " << axis << " out of range for rank " << rank;
-
-  auto elemType = resTy.getElementType();
-  if (!(elemType.isBF16() || elemType.isF16()))
-    return emitError() << "only bf16 or f16 is supported for now, got "
-                       << elemType;
+    return op->emitError() << "axis " << axis << " out of range for rank "
+                           << rank;
 
   for (int i = 0; i < rank; ++i) {
     if (i == axis) {
       if (resShape[i] != srcShape[i] * 2)
-        return emitError() << "axis " << axis
-                           << " dimension must be 2x source dimension (src="
-                           << srcShape[i] << ", dst=" << resShape[i] << ")";
+        return op->emitError()
+               << "axis " << axis
+               << " dimension must be 2x source dimension (src=" << srcShape[i]
+               << ", dst=" << resShape[i] << ")";
     } else {
       if (resShape[i] != srcShape[i])
-        return emitError() << "dimension " << i
-                           << " mismatch (src=" << srcShape[i]
-                           << ", dst=" << resShape[i] << ", axis=" << axis
-                           << ")";
+        return op->emitError()
+               << "dimension " << i << " mismatch (src=" << srcShape[i]
+               << ", dst=" << resShape[i] << ", axis=" << axis << ")";
     }
   }
   return success();
diff --git a/third_party/amd/include/Dialect/TritonAMDGPU/IR/TritonAMDGPUOps.td b/third_party/amd/include/Dialect/TritonAMDGPU/IR/TritonAMDGPUOps.td
@@ -519,6 +519,67 @@ def TTG_UpcastMXFPOp : TT_AMDGPU_Op<"upcast_mxfp", [Pure]> {
   }];
 }
 
+//===----------------------------------------------------------------------===//
+// ScaledUpcastFp4Op
+//===----------------------------------------------------------------------===//
+
+def ScaledUpcastFp4Op : TT_AMDGPU_Op<"scaled_upcast_fp4", [Pure]> {
+  let summary = "Upcast fp4 and then multiply scale";
+
+  let description = [{
+    Upcast fp4 (e2m1) values packed as i8 values and multiply with the given
+    E8M0 scale encoded as BF16. This maps to `v_cvt_scalef32_*` intrinsics
+    on the AMD CDNA4 architecture.
+
+    The lower 4 bits of the i8s represent the first fp4 element, and the upper
+    4 bits the second fp4 element.
+
+    The `axis` attribute specifies the axis along which the fp4 elements are
+    packed.
+  }];
+
+  let arguments = (ins
+    RankedTensorOf<[I8]>:$input,
+    RankedTensorOf<[BF16]>:$scale,
+    I32Attr:$axis);
+  let results = (outs RankedTensorOf<[AnyTypeOf<[F16, BF16, F32]>]>:$output);
+
+  let assemblyFormat = [{
+    $input `scale` $scale attr-dict
+        `:` type($input) `,` type($scale) `->` type($output)
+  }];
+
+  let hasVerifier = 1;
+}
+
+//===----------------------------------------------------------------------===//
+// ScaledUpcastFp8Op
+//===----------------------------------------------------------------------===//
+
+def ScaledUpcastFp8Op : TT_AMDGPU_Op<"scaled_upcast_fp8", [
+    Pure,
+    Elementwise,
+    SameOperandsAndResultShape,
+    SameOperandsAndResultEncoding]> {
+  let summary = "Upcast Fp8 and then multiply scale";
+
+  let description = [{
+    Upcast fp8 (e4m3/e5m2) values and multiply with the given E8M0 scale
+    encoded as BF16. This maps to `v_cvt_scalef32_*` intrinsics
+    on the AMD CDNA4 architecture.
+  }];
+
+  let arguments = (ins
+    RankedTensorOf<[AnyTypeOf<[F8E4M3FN, F8E5M2]>]>:$input,
+    RankedTensorOf<[BF16]>:$scale);
+  let results = (outs RankedTensorOf<[AnyTypeOf<[F16, BF16, F32]>]>:$output);
+
+  let assemblyFormat = [{
+    $input `scale` $scale attr-dict
+        `:` type($input) `,` type($scale) `->` type($output)
+  }];
+}
+
 //===----------------------------------------------------------------------===//
 // InThreadTransposeOp
 //===----------------------------------------------------------------------===//
diff --git a/third_party/amd/include/TritonAMDGPUToLLVM/PatternTritonAMDGPUToLLVM.h b/third_party/amd/include/TritonAMDGPUToLLVM/PatternTritonAMDGPUToLLVM.h
@@ -15,6 +15,10 @@ void populateConcatOpToLLVMPatterns(mlir::LLVMTypeConverter &typeConverter,
                                     mlir::RewritePatternSet &patterns,
                                     mlir::PatternBenefit benefit);
 
+void populateScaledUpcastOpToLLVMPatterns(
+    mlir::LLVMTypeConverter &typeConverter, mlir::RewritePatternSet &patterns,
+    mlir::PatternBenefit benefit);
+
 } // namespace mlir::triton::AMD
 
 #endif // TRITON_THIRD_PARTY_AMD_INCLUDE_TRITONAMDGPUTOLLVM_PATTERNTRITONAMDGPUTOLLVM_H_
diff --git a/third_party/amd/lib/Dialect/TritonAMDGPU/IR/Dialect.cpp b/third_party/amd/lib/Dialect/TritonAMDGPU/IR/Dialect.cpp
@@ -22,6 +22,7 @@
  */
 
 #include "triton/Dialect/Triton/IR/Dialect.h"
+#include "mlir/IR/BuiltinTypes.h"
 #include "mlir/IR/DialectImplementation.h"
 #include "mlir/IR/OpImplementation.h"
 #include "third_party/amd/include/Utils/Utility.h"
@@ -418,6 +419,19 @@ InThreadTransposeOp::deduceOutputLayout(ArrayRef<int64_t> shape,
   return transposedLL;
 }
 
+LogicalResult ScaledUpcastFp4Op::verify() {
+  RankedTensorType inputTy = getInput().getType();
+  RankedTensorType outputTy = getOutput().getType();
+  RankedTensorType scaleTy = getScale().getType();
+  auto axis = getAxis();
+
+  if (outputTy.getShape() != scaleTy.getShape())
+    return emitError() << "scale and output should have the same shape";
+
+  // Reuse Fp4ToFpOp's verifier to check types of input and output
+  return triton::gpu::Fp4ToFpOp::verifyFp4ToFp(*this, inputTy, outputTy, axis);
+}
+
 LogicalResult ConcatOp::verify() {
   auto sources = getSources();
   auto result = getResult();
diff --git a/third_party/amd/lib/TritonAMDGPUDialectToLLVM/CMakeLists.txt b/third_party/amd/lib/TritonAMDGPUDialectToLLVM/CMakeLists.txt
@@ -4,6 +4,7 @@ add_triton_library(TritonAMDGPUDialectToLLVM
     InThreadTransposeOpToTTG.cpp
     ConcatOpToLLVM.cpp
     Utility.cpp
+    ScaledUpcastToLLVM.cpp
 
     DEPENDS
     TritonAMDGPUIR
diff --git a/third_party/amd/lib/TritonAMDGPUDialectToLLVM/ScaledUpcastToLLVM.cpp b/third_party/amd/lib/TritonAMDGPUDialectToLLVM/ScaledUpcastToLLVM.cpp
@@ -0,0 +1,111 @@
+#include "Dialect/TritonAMDGPU/IR/Dialect.h"
+#include "TritonAMDGPUToLLVM/PatternTritonAMDGPUToLLVM.h"
+#include "mlir/Conversion/LLVMCommon/Pattern.h"
+#include "mlir/Dialect/LLVMIR/ROCDLDialect.h"
+#include "third_party/amd/lib/TritonAMDGPUToLLVM/Utility.h"
+#include "triton/Conversion/TritonGPUToLLVM/Utility.h"
+
+using namespace mlir;
+using namespace mlir::triton;
+using mlir::LLVM::AMD::upcast4xMxfp8_HW;
+using mlir::LLVM::AMD::upcast8xMxfp4_HW;
+
+namespace {
+struct ScaledUpcastFp4OpPattern
+    : ConvertOpToLLVMPattern<amdgpu::ScaledUpcastFp4Op> {
+  using ConvertOpToLLVMPattern::ConvertOpToLLVMPattern;
+
+  LogicalResult
+  matchAndRewrite(amdgpu::ScaledUpcastFp4Op upcastOp, OpAdaptor adaptor,
+                  ConversionPatternRewriter &rewriter) const override {
+    auto loc = upcastOp.getLoc();
+    auto elemType = upcastOp.getType().getElementType();
+
+    auto inputVals = unpackLLElements(loc, adaptor.getInput(), rewriter);
+    auto scaleVals = unpackLLElements(loc, adaptor.getScale(), rewriter);
+
+    assert(inputVals.size() % 4 == 0);
+    SmallVector<Value> results;
+    results.reserve(inputVals.size() * 2);
+
+    auto b = TritonLLVMOpBuilder(loc, rewriter);
+    for (int i = 0; i < inputVals.size(); i += 4) {
+      SmallVector<Value, 4> v4i32 =
+          elemType.isF16() ? upcast8xMxfp4_HW<ROCDL::CvtScaleF32PkF16Fp4Op>(
+                                 rewriter, loc, inputVals, i, scaleVals[i * 2],
+                                 /*useShiftedScale=*/true)
+                           : upcast8xMxfp4_HW<ROCDL::CvtScaleF32PkBf16Fp4Op>(
+                                 rewriter, loc, inputVals, i, scaleVals[i * 2],
+                                 /*useShiftedScale=*/true);
+      for (int j = 0; j < 4; j++) {
+        Value elements = b.bitcast(v4i32[j], vec_ty(elemType, 2));
+        results.push_back(b.extract_element(elements, b.i32_val(0)));
+        results.push_back(b.extract_element(elements, b.i32_val(1)));
+      }
+    }
+
+    Value result = packLLElements(loc, getTypeConverter(), results, rewriter,
+                                  upcastOp.getType());
+    rewriter.replaceOp(upcastOp, result);
+    return success();
+  }
+};
+
+struct ScaledUpcastFp8OpPattern
+    : ConvertOpToLLVMPattern<amdgpu::ScaledUpcastFp8Op> {
+  using ConvertOpToLLVMPattern::ConvertOpToLLVMPattern;
+
+  LogicalResult
+  matchAndRewrite(amdgpu::ScaledUpcastFp8Op upcastOp, OpAdaptor adaptor,
+                  ConversionPatternRewriter &rewriter) const override {
+    auto loc = upcastOp.getLoc();
+    auto elemType = upcastOp.getType().getElementType();
+    auto fp8ElemType = upcastOp.getInput().getType().getElementType();
+
+    auto inputVals = unpackLLElements(loc, adaptor.getInput(), rewriter);
+    auto scaleVals = unpackLLElements(loc, adaptor.getScale(), rewriter);
+
+    assert(inputVals.size() % 4 == 0);
+    assert(inputVals.size() == scaleVals.size());
+    SmallVector<Value> results;
+    results.reserve(inputVals.size());
+
+    auto b = TritonLLVMOpBuilder(loc, rewriter);
+    for (int i = 0; i < inputVals.size(); i += 4) {
+      SmallVector<Value, 2> v2i32 =
+          elemType.isF16()
+              ? (isa<Float8E4M3FNType>(fp8ElemType)
+                     ? upcast4xMxfp8_HW<ROCDL::CvtScaleF32PkF16Fp8Op>(
+                           rewriter, loc, inputVals, i, scaleVals[i],
+                           /*useShiftedScale=*/true)
+                     : upcast4xMxfp8_HW<ROCDL::CvtScaleF32PkF16Bf8Op>(
+                           rewriter, loc, inputVals, i, scaleVals[i],
+                           /*useShiftedScale=*/true))
+              : (isa<Float8E4M3FNType>(fp8ElemType)
+                     ? upcast4xMxfp8_HW<ROCDL::CvtScaleF32PkBf16Fp8Op>(
+                           rewriter, loc, inputVals, i, scaleVals[i],
+                           /*useShiftedScale=*/true)
+                     : upcast4xMxfp8_HW<ROCDL::CvtScaleF32PkBf16Bf8Op>(
+                           rewriter, loc, inputVals, i, scaleVals[i],
+                           /*useShiftedScale=*/true));
+      for (int j = 0; j < 2; j++) {
+        Value elements = b.bitcast(v2i32[j], vec_ty(elemType, 2));
+        results.push_back(b.extract_element(elements, b.i32_val(0)));
+        results.push_back(b.extract_element(elements, b.i32_val(1)));
+      }
+    }
+
+    Value result = packLLElements(loc, getTypeConverter(), results, rewriter,
+                                  upcastOp.getType());
+    rewriter.replaceOp(upcastOp, result);
+    return success();
+  }
+};
+} // anonymous namespace
+
+void mlir::triton::AMD::populateScaledUpcastOpToLLVMPatterns(
+    LLVMTypeConverter &typeConverter, RewritePatternSet &patterns,
+    PatternBenefit benefit) {
+  patterns.add<ScaledUpcastFp4OpPattern>(typeConverter, benefit);
+  patterns.add<ScaledUpcastFp8OpPattern>(typeConverter, benefit);
+}
diff --git a/third_party/amd/lib/TritonAMDGPUDialectToLLVM/TritonAMDGPUToLLVMPatterns.cpp b/third_party/amd/lib/TritonAMDGPUDialectToLLVM/TritonAMDGPUToLLVMPatterns.cpp
@@ -8,5 +8,6 @@ void populateTritonAMDGPUToLLVMPatterns(LLVMTypeConverter &typeConverter,
   populateExtractSliceOpToLLVMPatterns(typeConverter, patterns, benefit);
   populateInThreadTransposeOpToTTGPatterns(patterns, benefit);
   populateConcatOpToLLVMPatterns(typeConverter, patterns, benefit);
+  populateScaledUpcastOpToLLVMPatterns(typeConverter, patterns, benefit);
 }
 } // namespace mlir::triton::AMD
diff --git a/third_party/amd/lib/TritonAMDGPUToLLVM/Fp4ToFpOpToLLVM.cpp b/third_party/amd/lib/TritonAMDGPUToLLVM/Fp4ToFpOpToLLVM.cpp
@@ -2,18 +2,11 @@
 
 #include "Utility.h"
 #include "mlir/Conversion/LLVMCommon/Pattern.h"
-#include "mlir/Dialect/LLVMIR/LLVMTypes.h"
-#include "mlir/IR/BuiltinOps.h"
 #include "mlir/IR/TypeUtilities.h"
 #include "mlir/IR/ValueRange.h"
 #include "mlir/Transforms/DialectConversion.h"
 #include "triton/Conversion/TritonGPUToLLVM/Utility.h"
-#include "triton/Dialect/Triton/IR/Dialect.h"
-#include "triton/Dialect/TritonGPU/IR/Attributes.h"
-#include "llvm/ADT/STLExtras.h"
 #include "llvm/ADT/SmallVector.h"
-#include "llvm/Support/Debug.h"
-#include <array>
 
 using namespace mlir;
 using namespace mlir::triton;
diff --git a/third_party/amd/lib/TritonAMDGPUToLLVM/TritonGPUToLLVM.cpp b/third_party/amd/lib/TritonAMDGPUToLLVM/TritonGPUToLLVM.cpp
@@ -19,7 +19,6 @@
 #include "third_party/amd/include/Analysis/AxisInfoExt.h"
 #include "third_party/amd/include/Dialect/TritonAMDGPU/IR/Dialect.h"
 #include "triton/Analysis/Allocation.h"
-#include "triton/Analysis/AxisInfo.h"
 #include "triton/Analysis/Membar.h"
 #include "triton/Conversion/TritonGPUToLLVM/PatternTritonGPUOpToLLVM.h"
 #include "triton/Conversion/TritonGPUToLLVM/TypeConverter.h"
diff --git a/third_party/amd/lib/TritonAMDGPUToLLVM/Utility.h b/third_party/amd/lib/TritonAMDGPUToLLVM/Utility.h
diff --git a/third_party/amd/lib/TritonAMDGPUTransforms/AccelerateAMDMatmul.cpp b/third_party/amd/lib/TritonAMDGPUTransforms/AccelerateAMDMatmul.cpp

Original file line number	Diff line number	Diff line change
`@@ -8,5 +8,6 @@ void populateTritonAMDGPUToLLVMPatterns(LLVMTypeConverter &typeConverter,`
`8`	`8`	`populateExtractSliceOpToLLVMPatterns(typeConverter, patterns, benefit);`
`9`	`9`	`populateInThreadTransposeOpToTTGPatterns(patterns, benefit);`
`10`	`10`	`populateConcatOpToLLVMPatterns(typeConverter, patterns, benefit);`
	`11`	`+ populateScaledUpcastOpToLLVMPatterns(typeConverter, patterns, benefit);`
`11`	`12`	`}`
`12`	`13`	`} // namespace mlir::triton::AMD`