Merge commit '4b9efc5c476e2a60aa45b339eded566c56690aa1'

Author: whitneywhtsang

.github/workflows/integration-tests-nvidia.yml

@@ -10,7 +10,7 @@ on:
 jobs:
   integration-tests-nvidia:
     runs-on: ${{ matrix.runner }}
-    timeout-minutes: 30
+    timeout-minutes: 60
     strategy:
       matrix:
         runner: ${{ fromJson(inputs.matrix) }}
@@ -94,6 +94,8 @@ jobs:
           if [ "${{ matrix.runner[0] }}" == "nvidia-gb200" ]; then
             source /venv/bin/activate
           fi
+          nproc
+          nvidia-smi
           echo "PATH is '$PATH'"
           ccache --zero-stats
           make dev-install

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

@@ -29,6 +29,7 @@ include "triton/Dialect/Triton/IR/TritonTypes.td"
 include "triton/Dialect/Triton/IR/TritonAttrDefs.td"
 include "triton/Dialect/Triton/IR/TritonInterfaces.td"
 include "triton/Dialect/Triton/IR/TritonOpInterfaces.td"
+include "triton/Dialect/TritonGPU/IR/TritonGPUAttrDefs.td"
 include "triton/Dialect/TritonGPU/IR/TritonGPUTypes.td"
 include "triton/Dialect/TritonGPU/IR/TritonGPUTypeInterfaces.td"
 include "mlir/IR/OpBase.td"
@@ -584,6 +585,12 @@ def TTNG_TMEMStoreOp : TTNG_Op<"tmem_store"> {
   );
   let results = (outs Optional<TTG_AsyncToken>:$token);
 
+  let builders = [
+    OpBuilder<(ins "Value":$dst, "Value":$src, "Value":$pred), [{
+      build($_builder, $_state, Type(), dst, Value(), src, pred);
+    }]>
+  ];
+
   let assemblyFormat = [{
     $src `,` $dst `` custom<Token>($dep, type($token)) `,` $pred
     attr-dict `:` type($src) `->` qualified(type($dst))

include/triton/Dialect/TritonNvidiaGPU/Transforms/Passes.h

@@ -64,6 +64,8 @@ std::unique_ptr<Pass> createTritonNvidiaGPUOptimizeDescriptorEncodingPass();
 
 std::unique_ptr<Pass> createTritonNvidiaGPUOptimizeTMemLayoutsPass();
 
+std::unique_ptr<Pass> createTritonNvidiaGPUInterleaveTMemPass();
+
 /// Generate the code for registering passes.
 #define GEN_PASS_REGISTRATION
 #define GEN_PASS_DECL_TRITONNVIDIAGPULEGALIZETMALAYOUTS

include/triton/Dialect/TritonNvidiaGPU/Transforms/Passes.td

@@ -143,6 +143,16 @@ def TritonNvidiaGPUOptimizeTMemLayoutsPass : Pass<"triton-nvidia-optimize-tmem-l
                            "mlir::triton::TritonDialect"];
 }
 
