Merge OpenAI Triton commit 690f690 (#4914)

This PR change the Triton base from
4bcdbde to
690f690 (Aug 1).
Pass rate: 98.86%->98.85% (#4916)

Author: whitneywhtsang

.github/CODEOWNERS

@@ -46,3 +46,13 @@ lib/Dialect/TritonGPU/Transforms/TritonGPUConversion.cpp @ptillet
 # third_party
 # -----------
 third_party/amd/ @antiagainst @zhanglx13
+
+# -----------
+# gluon
+# -----------
+python/triton/experimental/gluon/ @peterbell10
+python/src/gluon_ir.cc @peterbell10
+python/test/gluon @peterbell10
+test/Gluon @peterbell10
+include/triton/Dialect/Gluon @peterbell10
+lib/Dialect/Gluon @peterbell10

.github/workflows/integration-tests-amd.yml

@@ -117,6 +117,10 @@ jobs:
           if [ ! -d "${INSTRUMENTATION_LIB_DIR}" ]; then
             echo "Could not find '${INSTRUMENTATION_LIB_DIR}'" ; exit -1
           fi
+
+          # Test gluon
+          pytest --capture=tee-sys -rfs -n 8 python/test/gluon/
+
           pytest --capture=tee-sys -rfs python/tutorials/06-fused-attention.py
           pytest --capture=tee-sys -rfs third_party/amd/python/test/test_extract_slice_concat_op.py
           TRITON_ALWAYS_COMPILE=1 pytest --capture=tee-sys -rfs third_party/amd/python/test/test_scalarize_packed_fops.py

include/triton/Conversion/TritonGPUToLLVM/TargetInfoBase.h

@@ -4,6 +4,7 @@
 #include "triton/Conversion/MLIRTypes.h"
 
 namespace mlir::triton {
+enum class ProgramIDDim : uint32_t;
 
 class TargetInfoBase {
 public:
@@ -48,7 +49,7 @@ class TargetInfoBase {
                            Value i) const = 0;
 
   virtual Value programId(RewriterBase &rewriter, Location loc,
-                          ModuleOp moduleOp, int axis) const = 0;
+                          ModuleOp moduleOp, ProgramIDDim axis) const = 0;
 
   virtual bool warpReduce(RewriterBase &rewriter, Location loc,
                           SmallVector<Value> &acc, triton::ReduceOp op,

include/triton/Conversion/TritonGPUToLLVM/Utility.h

@@ -614,6 +614,9 @@ inline bool isCanonicalIndex(unsigned index, unsigned freeVarMask) {
 // group code isolated from above by invoking this function.
 void makeAllWarpGroupsIsolatedFromAbove(Operation *op);
 
+// Set the correct loop annotation on LLVM branch ops.
+void fixUpLoopAnnotation(ModuleOp mod);
+
 /// Converts ConverLayoutOp to llvm using padded pattern.
 /// This pattern adds unused memory locations after every rows of tensor fastest
 /// changing dimension:

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

@@ -90,4 +90,15 @@ 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/TritonInstrument/IR/TritonInstrumentAttrDefs.td

@@ -6,8 +6,8 @@ include "mlir/IR/EnumAttr.td"
 def TT_MemTypeAttr : I32EnumAttr<
     "MemType", "",
     [
-        I32EnumAttrCase<"SHARED", 0, "shared">,
-        I32EnumAttrCase<"TENSOR", 1, "tensor">,
+        I32EnumAttrCase<"SHARED_MEM", 0, "shared_mem">,
+        I32EnumAttrCase<"TENSOR_MEM", 1, "tensor_mem">,
     ]> {
     let cppNamespace = "::mlir::triton::instrument";
 }

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

@@ -47,7 +47,7 @@ def TTI_ExperimentalBufferPointersOp : TTI_Op<"experimental_buffer_pointers", [P
 }
 
 
-def TTI_ExperimentalCheckOutstandingWritesOp : TTI_Op<"experimental_check_outstanding_writes", [MemoryEffects<[MemWrite<GlobalMemory>]>]> {
+def TTI_ExperimentalCheckWriteStateOp : TTI_Op<"experimental_check_write_state", [MemoryEffects<[MemWrite<GlobalMemory>]>]> {
   let summary = "check if there are outstanding writes to a buffer guarded by a mbar";
   let description = [{
     Check if the writeState tensor has non-zero value associated with the buffer.
@@ -76,7 +76,7 @@ def TTI_ExperimentalCheckOutstandingWritesOp : TTI_Op<"experimental_check_outsta
 }
 
 
-def TTI_ExperimentalCheckOutstandingReadsOp : TTI_Op<"experimental_check_outstanding_reads", [MemoryEffects<[MemWrite<GlobalMemory>]>]> {
+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.
@@ -95,8 +95,8 @@ def TTI_ExperimentalCheckOutstandingReadsOp : TTI_Op<"experimental_check_outstan
 }
 
 
-def TTI_ExperimentalMarkAsWriteOp : TTI_Op<"experimental_mark_as_write", [MemoryEffects<[MemWrite<GlobalMemory>]>]> {
-  let summary = "mark a buffer as being written to using mbar as a guard";
+def TTI_ExperimentalSetWriteStateOp : TTI_Op<"experimental_set_write_state", [MemoryEffects<[MemWrite<GlobalMemory>]>]> {
+  let summary = "mark a buffer as being written in writeState tensor";
   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
@@ -146,7 +146,7 @@ def TTI_ExperimentalCommitWriteWithBarrierOp : TTI_Op<"experimental_commit_write
 }
 
 
-def TTI_ExperimentalMarkAsReadOp : TTI_Op<"experimental_mark_as_read", [MemoryEffects<[MemWrite<GlobalMemory>]>]> {
+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.
@@ -226,72 +226,70 @@ def TTI_ExperimentalCheckBarrierWritesClearedOp : TTI_Op<"experimental_check_bar
 }
 
 
-// TODO: Potentially resolve the naming/functionality clash with commit_write_with_barrier
-def TTI_ExperimentalStageWriteForCommitOp : TTI_Op<"experimental_stage_write_for_commit", [MemoryEffects<[MemWrite<GlobalMemory>]>]> {
-  let summary = "Preapre to an async copy of a buffer. Staged until commit_group is called.";
+def TTI_ExperimentalStageAccessForCommitOp : TTI_Op<"experimental_stage_access_for_commit", [MemoryEffects<[MemWrite<GlobalMemory>]>]> {
+  let summary = "";
   let description = [{
-    Preapre to an async copy of a buffer. Staged until commit_group is called. The implementation will write `-1` to the
-    `write_commits` tensor under the indices corresponding to the buffer.
+    For operations that use `outstanding` to track the number of outstanding commits (rather than mbarriers),
+    mark the buffer as being accessed, but not commited yet, by marking it with `-1`.
   }];
   let arguments = (ins
     TTG_MemDescType:$buf,
     TT_Tensor:$buffers,
-    TT_PtrLike:$writeCommits,
-    TypeAttr:$writeCommitsType,
+    TT_PtrLike:$outstandingCommits,
+    TypeAttr:$outstandingCommitsType,
     Optional<I1>:$pred
   );
   let assemblyFormat = [{
-    $buf `{` $buffers `,` $writeCommits `(` $writeCommitsType `)` `}` (`,` $pred^)? attr-dict `:` type($buf) `,` type($buffers) `,` type($writeCommits)
+    $buf `{` $buffers `,` $outstandingCommits `(` $outstandingCommitsType `)` `}` (`,` $pred^)? attr-dict `:` type($buf) `,` type($buffers) `,` type($outstandingCommits)
   }];
-  // let hasVerifier = 1;
+  let hasVerifier = 1;
 }
 
-def TTI_ExperimentalCommitWritesOp : TTI_Op<"experimental_commit_writes", [MemoryEffects<[MemWrite<GlobalMemory>]>]> {
-  let summary = "Commit all the staged writes for all the buffers.";
+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 writes for all the buffers.
+    Commit all the staged accesses for all the buffers.
   }];
   let arguments = (ins
-    TT_PtrLike:$writeCommits,
-    TypeAttr:$writeCommitsType,
+    TT_PtrLike:$outstandingCommits,
+    TypeAttr:$outstandingCommitsType,
     Optional<I1>:$pred);
   let assemblyFormat = [{
-    `{` $writeCommits `(` $writeCommitsType `)` `}` (`,` $pred^)? attr-dict `:` type($writeCommits)
+    `{` $outstandingCommits `(` $outstandingCommitsType `)` `}` (`,` $pred^)? attr-dict `:` type($outstandingCommits)
   }];
-  // let hasVerifier = 1;
 }
 
-def TTI_ExperimentalClearWriteCommitsOp : TTI_Op<"experimental_clear_write_commits", [MemoryEffects<[MemWrite<GlobalMemory>]>]> {
-  let summary = "Clear all the write commits more distant than `outstandingNum.";
+def TTI_ExperimentalClearOutstandingCommitsOp : TTI_Op<"experimental_clear_outstanding_commits", [MemoryEffects<[MemWrite<GlobalMemory>]>]> {
+  let summary = "Clear all the outstanding commits more distant than `outstandingNum.";
   let description = [{
-    Clear all the write commits more distant than `outstandingNum` from the current thread.
+    Clear all the outstanding commits more distant than `outstandingNum` from the current thread.
   }];
   let arguments = (ins
-    TT_PtrLike:$writeCommits,
-    TypeAttr:$writeCommitsType,
+    TT_PtrLike:$outstandingCommits,
+    TypeAttr:$outstandingCommitsType,
     I32Attr:$outstandingNum,
     Optional<I1>:$pred);
   let assemblyFormat = [{
-    `{` $writeCommits `(` $writeCommitsType `)` `}` `,` $outstandingNum (`,` $pred^)? attr-dict `:` type($writeCommits)
+    `{` $outstandingCommits `(` $outstandingCommitsType `)` `}` `,` $outstandingNum (`,` $pred^)? attr-dict `:` type($outstandingCommits)
   }];
-  // let hasVerifier = 1;
 }
 
-def TTI_ExperimentalCheckWriteCommitOp : TTI_Op<"experimental_check_write_commit", [MemoryEffects<[MemWrite<GlobalMemory>]>]> {
-  let summary = "Check if the buffer has an outstanding write commit.";
+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 write commit.
+    Check if the buffer has an outstanding commit.
   }];
   let arguments = (ins
     TTG_MemDescType:$buf,
     TT_Tensor:$buffers,
-    TT_PtrLike:$writeCommits,
-    TypeAttr:$writeCommitsType,
+    TT_PtrLike:$outstandingCommits,
+    TypeAttr:$outstandingCommitsType,
+    StrAttr:$pendingAccessType,
     Optional<I1>:$pred);
   let assemblyFormat = [{
-    $buf `{` $buffers `,` $writeCommits `(` $writeCommitsType `)` `}` (`,` $pred^)? attr-dict `:` type($buf) `,` type($buffers) `,` type($writeCommits)
+    $buf `{` $buffers `,` $outstandingCommits `(` $outstandingCommitsType `)` `}` `,` $pendingAccessType (`,` $pred^)? attr-dict `:` type($buf) `,` type($buffers) `,` type($outstandingCommits)
   }];
-  // let hasVerifier = 1;
+  let hasVerifier = 1;
 }
 
 #endif // TRITONINSTRUMENT_OPS

include/triton/Dialect/TritonNvidiaGPU/IR/TritonNvidiaGPUOps.td

@@ -123,6 +123,7 @@ def TTNG_WarpGroupDotWaitOp : TTNG_Op<"warp_group_dot_wait", [DeclareOpInterface
   }];
 
   let assemblyFormat = "$inputs attr-dict `:` type($inputs)";
+  let hasVerifier = 1;
 }
 
 def TTNG_InitBarrierOp : TTNG_Op<"init_barrier"> {

lib/Conversion/TritonGPUToLLVM/AssertOpToLLVM.cpp

@@ -63,6 +63,9 @@ struct AssertOpConversion : public ConvertOpToLLVMPattern<triton::AssertOp> {
     while (auto callLoc = dyn_cast<CallSiteLoc>(loc))
       loc = callLoc.getCallee();
 
+    while (auto nameLoc = dyn_cast<NameLoc>(loc))
+      loc = nameLoc.getChildLoc();
+
     if (auto fileLineColLoc = dyn_cast<FileLineColLoc>(loc)) {
       file = fileLineColLoc.getFilename();
       line = fileLineColLoc.getLine();

lib/Conversion/TritonGPUToLLVM/PrintOpToLLVM.cpp

@@ -26,11 +26,10 @@ struct PrintOpConversion : public ConvertOpToLLVMPattern<triton::PrintOp> {
                   ConversionPatternRewriter &rewriter) const override {
     auto loc = op->getLoc();
 
-    auto getPid = [&](int axis) {
-      return targetInfo.programId(rewriter, loc,
-                                  op->getParentOfType<ModuleOp>(), axis);
-    };
-    std::array<Value, 3> pid = {getPid(0), getPid(1), getPid(2)};
+    std::array<Value, 3> pid;
+    auto module = op->getParentOfType<ModuleOp>();
+    for (auto axis : {ProgramIDDim::X, ProgramIDDim::Y, ProgramIDDim::Z})
+      pid[(int)axis] = targetInfo.programId(rewriter, loc, module, axis);
 
     // Simple printf of a string without any tensors.
     if (op.getNumOperands() == 0) {