Merge commit 'abd3bb0679dc12bad4e5e2d560211f8034b2f00b'

Author: whitneywhtsang

.gitignore

@@ -10,6 +10,7 @@ llvm-project-*/
 dist/
 triton*.egg-info/
 *.whl
+python/triton_kernels/triton*.egg-info/
 
 python/triton/_C/*.pyd
 python/triton/_C/*.so

bin/RegisterTritonDialects.h

@@ -61,7 +61,7 @@ inline void registerTritonDialects(mlir::DialectRegistry &registry) {
   mlir::registerAllPasses();
   mlir::triton::registerTritonPasses();
   mlir::triton::gpu::registerTritonGPUPasses();
-  mlir::registerTritonNvidiaGPUPasses();
+  mlir::triton::nvidia_gpu::registerTritonNvidiaGPUPasses();
   mlir::test::intel::registerTestAxisInfoPass();
   mlir::test::registerTestAliasPass();
   mlir::test::registerTestAlignmentPass();

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

@@ -38,38 +38,19 @@ struct ClusterInfo {
   int clusterDimZ;
 };
 
-} // namespace nvidia_gpu
-} // namespace triton
-} // namespace mlir
-
-namespace mlir {
-
 std::unique_ptr<Pass> createTritonNvidiaGPUPlanCTAPass(
     mlir::triton::nvidia_gpu::ClusterInfo *clusterInfo = nullptr);
 
-std::unique_ptr<Pass>
-createTritonNvidiaGPUFenceInsertionPass(int computeCapability = 90);
-
-std::unique_ptr<Pass> createTritonNvidiaGPUTMALoweringPass();
-
-std::unique_ptr<Pass> createTensorMemoryAllocationPass();
-
-std::unique_ptr<Pass> createTritonNvidiaGPUMMALoweringPass();
-
-std::unique_ptr<Pass> createTritonNvidiaGPUPromoteLHSToTMemPass();
-
-std::unique_ptr<Pass> createTritonNvidiaGPURemoveTMEMTokensPass();
-
-std::unique_ptr<Pass> createTritonNvidiaGPUOptimizeDescriptorEncodingPass();
-
-std::unique_ptr<Pass> createTritonNvidiaGPUOptimizeTMemLayoutsPass();
-
-std::unique_ptr<Pass> createTritonNvidiaGPUInterleaveTMemPass();
+#define GEN_PASS_DECL
+#include "triton/Dialect/TritonNvidiaGPU/Transforms/Passes.h.inc"
 
 /// Generate the code for registering passes.
 #define GEN_PASS_REGISTRATION
 #define GEN_PASS_DECL_TRITONNVIDIAGPULEGALIZETMALAYOUTS
 #include "triton/Dialect/TritonNvidiaGPU/Transforms/Passes.h.inc"
 
+} // namespace nvidia_gpu
+} // namespace triton
 } // namespace mlir
+
 #endif // TRITON_DIALECT_TRITONNVIDIAGPU_TRANSFORMS_PASSES_H_

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

@@ -32,7 +32,7 @@ def TritonGPUPlanCTAPass : Pass<"triton-nvidia-gpu-plan-cta", "mlir::ModuleOp">
     and StoreLikeOps operations.
   }];
 
-  let constructor = "mlir::createTritonNvidiaGPUPlanCTAPass()";
+  let constructor = "mlir::triton::nvidia_gpu::createTritonNvidiaGPUPlanCTAPass()";
 
   let dependentDialects = [
     "mlir::triton::gpu::TritonGPUDialect",
@@ -48,8 +48,6 @@ def TritonGPUFenceInsertion : Pass<"triton-nvidia-gpu-fence-insertion", "mlir::M
     properly ordered across generic and async operations.
   }];
 
-  let constructor = "mlir::createTritonNvidiaGPUFenceInsertionPass()";
-
   let dependentDialects = [
     "mlir::triton::gpu::TritonGPUDialect",
     "mlir::triton::nvidia_gpu::TritonNvidiaGPUDialect"
@@ -69,22 +67,18 @@ def TritonNvidiaGPUTMALoweringPass : Pass<"triton-nvidia-tma-lowering", "mlir::M
     Lower Triton experimental descriptor load to TMA load/store operations in TritonNvidiaGPUDialect.
   }];
 
-  let constructor = "mlir::createTritonNvidiaGPUTMALoweringPass()";
-
   let dependentDialects = [
     "mlir::triton::nvidia_gpu::TritonNvidiaGPUDialect"
   ];
 }
 
-def TritionTensorMemoryAllocationPass : Pass<"triton-tensor-memory-allocation", "mlir::ModuleOp"> {
+def TritonTensorMemoryAllocationPass : Pass<"triton-tensor-memory-allocation", "mlir::ModuleOp"> {
   let summary = "Assign tensor memory allocation";
 
   let description = [{
     Decide on tensor memory allocation and assign attributes to each allocation.
   }];
 
-  let constructor = "mlir::createTensorMemoryAllocationPass()";
-
   let dependentDialects = [
     "mlir::triton::nvidia_gpu::TritonNvidiaGPUDialect"
   ];
@@ -97,8 +91,6 @@ def TritonNvidiaGPUMMALoweringPass : Pass<"triton-nvidia-mma-lowering", "mlir::M
     Lower MMA ops to prepare for conversion to LLVM.
   }];
 
-  let constructor = "mlir::createTritonNvidiaGPUMMALoweringPass()";
-
   let dependentDialects = [
     "mlir::triton::nvidia_gpu::TritonNvidiaGPUDialect"
   ];
@@ -111,8 +103,6 @@ def TritonNvidiaGPUPromoteLHSToTMemPass : Pass<"tritongpu-promote-lhs-to-tmem",
     Promote LHS operand of MMAv5 op to Tensor Memory.
   }];
 
-  let constructor = "mlir::createTritonNvidiaGPUPromoteLHSToTMemPass()";
-
   let dependentDialects = ["mlir::triton::gpu::TritonGPUDialect",
                            "mlir::triton::nvidia_gpu::TritonNvidiaGPUDialect",
                            "mlir::triton::TritonDialect"];

lib/Conversion/TritonToTritonGPU/TritonGPUConversion.cpp

@@ -101,6 +101,15 @@ TritonGPUConversionTarget::TritonGPUConversionTarget(
       return true;
     return false;
   });
+  addDynamicallyLegalOp<triton::FuncOp>([](triton::FuncOp funcOp) -> bool {
+    for (auto arg : funcOp.getArguments()) {
+      if (auto tensor = dyn_cast<RankedTensorType>(arg.getType())) {
+        if (!tensor.getEncoding())
+          return false;
+      }
+    }
+    return true;
+  });
 }
 
 bool TritonGPUConversionTarget::isDynamicallyLegal(

lib/Conversion/TritonToTritonGPU/TritonToTritonGPUPass.cpp

@@ -481,14 +481,17 @@ class TritonFuncOpPattern : public OpConversionPattern<triton::FuncOp> {
   matchAndRewrite(triton::FuncOp op, OpAdaptor adaptor,
                   ConversionPatternRewriter &rewriter) const override {
     auto converter = getTypeConverter();
+    TypeConverter::SignatureConversion result(op.getNumArguments());
     auto newOp = rewriter.replaceOpWithNewOp<triton::FuncOp>(
         op, op.getName(), op.getFunctionType());
     addNamedAttrs(newOp, adaptor.getAttributes());
     rewriter.inlineRegionBefore(op.getBody(), newOp.getBody(),
                                 newOp.getBody().end());
-    if (failed(rewriter.convertRegionTypes(&newOp.getBody(), *converter)))
-      return failure();
-
+    // Convert just the entry block. The remaining unstructured control flow is
+    // converted by br patterns.
+    if (!newOp.getBody().empty())
+      rewriter.applySignatureConversion(&newOp.getBody().front(), result,
+                                        converter);
     return success();
   }
 };

lib/Dialect/Triton/Transforms/LoopAwareCSE.cpp

@@ -43,7 +43,7 @@ struct LoopCSEDriver {
   bool areEqualInLoop(Value a, Value b);
 
   scf::ForOp loop;
-  ValueEquivalence equalValues;
+  SmallVector<std::pair<int, int>> argStack;
 };
 } // namespace
 
@@ -52,14 +52,15 @@ bool LoopCSEDriver::areIterArgsEqual(int i, int j) {
     return true;
   if (loop.getInitArgs()[i] != loop.getInitArgs()[j])
     return false;
+  if (llvm::is_contained(argStack, std::make_pair(i, j)))
+    return true;
   BlockArgument aArg = loop.getRegionIterArg(i);
   BlockArgument bArg = loop.getRegionIterArg(j);
   // First, assume the arguments are equal. This is how recursion is broken.
-  equalValues.setKnownEquivalence(aArg, bArg, true);
+  argStack.push_back({i, j});
   bool result =
       areEqualInLoop(loop.getYieldedValues()[i], loop.getYieldedValues()[j]);
-  // Now update the equivalence based on the actual result.
-  equalValues.setKnownEquivalence(aArg, bArg, result);
+  argStack.pop_back();
   return result;
 }
 
@@ -83,14 +84,10 @@ bool LoopCSEDriver::areEqualInLoop(Value a, Value b) {
   if (a == loop.getInductionVar() || b == loop.getInductionVar())
     return false;
 
-  if (std::optional<bool> eq = equalValues.getKnownEquivalence(a, b))
-    return *eq;
-
   if (auto aArg = dyn_cast<BlockArgument>(a)) {
     auto bArg = cast<BlockArgument>(b);
     bool result =
         areIterArgsEqual(aArg.getArgNumber() - 1, bArg.getArgNumber() - 1);
-    equalValues.setKnownEquivalence(a, b, result);
     return result;
   }
 
@@ -107,9 +104,7 @@ bool LoopCSEDriver::areEqualInLoop(Value a, Value b) {
   bool result = OperationEquivalence::isEquivalentTo(
       aDef, bDef,
       [&](Value a, Value b) { return success(areEqualInLoop(a, b)); },
-      [&](Value a, Value b) { equalValues.setKnownEquivalence(a, b, true); },
-      OperationEquivalence::IgnoreLocations);
-  equalValues.setKnownEquivalence(a, b, result);
+      /*markEquivalent=*/nullptr, OperationEquivalence::IgnoreLocations);
   return result;
 }
 

