Merge OpenAI Triton commit 51021fb (#5238)

This PR change the Triton base from
22b1a44 to
51021fb (Sep 29).
Pass rate: 92.67%->92.74%

Author: whitneywhtsang

bin/RegisterTritonDialects.h

@@ -134,7 +134,8 @@ inline void registerTritonDialects(mlir::DialectRegistry &registry) {
   mlir::registerTritonAMDGPUHoistLayoutConversions();
   mlir::registerTritonAMDGPUReorderInstructions();
   mlir::registerTritonAMDGPUBlockPingpong();
-  mlir::registerTritonAMDGPUStreamPipeline();
+  mlir::registerTritonAMDGPUPipeline();
+  mlir::registerTritonAMDGPUScheduleLoops();
   mlir::registerTritonAMDGPUCanonicalizePointers();
   mlir::registerTritonAMDGPUConvertToBufferOps();
   mlir::registerTritonAMDGPUInThreadTranspose();

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

@@ -182,6 +182,8 @@ getLastUseOfPipelinedOp(ArrayRef<Operation *> ops, scf::ForOp forOp,
                         CoarseSchedule &schedule,
                         std::function<bool(Operation *)> filterUse = nullptr);
 
+// Clean up attributes passing over schedules across stages in pipelining
+void removePipeliningAttributes(ModuleOp moduleOp);
 } // namespace triton
 } // namespace mlir
 

include/triton/Dialect/TritonInstrument/IR/TritonInstrument.md

@@ -31,6 +31,8 @@ All types are generated on-demand (per partition) based on:
 - readVisibility (scratch, <B x 64 x i64>): Per-buffer, per-thread lanes. Each lane stores a 64-bit mask of other threads whose reads are visible to that lane’s thread
 - writeTracking (scratch, <B x K x i8>): Map buffers → barriers tracking writes (boolean stored in i8)
 - readTracking (scratch, <B x K x i64>): Map buffers → barriers tracking reads (bitmask of threads)
+- barrierStates (scratch, <K x i32>): Packed barrier metadata. Bit 0 stores the current phase, bits [1..8] the initial arrival count, bits [9..16] the current arrival count. The verifier checks underflow before updating, and flips the phase when the current count reaches zero.
+- waiting (scratch, <K x i32>): Per-barrier bitfield describing waiting threads. Each base thread gets two bits: bit (2 * thread + 0) is the waiting flag, bit (2 * thread + 1) stores the phase the thread is waiting on.
 - outstandingCommits (scratch, <B x 16 x i8>): Per-buffer, per-base-thread commit counters for cp.async and wgmma
 
 ## Visibility and legality rules
@@ -53,6 +55,20 @@ ConSan separates “tracking” from “visibility transfer”:
 - At arrive/commit sites (e.g., tc commit, arrive on mbarrier): ConSan emits the track ops for both reads and writes.
 - At waits: experimental_transfer_visible_reads / experimental_transfer_visible_writes propagates tracked visibility from the barrier back into the waiting thread’s visibility, and this transfer is repeated to peer threads (base, TMA, TC) to keep the three classes consistent.
 
+### Barrier phase/count tracking
+
+- experimental_init_barrier_state(barrier, count, barrierStates) initializes the per-barrier state with phase = 0 and both initial/current arrival counts = `count`.
+- experimental_verify_barrier_arrive(barrier, count, barrierStates) checks that subtracting `count` from the current arrival count would not underflow. The codegen emits an assert if it would.
+- experimental_update_barrier_state(barrier, count, barrierStates) applies the arrive: subtracts `count`, flips the phase when the count reaches zero, and reloads the current count from the initial count.
+
+### Deadlock detection
+
+ConSan records which phase each thread is waiting on:
+
+- experimental_set_waiting(barrier, baseThread, phase, barriers, waiting) sets the waiting flag for `baseThread` and stores the requested `phase`. The flag/phase bits share the waiting bitfield (two bits per base thread).
+- experimental_check_all_active_waiting(activeMask, barriers, waiting, barrierStates) filters waiting threads to those whose stored phase matches the current barrier phase. If all active threads are waiting on matching phases, it raises a deadlock assert.
+- experimental_clear_waiting(barrier, baseThread, barriers, waiting) clears the waiting bits for `baseThread`. Each wait clears its own state after the wait completes.
+
 ## Commit-count–based synchronization
 
 Some hardware ops synchronize via “number of outstanding commits” rather than mbarriers.

include/triton/Dialect/TritonInstrument/IR/TritonInstrumentOps.td

