Merge branch 'main' into llvm-head

Author: chsigg

include/triton/Dialect/Triton/IR/Dialect.h

@@ -72,10 +72,12 @@ class DialectInferLayoutInterface
 
   virtual LogicalResult
   inferJoinOpEncoding(Attribute srcEnc, Attribute &dstEnc,
+                      ArrayRef<int64_t> shape,
                       std::optional<Location> loc) const = 0;
 
   virtual LogicalResult
   inferSplitOpEncoding(Attribute srcEnc, Attribute &dstEnc,
+                       ArrayRef<int64_t> shape,
                        std::optional<Location> loc) const = 0;
 
   // Verify that the encoding are compatible to be used together in a dot

include/triton/Dialect/TritonGPU/IR/TritonGPUAttrDefs.td

@@ -620,7 +620,9 @@ def LinearEncodingAttr : DistributedEncoding<"LinearEncoding", "linear_encoding"
     unsigned getTotalElemsPerThread(ArrayRef<int64_t> shape) const;
     SmallVector<unsigned> getElemsPerThread(ArrayRef<int64_t> shape) const;
 
+    SmallVector<unsigned int> getContig(const char *, SmallVector<unsigned int>) const;
     SmallVector<unsigned> getContigPerThread() const;
+    SmallVector<unsigned> getContigPerWarp() const;
     SmallVector<unsigned> getOrder() const;
 
     // Generalizes get{Warp,Thread,CTA}Order to linear layouts.

include/triton/Dialect/TritonGPU/Transforms/Utility.h

@@ -49,6 +49,9 @@ unsigned
 getNumElementsPerThread(Operation *op, SmallVector<unsigned> order,
                         triton::ModuleAxisInfoAnalysis &axisInfoAnalysis);
 
+// Returns whether the op is a "view op", i.e. doesn't move any data
+bool isView(Operation *op);
+
 /* Dump Triton IR in graphviz dot format.
  *
  * You can override `onValue` and `onOperation` in a subclass to mark

include/triton/Tools/LayoutUtils.h

@@ -87,6 +87,11 @@ LinearLayout ensureLayoutNotSmallerThan(
 // are "dim0", "dim1", etc.
 SmallVector<StringAttr> standardOutDimNames(MLIRContext *ctx, int rank);
 
+// Return a vector of the standard out dimension name/value pairs, i.e.
+// ("dim0", dstShape[0]), ("dim1", dstShape[1]), etc.
+SmallVector<std::pair<StringAttr, int32_t>>
+standardOutDimPairs(MLIRContext *ctx, ArrayRef<int64_t> dstShape);
+
 // Return an identity mapping from `inDimName` to the standard out dimensions,
 // with the dimensions sized according to the shape. The bases are sorted
 // according to `order`, with the most minor dimension first.

lib/Analysis/Membar.cpp

@@ -60,8 +60,9 @@ void MembarAnalysis::resolve(FunctionOpInterface funcOp,
       // the outputBlockInfo, we skip the successors
       continue;
     }
-    // Update the current block
-    outputBlockInfoMap[block].join(inputBlockInfo);
+    // Update the current block. The block transfer function is not monotonic,
+    // so overwrite the output state entirely.
+    outputBlockInfoMap[block] = inputBlockInfo;
     // Update the successors
     for (auto *successor : successors) {
       inputBlockInfoMap[successor].join(outputBlockInfoMap[block]);

lib/Analysis/Utility.cpp

@@ -638,7 +638,11 @@ bool supportMMA(triton::DotOp op, int version) {
       return false;
     if (op.getType().getRank() != 2)
       return false;
-    if (!(numWarps % 4 == 0 && retShapePerCTA[rank - 2] % 64 == 0 &&
+    if (numWarps != 4 && numWarps != 8) {
+      // Currently only support numWarps 4 or 8 for TMEM load and store.
+      return false;
+    }
+    if (!(retShapePerCTA[rank - 2] % 64 == 0 &&
           retShapePerCTA[rank - 1] % 8 == 0))
       return false;
     return true;

lib/Dialect/Triton/IR/Ops.cpp

@@ -1046,7 +1046,7 @@ JoinOp::inferReturnTypes(MLIRContext *context, std::optional<Location> location,
   Attribute retEnc;
   if (srcEnc) {
     if (cast<DialectInferLayoutInterface>(&srcEnc.getDialect())
-            ->inferJoinOpEncoding(srcEnc, retEnc, location)
+            ->inferJoinOpEncoding(srcEnc, retEnc, srcTy.getShape(), location)
             .failed()) {
       return failure();
     }
@@ -1079,7 +1079,7 @@ LogicalResult SplitOp::inferReturnTypes(
   Attribute retEnc;
   if (srcEnc) {
     if (cast<DialectInferLayoutInterface>(&srcEnc.getDialect())
-            ->inferSplitOpEncoding(srcEnc, retEnc, location)
+            ->inferSplitOpEncoding(srcEnc, retEnc, srcTy.getShape(), location)
             .failed()) {
       return failure();
     }