Merge OpenAI commit 1b27b93 (#5202)

This PR change the Triton base from
e15cb57
to 1b27b93 (Sep 22).
Pass rate: 97.22%

Author: anmyachev

.github/workflows/wheels.yml

@@ -88,8 +88,8 @@ jobs:
             export CIBW_MANYLINUX_AARCH64_IMAGE="quay.io/pypa/manylinux_2_28_${{ matrix.config.arch }}:latest"
           fi
 
-          export CIBW_BUILD="cp3{9,10,11,12,13,13t,14,14t}-manylinux_${{ matrix.config.arch }}"
-          export CIBW_SKIP="cp{35,36,37,38}-*"
+          export CIBW_BUILD="cp3{10,11,12,13,13t,14,14t}-manylinux_${{ matrix.config.arch }}"
+          export CIBW_SKIP="cp{35,36,37,38,39}-*"
           export CIBW_ENABLE=cpython-freethreading
           python3 -m cibuildwheel . --output-dir wheelhouse
 

README.md

@@ -38,7 +38,7 @@ You can install the latest stable release of Triton from pip:
 pip install triton
 ```
 
-Binary wheels are available for CPython 3.10-3.13.
+Binary wheels are available for CPython 3.10-3.14.
 
 # Install from source
 
@@ -262,6 +262,15 @@ export TRITON_OVERRIDE_DIR=<override_dir>
 # Step 4: Run the kernel again to see the overridden result
 ```
 
+**Compiler Pipeline Inspection Steps**
+To introspect the pipeline `add_stages`, before running your kernels, simply set
+the add_stages_inspection_hook like so:
+
+```python
+def inspect_stages(_self, stages, options, language, capability):
+    # inspect or modify add_stages here
+triton.knobs.runtime.add_stages_inspection_hook = inspect_stages
+```
 
 # Changelog
 

docs/getting-started/installation.rst

@@ -14,7 +14,7 @@ You can install the latest stable release of Triton from pip:
 
       pip install triton
 
-Binary wheels are available for CPython 3.10-3.13.
+Binary wheels are available for CPython 3.10-3.14.
 
 -----------
 From Source

lib/Conversion/TritonGPUToLLVM/ViewOpToLLVM.cpp

@@ -3,6 +3,7 @@
 #include "triton/Conversion/TritonGPUToLLVM/Utility.h"
 #include "triton/Dialect/TritonGPU/IR/Attributes.h"
 #include "triton/Dialect/TritonGPU/IR/Types.h"
+#include "triton/Dialect/TritonNvidiaGPU/IR/Dialect.h"
 #include "triton/Tools/LayoutUtils.h"
 
 using namespace mlir;

lib/Dialect/TritonGPU/IR/Dialect.cpp

@@ -1670,7 +1670,7 @@ void SharedLinearEncodingAttr::print(AsmPrinter &printer) const {
         layout.sublayout({kOffset}, llvm::to_vector(layout.getOutDimNames()));
   }
   printLinearLayout(printer, layout);
-  printer << "}, alignment = " << getAlignment() << "}>";
+  printer << "}, alignment = " << getAlignment() << ">";
 }
 
 Attribute SharedLinearEncodingAttr::parse(AsmParser &parser, Type type) {
@@ -2701,19 +2701,17 @@ struct TritonGPUInferLayoutInterface
     }
 
     if (auto enc = dyn_cast<NVMMASharedEncodingAttr>(operandEncoding)) {
-      if (failed(checkRank(enc.getRank())))
-        return failure();
-      if (order != ArrayRef<int32_t>({1, 0})) {
-        return emitOptionalError(
-            loc, "NVMMSharedEncoding can only be transposed in 2D");
-      }
+      if (order == ArrayRef<int32_t>({1, 0})) {
+        if (failed(checkRank(enc.getRank())))
+          return failure();
 
-      CTALayoutAttr ctaLayout =
-          permuteCTALayout(ctx, enc.getCTALayout(), order);
-      resultEncoding = NVMMASharedEncodingAttr::get(
-          ctx, enc.getSwizzlingByteWidth(), !enc.getTransposed(),
-          enc.getElementBitWidth(), enc.getFp4Padded(), ctaLayout);
-      return success();
+        CTALayoutAttr ctaLayout =
+            permuteCTALayout(ctx, enc.getCTALayout(), order);
+        resultEncoding = NVMMASharedEncodingAttr::get(
+            ctx, enc.getSwizzlingByteWidth(), !enc.getTransposed(),
+            enc.getElementBitWidth(), enc.getFp4Padded(), ctaLayout);
+        return success();
+      }
     }
 
     if (auto enc = dyn_cast<BlockedEncodingAttr>(operandEncoding)) {
@@ -2729,20 +2727,25 @@ struct TritonGPUInferLayoutInterface
           applyPermutation(invOrderUnsigned, enc.getOrder()), ctaLayout);
       return success();
     }
