[CONSAN] Adding support for aliasing (#8939)

This change adds support for buffer aliasing in ConSan. This is achieved
by using BufferRegion analysis to create AliasingMatrix that marks which
buffers may overlap with each other. This information is then used
during the visibility checks: instead of verifying that the buffer being
accessed is visible, we load a row from the aliasing matrix, describing
all the buffers that may alias buffer being access (including the buffer
in question). This set of buffers is then checked for having visibility
into potential outstanding reads/writes.

There is still follow-up change coming: moving BufferRegion analysis to
TritonInstrument. It was not done in this PR to make review easier.

Author: pawelszczerbuk

include/triton/Analysis/BufferRegion.h

@@ -162,6 +162,8 @@ class BufferRegionAnalysis : public dataflow::SparseForwardDataFlowAnalysis<
 private:
   // Global registry of all regions
   std::set<BufferRegion> usedBufferRegions[NUM_REGION_TYPES];
+
+  static void verifyOpIsSupported(Operation *op);
 };
 
 } // namespace mlir::triton

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

@@ -117,25 +117,27 @@ class FunctionBuilder {
   // from the visibility bitmask. We know this is safe because there cannot be
   // outstanding writes to this buffer at this point.
   void createSetWriteVisibilityCall(ImplicitLocOpBuilder &b, Value buf,
-                                    uint64_t threadMask, Value pred,
-                                    MemType memType, Operation *insertPoint);
+                                    uint32_t length, uint64_t threadMask,
+                                    Value pred, MemType memType,
+                                    Operation *insertPoint);
   // setReadVisibility: add the threads set in threadMask to the buffer's read
   // visibility bitmask.
   void createSetReadVisibilityCall(ImplicitLocOpBuilder &b, Value buf,
-                                   uint64_t threadMask, Value pred,
-                                   MemType memType, Operation *insertPoint);
+                                   uint32_t length, uint64_t threadMask,
+                                   Value pred, MemType memType,
+                                   Operation *insertPoint);
   // clearWriteTracking: clear all the information about threads writing to a
   // buffer.
   void createClearWriteTrackingCall(ImplicitLocOpBuilder &b, Value buf,
-                                    Value pred, MemType memType,
-                                    Operation *insertPoint);
+                                    uint32_t length, Value pred,
+                                    MemType memType, Operation *insertPoint);
   // clearReadVisibility: clear the read visibility for a buffer.
   void createClearReadVisibilityCall(ImplicitLocOpBuilder &b, Value buf,
-                                     Value pred, MemType memType,
-                                     Operation *insertPoint);
+                                     uint32_t length, Value pred,
+                                     MemType memType, Operation *insertPoint);
   // clearReadTracking: clear the read tracking for a buffer.
   void createClearReadTrackingCall(ImplicitLocOpBuilder &b, Value buf,
-                                   Value pred, MemType memType,
+                                   uint32_t length, Value pred, MemType memType,
                                    Operation *insertPoint);
   // trackVisibleWrites: snapshot buffers currently visible to the thread into
   // the tracking table for a barrier.