+def TritonNvidiaGPUInterleaveTMemPass : Pass<"triton-nvidia-interleave-tmem", "mlir::ModuleOp"> {
+  let summary = "Interleave TMEM loads/stores.";
+
+  let description = [{
+    The `triton-nvidia-interleave-tmem` pass attempts to sink TMEM loads and
+    hoist TMEM stores, and potentially interleave them, to reduce register
+    pressure.
+  }];
+}
+
 def TritonNvidiaGPURemoveTMEMTokensPass : Pass<"triton-nvidia-gpu-remove-tmem-tokens", "mlir::ModuleOp"> {
   let summary = "remove TMEM tokens";
 

include/triton/Dialect/TritonNvidiaGPU/Transforms/TMAUtilities.h

@@ -39,14 +39,30 @@ triton::gpu::SharedEncodingTrait
 getEncodingFromDescriptor(Operation *op, RankedTensorType tensorType,
                           Value desc);
 
-int64_t getTMAContigDim(Attribute encoding, ArrayRef<int64_t> shape);
+SmallVector<int64_t> getTMABlockShape(ArrayRef<int64_t> shapePerCTA,
+                                      int elementBitWidth, int swizzleBytes,
+                                      bool fp4Padded, bool transposed,
+                                      bool packedSize);
+
+inline SmallVector<int64_t> getTMABlockShape(Attribute encoding,
+                                             ArrayRef<int64_t> shapePerCTA,
+                                             bool packedSize) {
+  auto mmaEnc = cast<gpu::NVMMASharedEncodingAttr>(encoding);
+  return getTMABlockShape(shapePerCTA, mmaEnc.getElementBitWidth(),
+                          mmaEnc.getSwizzlingByteWidth(), mmaEnc.getFp4Padded(),
+                          mmaEnc.getTransposed(), packedSize);
+}
 
-inline int64_t getTMAContigDim(RankedTensorType tensorType) {
-  return getTMAContigDim(tensorType.getEncoding(), tensorType.getShape());
+inline SmallVector<int64_t> getTMABlockShape(RankedTensorType ty,
+                                             bool packedSize) {
+  auto shapePerCTA = gpu::getShapePerCTA(ty);
+  return getTMABlockShape(ty.getEncoding(), shapePerCTA, packedSize);
 }
 
-inline int64_t getTMAContigDim(gpu::MemDescType memDescType) {
-  return getTMAContigDim(memDescType.getEncoding(), memDescType.getShape());
+inline SmallVector<int64_t> getTMABlockShape(triton::gpu::MemDescType ty,
+                                             bool packedSize) {
+  auto shapePerCTA = gpu::getShapePerCTA(ty);
+  return getTMABlockShape(ty.getEncoding(), shapePerCTA, packedSize);
 }
 
 std::optional<int> getTMASwizzleMode(Operation *op, TensorDescType ty);
@@ -74,16 +90,18 @@ mlir::LogicalResult createTMADesc(mlir::Value tmaPtr,
 
   int paddingScale = fp4Padded ? 2 : 1;
   auto shapePerCTA = gpu::getShapePerCTA(encoding, op.getTensorShape());
-  int32_t contig_dim_size = getTMAContigDim(encoding, op.getTensorShape());
+  auto blockShape =
+      getTMABlockShape(encoding, shapePerCTA, /*packedSize=*/false);
+  auto contigDimSize = blockShape.back();
 
   llvm::SmallVector<Value> boxDim;
-  if (fp4Padded && contig_dim_size != 128) {
+  if (fp4Padded && contigDimSize != 128) {
     return op->emitError(
         "FP4 padded loads require 128 elements or more in the last dim");
   }
-  boxDim.push_back(mkI32Constant(contig_dim_size));
+  boxDim.push_back(mkI32Constant(contigDimSize));
   for (int k = shapePerCTA.size() - 2; k >= 0; --k)
-    boxDim.push_back(mkI32Constant(shapePerCTA[k]));
+    boxDim.push_back(mkI32Constant(blockShape[k]));
 
   unsigned swizzleBytes = mmaEncoding ? mmaEncoding.getSwizzlingByteWidth() : 0;
   if (!mmaEncoding) {

include/triton/Tools/LayoutUtils.h

@@ -83,6 +83,17 @@ LinearLayout ensureLayoutNotSmallerThan(
     const LinearLayout &layout,
     const llvm::SmallDenseMap<StringAttr, int64_t> &shape);
 
+inline LinearLayout
+ensureLayoutNotSmallerThan(const LinearLayout &layout,
+                           const llvm::ArrayRef<StringAttr> dimNames,
+                           const llvm::ArrayRef<int64_t> shape) {
+  llvm::SmallDenseMap<StringAttr, int64_t> namedDims;
+  for (auto [dimName, length] : llvm::zip_equal(dimNames, shape))
+    namedDims[dimName] = length;
+  assert(namedDims.size() == shape.size() && "duplicate dimension names given");
+  return ensureLayoutNotSmallerThan(layout, namedDims);
+}
+
 // Return a vector of the standard out dimension names for tensor layouts. These
 // are "dim0", "dim1", etc.
 SmallVector<StringAttr> standardOutDimNames(MLIRContext *ctx, int rank);

include/triton/Tools/LinearLayout.h

@@ -325,23 +325,32 @@ class LinearLayout {
       bases;
 
   llvm::MapVector<StringAttr, int32_t /*size*/> outDims;
-  bool surjective;
+  bool surjective = true;
 
 public:
   using BasesT = decltype(bases);
 
+  LinearLayout() = default;
+
   // The 0-dimensional layout that maps everything to 0.  This is useful as a
   // starting point when doing something like
   //
   //   LinearLayout ret = LinearLayout::empty();
   //   for (...) ret *= ...;
   //   return ret;
-  static LinearLayout empty() { return LinearLayout(BasesT{}, {}); }
+  static LinearLayout empty() { return {}; }
+
+  // Creates a 1D -> 1D layout that's the function L(x) = stride * x
+  // for x in [0, size).
+  static LinearLayout strided1D(int32_t size, int32_t stride, StringAttr inDim,
+                                StringAttr outDim);
 
   // Creates a 1D -> 1D layout that's the identity function, i.e. L(x) = x
   // for x in [0, size).
   static LinearLayout identity1D(int32_t size, StringAttr inDim,
-                                 StringAttr outDim);
+                                 StringAttr outDim) {
+    return strided1D(size, /*stride=*/1, inDim, outDim);
+  }
 
   // Creates a 1D -> 1D layout that maps every input value to 0, i.e. L(x) = 0
   // for x in [0, size). By default this creates a surjective layout where