[ConSan] Support for WGMMA. Checks on non-async shmem and tmem accesses (#7712)

This PR covers:
* Adding support for wgmma checks (ops marked as read until wgmma_wait)
* Adding checks on any shmem and tmem access - this finalizes the basic
support for pipelined kernels
* Slight refactor of IR - op and attribute names

Author: pawelszczerbuk

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

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();