+    // Generic case
+    auto padded = dyn_cast<PaddedSharedEncodingAttr>(operandEncoding);
 
-    if (auto enc = dyn_cast<PaddedSharedEncodingAttr>(operandEncoding)) {
-      if (failed(checkRank(enc.getRank())))
-        return failure();
-      const auto &transLL =
-          transposeLinearLayout(enc.getLinearComponent(), order);
-      resultEncoding = PaddedSharedEncodingAttr::get(
-          ctx, enc.getIntervals(), enc.getPaddings(), transLL);
-      return success();
-    }
-
-    auto ll = toLinearLayout(shape, operandEncoding);
+    auto ll = padded ? padded.getLinearComponent()
+                     : toLinearLayout(shape, operandEncoding);
+    if (failed(checkRank(ll.getNumOutDims())))
+      return failure();
     auto transposedLl = transposeLinearLayout(ll, order);
-    resultEncoding = LinearEncodingAttr::get(ctx, std::move(transposedLl));
+    if (isa<DistributedEncodingTrait>(operandEncoding)) {
+      resultEncoding = LinearEncodingAttr::get(ctx, std::move(transposedLl));
+    } else if (padded) {
+      resultEncoding = PaddedSharedEncodingAttr::get(ctx, padded.getIntervals(),
+                                                     padded.getPaddings(),
+                                                     std::move(transposedLl));
+    } else {
+      auto shared = cast<SharedEncodingTrait>(operandEncoding);
+      resultEncoding = SharedLinearEncodingAttr::get(
+          ctx, std::move(transposedLl), shared.getAlignment());
+    }
     return success();
   }
 

lib/Dialect/TritonGPU/IR/Ops.cpp

