Lowering Fp4ToFP to LLVM (#3607)

leonling-ll · web-flow · commit 580fda2b069e · 2025-03-14T16:06:43.000+08:00
This PR is splitted as the second part of #3538. It lowers `Fp4ToFPOp` to LLVM and remove `UpcastMXFPOp`. CI depends on #3606.
diff --git a/include/triton/Tools/Sys/GetEnv.hpp b/include/triton/Tools/Sys/GetEnv.hpp
@@ -49,7 +49,6 @@ inline const std::set<std::string> CACHE_INVALIDATING_ENV_VARS = {
     "TRITON_INTEL_ENABLE_INSTR_SCHED",
     "TRITON_INTEL_RAISE_BLOCK_POINTER",
     "TRITON_INTEL_REDUCE_TRANSPOSE",
-    "TRITON_INTEL_DECOMPOSE_SCALED_BLOCKED",
     // clang-format on
 };
 
diff --git a/third_party/intel/include/Dialect/TritonIntelGPU/IR/TritonIntelGPUOps.td b/third_party/intel/include/Dialect/TritonIntelGPU/IR/TritonIntelGPUOps.td
@@ -202,33 +202,4 @@ def TTIG_SubGroupTransposeOp
   let hasVerifier = 1;
 }
 
-// The same as ttg.upcast_mxfp, but we want Dot Layout from Dpas layout for input tensor
-def TTIG_UpcastMXFPOp : TTIG_Op<"upcast_mxfp", [Pure]> {
-  let summary = "Convert an mxfp tensor to bf16/fp16";
-
-  let hasVerifier = 1;
-
-  let description = [{
-    Compute the bf16 encoded in the given mxfp number as per
-    https://www.opencompute.org/documents/ocp-microscaling-formats-mx-v1-0-spec-final-pdf
-  }];
-  let arguments = (
-    ins
-    TT_Tensor:$src,
-    TT_Tensor:$scale,
-    TT_ScaleDotElemTypeAttr:$fp_type,
-    BoolAttr:$fastMath
-  );
-  let results = (outs TT_Tensor:$result);
-
-  let assemblyFormat = [{
-    $src `,` $scale  `fp_type` `=` $fp_type attr-dict `:` type($src) `,` type($scale) `->` type($result)
-  }];
-
-  let extraClassDeclaration = [{
-    static RankedTensorType deduceOutputType(
-        TypedValue<RankedTensorType> inputTensor, ScaleDotElemType inputElemType, Type outputElemType);
-  }];
-}
-
 #endif
diff --git a/third_party/intel/lib/Dialect/TritonIntelGPU/IR/Ops.cpp b/third_party/intel/lib/Dialect/TritonIntelGPU/IR/Ops.cpp
@@ -209,129 +209,4 @@ LogicalResult SubGroupTransposeOp::verify() {
   return success();
 }
 
