Reapply "[NVIDIA] Add native MXFP FP8 scaled_dot for SM120 (#7918)" (#8029) (#8129)

Add a fix for the crash happening on BW. 

Still doing some perf test

Author: ThomasRaoux

include/triton/Dialect/Triton/IR/TritonOpInterfaces.td

@@ -49,7 +49,12 @@ def DotOpInterface : OpInterface<"DotOpInterface"> {
       /*retType=*/"::mlir::Value",
       /*methodName=*/"getB",
       /*args=*/(ins)>,
-  InterfaceMethod<
+    InterfaceMethod<
+      /*desc=*/"Get the output tensor",
+      /*retType=*/"::mlir::Value",
+      /*methodName=*/"getD",
+      /*args=*/(ins)>,
+    InterfaceMethod<
       /*desc=*/"Verify the dimensions of the A and B DotOp operands.",
       /*retType=*/"bool",
       /*methodName=*/"verifyDims",
@@ -64,6 +69,7 @@ def DotOpInterface : OpInterface<"DotOpInterface"> {
         auto aTy = cast<ShapedType>($_op.getA().getType());
         auto bTy = cast<ShapedType>($_op.getB().getType());
         auto cTy = cast<ShapedType>($_op->getOperand(2).getType());
+        auto dTy = cast<ShapedType>($_op.getD().getType());
         auto aShape = aTy.getShape();
         auto bShape = bTy.getShape();
         auto cShape = cTy.getShape();

include/triton/Dialect/TritonGPU/IR/LinearLayoutConversions.h

@@ -135,6 +135,13 @@ LinearLayout chooseScaledMfmaScaleLayout(MLIRContext *ctx, int dotOperandIdx,
                                          ArrayRef<unsigned> tilesPerWarp,
                                          ArrayRef<unsigned> warpsPerCTA);
 
+LinearLayout getSM120DotScaledScaleLayout(MLIRContext *ctx, int dotOperandIdx,
+                                          ArrayRef<int64_t> dotOperandShape,
+                                          ArrayRef<unsigned> tilesPerWarp,
+                                          ArrayRef<unsigned> warpsPerCTA,
+                                          unsigned instrM, unsigned instrN,
+                                          CTALayoutAttr ctaLayoutAttr);
+
 // Create LinearLayout for nvidia mma tile.
 LinearLayout nvidiaMmaTile(MLIRContext *ctx, ArrayRef<unsigned> tileShape,
                            unsigned kWidth, ArrayRef<unsigned> order,

lib/Dialect/TritonGPU/IR/LinearLayoutConversions.cpp

@@ -1401,6 +1401,83 @@ LinearLayout chooseDsReadB64TrLayout(Attribute enc, ArrayRef<int64_t> shape,
   return chooseDotDsReadB64TrLayout(dot, shape, elemBitWidth);
 }
 
+// Warp-level block scaling (sm_120, m16n8k32)
+// Reference: NVIDIA PTX ISA "Warp-level block scaling"
+// https://docs.nvidia.com/cuda/parallel-thread-execution/#warp-level-block-scaling
+//
+// Semantics:
+//   D = (A * SF_A) * (B * SF_B) + C
+//   scale_vec::1X  -> SF_A shape Mx1 (per-row),   SF_B shape 1xN (per-col)
+//
+// Providers (within each warp quad of 4 lanes):
+//   - A scales are provided by a lane-pair selected by thread-id-a ∈ {0,1}
+//       (0 => lanes {0,1}, 1 => lanes {2,3} in the quad).
+//   - B scales are provided by a single lane selected by thread-id-b ∈
+//   {0,1,2,3}.
+//
+// Byte selectors (which subfield of the 32-bit metadata is used):
+//   - 1X: 1 byte  => byte-id ∈ {0,1,2,3}
+//
+// Implementation notes:
+//   - We support only scale_vec::1X for now.
+//   - We choose a fixed provider for A (thread-id-a = 0) and B (thread-id-b =
+//   0)
+//   - In this implementation, each lane in a quad has the same scale factor.
+LinearLayout getSM120DotScaledScaleLayout(
+    MLIRContext *ctx, int dotOperandIdx, ArrayRef<int64_t> dotOperandShape,
+    ArrayRef<unsigned> tilesPerWarp, ArrayRef<unsigned> warpsPerCTA,
+    unsigned mmaInstrM, unsigned mmaInstrN, CTALayoutAttr ctaLayoutAttr) {
+  unsigned rank = dotOperandShape.size();
+  auto outDims = standardOutDimNames(ctx, rank);
+
+  StringAttr kRegister = StringAttr::get(ctx, "register");
+  StringAttr kLane = StringAttr::get(ctx, "lane");
+  StringAttr kWarp = StringAttr::get(ctx, "warp");
+
+  const unsigned mIndex = 0;
+  const unsigned nIndex = 1;
+  const int instrM = mmaInstrM;
+  const int instrN = mmaInstrN;
+  const int kSize = dotOperandShape[1];
+  const int mWarps = warpsPerCTA[mIndex];
+  const int nWarps = warpsPerCTA[nIndex];
+  const int totalWarps = mWarps * nWarps;
+  const unsigned mRep_warp = tilesPerWarp[mIndex];
+  const unsigned nRep_warp = tilesPerWarp[nIndex];
+  const unsigned kRep = std::min<unsigned>(kSize, 2);
+
+  std::vector<std::vector<int32_t>> registerBase;
+  std::vector<std::vector<int32_t>> laneBase;
+  std::vector<std::vector<int32_t>> warpBase;
+  if (dotOperandIdx == 0) { // per-row A-scale
+    laneBase = {{0, 8}, {0, 0}, {0, 1}, {0, 2}, {0, 4}};
+    for (int offset = instrM * mWarps; offset < instrM * mWarps * mRep_warp;
+         offset <<= 1)
+      registerBase.push_back({0, offset});
+    for (int w = mWarps; w < totalWarps; w <<= 1)
+      warpBase.push_back({0, 0});
+    for (int offset = instrM; offset < instrM * mWarps; offset <<= 1)
+      warpBase.push_back({0, offset});
+  } else { // per-col B-scale
+    laneBase = {{0, 0}, {0, 0}, {0, 1}, {0, 2}, {0, 4}};
+    if (nRep_warp > 1)
+      registerBase.push_back({0, nWarps * instrN});
+    for (int k = 1; k < kRep; k += 1)
+      registerBase.push_back({1 << (k - 1), 0});
+    for (int offset = instrN; offset < instrN * nWarps; offset <<= 1)
+      warpBase.push_back({0, offset});
+    for (int w = nWarps; w < totalWarps; w <<= 1)
+      warpBase.push_back({0, 0});
+  }
+
+  const unsigned kIdx = (dotOperandShape[0] == 1) ? 0 : 1;
+  const unsigned mnIdx = 1 - kIdx;
+  LinearLayout ctaLayout(
+      {{kRegister, registerBase}, {kLane, laneBase}, {kWarp, warpBase}},
+      {outDims[kIdx], outDims[mnIdx]});
+  return combineCtaCgaWithShape(ctaLayout, ctaLayoutAttr, dotOperandShape);
+}
+
 LinearLayout chooseScaledMfmaScaleLayout(MLIRContext *ctx, int dotOperandIdx,
                                          ArrayRef<int64_t> dotOperandShape,
                                          unsigned mfmaMDim,