[ConSan] Add support for mmav5, checks for buffer reads (#7433)

Main functionality included in this change is a support for tracking and
checking outstanding reads.
With this we can catch issues like insufficient multibuffering of the
mma operands - added test confirming the functionality works.
This change also features slight refactor of the IR - instead of having
ops per the TritonGPU op that we want to instrument, I have introduced
more modular ops that do individual checks/aux data updates. This should
make adding further instrumentation much easier.

Author: pawelszczerbuk

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

@@ -411,12 +411,6 @@ def TTG_GlobalScratchAllocOp : TTG_Op<"global_scratch_alloc"> {
   );
   let results = (outs Arg<TT_Ptr, "", [MemAlloc<GlobalMemory>]>:$result);
 
-  let builders = [
-    OpBuilder<(ins "Type":$result, "int32_t":$nbytes, "int32_t":$alignment),
-              [{ build($_builder, $_state, result,
-                       $_builder.getI32IntegerAttr(nbytes), $_builder.getI32IntegerAttr(alignment)); }]>
-  ];
-
   let assemblyFormat = [{attr-dict `:` qualified(type($result))}];
 }
 

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

@@ -7,11 +7,32 @@ include "triton/Dialect/Triton/IR/TritonTypes.td"
 include "mlir/IR/OpBase.td"
 include "mlir/Interfaces/SideEffectInterfaces.td"
 
+//
+// Interfaces
+//
+def GlobalMemory : Resource<"::mlir::triton::GlobalMemory">;
+
+//
+// Ops
+//
+
 class TTI_Op<string mnemonic, list<Trait> traits = []> :
     Op<TritonInstrument_Dialect, mnemonic, traits> {
 }
 
-// Define an array of pointers to shared memory buffers
+def TTI_ExperimentalAssertInThreadOp : TTI_Op<"experimental_assert_in_thread", [MemoryEffects<[MemWrite<GlobalMemory>]>]> {
+  let summary = "assert the condition within the current thread";
+  let description = [{
+    Assert that the condition is true given all the values are available in the current thread.
+    If the condition is false, the message is printed, and the program is aborted.
+    If check_any is true, any of the values in the condition must be true. Otherwise, all the
+    values in the condition must be true.
+  }];
+  let arguments = (ins I1Tensor:$condition, StrAttr:$message, BoolAttr:$check_any);
+  let assemblyFormat = "$condition `,` $message attr-dict `:` type($condition)";
+}
+
+
 def TTI_ExperimentalSharedBufferPointersOp : TTI_Op<"experimental_shared_buffer_pointers", [Pure]> {
   let summary = "definte an array of pointers to shared memory buffers";
   let description = [{
@@ -24,59 +45,119 @@ def TTI_ExperimentalSharedBufferPointersOp : TTI_Op<"experimental_shared_buffer_
   }];
 }
 
-// Check if writing to a buffer guarded by a mbar is valid
-def TTI_ExperimentalCheckAsyncWriteWithMbarSharedOp : TTI_Op<"experimental_check_async_write_with_mbar_shared", [Pure]> {
-  let summary = "check if writing to a buffer guarded by a mbar is valid";
+
+def TTI_ExperimentalCheckOutstandingWritesOp : TTI_Op<"experimental_check_outstanding_writes", [MemoryEffects<[MemWrite<GlobalMemory>]>]> {
+  let summary = "check if there are outstanding writes to a buffer guarded by a mbar";
   let description = [{
-    Check if writing to a shared memory buffer guarded by a mbar is valid.
-    Update the buffer state and assert if the buffer is being read or written.
+    Check if there are outstanding writes to a buffer guarded by a mbar.
   }];
   let arguments = (ins
-    TTG_MemDescType:$buffer,
-    TTG_MemDescType:$mbar,
+    TTG_MemDescType:$buf,
     TT_Tensor:$buffers,
-    TT_Tensor:$states,
-    TT_Tensor:$barriers
+    TT_PtrLike:$writeBars,
+    TypeAttr:$writeBarsType,
+    Optional<I1>:$pred
   );
-  let results = (outs
-    TT_Tensor:$outStates,
-    TT_Tensor:$outBarriers
+  let assemblyFormat = [{
+    $buf `{` $buffers `,` $writeBars `(` $writeBarsType `)` `}` (`,` $pred^)? attr-dict `:` type($buf) `,` type($buffers) `,` type($writeBars)
+  }];
+  let hasVerifier = 1;
+}
+
+
+def TTI_ExperimentalCheckOutstandingReadsOp : TTI_Op<"experimental_check_outstanding_reads", [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.
+  }];
+  let arguments = (ins
+    TTG_MemDescType:$buf,
+    TT_Tensor:$buffers,
+    TT_PtrLike:$readBars,
+    TypeAttr:$readBarsType,
+    Optional<I1>:$pred
   );
   let assemblyFormat = [{
-    $buffer `,` $mbar `{` $buffers `,` $states `,` $barriers `}` attr-dict `:` type($buffer) `,` type($mbar) `,` type($buffers) `,` type($states) `,` type($barriers) `->` type($outStates) `,` type($outBarriers)
+    $buf `{` $buffers `,` $readBars `(` $readBarsType `)` `}` (`,` $pred^)? attr-dict `:` type($buf) `,` type($buffers) `,` type($readBars)
   }];
-  let builders = [
-    OpBuilder<(ins "Value":$buffer, "Value":$mbar, "Value":$buffers, "Value":$states, "Value":$barriers),[{
-      build($_builder, $_state, {states.getType(), barriers.getType()}, buffer, mbar, buffers, states, barriers);
-    }]>
-  ];
+  let hasVerifier = 1;
 }
 