@@ -12,16 +12,18 @@ include "triton/Dialect/TritonInstrument/IR/TritonInstrumentAttrDefs.td"
 // ConSan keeps auxilary data requied for tracking memory accesses in tensors.
 // These tensors are stored as a distributed tensor or in global scratch memory.
 //
-// Name            | Storage | Rank/Type       | Description
-// ----------------|---------|-----------------|------------
-// buffers         | tensor  | <B x i64>       | Base pointers of all (sub)buffers
-// barriers        | tensor  | <K x i64>       | Pointers to all individual mbarriers
-// writeVisibility | scratch | <B x i64>       | Per-buffer thread-visibility bitmask (bit i => thread i visible)
-// readVisibility  | scratch | <B x T x i64>   | Per-buffer, per-thread visibility lanes (row-updated; values are bitmasks)
-// writeTracking   | scratch | <B x K x i8>    | Map buffers -> barriers that track writes
-// readTracking    | scratch | <B x K x i64>   | Map buffers -> barriers that track reads
+// Name              | Storage | Rank/Type       | Description
+// ------------------|---------|-----------------|------------
+// buffers           | tensor  | <B x i64>       | Base pointers of all (sub)buffers
+// barriers          | tensor  | <K x i64>       | Pointers to all individual mbarriers
+// barrierStates     | scratch | <K x i32>       | Packed barrier phase (bit 0) and arrival counts (bits[1..8] init, [9..16] current)
+// waiting           | scratch | <K x i32>       | Two bits per thread: waiting flag bit (LSB), stored phase bit (bit 1)
+// writeVisibility   | scratch | <B x i64>       | Per-buffer thread-visibility bitmask (bit i => thread i visible)
+// readVisibility    | scratch | <B x T x i64>   | Per-buffer, per-thread visibility lanes (row-updated; values are bitmasks)
+// writeTracking     | scratch | <B x K x i8>    | Map buffers -> barriers that track writes
+// readTracking      | scratch | <B x K x i64>   | Map buffers -> barriers that track reads
 // outstandingCommits
-//   (async/wgmma) | scratch | <B x T x i8>    | Number of outstanding commits per buffer/thread (2D replaces prior 1D)
+//   (async/wgmma)   | scratch | <B x T x i8>    | Number of outstanding commits per buffer/thread (2D replaces prior 1D)
 
 //
 // Interfaces
@@ -62,6 +64,9 @@ def TTI_ExperimentalBufferPointersOp : TTI_Op<"experimental_buffer_pointers", [P
 }
 
 