@@ -531,40 +531,44 @@ static LogicalResult inferMemDescReshapeOpEncoding(ArrayRef<int64_t> srcShape,
                                                    Attribute srcEnc,
                                                    ArrayRef<int64_t> dstShape,
                                                    Attribute &dstEnc) {
+  // TODO Delete this once SharedLinearEncodingAttr is more widely supported.
   if (auto mmaEncoding = dyn_cast<NVMMASharedEncodingAttr>(srcEnc)) {
-    // TODO: supporting reshape of CTA layouts is non-trivial.
-    if (getNumCTAs(mmaEncoding) > 1)
-      return failure();
-    int innerDimDst =
-        mmaEncoding.getTransposed() ? dstShape.front() : dstShape.back();
-    int innerDimSrc =
-        mmaEncoding.getTransposed() ? srcShape.front() : srcShape.back();
-    // For now disallow reshape of the inner dimension.
-    if (innerDimDst != innerDimSrc)
-      return failure();
     auto *ctx = srcEnc.getContext();
-
-    // CTALayout can be all 1's because we bailed on multi-CTA layouts above.
-    auto CTALayout = CTALayoutAttr::get(
-        ctx,
-        /*CTAsPerCGA=*/SmallVector<unsigned>(dstShape.size(), 1),
-        /*CTASplitNum=*/SmallVector<unsigned>(dstShape.size(), 1),
-        /*CTAOrder=*/llvm::to_vector(llvm::seq<unsigned>(dstShape.size())));
-    dstEnc = NVMMASharedEncodingAttr::get(
-        ctx, mmaEncoding.getSwizzlingByteWidth(), mmaEncoding.getTransposed(),
-        mmaEncoding.getElementBitWidth(), mmaEncoding.getFp4Padded(),
-        CTALayout);
-    // Big guns, check linear layouts are equivalent
-    // We disallow reshaping memdesc_subslice in the verifier
-    // so allocShape == shape
-    auto srcLL = toLinearLayout(srcShape, srcEnc);
-    auto dstLL = toLinearLayout(dstShape, dstEnc);
-    if (reshapeLayout(ctx, srcLL, dstShape) != dstLL) {
-      return failure();
+    if (getNumCTAs(mmaEncoding) == 1) {
+      int innerDimDst =
+          mmaEncoding.getTransposed() ? dstShape.front() : dstShape.back();
+      int innerDimSrc =
+          mmaEncoding.getTransposed() ? srcShape.front() : srcShape.back();
+      // We can keep an NVMMAShared encoding only if the innermost dimension is
+      // preserved. Otherwise fall back to the generic shared-linear encoding
+      // logic below.
+      if (innerDimDst == innerDimSrc) {
+        auto CTALayout = CTALayoutAttr::get(
+            ctx,
+            /*CTAsPerCGA=*/SmallVector<unsigned>(dstShape.size(), 1),
+            /*CTASplitNum=*/SmallVector<unsigned>(dstShape.size(), 1),
+            /*CTAOrder=*/llvm::to_vector(llvm::seq<unsigned>(dstShape.size())));
+        auto candidateEncoding = NVMMASharedEncodingAttr::get(
+            ctx, mmaEncoding.getSwizzlingByteWidth(),
+            mmaEncoding.getTransposed(), mmaEncoding.getElementBitWidth(),
+            mmaEncoding.getFp4Padded(), CTALayout);
+        auto srcLL = toLinearLayout(srcShape, srcEnc);
+        auto dstLL = toLinearLayout(dstShape, candidateEncoding);
+        if (reshapeLayout(ctx, srcLL, dstShape) == dstLL) {
+          dstEnc = candidateEncoding;
+          return success();
+        }
+      }
     }
-    return success();
   }
-  return failure();
+
+  // Generic LL case
+  auto sharedEnc = cast<SharedEncodingTrait>(srcEnc);
+  auto *ctx = srcEnc.getContext();
+  auto srcLL = toLinearLayout(srcShape, srcEnc);
+  auto dstLL = reshapeLayout(ctx, srcLL, dstShape);
+  dstEnc = SharedLinearEncodingAttr::get(ctx, dstLL, sharedEnc.getAlignment());
+  return success();
 }
 
 LogicalResult MemDescReshapeOp::inferReturnTypes(

lib/Dialect/TritonGPU/Transforms/OptimizeDotOperands.cpp

@@ -257,7 +257,24 @@ class UseShmemForScales
     if (!isTmemCopyCompatible(localLoad.getSrc().getType(), usesTMAload))
       return failure();
 
-    opOperand.assign(localLoad.getSrc());
+    PatternRewriter::InsertionGuard guard(rewriter);
+    rewriter.setInsertionPoint(tmemAlloc);
+
+    Value shared = localLoad.getSrc();
+
+    Value reshaped5D = rewriter.create<MemDescReshapeOp>(
+        reshapeOp5D.getLoc(), shared, reshape5DShape);
+    SmallVector<int32_t> transposeOrder32(transposeOrder.begin(),
+                                          transposeOrder.end());
+    Value transposed = rewriter.create<MemDescTransOp>(
+        transOp.getLoc(), reshaped5D, transposeOrder32);
+    SmallVector<int64_t> scale2DShapeVec(scale2DShape.begin(),
+                                         scale2DShape.end());
+    Value reshaped2D = rewriter.create<MemDescReshapeOp>(
+        reshapeOp2D.getLoc(), transposed, scale2DShapeVec);
+
+    opOperand.assign(reshaped2D);
+    rewriter.eraseOp(tmemAlloc);
     return success();
   }
 

lib/Dialect/TritonNvidiaGPU/IR/Ops.cpp

@@ -680,14 +680,20 @@ LogicalResult TMEMCopyOp::verify() {
           getSrc().getType().getMemorySpace()))
     return emitOpError("The source must be a shared memory buffer");
 
