[Backend] Bump LLVM to 2eb709b95d8f (triton-lang#9264)

AlexAUT · neildhar · antiagainst · paultrojahnamd · commit 4a63ec93c871 · 2026-03-19T19:48:28.000+01:00
Bumps LLVM to `2eb709b95d8f` to consume 2 fixes related to `gfx1250`: - llvm/llvm-project#176206 - llvm/llvm-project#176355 --------- Co-authored-by: neildhar <neildhar@users.noreply.github.com> Co-authored-by: Lei Zhang <antiagainst@gmail.com>
diff --git a/cmake/llvm-hash.txt b/cmake/llvm-hash.txt
@@ -1 +1 @@
-0729a74e66aeeb7a9839d80bfd64fc49b2e69f52
+2eb709b95d8f521aa15401e159fac0729d56a677
diff --git a/include/triton/Analysis/Alias.h b/include/triton/Analysis/Alias.h
@@ -92,6 +92,7 @@ class SharedMemoryAliasAnalysis
 
   void visitNonControlFlowArguments(
       Operation *op, const RegionSuccessor &successor,
+      ValueRange successorInputs,
       ArrayRef<dataflow::Lattice<AliasInfo> *> argLattices,
       unsigned firstIndex) override;
 };
diff --git a/include/triton/Dialect/TritonGPU/IR/TritonGPUOps.td b/include/triton/Dialect/TritonGPU/IR/TritonGPUOps.td
@@ -454,7 +454,7 @@ def TTG_GlobalScratchAllocOp : TTG_Op<"global_scratch_alloc"> {
 
 def TTG_WarpSpecializeOp : TTG_Op<"warp_specialize", [
   RecursiveMemoryEffects, RecursivelySpeculatable, AsyncRegions,
-  DeclareOpInterfaceMethods<RegionBranchOpInterface>
+  DeclareOpInterfaceMethods<RegionBranchOpInterface, ["getSuccessorInputs"]>
 ]> {
   let summary = "asynchronously execute code on multiple warpgroups";
   let description = [{
@@ -527,10 +527,13 @@ def TTG_WarpSpecializeOp : TTG_Op<"warp_specialize", [
   let hasCanonicalizeMethod = 1;
 }
 
-def TTG_WarpSpecializePartitionsOp : TTG_Op<"warp_specialize.partitions", [
-  IsolatedFromAbove, RecursiveMemoryEffects, RecursivelySpeculatable,
-  Terminator, HasParent<"WarpSpecializeOp">
-]> {
+def TTG_WarpSpecializePartitionsOp
+    : TTG_Op<"warp_specialize.partitions",
+             [IsolatedFromAbove, RecursiveMemoryEffects,
+              RecursivelySpeculatable, Terminator,
+              HasParent<"WarpSpecializeOp">,
+              DeclareOpInterfaceMethods<
+                  RegionBranchOpInterface, ["getSuccessorInputs"]>]> {
   let summary = "container op for `ttg.warp_specialize`";
   let description = [{
     Because MLIR requires entire operations be isolated from above, this op
diff --git a/lib/Analysis/Alias.cpp b/lib/Analysis/Alias.cpp
@@ -60,12 +60,13 @@ LogicalResult SharedMemoryAliasAnalysis::visitOperation(
 
 void SharedMemoryAliasAnalysis::visitNonControlFlowArguments(
     Operation *op, const RegionSuccessor &successor,
+    ValueRange successorInputs,
     ArrayRef<dataflow::Lattice<AliasInfo> *> argLattices, unsigned firstIndex) {
   auto wsOp = dyn_cast<triton::gpu::WarpSpecializePartitionsOp>(op);
   if (!wsOp) {
     setAllToEntryStates(argLattices.take_front(firstIndex));
     setAllToEntryStates(argLattices.drop_front(
-        firstIndex + successor.getSuccessorInputs().size()));
+        firstIndex + successorInputs.size()));
     return;
   }
 
diff --git a/lib/Analysis/AxisInfo.cpp b/lib/Analysis/AxisInfo.cpp
@@ -146,6 +146,7 @@ class AxisInfoAnalysis : public dataflow::SparseForwardDataFlowAnalysis<
 
   void visitNonControlFlowArguments(
       Operation *op, const RegionSuccessor &successor,
+      ValueRange successorInputs,
       ArrayRef<dataflow::Lattice<AxisInfo> *> argLattices,
       unsigned firstIndex) override {
     if (auto forOp = dyn_cast<scf::ForOp>(op)) {
@@ -154,8 +155,8 @@ class AxisInfoAnalysis : public dataflow::SparseForwardDataFlowAnalysis<
       visitWarpSpecializeExplicitCaptures(ws, successor, argLattices);
     } else {
       setAllToEntryStates(argLattices.take_front(firstIndex));
-      setAllToEntryStates(argLattices.drop_front(
-          firstIndex + successor.getSuccessorInputs().size()));
+      setAllToEntryStates(
+          argLattices.drop_front(firstIndex + successorInputs.size()));
     }
   }
 
diff --git a/lib/Dialect/TritonGPU/IR/Ops.cpp b/lib/Dialect/TritonGPU/IR/Ops.cpp
@@ -929,9 +929,30 @@ void WarpSpecializeOp::getSuccessorRegions(
     return;
   }
   // And the default region branches transparently back to the parent.
-  assert(src.getTerminatorPredecessorOrNull()->getParentRegion() ==
-         &getDefaultRegion());
-  successors.push_back(RegionSuccessor(getOperation(), getResults()));
+  if (src.getTerminatorPredecessorOrNull()->getParentRegion() ==
+      &getDefaultRegion())
+    successors.push_back(RegionSuccessor::parent());
+}
+
+ValueRange WarpSpecializeOp::getSuccessorInputs(RegionSuccessor successor) {
+  // When returning to parent, the successor inputs are the op results.
+  return successor.isParent() ? getResults() : ValueRange();
+}
+
+void WarpSpecializePartitionsOp::getSuccessorRegions(
+    RegionBranchPoint src, SmallVectorImpl<RegionSuccessor> &successors) {
+  // The parent branches to each of the partition regions, but nothing flows out
+  // of the partition regions.
+  if (src.isParent())
+    for (Region &region : getPartitionRegions())
+      successors.emplace_back(&region);
+}
+
+ValueRange
+WarpSpecializePartitionsOp::getSuccessorInputs(RegionSuccessor successor) {
+  // The successor inputs are the block arguments of the partition region.
+  Region *region = successor.getSuccessor();
+  return region ? region->getArguments() : ValueRange();
 }
 
 LogicalResult WarpSpecializeOp::verify() {
diff --git a/lib/Dialect/TritonGPU/Transforms/WarpSpecialization/OptimizePartitionWarps.cpp b/lib/Dialect/TritonGPU/Transforms/WarpSpecialization/OptimizePartitionWarps.cpp
@@ -305,7 +305,7 @@ void OptimizePartitionWarps::runOnOperation() {
     // The module must be directly nested under the current op for `runPipeline`
     // to work.
     getOperation().push_back(container);
-    auto remove = llvm::make_scope_exit([&] { container->remove(); });
+    llvm::scope_exit remove([&] { container->remove(); });
     return runPipeline(pm, container);
   };
 
diff --git a/lib/Instrumentation/PrintLoadStoreMemSpaces.cpp b/lib/Instrumentation/PrintLoadStoreMemSpaces.cpp
@@ -1,7 +1,7 @@
 #include "llvm/IR/Module.h"
 #include "llvm/IR/PassManager.h"
 #include "llvm/Passes/PassBuilder.h"
-#include "llvm/Passes/PassPlugin.h"
+#include "llvm/Plugins/PassPlugin.h"
 #include <map>
 
 using namespace llvm;
diff --git a/python/src/llvm.cc b/python/src/llvm.cc
@@ -17,8 +17,8 @@
 #include "llvm/Pass.h"
 #include "llvm/Passes/OptimizationLevel.h"
 #include "llvm/Passes/PassBuilder.h"
-#include "llvm/Passes/PassPlugin.h"
 #include "llvm/Passes/StandardInstrumentations.h"
+#include "llvm/Plugins/PassPlugin.h"
 #include "llvm/Support/CodeGen.h"
 #include "llvm/Support/Signals.h"
 #include "llvm/Support/SourceMgr.h"
diff --git a/test/lib/Instrumentation/GPUHello.cpp b/test/lib/Instrumentation/GPUHello.cpp
@@ -3,7 +3,7 @@
 #include "llvm/IR/PassManager.h"
 #include "llvm/Pass.h"
 #include "llvm/Passes/PassBuilder.h"
-#include "llvm/Passes/PassPlugin.h"
+#include "llvm/Plugins/PassPlugin.h"
 #include "llvm/Support/raw_ostream.h"
 #include <iostream>
 #include <vector>
diff --git a/third_party/amd/lib/Analysis/RangeAnalysis.cpp b/third_party/amd/lib/Analysis/RangeAnalysis.cpp
@@ -711,24 +711,22 @@ void TritonIntegerRangeAnalysis::visitRegionSuccessors(
 
     unsigned firstIndex = 0;
     if (inputs.size() != lattices.size()) {
-      if (!point->isBlockStart()) {
+      if (successor.isParent()) {
         if (!inputs.empty()) {
           firstIndex = cast<OpResult>(inputs.front()).getResultNumber();
         }
         visitNonControlFlowArguments(
-            branch,
-            RegionSuccessor(
-                branch, branch->getResults().slice(firstIndex, inputs.size())),
-            lattices, firstIndex);
+            branch, successor,
+            branch->getResults().slice(firstIndex, inputs.size()), lattices,
+            firstIndex);
       } else {
         if (!inputs.empty()) {
           firstIndex = cast<BlockArgument>(inputs.front()).getArgNumber();
         }
-        Region *region = point->getBlock()->getParent();
         visitNonControlFlowArguments(
-            branch,
-            RegionSuccessor(region, region->getArguments().slice(
-                                        firstIndex, inputs.size())),
+            branch, successor,
+            successor.getSuccessor()->getArguments().slice(firstIndex,
+                                                           inputs.size()),
             lattices, firstIndex);
       }
     }
@@ -737,7 +735,14 @@ void TritonIntegerRangeAnalysis::visitRegionSuccessors(
          llvm::zip(*operands, ArrayRef(lattices).drop_front(firstIndex))) {
       std::pair loopArgLat = {loop, argLat};
       // If we've "run the loop" #tripcount times, stop propagating.
-      if (loop && loopVisits[loopArgLat] >= loopTripCounts[loop])
+      bool reachedTripCount =
+          loop && loopVisits[loopArgLat] >= loopTripCounts[loop];
+      // However, if trip count is 0, we still need to initialize loop-carried
+      // values from the initial iter_args (so loop results equal initial
+      // values).
+      bool needsZeroTripInit = loop && loopTripCounts[loop] == 0 &&
+                               argLat->getValue().isUninitialized();
+      if (reachedTripCount && !needsZeroTripInit)
         continue;
 
       ChangeResult changed;
@@ -765,7 +770,9 @@ void TritonIntegerRangeAnalysis::visitRegionSuccessors(
       // lattice because otherwise we will over count the number of visits
       // (since not all iter_arg lattices are updated/propagated on each
       // visit).
-      if (loop && changed == ChangeResult::Change)
+      // For initial iterations of zero trip count loops we do not increment
+      // the visit count to avoid overcounting.
+      if (loop && changed == ChangeResult::Change && !needsZeroTripInit)
         ++loopVisits[loopArgLat];
     }
   }
diff --git a/third_party/nvidia/hopper/lib/Transforms/WarpSpecialization/TaskIdPropagation.cpp b/third_party/nvidia/hopper/lib/Transforms/WarpSpecialization/TaskIdPropagation.cpp
@@ -193,6 +193,9 @@ void TaskIdBackwardPropagation::visitCallOperand(OpOperand &operand) {
       "Should not have any call operands in the IR after inlining.");
 }
 
+void TaskIdBackwardPropagation::visitNonControlFlowArguments(
+    RegionSuccessor &successor, ArrayRef<BlockArgument> arguments) {}
+
 void TaskIdBackwardPropagation::setToExitState(TaskIdLattice *lattice) {}
 
 } // namespace mlir::triton::gpu
diff --git a/third_party/nvidia/hopper/lib/Transforms/WarpSpecialization/TaskIdPropagation.h b/third_party/nvidia/hopper/lib/Transforms/WarpSpecialization/TaskIdPropagation.h
@@ -87,6 +87,9 @@ class TaskIdBackwardPropagation
 
   void visitCallOperand(OpOperand &operand) override;
 
+  void visitNonControlFlowArguments(RegionSuccessor &successor,
+                                    ArrayRef<BlockArgument> arguments) override;
+
   void setToExitState(TaskIdLattice *lattice) override;
 
   void propagateToYield(scf::YieldOp yieldOp, SmallVector<TaskId> &lattices);

Original file line number	Diff line number	Diff line change
`@@ -1 +1 @@`
`1`		`-0729a74e66aeeb7a9839d80bfd64fc49b2e69f52`
	`1`	`+2eb709b95d8f521aa15401e159fac0729d56a677`