+// ===== Critical section lock ops =====
+
+
 def TTI_ExperimentalLockAcquireOp : TTI_Op<"experimental_lock_acquire", [MemoryEffects<[MemWrite<GlobalMemory>]>]> {
   let summary = "Acquire a lock.";
   let description = [{
@@ -86,6 +91,9 @@ def TTI_ExperimentalLockReleaseOp : TTI_Op<"experimental_lock_release", [MemoryE
 }
 
 
+// ===== Barrier based synchronization ops =====
+
+
 def TTI_ExperimentalSetWriteVisibilityOp : TTI_Op<"experimental_set_write_visibility", [MemoryEffects<[MemWrite<GlobalMemory>]>]> {
   let summary = "set the write visibility for a buffer";
   let description = [{
@@ -315,6 +323,9 @@ def TTI_ExperimentalVerifyReadVisibilityOp : TTI_Op<"experimental_verify_read_vi
 }
 
 
+// ===== Commit-count–based synchronization ops =====
+
+
 def TTI_ExperimentalStageAccessForCommitOp : TTI_Op<"experimental_stage_access_for_commit", [MemoryEffects<[MemWrite<GlobalMemory>]>]> {
   let summary = "";
   let description = [{
@@ -409,4 +420,155 @@ def TTI_ExperimentalCheckOutstandingCommitsOp : TTI_Op<"experimental_check_outst
   let hasVerifier = 1;
 }
 
+
+// ===== Barrier state ops =====
+
+def TTI_ExperimentalInitBarrierStateOp : TTI_Op<"experimental_init_barrier_state", [MemoryEffects<[MemWrite<GlobalMemory>]>]> {
+  let summary = "initialize the auxiliary barrier state";
+  let description = [{
+    Initialize the tracked barrier state to phase 0 and set both the initial and current arrival counts.
+  }];
+  let arguments = (ins
+    TTG_MemDescType:$mbar,
+    I32Attr:$count,
+    TT_Tensor:$barriers,
+    TT_PtrLike:$states,
+    TypeAttr:$statesType
+  );
+  let assemblyFormat = [{
+    $mbar `,` $count `{` $barriers `,` $states `(` $statesType `)` `}` attr-dict `:` type($mbar) `,` type($barriers) `,` type($states)
+  }];
+}
+
+def TTI_ExperimentalVerifyBarrierArriveOp : TTI_Op<"experimental_verify_barrier_arrive", [MemoryEffects<[MemWrite<GlobalMemory>]>]> {
+  let summary = "verify an arrive count against the tracked barrier state";
+  let description = [{
+    Check that applying the arrive count would not drive the tracked current count negative. Triggers an assertion on failure.
+  }];
+  let arguments = (ins
+    TTG_MemDescType:$mbar,
+    I32Attr:$count,
+    TT_Tensor:$barriers,
+    TT_PtrLike:$states,
+    TypeAttr:$statesType,
+    Optional<I1>:$pred
+  );
+  let assemblyFormat = [{
+    $mbar `,` $count `{` $barriers `,` $states `(` $statesType `)` `}` (`,` $pred^)? attr-dict `:` type($mbar) `,` type($barriers) `,` type($states)
+  }];
+}
+
+def TTI_ExperimentalUpdateBarrierStateOp : TTI_Op<"experimental_update_barrier_state", [MemoryEffects<[MemWrite<GlobalMemory>]>]> {
+  let summary = "update the auxiliary barrier state after a verified arrive";
+  let description = [{
+    Apply an arrive count to the tracked barrier state, toggling the phase when the count reaches zero and reloading the current count from the initial count.
+  }];
+  let arguments = (ins
+    TTG_MemDescType:$mbar,
+    I32Attr:$count,
+    TT_Tensor:$barriers,
+    TT_PtrLike:$states,
+    TypeAttr:$statesType,
+    Optional<I1>:$pred
+  );
+  let assemblyFormat = [{
+    $mbar `,` $count `{` $barriers `,` $states `(` $statesType `)` `}` (`,` $pred^)? attr-dict `:` type($mbar) `,` type($barriers) `,` type($states)
+  }];
+}
+
+// ===== Deadlock detection ops =====
+
+def TTI_ExperimentalSetWaitingOp : TTI_Op<"experimental_set_waiting", [MemoryEffects<[MemWrite<GlobalMemory>]>]> {
+  let summary = "Mark the base thread as waiting on the given barrier phase";
+  let description = [{
+    For the barrier row matching mbar, set the waiting flag for baseThread and record the barrier phase being waited on.
+  }];
+  let arguments = (ins
+    TTG_MemDescType:$mbar,
+    I32Attr:$baseThread,
+    I32:$phase,
+    TT_Tensor:$barriers,
+    TT_PtrLike:$waiting,
+    TypeAttr:$waitingType,
+    Optional<I1>:$pred
+  );
+  let assemblyFormat = [{
+    $mbar `,` $baseThread `,` $phase `{` $barriers `,` $waiting `(` $waitingType `)` `}` (`,` $pred^)? attr-dict `:` type($mbar) `,` type($barriers) `,` type($waiting)
+  }];
+}
+
+def TTI_ExperimentalCheckAllActiveWaitingOp : TTI_Op<"experimental_check_all_active_waiting", [MemoryEffects<[MemWrite<GlobalMemory>]>]> {
+  let summary = "Assert that not all active threads are waiting on matching phases";
+  let description = [{
+    Filter waiting threads to those whose recorded phase matches the current barrier phase, OR-reduce across barriers, and assert that (waitingMask & activeMask) != activeMask.
+  }];
+  let arguments = (ins
+    I32Attr:$activeMask,
+    TT_Tensor:$barriers,
+    TT_PtrLike:$waiting,
+    TypeAttr:$waitingType,
+    TT_PtrLike:$barrierStates,
+    TypeAttr:$barrierStatesType,
+    Optional<I1>:$pred
+  );
+  let assemblyFormat = [{
+    $activeMask `,` $barriers `,` $waiting `(` $waitingType `)` `,` $barrierStates `(` $barrierStatesType `)` (`,` $pred^)? attr-dict `:` type($barriers) `,` type($waiting) `,` type($barrierStates)
+  }];
+}
+
+def TTI_ExperimentalClearWaitingOp : TTI_Op<"experimental_clear_waiting", [MemoryEffects<[MemWrite<GlobalMemory>]>]> {
+  let summary = "Clear the waiting state for the given base thread";
+  let description = [{
+    For the barrier row matching mbar, clear both the waiting flag and stored phase for baseThread.
+  }];
+  let arguments = (ins
+    TTG_MemDescType:$mbar,
+    I32Attr:$baseThread,
+    TT_Tensor:$barriers,
+    TT_PtrLike:$waiting,
+    TypeAttr:$waitingType,
+    Optional<I1>:$pred
+  );
+  let assemblyFormat = [{
+    $mbar `,` $baseThread `{` $barriers `,` $waiting `(` $waitingType `)` `}` (`,` $pred^)? attr-dict `:` type($mbar) `,` type($barriers) `,` type($waiting)
+  }];
+}
+
+
+// ===== Visibility replication ops =====
+
+def TTI_ExperimentalCopyWriteVisibilityOp : TTI_Op<"experimental_copy_write_visibility", [MemoryEffects<[MemWrite<GlobalMemory>]>]> {
+  let summary = "replicate write visibility from one thread to others";
+  let description = [{
+    Copy the write-visibility bit of sourceThread to every thread listed in destMask for all buffers. Destination bits are overwritten.
+  }];
+  let arguments = (ins
+    I32Attr:$sourceThread,
+    I64Attr:$destMask,
+    TT_PtrLike:$writeVisibility,
+    TypeAttr:$writeVisibilityType,
+    Optional<I1>:$pred
+  );
+  let assemblyFormat = [{
+    $sourceThread `,` $destMask `{` $writeVisibility `(` $writeVisibilityType `)` `}` (`,` $pred^)? attr-dict `:` type($writeVisibility)
+  }];
+}
+
+def TTI_ExperimentalCopyReadVisibilityOp : TTI_Op<"experimental_copy_read_visibility", [MemoryEffects<[MemWrite<GlobalMemory>]>]> {
+  let summary = "replicate read visibility rows from one thread to others";
+  let description = [{
+    Copy the read-visibility row of sourceThread to every thread listed in destMask for all buffers. Destination rows are overwritten.
+  }];
+  let arguments = (ins
+    I32Attr:$sourceThread,
+    I64Attr:$destMask,
+    TT_PtrLike:$readVisibility,
+    TypeAttr:$readVisibilityType,
+    Optional<I1>:$pred
+  );
+  let assemblyFormat = [{
+    $sourceThread `,` $destMask `{` $readVisibility `(` $readVisibilityType `)` `}` (`,` $pred^)? attr-dict `:` type($readVisibility)
+  }];
+}
+
 #endif // TRITONINSTRUMENT_OPS

include/triton/Dialect/TritonInstrument/IR/Utility.h

@@ -19,7 +19,8 @@ Operation *createLoadScratchMemory(OpBuilder &b, Location loc, Value alloc,
 Value expandOuterSlicedDim(OpBuilder &b, Location loc, Value tensor);
 TypedValue<RankedTensorType> createConstIntTensor(OpBuilder &builder,
                                                   Location loc, int64_t val,
-                                                  RankedTensorType tensorType);
+                                                  RankedTensorType tensorType,
+                                                  bool isSigned = false);
 FuncOp getEntryPoint(ModuleOp module);
 gpu::DistributedEncodingTrait
 getSingleDimSliceEncoding(gpu::BlockedEncodingAttr encoding, int dim);
@@ -43,8 +44,11 @@ struct AuxDataMap {
     Region *getEnclosingParitionOrFunctionRegion(Operation *op);
   };
 
+  // Please see TritonInstrumentOps.td for more information on the auxiliary
+  // data structures.
   RegionToValueMap buffers[numMemTypes];
   RegionToValueMap barriers;
+  RegionToValueMap barrierStates;
 
   RegionToValueMap writeVisibility[numMemTypes];
   RegionToValueMap writeTracking[numMemTypes];
@@ -53,6 +57,7 @@ struct AuxDataMap {
   RegionToValueMap asyncCpCommits;
   RegionToValueMap wgmmaCommits;
   RegionToValueMap lock;
+  RegionToValueMap waiting;
 
   void populateAndPassToWarpSpecialize(ModuleOp module);
 

include/triton/Tools/Sys/GetEnv.hpp

@@ -32,8 +32,6 @@ inline const std::set<std::string> CACHE_INVALIDATING_ENV_VARS = {
     "TRITON_DEFAULT_FP_FUSION",
     "TRITON_DISABLE_LINE_INFO",
     "TRITON_ENABLE_LLVM_DEBUG",
-    "TRITON_HIP_GLOBAL_PREFETCH",
-    "TRITON_HIP_LOCAL_PREFETCH",
     "TRITON_HIP_USE_ASYNC_COPY",
     "TRITON_HIP_USE_BLOCK_PINGPONG",
     "TRITON_HIP_USE_IN_THREAD_TRANSPOSE",