-def TTI_ExperimentalCheckWaitMbarOp : TTI_Op<"experimental_check_wait_mbar", [Pure]> {
-  let summary = "check if waiting on a mbar is valid and update the barrier state";
+
+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";
   let description = [{
-    Check if waiting on a mbar is valid and update the barrier state.
+    Mark a buffer as being written to using mbar as a guard.
   }];
   let arguments = (ins
+    TTG_MemDescType:$buf,
     TTG_MemDescType:$mbar,
+    TT_Tensor:$buffers,
+    TT_PtrLike:$writeBars,
+    TypeAttr:$writeBarsType,
+    Optional<I1>:$pred
+  );
+  let assemblyFormat = [{
+    $buf `,` $mbar `{` $buffers `,` $writeBars `(` $writeBarsType `)` `}` (`,` $pred^)? attr-dict `:` type($buf) `,` type($mbar) `,` type($buffers) `,` type($writeBars)
+  }];
+  let hasVerifier = 1;
+}
+
+
+def TTI_ExperimentalMarkAsReadOp : TTI_Op<"experimental_mark_as_read", [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.
+  }];
+  let arguments = (ins
+    TTG_MemDescType:$buf,
+    TTG_MemDescType:$mbar,
+    TT_Tensor:$buffers,
     TT_Tensor:$barriers,
-    TT_Tensor:$states
-    );
+    TT_PtrLike:$readBars,
+    TypeAttr:$readBarsType,
+    Optional<I1>:$pred
+  );
+  let assemblyFormat = [{
+    $buf `,` $mbar `{` $buffers `,` $barriers `,` $readBars `(` $readBarsType `)` `}` (`,` $pred^)? attr-dict `:` type($buf) `,` type($mbar) `,` type($buffers) `,` type($barriers) `,` type($readBars)
+  }];
+  let hasVerifier = 1;
+}
 
-  let results = (outs
-    TT_Tensor:$outStates,
-    TT_Tensor:$outBarriers);
 
+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.
+  }];
+  let arguments = (ins
+    TTG_MemDescType:$mbar,
+    TT_PtrLike:$writeBars,
+    TypeAttr:$writeBarsType,
+    Optional<I1>:$pred
+  );
   let assemblyFormat = [{
-    $mbar `{` $states `,` $barriers `}` attr-dict `:` type($mbar) `,` type($states) `,` type($barriers) `->` type($outStates) `,` type($outBarriers)
+    $mbar `{` $writeBars `(` $writeBarsType `)` `}` (`,` $pred^)? attr-dict `:` type($mbar) `,` type($writeBars)
   }];
+}
 
-  let builders = [
-    OpBuilder<(ins "Value":$mbar, "Value":$barriers, "Value":$states),
-    [{
-      build($_builder, $_state, {states.getType(), barriers.getType()}, mbar, barriers, states);
-    }]>];
 
+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.
+  }];
+  let arguments = (ins
+    TTG_MemDescType:$mbar,
+    TT_Tensor:$barriers,
+    TT_PtrLike:$readBars,
+    TypeAttr:$readBarsType,
+    Optional<I1>:$pred
+  );
+  let assemblyFormat = [{
+    $mbar `{` $barriers `,` $readBars `(` $readBarsType `)` `}` (`,` $pred^)? attr-dict `:` type($mbar) `,` type($barriers) `,` type($readBars)
+  }];
+  let hasVerifier = 1;
 }
 
 #endif // TRITONINSTRUMENT_OPS

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

@@ -0,0 +1,14 @@
+#include "triton/Dialect/TritonInstrument/IR/Dialect.h"
+
+namespace mlir::triton::instrument {
+
+Operation *createStoreScratchMemory(OpBuilder &b, Location loc, Value alloc,
+                                    Value tensor, RankedTensorType tensorType);
+Operation *createLoadScratchMemory(OpBuilder &b, Location loc, Value alloc,
+                                   RankedTensorType tensorType);
+Value expandOuterSlicedDim(OpBuilder &b, Location loc, Value tensor);
+TypedValue<RankedTensorType> createConstIntTensor(OpBuilder &builder,
+                                                  Location loc, int val,
+                                                  RankedTensorType tensorType);
+
+} // namespace mlir::triton::instrument