-LogicalResult UpcastMXFPOp::verify() {
-  auto fpType = getFpType();
-
-  auto xTy = getSrc().getType();
-  auto scaleTy = getScale().getType();
-  Builder b(getContext());
-  if (xTy.getElementType() != b.getBF16Type() &&
-      xTy.getElementType() != b.getF16Type() &&
-      xTy.getElementType() != b.getI8Type()) {
-    return emitOpError(
-        "element type of the first operand must be bf16/fp16 or i8");
-  }
-
-  if (scaleTy.getElementType() != b.getI8Type()) {
-    return emitOpError("element type of the second operand must be uint8");
-  }
-
-  auto xShape = xTy.getShape();
-  auto scaleShape = scaleTy.getShape();
-
-  if (xShape.size() != scaleShape.size() || xShape.size() < 2) {
-    return emitOpError(
-        "operands must have the same number of dimensions, at least 2");
-  }
-
-  if (!(fpType == ScaleDotElemType::E2M1 || fpType == ScaleDotElemType::E4M3 ||
-        fpType == ScaleDotElemType::E5M2)) {
-    return emitOpError("NYI: fpType must be E2M1, E4M3, or E5M2");
-  }
-
-  auto layoutX = xTy.getEncoding();
-  auto layoutScale = scaleTy.getEncoding();
-  if (bool(layoutX) != bool(layoutScale)) {
-    return emitOpError(
-        "Expected either both or neither operands to have an encoding");
-  }
-  // Nothing to check if no encoding. This is used to infer the return type in
-  // AccelerateMatmul.cpp
-  if (!layoutX) {
-    return success();
-  }
-
-  auto dotEncoding = dyn_cast<DotOperandEncodingAttr>(layoutX);
-  if (!dotEncoding) {
-    return emitOpError("Expected a DotOperandEncodingAttr for values");
-  }
-  if (!isa<BlockedEncodingAttr, LinearEncodingAttr>(layoutScale)) {
-    return emitOpError(
-        "Expected a BlockOperandEncoding or LinearOperandEncoding "
-        "for scales");
-  }
-
-  // Change to support fp8 types
-  const auto elemsPacked = fpType == ScaleDotElemType::E2M1 ? 2 : 1;
-  // Figure out the K dimension for the input A/B. For A/B scale, the K
-  // dimension is always the last dimension.
-  const int opIdx = dotEncoding.getOpIdx();
-  const bool hasBatch = xShape.size() == 3;
-  const int kIdx = (opIdx == 0 ? 1 : 0) + hasBatch;
-
-  if (xShape[kIdx] != (32 / elemsPacked) * scaleShape.back()) {
-    return emitOpError("K dimension of first operand must be 16 times "
-                       "larger than last/K dimension of the second operand");
-  }
-
-  // Check other dimensions match too. For input A/B, we need to figure out the
-  // index for the M/N dimension. For scale, it's always {(batch), M/N, K}.
-  const int mnIdx = (opIdx == 0 ? 0 : 1) + hasBatch;
-  if (hasBatch && xShape[0] != scaleShape[0])
-    return emitOpError("batch dimension must match between operands");
-  if (xShape[mnIdx] != scaleShape[hasBatch]) {
-    return emitOpError("M/N dimension must match between operands");
-  }
-
-  return success();
-}
-
-RankedTensorType
-UpcastMXFPOp::deduceOutputType(TypedValue<RankedTensorType> inputTensor,
-                               ScaleDotElemType inputElemType,
-                               Type outputElemType) {
-  MLIRContext *ctx = inputTensor.getContext();
-  auto xTy = inputTensor.getType();
-  if (inputElemType != ScaleDotElemType::E2M1)
-    return xTy;
-
-  auto xShape = xTy.getShape();
-  auto newShape = llvm::to_vector(xShape);
-  auto encoding = xTy.getEncoding();
-  if (!encoding) {
-    newShape.back() *= 2;
-    return RankedTensorType::get(xShape, outputElemType);
-  }
-
-  auto oldEncoding = cast<DotOperandEncodingAttr>(encoding);
-  const int opIdx = oldEncoding.getOpIdx();
-  // Note: For Intel the dot operands layout's kWidth parameter must match
-  // the parent's DPAS layout opsPerChannel so we need to materialize a
-  // new DPAS layout.
-  auto dpasEncoding = cast<intel::DpasEncodingAttr>(oldEncoding.getParent());
-  unsigned opsPerChannel =
-      intel::DpasEncodingAttr::getOpsPerChannel(outputElemType);
-  // e2m1 is packed 2 elements per int8, we must handle continuous 2
-  // elements when upcasting to bf16
-  if (xTy.getElementType() == IntegerType::get(ctx, 8))
-    opsPerChannel *= 2;
-  auto newDpasEncoding = intel::DpasEncodingAttr::get(
-      ctx, dpasEncoding.getRepeatCount(), dpasEncoding.getSystolicDepth(),
-      dpasEncoding.getExecutionSize(), opsPerChannel,
-      dpasEncoding.getWarpsPerCTA(), dpasEncoding.getRepCluster(),
-      product<unsigned>(dpasEncoding.getThreadsPerWarp()));
-
-  // Operand A is packed to i16 for scalar type < 16 bits.
-  int kWidth =
-      (opIdx == 0) && (opsPerChannel != 1) ? opsPerChannel / 2 : opsPerChannel;
-
-  Attribute newVEncoding =
-      DotOperandEncodingAttr::get(ctx, opIdx, newDpasEncoding, kWidth);
-
-  const bool hasBatch = xShape.size() == 3;
-  const int kIdx = (opIdx == 0 ? 1 : 0) + hasBatch;
-  newShape[kIdx] *= 2;
-  return RankedTensorType::get(newShape, outputElemType, newVEncoding);
-}
-
 } // namespace mlir::triton::gpu::intel
