Merge OpenAI Triton commit d412906 (#5045)

This PR change the Triton base from
818e892 to
d412906 (Aug 26).
Pass rate: 98.78%->98.74%

Author: whitneywhtsang

Makefile

@@ -54,7 +54,7 @@ test-distributed: all
 .PHONY: test-gluon
 test-gluon: all
 	$(PYTEST) -s -n $(NUM_PROCS) python/test/gluon
-	$(PYTEST) -vs python/tutorials/gluon/01-attention-forward.py
+	$(PYTEST) -vs python/examples/gluon/01-attention-forward.py
 
 .PHONY: test-regression
 test-regression: all

include/triton/Conversion/TritonGPUToLLVM/Utility.h

@@ -377,6 +377,8 @@ class SharedMemoryObject {
 
   Value getShmemOffset(Location loc, RewriterBase &rewriter,
                        triton::gpu::MemDescType srcTy) const;
+  Value getShmemAffineBase(Location loc, RewriterBase &rewriter,
+                           triton::gpu::MemDescType srcTy) const;
 
   // TODO(Keren): deprecate the method once AMD backend has cleaned up
   Value getCSwizzleOffset(int dim) const {

include/triton/Dialect/TritonNvidiaGPU/IR/TritonNvidiaGPUOps.td

@@ -701,8 +701,13 @@ def TTNG_TMEMCopyOp : TTNG_Op<"tmem_copy"> {
     for the completion of the copy before MMA, since tcgen05.cp followed by tcgen05.mma is guaranteed to
     execute in that order.
 
-    This op lowers to the PTX instruction tcgen05.cp. Right now, we only support 1CTA and the warpx4.32x128b
-    variant of the instruction. Each 32x128b block in SMEM is duplicated over 4 warps and stored into 128 rows
+    This op lowers to the PTX instruction tcgen05.cp. This supports writing either to scales tmem layout as well as default tmem layout.
+    Currently the semantic is different when writing to tmem scale layout.
+
+    In case of default layout the copy doesn't change the logical elements between the source and destination memdesc.
+
+    In case of scale layout:
+    Each 32x128b block in SMEM is duplicated over 4 warps and stored into 128 rows
     and 4 columns of TMEM. The primary use case of this op is to copy blocked scales from SMEM to TMEM.
 
     The shape of the input SMEM can be flexibily chosen depending on use cases. In the simplest case (e.g. unit test),
@@ -741,7 +746,7 @@ def TTNG_TMEMCopyOp : TTNG_Op<"tmem_copy"> {
     Optional<TTG_MemDescType>:$barrier
   );
 
-  let assemblyFormat = [{$src `,` $dst `,` $barrier attr-dict `:` functional-type(operands, results)}];
+  let assemblyFormat = [{$src `,` $dst (`,` $barrier^)? attr-dict `:` qualified(type(operands))}];
   let hasVerifier = 1;
 }
 

lib/Conversion/TritonGPUToLLVM/Utility.cpp

@@ -1209,6 +1209,14 @@ Value SharedMemoryObject::getShmemOffset(Location loc, RewriterBase &rewriter,
   return offset;
 }
 
+Value SharedMemoryObject::getShmemAffineBase(
+    Location loc, RewriterBase &rewriter,
+    triton::gpu::MemDescType srcTy) const {
+  auto b = TritonLLVMOpBuilder(loc, rewriter);
+  Value offset = getShmemOffset(loc, rewriter, srcTy);
+  return b.gep(base.getType(), baseElemType, base, offset);
+}
+
 Value getStructFromSharedMemoryObject(Location loc,
                                       const SharedMemoryObject &smemObj,
                                       RewriterBase &rewriter) {

lib/Conversion/TritonGPUToLLVM/ViewOpToLLVM.cpp

@@ -561,8 +561,8 @@ struct MemDescReinterpretOpConversion
 
     auto smemObj =
         getSharedMemoryObjectFromStruct(loc, adaptor.getSrc(), srcElemTy, b);
-    SharedMemoryObject newObj(smemObj.getBase(), dstElemTy, dstTy.getRank(),
-                              loc, b);
+    Value newBase = smemObj.getShmemAffineBase(loc, b, srcTy);
+    SharedMemoryObject newObj(newBase, dstElemTy, dstTy.getRank(), loc, b);
     b.replaceOp(op, getStructFromSharedMemoryObject(loc, newObj, b));
     return success();
   }

lib/Dialect/TritonInstrument/Transforms/ConcurrencySanitizer.cpp

@@ -552,6 +552,27 @@ class ConcurrencySanitizerPass
         }
       }
       if (auto commitOp = dyn_cast<ttng::TCGen5CommitOp>(op)) {
+        // Workaround: scan towards the beginning of the current block looking
+        // for mmav5s and mark their operands as reads guarded by the barrier.
+        Operation *prevOp = op->getPrevNode();
+        while (prevOp) {
+          auto setBarrier = [&](TypedValue<ttg::MemDescType> buf) {
+            MemType memType = MemType::TENSOR_MEM;
+            if (isa<ttg::SharedEncodingTrait>(buf.getType().getEncoding())) {
+              memType = MemType::SHARED_MEM;
+            }
+            b.create<tti::ExperimentalSetReadBarrierOp>(
+                buf, commitOp.getBarrier(), buffersTensor[(int)memType],
+                barriers, readBarriersAlloc[(int)memType],
+                readBarriersType[(int)memType], commitOp.getPred());
+          };
+          if (auto mmav5Op = dyn_cast<ttng::TCGen5MMAOp>(prevOp)) {
+            setBarrier(mmav5Op.getA());
+            setBarrier(mmav5Op.getB());
+          }
+          prevOp = prevOp->getPrevNode();
+        }
+
         b.create<tti::ExperimentalCommitWriteWithBarrierOp>(
             commitOp.getBarrier(), barriers,
             writeBarriersAlloc[(int)MemType::TENSOR_MEM],

lib/Dialect/TritonNvidiaGPU/IR/Ops.cpp

@@ -588,9 +588,6 @@ LogicalResult TMEMCopyOp::verify() {
   if (!isa<triton::gpu::SharedMemorySpaceAttr>(
           getSrc().getType().getMemorySpace()))
     return emitOpError("The source must be a shared memory buffer");
-  if (!isa<TensorMemoryEncodingAttr, TensorMemoryScalesEncodingAttr>(
-          getDst().getType().getEncoding()))
-    return emitOpError("The destination must be a tensor memory buffer.");
 
   if (getBarrier() && !isa<triton::gpu::SharedMemorySpaceAttr>(
                           getBarrier().getType().getMemorySpace())) {
@@ -599,19 +596,41 @@ LogicalResult TMEMCopyOp::verify() {
   if (!getDst().getType().getMutableMemory()) {
     return emitOpError("Cannot copy into an immutable alloc");
   }
-
   auto srcTy = cast<triton::gpu::MemDescType>(getSrc().getType());
   auto sharedEnc =
-      cast<triton::gpu::NVMMASharedEncodingAttr>(srcTy.getEncoding());
-
-  if (!sharedEnc || sharedEnc.getTransposed() || sharedEnc.getFp4Padded() ||
-      sharedEnc.getSwizzlingByteWidth() != 0)
-    return emitOpError("The source should not have swizzling applied for now");
-
-  if (!triton::gpu::isInnermostContiguous(srcTy, 512)) {
-    return emitOpError("The source must be in a row-major order.");
+      dyn_cast<triton::gpu::NVMMASharedEncodingAttr>(srcTy.getEncoding());
+  if (!sharedEnc) {
+    return emitOpError("Source must have nvmma layout.");
+  }
+  if (sharedEnc.getTransposed() || sharedEnc.getFp4Padded())
+    return emitOpError("The source should not be transposed or passed");
+  if (isa<TensorMemoryScalesEncodingAttr>(getDst().getType().getEncoding())) {
+    if (sharedEnc.getSwizzlingByteWidth() != 0) {
+      return emitOpError("The source should not be swizzled for now");
+    }
+    if (!triton::gpu::isInnermostContiguous(srcTy, 512)) {
+      return emitOpError("The source must be in a row-major order.");
+    }
+  } else {
+    if (getSrc().getType().getShape() != getDst().getType().getShape()) {
+      return emitOpError(
+          "The source and destination must have the same shape.");
+    }
+    auto tmemEnc = dyn_cast<triton::nvidia_gpu::TensorMemoryEncodingAttr>(
+        getDst().getType().getEncoding());
+    if (!tmemEnc) {
+      return emitOpError("Incorrect tmem layout.");
+    }
+    if (tmemEnc.getBlockM() != 128) {
+      return emitOpError("Tmem layout ahouls have M=128.");
+    }
+    if (sharedEnc.getSwizzlingByteWidth() == 0) {
+      return emitOpError("Source layout should be swizzled.");
+    }
+    if (srcTy.getElementType().getIntOrFloatBitWidth() != 32) {
+      return emitOpError("Source element type should be 32-bit.");
+    }
   }
-
   // Given that we want to support flexible input SMEM shapes, kinds of shape
   // checking we can do here are limited. For simplicity, shape checking is
   // omitted.

pyproject.toml

@@ -35,3 +35,6 @@ line-length = 120
 
 [tool.ruff.lint]
 ignore = ["E501", "E701", "E731", "E741"]
+
+[tool.ruff.lint.per-file-ignores]
+"__init__.py" = ["F401"]

python/src/gluon_ir.cc

@@ -403,6 +403,14 @@ void init_gluon_ir(py::module &&m) {
                  ctx, block[0], block[1], unpacked, ctaSplitNum[0],
                  ctaSplitNum[1]);
            })
+      .def("get_tensor_memory_scales_layout",
+           [](GluonOpBuilder &self,
+              std::vector<unsigned> &ctaSplitNum) -> Attribute {
+             auto ctx = self.getContext();
+             assert(ctaSplitNum.size() == 2);
+             return self.getChecked<ttng::TensorMemoryScalesEncodingAttr>(
+                 ctx, ctaSplitNum[0], ctaSplitNum[1]);
+           })
       .def("get_gluon_layout_from_tensor",
            [](GluonOpBuilder &self, Value tensor) -> py::object {
              auto ty = dyn_cast<RankedTensorType>(tensor.getType());
@@ -548,6 +556,10 @@ void init_gluon_ir(py::module &&m) {
            [](GluonOpBuilder &self, Type resultTy, Value memDesc) -> Value {
              return self.create<ttng::TMEMLoadOp>(resultTy, memDesc);
            })
+      .def("create_tmem_copy",
+           [](GluonOpBuilder &self, Value src, Value dst) {
+             self.create<ttng::TMEMCopyOp>(src, dst, /*barrier=*/Value());
+           })
       .def("create_tmem_subslice",
            [](GluonOpBuilder &self, Type resultTy, Value memDesc,
               int N) -> Value {
@@ -585,6 +597,10 @@ void init_gluon_ir(py::module &&m) {
                                             pred, two_ctas, mbarriers,
                                             mbarrier_preds);
            })
+      .def("create_tcgen05_cp",
+           [](GluonOpBuilder &self, Value src, Value dst) {
+             self.create<ttng::TMEMCopyOp>(src, dst, Value());
+           })
       .def("create_tcgen05_commit",
            [](GluonOpBuilder &self, Value &barrier) {
              self.create<ttng::TCGen5CommitOp>(barrier);

python/test/gluon/test_consan.py

@@ -228,7 +228,8 @@ def tcgen5_mma_kernel(input_desc, XBLOCK: ttgl.constexpr, FAILURE: ttgl.constexp
     mbarrier.init(bar.index(0), count=1)
     mbarrier.init(bar.index(1), count=1)
 
-    blackwell.tcgen05_mma(smemA, smemB.permute([1, 0]), acc, mbarriers=[bar.index(0)])
+    blackwell.tcgen05_mma(smemA, smemB.permute([1, 0]), acc)
+    blackwell.tcgen05_commit(bar.index(0))
 
     if not FAILURE:
         mbarrier.wait(bar.index(0), 0)
@@ -285,32 +286,6 @@ def alloc_fn(size: int, alignment: int, stream: Optional[int]):
     tcgen5_mma_kernel[(1, )](input_desc, XBLOCK, FAILURE=FAILURE, MEM_ACCESS_KIND=MEM_ACCESS_KIND, num_warps=4)
 
 
-@gluon.jit
-def tcgen5_mma_multibar_kernel(input_desc, XBLOCK: ttgl.constexpr, BUF_IDX: ttgl.constexpr, BAR_IDX: ttgl.constexpr):
-    acc_layout: ttgl.constexpr = blackwell.TensorMemoryLayout([XBLOCK, XBLOCK], unpacked=True, cta_split_num=[1, 1])
-    blocked_layout: ttgl.constexpr = ttgl.BlockedLayout(size_per_thread=[1, XBLOCK], threads_per_warp=[32, 1],
-                                                        warps_per_cta=[4, 1], order=[0, 1])
-    smemA = ttgl.allocate_shared_memory(ttgl.float16, [XBLOCK, XBLOCK], input_desc.layout)
-    smemB = ttgl.allocate_shared_memory(ttgl.float16, [XBLOCK, XBLOCK], input_desc.layout)
-    bar = ttgl.allocate_shared_memory(ttgl.int64, [4, 1], mbarrier.MBarrierLayout())
-    acc = blackwell.allocate_tensor_memory(ttgl.float32, [2, XBLOCK, XBLOCK], acc_layout)
-    for i in range(4):
-        mbarrier.init(bar.index(i), count=1)
-
-    blackwell.tcgen05_mma(smemA, smemB.permute([1, 0]), acc.index(0), mbarriers=[bar.index(0),
-                                                                                 bar.index(1)],
-                          mbarrier_preds=[False, True])
-    blackwell.tcgen05_mma(smemA, smemB.permute([1, 0]), acc.index(1), mbarriers=[bar.index(2)])
-    blackwell.tcgen05_commit(bar.index(3))
-
-    mbarrier.wait(bar.index(BAR_IDX), 0)
-
-    acc.index(BUF_IDX).store(ttgl.full([XBLOCK, XBLOCK], 42, ttgl.float32, blocked_layout))
-
-    for i in range(4):
-        mbarrier.invalidate(bar.index(i))
-
-
 @gluon.jit
 def warpgroup_mma_kernel(input, XBLOCK: ttgl.constexpr, FAILURE: ttgl.constexpr):
     smem_layout: ttgl.constexpr = ttgl.NVMMASharedLayout(swizzle_byte_width=128, element_bitwidth=16, rank=2)
@@ -405,6 +380,32 @@ def alloc_fn(size: int, alignment: int, stream: Optional[int]):
     warpgroup_mma_kernel[(1, )](input, XBLOCK, FAILURE=FAILURE)
 
 
+@gluon.jit
+def tcgen5_mma_multibar_kernel(input_desc, XBLOCK: ttgl.constexpr, BUF_IDX: ttgl.constexpr, BAR_IDX: ttgl.constexpr):
+    acc_layout: ttgl.constexpr = blackwell.TensorMemoryLayout([XBLOCK, XBLOCK], unpacked=True, cta_split_num=[1, 1])
+    blocked_layout: ttgl.constexpr = ttgl.BlockedLayout(size_per_thread=[1, XBLOCK], threads_per_warp=[32, 1],
+                                                        warps_per_cta=[4, 1], order=[0, 1])
+    smemA = ttgl.allocate_shared_memory(ttgl.float16, [XBLOCK, XBLOCK], input_desc.layout)
+    smemB = ttgl.allocate_shared_memory(ttgl.float16, [XBLOCK, XBLOCK], input_desc.layout)
+    bar = ttgl.allocate_shared_memory(ttgl.int64, [4, 1], mbarrier.MBarrierLayout())
+    acc = blackwell.allocate_tensor_memory(ttgl.float32, [2, XBLOCK, XBLOCK], acc_layout)
+    for i in range(4):
+        mbarrier.init(bar.index(i), count=1)
+
+    blackwell.tcgen05_mma(smemA, smemB.permute([1, 0]), acc.index(0), mbarriers=[bar.index(0),
+                                                                                 bar.index(1)],
+                          mbarrier_preds=[False, True])
+    blackwell.tcgen05_mma(smemA, smemB.permute([1, 0]), acc.index(1), mbarriers=[bar.index(2)])
+    blackwell.tcgen05_commit(bar.index(3))
+
+    mbarrier.wait(bar.index(BAR_IDX), 0)
+
+    acc.index(BUF_IDX).store(ttgl.full([XBLOCK, XBLOCK], 42, ttgl.float32, blocked_layout))
+
+    for i in range(4):
+        mbarrier.invalidate(bar.index(i))
+
+
 @pytest.mark.skipif(not is_cuda() or torch.cuda.get_device_capability()[0] < 10, reason="Requires blackwell or newer")
 @pytest.mark.parametrize("BUF_IDX", [0, 1])
 @pytest.mark.parametrize("BAR_IDX", [0, 1, 2, 3])