Merge OpenAI Triton commit 2fafd63 (#5026)

This PR change the Triton base from
83fbc0e to
2fafd63 (Aug 20).
Pass rate: 98.85%->98.84%

Author: whitneywhtsang

bin/CMakeLists.txt

@@ -1,5 +1,3 @@
-get_property(dialect_libs GLOBAL PROPERTY MLIR_DIALECT_LIBS)
-get_property(conversion_libs GLOBAL PROPERTY MLIR_CONVERSION_LIBS)
 get_property(triton_libs GLOBAL PROPERTY TRITON_LIBS)
 
 add_llvm_executable(triton-opt triton-opt.cpp PARTIAL_SOURCES_INTENDED)
@@ -8,8 +6,6 @@ add_llvm_executable(triton-opt triton-opt.cpp PARTIAL_SOURCES_INTENDED)
 llvm_update_compile_flags(triton-opt)
 target_link_libraries(triton-opt PRIVATE
   TritonIntelLLVMIR
-  ${dialect_libs}
-  ${conversion_libs}
   ${triton_libs}
   # tests
   TritonTestAnalysis
@@ -19,6 +15,8 @@ target_link_libraries(triton-opt PRIVATE
   # MLIR core
   MLIROptLib
   MLIRPass
+  MLIRRegisterAllDialects
+  MLIRRegisterAllPasses
   MLIRTransforms
 )
 
@@ -29,8 +27,6 @@ mlir_check_all_link_libraries(triton-reduce)
 
 llvm_update_compile_flags(triton-reduce)
 target_link_libraries(triton-reduce PRIVATE
-  ${dialect_libs}
-  ${conversion_libs}
   ${triton_libs}
   # tests
   TritonTestAnalysis
@@ -40,6 +36,8 @@ target_link_libraries(triton-reduce PRIVATE
   # MLIR core
   MLIRReduceLib
   MLIRPass
+  MLIRRegisterAllDialects
+  MLIRRegisterAllPasses
   MLIRTransforms
 )
 
@@ -49,8 +47,6 @@ add_llvm_executable(triton-lsp triton-lsp.cpp PARTIAL_SOURCES_INTENDED)
 
 llvm_update_compile_flags(triton-lsp)
 target_link_libraries(triton-lsp PRIVATE
-  ${dialect_libs}
-  ${conversion_libs}
   ${triton_libs}
   # tests
   TritonTestAnalysis
@@ -60,6 +56,8 @@ target_link_libraries(triton-lsp PRIVATE
   # MLIR core
   MLIRLspServerLib
   MLIRPass
+  MLIRRegisterAllDialects
+  MLIRRegisterAllPasses
   MLIRTransforms
 )
 
@@ -91,10 +89,11 @@ export_executable_symbols_for_plugins(triton-llvm-opt)
 add_llvm_executable(triton-tensor-layout triton-tensor-layout.cpp PARTIAL_SOURCES_INTENDED)
 target_link_libraries(triton-tensor-layout PRIVATE
   ${triton_libs}
-  ${conversion_libs}
-  ${dialect_libs}
   TritonTestAnalysis
   TritonTestDialect
   TritonTestProton
   TritonAMDGPUTestAnalysis
+  MLIRRegisterAllDialects
+  MLIRRegisterAllPasses
+  MLIRTransforms
   )

cmake/llvm-hash.txt

@@ -1 +1 @@
-570885128351868c1308bb22e8ca351d318bc4a1
+bc773632355b3cebde350b0341624e88be40b744

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

@@ -90,15 +90,4 @@ def TritonLoopAwareCSE : Pass<"triton-loop-aware-cse", "mlir::ModuleOp"> {
   }];
 }
 
-def TritonSCFToCF : Pass</*cli-arg*/"triton-scf-to-cf", /*Op*/"mlir::ModuleOp"> {
-  let summary = "MLIR's SCF To CF plus some extra attributes propagation.";
-  let description = [{
-    This pass uses MLIR's SCF To CF pass as base. Additionally, it propagates
-    some extra attributes to the converted CFG.
-    TODO: upstream the llvm loop attribute propagation and remove this pass.
-  }];
-
-  let dependentDialects = [];
-}
-
 #endif

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

@@ -100,12 +100,6 @@ def TTG_AsyncCopyGlobalToLocalOp : TTG_Op<"async_copy_global_to_local", [
     DefaultValuedAttr<BoolAttr, "false">:$isVolatile
   );
 
-  let builders = [
-      OpBuilder<(ins "Value":$src, "Value":$result,
-                     "triton::CacheModifier":$cache,
-                     "triton::EvictionPolicy":$evict, "bool":$isVolatile)>,
-  ];
-
   let results = (outs TTG_AsyncToken:$token);
 
   let extraClassDeclaration = [{
@@ -395,9 +389,6 @@ def TTG_MaskOp: TTG_Op<"mask",
     let arguments = (ins I1:$pred);
     let results = (outs Variadic<AnyType>:$result);
     let regions = (region SizedRegion<1>:$region);
-    let builders = [
-        OpBuilder<(ins "Value":$pred)>,
-    ];
 }
 
 def TTG_MaskReturnOp: TTG_Op<"mask.return",

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

@@ -8,6 +8,25 @@ include "mlir/IR/OpBase.td"
 include "mlir/Interfaces/SideEffectInterfaces.td"
 include "triton/Dialect/TritonInstrument/IR/TritonInstrumentAttrDefs.td"
 
+// Concurrency Sanitizer data structures:
+// ConSan keeps auxilary data requied for tracking memory accesses in tensors.
+// These tensors are stored as a distributed tensor or in global scratch memory.
+//
+// Tensor name   | Storage | Type     | Description
+// ------------- | ------- | -------- | -----------
+// buffers       | tensor  | <Bxi64>  | List of base pointers of all the buffers and sub-buffers in the program
+// barriers      | tensor  | <Mxi64>  | List of pointers to all individual mbarriers in the program
+// writeState    | scratch | <Mxi8>   | Marks which buffers are being written to.
+//               |         |          | Entries in this tensor are set when write operation is issued. Entries are bitfields, where:
+//               |         |          | - bit 0: 1 if the buffer is being written to
+//               |         |          | - bit 1: 1 if the write is *not* hwPipelined
+// writeBars     | scratch | <BxMxi8> | Which barriers track writes to which buffers.
+//               |         |          | Entries in this tensor are set when commit with barrier is called.
+// readBars      | scratch | <BxMxi8> | Which barriers track reads from which buffers.
+//               |         |          | Entries in this tensor are set when read operation with barrier is issued.
+// asyncCpCommits | scratch | <Bxi8>  | Tracks number of outstanding commits for buffers written with cp-async.
+// wgmmaCommits  | scratch | <Bxi8>  | Tracks number of outstanding commits for buffers being read by wgmma.
+
 //
 // Interfaces
 //
@@ -52,10 +71,6 @@ def TTI_ExperimentalCheckWriteStateOp : TTI_Op<"experimental_check_write_state",
   let description = [{
     Check if the writeState tensor has non-zero value associated with the buffer.
 
-    `writeState` is a tensor of 8b bitfields, where:
-    - bit 0: 1 if the buffer is being written to
-    - bit 1: 1 if the write is *not* hwPipelined
-
     If hwPipelined is true, shift the bitfield by 1 to check the second bit - this
     means that the error won't be triggered if another pipelined write is outstanding.
   }];
@@ -79,7 +94,7 @@ def TTI_ExperimentalCheckWriteStateOp : TTI_Op<"experimental_check_write_state",
 def TTI_ExperimentalCheckReadBarriersOp : TTI_Op<"experimental_check_read_barriers", [MemoryEffects<[MemWrite<GlobalMemory>]>]> {
   let summary = "check if there are outstanding reads from a buffer guarded by a mbar";
   let description = [{
-    Check if there are outstanding reads from a buffer guarded by a mbar.
+    Check if any of the entries in readBars in the row corresponding to the buffer is non-zero.
   }];
   let arguments = (ins
     TTG_MemDescType:$buf,
@@ -100,11 +115,7 @@ def TTI_ExperimentalSetWriteStateOp : TTI_Op<"experimental_set_write_state", [Me
   let description = [{
     Mark a buffer as being written to. It is not yet tracked by a barrier, until
     `commit_write_with_barrier` is called, at which point all the buffers being written
-    to are marked as tracked by the barrier.
-
-    `writeState` is a tensor of 8b bitfields, where:
-    - bit 0: 1 if the buffer is being written to
-    - bit 1: 1 if the write is *not* hwPipelined
+    to are marked as tracked by the barrier in writeBars tensor.
 
     If hwPipelined is true, the write won't trigger an error if another pipelined
     write is executed later without waiting for the barrier.
@@ -149,7 +160,7 @@ def TTI_ExperimentalCommitWriteWithBarrierOp : TTI_Op<"experimental_commit_write
 def TTI_ExperimentalSetReadBarrierOp : TTI_Op<"experimental_set_read_barrier", [MemoryEffects<[MemWrite<GlobalMemory>]>]> {
   let summary = "mark a buffer as being read from using mbar as a guard";
   let description = [{
-    Mark a buffer as being read from using mbar as a guard.
+    Set the entry under [buffer, mbar] in readBars tensor to 1, marking the buffer as tracked by the barrier.
   }];
   let arguments = (ins
     TTG_MemDescType:$buf,
@@ -170,7 +181,8 @@ def TTI_ExperimentalSetReadBarrierOp : TTI_Op<"experimental_set_read_barrier", [
 def TTI_ExperimentalClearWriteBarrierOp : TTI_Op<"experimental_clear_write_barrier", [MemoryEffects<[MemWrite<GlobalMemory>]>]> {
   let summary = "clear the write state for buffers being guarded by an mbar";
   let description = [{
-    Clear the write state for buffers being guarded by an mbar.
+    For each buffer that has [buffer, mbar] entry in writeBars tensor, set the corresponding entry in writeState tensor to 0.
+    Also, set the corresponding entry in writeBars tensor to 0.
   }];
   let arguments = (ins
     TTG_MemDescType:$mbar,
@@ -191,7 +203,7 @@ def TTI_ExperimentalClearWriteBarrierOp : TTI_Op<"experimental_clear_write_barri
 def TTI_ExperimentalClearReadBarrierOp : TTI_Op<"experimental_clear_read_barrier", [MemoryEffects<[MemWrite<GlobalMemory>]>]> {
   let summary = "clear the read state for buffers being guarded by an mbar";
   let description = [{
-    Clear the read state for buffers being guarded by an mbar.
+    Set all the entries in the column corresponding to the mbar in readBars tensor to 0.
   }];
   let arguments = (ins
     TTG_MemDescType:$mbar,
@@ -210,7 +222,7 @@ def TTI_ExperimentalClearReadBarrierOp : TTI_Op<"experimental_clear_read_barrier
 def TTI_ExperimentalCheckBarrierWritesClearedOp : TTI_Op<"experimental_check_barrier_writes_cleared", [MemoryEffects<[MemWrite<GlobalMemory>]>]> {
   let summary = "verify that the barrier is not used to track any writes";
   let description = [{
-    Verify that the barrier is not used to track any writes.
+    Verify that the column corresponding to the mbar in writeBars tensor is all 0.
   }];
   let arguments = (ins
     TTG_MemDescType:$mbar,
@@ -248,7 +260,8 @@ def TTI_ExperimentalStageAccessForCommitOp : TTI_Op<"experimental_stage_access_f
 def TTI_ExperimentalCommitAccessesOp : TTI_Op<"experimental_commit_accesses", [MemoryEffects<[MemWrite<GlobalMemory>]>]> {
   let summary = "Commit all the staged accesses for all the buffers.";
   let description = [{
-    Commit all the staged accesses for all the buffers.
+    Increment the value in outstandingCommits tensor for each entry greater than 0.
+    Change all the `-1` entries in outstandingCommits tensor to 1, signifying 1 outstanding commit.
   }];
   let arguments = (ins
     TT_PtrLike:$outstandingCommits,
@@ -277,7 +290,7 @@ def TTI_ExperimentalClearOutstandingCommitsOp : TTI_Op<"experimental_clear_outst
 def TTI_ExperimentalCheckOutstandingCommitsOp : TTI_Op<"experimental_check_outstanding_commits", [MemoryEffects<[MemWrite<GlobalMemory>]>]> {
   let summary = "Check if the buffer has an outstanding commit.";
   let description = [{
-    Check if the buffer has an outstanding commit.
+    Verify that the entry corresponding to the buffer in outstandingCommits tensor is 0.
   }];
   let arguments = (ins
     TTG_MemDescType:$buf,

lib/Dialect/Triton/Transforms/CMakeLists.txt

@@ -13,7 +13,6 @@ add_triton_library(TritonTransforms
   RewriteTensorDescriptorToPointer.cpp
   ArithTypeConversion.cpp
   FunctionTypeConversion.cpp
-  SCFToCF.cpp
 
   DEPENDS
   TritonTransformsIncGen

lib/Dialect/Triton/Transforms/RewriteTensorDescriptorToPointer.cpp

@@ -166,7 +166,7 @@ Value generateMaskFromOffsetRanges(OpBuilder &builder, const Location &loc,
 
     // Compare with lower bound
     Value lowerBound = builder.create<mlir::arith::ConstantIntOp>(
-        loc, 0, builder.getI64Type());
+        loc, builder.getI64Type(), 0);
     Value splatLowerBound = builder.create<triton::SplatOp>(
         loc, offsetWithRange.getType(), lowerBound);
     Value cmpLower = builder.create<arith::CmpIOp>(

lib/Dialect/Triton/Transforms/RewriteTensorPointer.cpp

@@ -135,7 +135,7 @@ struct RewritedInfo {
 
       // Compare with lower bound
       Value lowerBound = builder.create<mlir::arith::ConstantIntOp>(
-          loc, 0, builder.getI64Type());
+          loc, builder.getI64Type(), 0);
       Value splatLowerBound = builder.create<triton::SplatOp>(
           loc, offsetWithRange.getType(), lowerBound);
       Value cmpLower = builder.create<arith::CmpIOp>(