@@ -160,15 +162,15 @@ class FunctionBuilder {
   // verifyWriteVisibility: ensure the thread either sees the latest write or no
   // other thread is writing the buffer.
   void createVerifyWriteVisibilityCall(ImplicitLocOpBuilder &b, Value buf,
-                                       int thread, StringRef operandName,
-                                       Value pred, MemType memType,
-                                       Operation *insertPoint);
+                                       uint32_t length, int thread,
+                                       StringRef operandName, Value pred,
+                                       MemType memType, Operation *insertPoint);
   // verifyReadVisibility: ensure all reads from the buffer are visible to the
   // thread.
   void createVerifyReadVisibilityCall(ImplicitLocOpBuilder &b, Value buf,
-                                      int thread, StringRef operandName,
-                                      Value pred, MemType memType,
-                                      Operation *insertPoint);
+                                      uint32_t length, int thread,
+                                      StringRef operandName, Value pred,
+                                      MemType memType, Operation *insertPoint);
   // copyWriteVisibility: replicate the write visibility bit of sourceThread to
   // every destination thread in destMask.
   void createCopyWriteVisibilityCall(ImplicitLocOpBuilder &b, int sourceThread,
@@ -182,7 +184,8 @@ class FunctionBuilder {
   // stageAccessForCommit: mark the buffer as staged (value -1) in the
   // outstanding commit table for this thread.
   void createStageAccessForCommitCall(ImplicitLocOpBuilder &b, Value buf,
-                                      int thread, Value pred, MemType memType,
+                                      uint32_t length, int thread, Value pred,
+                                      MemType memType,
                                       CommitKind::Kind commitKind,
                                       Operation *insertPoint);
   // commitAccesses: convert staged entries to 1 and increment outstanding
@@ -207,7 +210,7 @@ class FunctionBuilder {
   // checkOutstandingCommits: assert that the outstanding commit row for the
   // buffer is zero before the access described by pendingAccessType.
   void createCheckOutstandingCommitsCall(ImplicitLocOpBuilder &b, Value buf,
-                                         int thread,
+                                         uint32_t length, int thread,
                                          StringRef pendingAccessType,
                                          Value pred, MemType memType,
                                          CommitKind::Kind commitKind,

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

@@ -34,30 +34,33 @@ def TTI_ExperimentalAssertInThreadOp : TTI_Op<"experimental_assert_in_thread", [
 }
 
 
-def TTI_ExperimentalBufferPointersOp : TTI_Op<"experimental_buffer_pointers", [Pure]> {
-  let summary = "definte an array of pointers to shared memory buffers";
+def TTI_ExperimentalBufferDescriptorsOp
+    : TTI_Op<"experimental_buffer_descriptors", [Pure]> {
+  let summary = "define an array of buffer descriptors";
   let description = [{
-    Create a tensor of pointers to shared memory buffers.
+    Create a tensor of buffer descriptors packing 32-bit pointer offsets and
+    32-bit lengths into 64-bit elements.
   }];
-  let arguments = (ins DenseI32ArrayAttr:$offsets, TT_MemTypeAttr:$memType);
+  let arguments = (ins DenseI32ArrayAttr:$offsets, DenseI32ArrayAttr:$lengths,
+                   TT_MemTypeAttr:$memType);
   let results = (outs TT_Tensor:$result);
   let assemblyFormat = [{
-    $offsets `,` $memType attr-dict `:` type($result)
+    $offsets `,` $lengths `,` $memType attr-dict `:` type($result)
   }];
 }
 
-def TTI_ExperimentalMemDescToI64Op : TTI_Op<"experimental_memdesc_to_i64", [Pure]> {
-  let summary = "Convert a memdesc into its base pointer as i64";
+def TTI_ExperimentalMemDescToI32Op : TTI_Op<"experimental_memdesc_to_i32", [Pure]> {
+  let summary = "Convert a memdesc into its base pointer as i32";
   let description = [{
-    Extract the base pointer from the given memdesc and return it as a 64-bit
+    Extract the base pointer from the given memdesc and return it as a 32-bit
     integer. This can be used to compare the memdesc against tensors of barrier
     pointers maintained by the concurrency sanitizer.
   }];
   let arguments = (ins TTG_MemDescType:$memdesc);
-  let results = (outs I64:$result);
+  let results = (outs I32:$result);
   let builders = [
     OpBuilder<(ins "Value":$memdesc), [{
-      build($_builder, $_state, $_builder.getI64Type(), memdesc);
+      build($_builder, $_state, $_builder.getI32Type(), memdesc);
     }]>
   ];
   let assemblyFormat = "$memdesc attr-dict `:` type($memdesc)";

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

@@ -1,6 +1,7 @@
 #ifndef TRITONINSTRUMENT_UTILITY_H
 #define TRITONINSTRUMENT_UTILITY_H
 
+#include "triton/Analysis/BufferRegion.h"
 #include "triton/Dialect/Triton/IR/Utility.h"
 #include "triton/Dialect/TritonGPU/IR/Attributes.h"
 #include "triton/Dialect/TritonInstrument/IR/Dialect.h"
@@ -74,15 +75,17 @@ struct AuxDataMap {
   RegionToValueMap readVisibility[numMemTypes];
   RegionToValueMap readTracking[numMemTypes];
   RegionToValueMap commits[CommitKind::NumCommitKinds];
+  RegionToValueMap aliasMatrices[numMemTypes];
   RegionToValueMap lock;
   RegionToValueMap waiting;
 
   void populateAndPassToWarpSpecialize(ModuleOp module);
 
 private:
-  void getBuffersAndBarriers(ModuleOp module,
-                             SmallVector<SmallVector<uint32_t>, 2> &bufValues,
-                             SmallVector<uint32_t> &barrierValues);
+  void getBuffersAndBarriers(
+      ModuleOp module,
+      SmallVector<SmallVector<triton::BufferRegion>, 2> &bufRegions,
+      SmallVector<triton::BufferRegion> &barrierRegions);
   void passToWarpSpecialize(triton::FuncOp func, ValueType value,
                             RegionToValueMap &map);
   void createInWarpSpecialize(

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

@@ -15,13 +15,22 @@ def MMAv5OpInterface : OpInterface<"MMAv5OpInterface"> {
     InterfaceMethod<"Return the A operand.",
                     "::mlir::TypedValue<::mlir::triton::gpu::MemDescType>",
                     "getA">,
+    InterfaceMethod<"Return the B operand.",
+                    "::mlir::TypedValue<::mlir::triton::gpu::MemDescType>",
+                    "getB">,
     InterfaceMethod<"Return the accumulator init flag.",
                     "::mlir::Value",
                     "useAccumulator">,
     InterfaceMethod<"Set the accumulator init flag.",
                     "void",
                     "setUseAccumulator",
                     (ins "::mlir::Value":$flag)>,
+    InterfaceMethod<"Return the completion barriers of this MMAv5 op.",
+                    "::mlir::ValueRange",
+                    "getCompletionBarriers">,
+    InterfaceMethod<"Return the completion barrier predicates of this MMAv5 op.",
+                    "::mlir::ValueRange",
+                    "getCompletionBarrierPreds">,
     InterfaceMethod<"Associate a new completion barrier to this MMAv5 op.",
                     "void",
                     "addCompletionBarrier",