diff --git a/third_party/intel/lib/TritonIntelGPUToLLVM/CMakeLists.txt b/third_party/intel/lib/TritonIntelGPUToLLVM/CMakeLists.txt
@@ -8,6 +8,7 @@ add_triton_library(TritonIntelGPUToLLVM
     DotOpToLLVM/FMA.cpp
     DotOpToLLVM.cpp
     ElementwiseOpToLLVM.cpp
+    Fp4ToFpOpToLLVM.cpp
     HistogramOpToLLVM.cpp
     LoadStoreOpToLLVM.cpp
     MakeRangeOpToLLVM.cpp
@@ -20,7 +21,6 @@ add_triton_library(TritonIntelGPUToLLVM
     TritonGPUToLLVM.cpp
     TritonOpsToLLVM.cpp
     TypeConverter.cpp
-    UpcastMXFPToLLVM.cpp
     Utility.cpp
 
     DEPENDS
diff --git a/third_party/intel/lib/TritonIntelGPUToLLVM/Fp4ToFpOpToLLVM.cpp b/third_party/intel/lib/TritonIntelGPUToLLVM/Fp4ToFpOpToLLVM.cpp
@@ -0,0 +1,124 @@
+#include "PatternTritonGPUOpToLLVM.h"
+
+#include "mlir/Conversion/LLVMCommon/Pattern.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 "llvm/ADT/STLExtras.h"
+#include "llvm/ADT/SmallVector.h"
+
+using namespace mlir;
+using namespace mlir::triton;
+using namespace mlir::triton::gpu;
+using namespace mlir::triton::gpu::intel;
+
+namespace {
+SmallVector<Value> convertMxfp4x2ToBf16x2(RewriterBase &rewriter, Location loc,
+                                          ArrayRef<Value> values) {
+  auto b = TritonLLVMOpBuilder(loc, rewriter);
+  SmallVector<Value> results;
+  for (auto v : values) {
+    auto em0 = b.and_(v, b.i8_val(0x7));
+    auto em1 = b.and_(v, b.i8_val(0x70));
+    Value v0 =
+        b.or_(b.shl(b.zext(i16_ty, em0), b.i16_val(6)),
+              b.shl(b.zext(i16_ty, b.and_(v, b.i8_val(0x8))), b.i16_val(12)));
+    Value v1 =
+        b.or_(b.shl(b.zext(i16_ty, em1), b.i16_val(2)),
+              b.shl(b.zext(i16_ty, b.and_(v, b.i8_val(0x80))), b.i16_val(8)));
+    // Three cases:
+    // 1) x is normal and non-zero: Correct bias
+    v0 = b.select(b.icmp_ne(b.and_(em0, b.i8_val(0x6)), b.i8_val(0)),
+                  b.add(v0, b.i16_val((127 - 1) << 7)), v0);
+    v1 = b.select(b.icmp_ne(b.and_(em1, b.i8_val(0x60)), b.i8_val(0)),
+                  b.add(v1, b.i16_val((127 - 1) << 7)), v1);
+    // 2) x is subnormal (x == 0bs001 where s is the sign): Map to +-0.5 in
+    // bf16
+    v0 = b.bitcast(
+        b.select(b.icmp_eq(em0, b.i8_val(0x1)),
+                 b.or_(b.i16_val(16128), b.and_(v0, b.i16_val(0x8000))), v0),
+        bf16_ty);
+    v1 = b.bitcast(
+        b.select(b.icmp_eq(em1, b.i8_val(0x10)),
+                 b.or_(b.i16_val(16128), b.and_(v1, b.i16_val(0x8000))), v1),
+        bf16_ty);
+    // 3) x is zero, nothing to do
+    results.push_back(v0);
+    results.push_back(v1);
+  }
+  return results;
+}
+
+SmallVector<Value> convertMxfp4x2ToFp16x2(RewriterBase &rewriter, Location loc,
+                                          ArrayRef<Value> values) {
+  auto b = TritonLLVMOpBuilder(loc, rewriter);
+  SmallVector<Value> results;
+  for (auto v : values) {
+    auto em0 = b.and_(v, b.i8_val(0x7));
+    auto em1 = b.and_(v, b.i8_val(0x70));
+    // FP16 bits: sign = 1, exponent = 5, mantissa = 10
+    Value v0 =
+        b.or_(b.shl(b.zext(i16_ty, em0), b.i16_val(10 - 1)),
+              b.shl(b.zext(i16_ty, b.and_(v, b.i8_val(0x8))), b.i16_val(12)));
+    Value v1 =
+        b.or_(b.shl(b.zext(i16_ty, em1), b.i16_val(10 - 1 - 4)),
+              b.shl(b.zext(i16_ty, b.and_(v, b.i8_val(0x80))), b.i16_val(8)));
+
+    // Three cases:
+    // 1) x is normal and non-zero: Correct bias
+    v0 = b.select(b.icmp_ne(b.and_(em0, b.i8_val(0x6)), b.i8_val(0)),
+                  b.add(v0, b.i16_val((15 - 1) << 10)), v0);
+    v1 = b.select(b.icmp_ne(b.and_(em1, b.i8_val(0x60)), b.i8_val(0)),
+                  b.add(v1, b.i16_val((15 - 1) << 10)), v1);
+
+    // 2) x is subnormal (x == 0bs001 where s is the sign): Map to fp16 +-0.5
+    v0 = b.bitcast(
+        b.select(b.icmp_eq(em0, b.i8_val(0x1)),
+                 b.or_(b.i16_val(0x3800), b.and_(v0, b.i16_val(0x8000))), v0),
+        f16_ty);
+    v1 = b.bitcast(
+        b.select(b.icmp_eq(em1, b.i8_val(0x10)),
+                 b.or_(b.i16_val(0x3800), b.and_(v1, b.i16_val(0x8000))), v1),
+        f16_ty);
+    // 3) x is zero, nothing to do
+    results.push_back(v0);
+    results.push_back(v1);
+  }
+  return results;
+}
+
+class Fp4ToFpOpPattern : public ConvertOpToLLVMPattern<Fp4ToFpOp> {
+public:
+  Fp4ToFpOpPattern(LLVMTypeConverter &typeConverter, PatternBenefit benefit)
+      : ConvertOpToLLVMPattern<Fp4ToFpOp>(typeConverter, benefit) {}
+
+  LogicalResult
+  matchAndRewrite(Fp4ToFpOp op, OpAdaptor adaptor,
+                  ConversionPatternRewriter &rewriter) const override {
+    Location loc = op.getLoc();
+    auto *ctx = op.getContext();
+    Type elemType = op.getType().getElementType();
+    assert(elemType == f16_ty || elemType == bf16_ty);
+    bool toFp16 = elemType == f16_ty;
+
+    SmallVector<Value> xVals =
+        unpackLLElements(loc, adaptor.getSrc(), rewriter);
+    xVals = toFp16 ? convertMxfp4x2ToFp16x2(rewriter, loc, xVals)
+                   : convertMxfp4x2ToBf16x2(rewriter, loc, xVals);
+
+    Value result =
+        packLLElements(loc, getTypeConverter(), xVals, rewriter, op.getType());
+    rewriter.replaceOp(op, result);
+    return success();
+  }
+};
+} // anonymous namespace
+
+void mlir::triton::intel::populateFp4ToFpToLLVMPatterns(
+    LLVMTypeConverter &typeConverter, RewritePatternSet &patterns,
+    PatternBenefit benefit) {
+  patterns.add<Fp4ToFpOpPattern>(typeConverter, benefit);
+}
diff --git a/third_party/intel/lib/TritonIntelGPUToLLVM/PatternTritonGPUOpToLLVM.h b/third_party/intel/lib/TritonIntelGPUToLLVM/PatternTritonGPUOpToLLVM.h
@@ -38,10 +38,9 @@ void populateElementwiseOpToLLVMPatterns(
     ModuleAxisInfoAnalysis &axisInfoAnalysis, const TargetInfoBase &targetInfo,
     PatternBenefit benefit);
 
-void populateUpcastMXFPToLLVMPatterns(LLVMTypeConverter &typeConverter,
-                                      RewritePatternSet &patterns,
-                                      const TargetInfo &targetInfo,
-                                      PatternBenefit benefit);
+void populateFp4ToFpToLLVMPatterns(LLVMTypeConverter &typeConverter,
+                                   RewritePatternSet &patterns,
+                                   PatternBenefit benefit);
 
 void populateBF16CastsLLVMPatterns(LLVMTypeConverter &typeConverter,
                                    RewritePatternSet &patterns,
diff --git a/third_party/intel/lib/TritonIntelGPUToLLVM/PipelineManager.h b/third_party/intel/lib/TritonIntelGPUToLLVM/PipelineManager.h
@@ -278,8 +278,7 @@ class TritonGPUToLLVMPipelineManager {
                                                        targetInfo, benefit);
       intel::populateMakeRangeOpToLLVMPattern(typeConverter, targetInfo,
                                               patterns, benefit);
-      intel::populateUpcastMXFPToLLVMPatterns(typeConverter, patterns,
-                                              targetInfo, benefit);
+      intel::populateFp4ToFpToLLVMPatterns(typeConverter, patterns, benefit);
     }
 
     intel::populateSPMDOpToLLVMPattern(typeConverter, patterns, targetInfo,
diff --git a/third_party/intel/lib/TritonIntelGPUToLLVM/UpcastMXFPToLLVM.cpp b/third_party/intel/lib/TritonIntelGPUToLLVM/UpcastMXFPToLLVM.cpp
diff --git a/third_party/intel/lib/TritonIntelGPUTransforms/AccelerateMatmul.cpp b/third_party/intel/lib/TritonIntelGPUTransforms/AccelerateMatmul.cpp
diff --git a/third_party/intel/lib/TritonIntelGPUTransforms/RemoveLayoutConversions.cpp b/third_party/intel/lib/TritonIntelGPUTransforms/RemoveLayoutConversions.cpp

Original file line number	Diff line number	Diff line change
`@@ -278,8 +278,7 @@ class TritonGPUToLLVMPipelineManager {`
`278`	`278`	`targetInfo, benefit);`
`279`	`279`	`intel::populateMakeRangeOpToLLVMPattern(typeConverter, targetInfo,`
`280`	`280`	`patterns, benefit);`
`281`		`- intel::populateUpcastMXFPToLLVMPatterns(typeConverter, patterns,`
`282`		`- targetInfo, benefit);`
	`281`	`+ intel::populateFp4ToFpToLLVMPatterns(typeConverter, patterns, benefit);`
`283`	`282`	`}`
`284`	`283`
`285`	`284`	`intel::populateSPMDOpToLLVMPattern(typeConverter, patterns, targetInfo,`