lib/Dialect/TritonNvidiaGPU/Transforms/FenceInsertion.cpp

@@ -13,24 +13,21 @@
 //
 //===----------------------------------------------------------------------===//
 
-using namespace mlir;
-namespace tt = ::mlir::triton;
-namespace ttg = ::mlir::triton::gpu;
-namespace ttng = ::mlir::triton::nvidia_gpu;
+namespace ttg = mlir::triton::gpu;
 
-#define GEN_PASS_CLASSES
-#include "triton/Dialect/TritonNvidiaGPU/Transforms/Passes.h.inc"
+namespace mlir {
+namespace triton {
+namespace nvidia_gpu {
 
-namespace {
+#define GEN_PASS_DEF_TRITONGPUFENCEINSERTION
+#include "triton/Dialect/TritonNvidiaGPU/Transforms/Passes.h.inc"
 
 struct FenceInsertionPass
-    : public TritonGPUFenceInsertionBase<FenceInsertionPass> {
+    : public impl::TritonGPUFenceInsertionBase<FenceInsertionPass> {
 
 public:
-  FenceInsertionPass() = default;
-  FenceInsertionPass(int computeCapability) {
-    this->computeCapability = computeCapability;
-  }
+  using impl::TritonGPUFenceInsertionBase<
+      FenceInsertionPass>::TritonGPUFenceInsertionBase;
   // TODO: support more general patterns to insert fences. eg. any op(generic)
   // to shared in use-def chain which refers by async proxy. We have generic(
   // convertlayout with sts/stmatix) + fence + async(wgmma) up to now
@@ -39,7 +36,7 @@ struct FenceInsertionPass
     if (computeCapability < 90)
       return;
     ModuleOp mod = getOperation();
-    mod.walk([&](tt::DotOpInterface dotOp) {
+    mod.walk([&](DotOpInterface dotOp) {
       Value a = dotOp.getA();
       Value b = dotOp.getB();
       bool aDependsOnShared = dependOnCopyRegToShared(a);
@@ -48,8 +45,8 @@ struct FenceInsertionPass
         return WalkResult::advance();
 
       OpBuilder builder(dotOp);
-      auto fence = builder.create<ttng::FenceAsyncSharedOp>(dotOp.getLoc(),
-                                                            /*bCluster=*/false);
+      auto fence = builder.create<FenceAsyncSharedOp>(dotOp.getLoc(),
+                                                      /*bCluster=*/false);
       // If there is all the dependencies are outside of the loop try to hoist
       // the fence.
       while (auto loopOp = fence->getParentOfType<LoopLikeOpInterface>()) {
@@ -63,8 +60,8 @@ struct FenceInsertionPass
       }
 
       // If the previous op is already a fence, this one isn't needed.
-      if (auto lastFence = dyn_cast_or_null<ttng::FenceAsyncSharedOp>(
-              fence->getPrevNode())) {
+      if (auto lastFence =
+              dyn_cast_or_null<FenceAsyncSharedOp>(fence->getPrevNode())) {
         if (lastFence.getBCluster() == fence.getBCluster())
           fence.erase();
       }
@@ -129,9 +126,7 @@ struct FenceInsertionPass
     return true;
   }
 };
-} // namespace
 
-std::unique_ptr<Pass>
-mlir::createTritonNvidiaGPUFenceInsertionPass(int computeCapability) {
-  return std::make_unique<FenceInsertionPass>(computeCapability);
-}
+} // namespace nvidia_gpu
+} // namespace triton
+} // namespace mlir

lib/Dialect/TritonNvidiaGPU/Transforms/InterleaveTMem.cpp

@@ -4,17 +4,17 @@
 #include "triton/Dialect/TritonNvidiaGPU/Transforms/Passes.h"
 #include "llvm/ADT/AddressRanges.h"
 
-namespace {
-
-using namespace mlir;
+namespace ttg = mlir::triton::gpu;
 
-namespace ttng = triton::nvidia_gpu;
-namespace ttg = triton::gpu;
-namespace tt = triton;
+namespace mlir {
+namespace triton {
+namespace nvidia_gpu {
 
-#define GEN_PASS_CLASSES
+#define GEN_PASS_DEF_TRITONNVIDIAGPUINTERLEAVETMEMPASS
 #include "triton/Dialect/TritonNvidiaGPU/Transforms/Passes.h.inc"
 
+namespace {
+
 // If we don't know the effects of the op, we add all possible effects.
 void addAllValuelessEffects(
     SmallVectorImpl<MemoryEffects::EffectInstance> &effects) {
@@ -77,7 +77,7 @@ std::pair<Value, AccessRange> findBufferAccess(Value a) {
 
   Operation *defOp = a.getDefiningOp();
   // Accessing the alloc accesses the whole buffer.
-  if (auto alloc = dyn_cast<ttng::TMEMAllocOp>(defOp)) {
+  if (auto alloc = dyn_cast<TMEMAllocOp>(defOp)) {
     AccessRange access;
     for (uint64_t dim : alloc.getType().getShape())
       access.ranges.push_back({{0, dim}});
@@ -128,7 +128,7 @@ std::pair<Value, AccessRange> findBufferAccess(Value a) {
   }
 
   // Subslice is a subview only on the N dimension.
-  if (auto subslice = dyn_cast<ttng::TMEMSubSliceOp>(defOp)) {
+  if (auto subslice = dyn_cast<TMEMSubSliceOp>(defOp)) {
     auto [alloc, parentAccess] = findBufferAccess(subslice.getSrc());
     if (!alloc)
       return {};
@@ -186,7 +186,7 @@ bool sinkOps(Value buffer, ArrayRef<Operation *> useChain) {
     }
     // Don't sink past barrier signals, since they may guard the liverange
     // of the buffer.
-    if (isa<ttng::ArriveBarrierOp>(next))
+    if (isa<ArriveBarrierOp>(next))
       break;
     if (!isMemoryEffectFree(next)) {
       SmallVector<MemoryEffects::EffectInstance> effects;
@@ -199,7 +199,7 @@ bool sinkOps(Value buffer, ArrayRef<Operation *> useChain) {
           dep = true;
           break;
         }
-        if (isa<ttng::TensorMemory>(effect.getResource()) &&
+        if (isa<TensorMemory>(effect.getResource()) &&
             (!effect.getValue() || tmemMayAlias(effect.getValue(), buffer))) {
           dep = true;
           break;
@@ -229,20 +229,22 @@ bool trySinkOp(Operation *op, Value buffer) {
   return sinkOps(buffer, useChain);
 }
 
+} // anonymous namespace
+
 struct TritonNvidiaGPUInterleaveTMemPass
-    : public TritonNvidiaGPUInterleaveTMemPassBase<
+    : public impl::TritonNvidiaGPUInterleaveTMemPassBase<
           TritonNvidiaGPUInterleaveTMemPass> {
-  using TritonNvidiaGPUInterleaveTMemPassBase::
-      TritonNvidiaGPUInterleaveTMemPassBase;
+  using impl::TritonNvidiaGPUInterleaveTMemPassBase<
+      TritonNvidiaGPUInterleaveTMemPass>::TritonNvidiaGPUInterleaveTMemPassBase;
 
   void runOnOperation() override {
     MLIRContext *context = &getContext();
     ModuleOp m = getOperation();
     SmallVector<std::pair<Operation *, Value>> opsToSink;
     m.walk([&](Operation *op) {
-      if (auto load = dyn_cast<ttng::TMEMLoadOp>(op))
+      if (auto load = dyn_cast<TMEMLoadOp>(op))
         opsToSink.emplace_back(load, load.getSrc());
-      else if (auto alloc = dyn_cast<ttng::TMEMAllocOp>(op))
+      else if (auto alloc = dyn_cast<TMEMAllocOp>(op))
         opsToSink.emplace_back(alloc, alloc.getResult());
     });
     for (auto [op, buffer] : opsToSink) {
@@ -253,8 +255,6 @@ struct TritonNvidiaGPUInterleaveTMemPass
   }
 };
 
-} // namespace
-
-std::unique_ptr<Pass> mlir::createTritonNvidiaGPUInterleaveTMemPass() {
-  return std::make_unique<TritonNvidiaGPUInterleaveTMemPass>();
-}
+} // namespace nvidia_gpu
+} // namespace triton
+} // namespace mlir