+  auto srcTy = cast<triton::gpu::MemDescType>(getSrc().getType());
+  auto dstTy = cast<triton::gpu::MemDescType>(getDst().getType());
+  if (srcTy.getShape() != dstTy.getShape())
+    return emitOpError("source shape ")
+           << srcTy.getShape() << " must match destination shape "
+           << dstTy.getShape();
+
   if (getBarrier() && !isa<triton::gpu::SharedMemorySpaceAttr>(
                           getBarrier().getType().getMemorySpace())) {
     return emitOpError("The optional barrier should be a shared memory buffer");
   }
   if (!getDst().getType().getMutableMemory()) {
     return emitOpError("Cannot copy into an immutable alloc");
   }
-  auto srcTy = cast<triton::gpu::MemDescType>(getSrc().getType());
   auto sharedEnc =
       dyn_cast<triton::gpu::SharedEncodingTrait>(srcTy.getEncoding());
   if (sharedEnc.getAlignment() < 16) {
@@ -700,21 +706,16 @@ LogicalResult TMEMCopyOp::verify() {
   if (numCTAs != 1)
     return emitOpError("NYI: Only one CTA is supported for now.");
 
+  // Fp4 we could lift if we needed
   auto nvmmaEnc =
       dyn_cast<triton::gpu::NVMMASharedEncodingAttr>(srcTy.getEncoding());
-  if (!nvmmaEnc) {
-    return emitOpError("Source must have nvmma layout.");
-  }
-  // Fp4 we could lift if we needed
-  if (nvmmaEnc.getTransposed() || nvmmaEnc.getFp4Padded())
+  if (nvmmaEnc && (nvmmaEnc.getTransposed() || nvmmaEnc.getFp4Padded())) {
     return emitOpError("The source should not be transposed or padded");
+  }
   if (isa<TensorMemoryScalesEncodingAttr>(getDst().getType().getEncoding())) {
-    if (nvmmaEnc.getSwizzlingByteWidth() != 0) {
+    if (nvmmaEnc && nvmmaEnc.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(
@@ -728,7 +729,7 @@ LogicalResult TMEMCopyOp::verify() {
     if (tmemEnc.getBlockM() != 128) {
       return emitOpError("Tmem layout ahouls have M=128.");
     }
-    if (nvmmaEnc.getSwizzlingByteWidth() == 0) {
+    if (nvmmaEnc && nvmmaEnc.getSwizzlingByteWidth() == 0) {
       return emitOpError("Source layout should be swizzled.");
     }
     // When we lift this, we should make sure we handle unpacked cleanly

python/src/gluon_ir.cc

@@ -98,6 +98,7 @@ struct GluonLayouts {
   py::handle NVMMADistributedLayout;
   py::handle NVMMASharedLayout;
   py::handle SwizzledSharedLayout;
+  py::handle SharedLinearLayout;
   py::handle AMDMFMALayout;
   py::handle AMDWMMALayout;
   py::handle PaddedSharedLayout;
@@ -119,6 +120,8 @@ struct GluonLayouts {
     NVMMASharedLayout = py::object(layouts.attr("NVMMASharedLayout")).release();
     SwizzledSharedLayout =
         py::object(layouts.attr("SwizzledSharedLayout")).release();
+    SharedLinearLayout =
+        py::object(layouts.attr("SharedLinearLayout")).release();
     AMDMFMALayout = py::object(amdLayouts.attr("AMDMFMALayout")).release();
     AMDWMMALayout = py::object(amdLayouts.attr("AMDWMMALayout")).release();
     PaddedSharedLayout =
@@ -203,6 +206,14 @@ py::object layoutToGluon(Attribute layout) {
         toStdVector(ctaLayout.getCTAsPerCGA()),
         toStdVector(ctaLayout.getCTASplitNum()),
         toStdVector(ctaLayout.getCTAOrder()));
+  } else if (auto sharedLl = dyn_cast<ttg::SharedLinearEncodingAttr>(layout)) {
+    const auto &ll = sharedLl.getLinearLayout();
+    auto ctx = layout.getContext();
+    auto kOffset = mlir::StringAttr::get(ctx, "offset");
+    auto kBlock = mlir::StringAttr::get(ctx, "block");
+    return layouts.SharedLinearLayout(
+        toStdVector(ll.getBases().lookup(kOffset)),
+        toStdVector(ll.getBases().lookup(kBlock)), sharedLl.getAlignment());
   } else if (auto autoEnc = dyn_cast<gluon::AutoEncodingAttr>(layout)) {
     return layouts.AutoLayout();
   } else if (auto amdMfma = dyn_cast<ttg::AMDMfmaEncodingAttr>(layout)) {
@@ -410,14 +421,13 @@ void init_gluon_ir(py::module &&m) {
       .def("get_shared_linear_layout",
            [](GluonOpBuilder &self, std::vector<std::vector<int>> &offsetBases,
               std::vector<std::vector<int>> &blockBases,
-              std::vector<int64_t> &shape, unsigned alignment) -> Attribute {
+              unsigned alignment) -> Attribute {
              auto ctx = self.getContext();
              auto kOffset = mlir::StringAttr::get(ctx, "offset");
              auto kBlock = mlir::StringAttr::get(ctx, "block");
+             auto outDims = tt::standardOutDimNames(ctx, offsetBases[0].size());
              auto ll = tt::LinearLayout(
-                 {{kOffset, offsetBases}, {kBlock, blockBases}},
-                 tt::standardOutDimPairs(ctx, shape),
-                 /*requireSurjective=*/true);
+                 {{kOffset, offsetBases}, {kBlock, blockBases}}, outDims);
              return self.getChecked<ttg::SharedLinearEncodingAttr>(ctx, ll,
                                                                    alignment);
            })

python/test/backend/test_stages_inspection.py

@@ -0,0 +1,53 @@
+import triton
+from triton import knobs
+
+import os
+import pathlib
+
+
+def test_inspection(monkeypatch, tmp_path: pathlib.Path):
+    stage_name = 'make_ttgir'
+    curr_repro_path = tmp_path / ("repro_prefix." + stage_name + ".repro.mlir")
+    repro_path = tmp_path / "repro_prefix"
+
+    monkeypatch.setenv("TRITON_ALWAYS_COMPILE", "1")
+    monkeypatch.setenv("TRITON_REPRODUCER_PATH", str(repro_path))
+
+    inspect_stages_hook_called = False
+    make_ttgir_wrapper_called = False
+
+    def inspect_stages_hook(self, stages, options, language, capability):
+        nonlocal inspect_stages_hook_called
+        inspect_stages_hook_called = True
+
+        def make_ttgir_wrapper(src, metadata, options, capability):
+            nonlocal make_ttgir_wrapper_called
+            make_ttgir_wrapper_called = True
+            return self.make_ttgir(src, metadata, options, capability)
+
+        stages["ttgir"] = lambda src, metadata: make_ttgir_wrapper(src, metadata, options, capability)
+
+    @triton.jit
+    def k1():
+        return
+
+    @triton.jit
+    def k2():
+        return
+
+    # Run once to get the clean/golden repro dump
+    k1[(1, )]()
+    assert not inspect_stages_hook_called and not make_ttgir_wrapper_called
+    assert os.path.exists(curr_repro_path)
+    golden_repro = curr_repro_path.read_text()
+    curr_repro_path.unlink()
+
+    # Setup hook and call again, check if hooks got called
+    knobs.runtime.add_stages_inspection_hook = inspect_stages_hook
+    k2[(1, )]()
+    assert inspect_stages_hook_called and make_ttgir_wrapper_called
+    assert os.path.exists(curr_repro_path)
+    hook_repro = curr_repro_path.read_text()
+
+    # Check that repros match
+    assert golden_repro.replace('k1', 'dummy') == hook_repro.replace('k2', 'dummy')