[ConSan] Concurrency Sanitizer - initial scaffolding and introduction of TritonInstrument dialect (#7157)

ConcurrencySanitizer is a mechanism that is intended to help with
debugging concurrency issues in Triton and Gluon. It consists of:
1. `TritonInstrumentDialect` (`tti`) intended to be used for all
instrumentation purposes. The reason of introducing new dialect is to
make it easy to move all the functionality away from triton if needed,
and to clearly distinguish between functional and instrumentation
opcodes in IR.
2. ConcurrencySanitizer pass that adds instrumentation to the IR during
lowering to LLVMIR
3. Set of TTG Ops creating and updating auxiliary data structures and
performing checks
This change is an initial implementation of base functionality to enable
checks in simple test and catch intentionally introduced issues.
Current version only supports modules with single function, as aux
tensors are passed around by value. This will be changed in the future
with tensors being passed through global scratch and multi function
support.

Author: pawelszczerbuk

bin/RegisterTritonDialects.h

@@ -6,6 +6,7 @@
 #include "third_party/proton/dialect/include/Dialect/Proton/IR/Dialect.h"
 #include "triton/Dialect/Triton/IR/Dialect.h"
 #include "triton/Dialect/TritonGPU/IR/Dialect.h"
+#include "triton/Dialect/TritonInstrument/IR/Dialect.h"
 #include "triton/Dialect/TritonNvidiaGPU/IR/Dialect.h"
 
 // Below headers will allow registration to ROCm passes
@@ -15,6 +16,7 @@
 
 #include "triton/Dialect/Triton/Transforms/Passes.h"
 #include "triton/Dialect/TritonGPU/Transforms/Passes.h"
+#include "triton/Dialect/TritonInstrument/Transforms/Passes.h"
 #include "triton/Dialect/TritonNvidiaGPU/Transforms/Passes.h"
 
 #include "nvidia/hopper/include/Transforms/Passes.h"
@@ -47,6 +49,7 @@ inline void registerTritonDialects(mlir::DialectRegistry &registry) {
   mlir::triton::registerTritonPasses();
   mlir::triton::gpu::registerTritonGPUPasses();
   mlir::triton::nvidia_gpu::registerTritonNvidiaGPUPasses();
+  mlir::triton::instrument::registerTritonInstrumentPasses();
   mlir::test::registerTestAliasPass();
   mlir::test::registerTestAlignmentPass();
   mlir::test::registerAMDTestAlignmentPass();
@@ -95,9 +98,10 @@ inline void registerTritonDialects(mlir::DialectRegistry &registry) {
   registry.insert<
       mlir::triton::TritonDialect, mlir::cf::ControlFlowDialect,
       mlir::triton::nvidia_gpu::TritonNvidiaGPUDialect,
-      mlir::triton::gpu::TritonGPUDialect, mlir::math::MathDialect,
-      mlir::arith::ArithDialect, mlir::scf::SCFDialect, mlir::gpu::GPUDialect,
-      mlir::LLVM::LLVMDialect, mlir::NVVM::NVVMDialect,
+      mlir::triton::gpu::TritonGPUDialect,
+      mlir::triton::instrument::TritonInstrumentDialect,
+      mlir::math::MathDialect, mlir::arith::ArithDialect, mlir::scf::SCFDialect,
+      mlir::gpu::GPUDialect, mlir::LLVM::LLVMDialect, mlir::NVVM::NVVMDialect,
       mlir::triton::nvgpu::NVGPUDialect, mlir::triton::nvws::NVWSDialect,
       mlir::triton::amdgpu::TritonAMDGPUDialect,
       mlir::triton::proton::ProtonDialect, mlir::ROCDL::ROCDLDialect>();

include/triton/Conversion/TritonGPUToLLVM/PatternTritonGPUOpToLLVM.h

@@ -102,6 +102,11 @@ void populatePrintOpToLLVMPattern(LLVMTypeConverter &typeConverter,
                                   const TargetInfoBase &targetInfo,
                                   PatternBenefit benefit);
 
+void populateInstrumentationToLLVMPatterns(LLVMTypeConverter &typeConverter,
+                                           const TargetInfoBase &targetInfo,
+                                           RewritePatternSet &patterns,
+                                           PatternBenefit benefit);
+
 } // namespace triton
 } // namespace mlir
 

include/triton/Dialect/CMakeLists.txt

@@ -1,3 +1,4 @@
 add_subdirectory(Triton)
 add_subdirectory(TritonGPU)
 add_subdirectory(TritonNvidiaGPU)
+add_subdirectory(TritonInstrument)

include/triton/Dialect/TritonInstrument/CMakeLists.txt

@@ -0,0 +1,2 @@
+add_subdirectory(IR)
+add_subdirectory(Transforms)

include/triton/Dialect/TritonInstrument/IR/CMakeLists.txt

@@ -0,0 +1,13 @@
+set(MLIR_BINARY_DIR ${CMAKE_BINARY_DIR})
+
+set(LLVM_TARGET_DEFINITIONS TritonInstrumentDialect.td)
+mlir_tablegen(Dialect.h.inc -gen-dialect-decls -dialect=tti)
+mlir_tablegen(Dialect.cpp.inc -gen-dialect-defs -dialect=tti)
+add_mlir_doc(TritonInstrumentDialect TritonInstrumentDialect dialects/ -gen-dialect-doc)
+
+set(LLVM_TARGET_DEFINITIONS TritonInstrumentOps.td)
+mlir_tablegen(Ops.h.inc -gen-op-decls)
+mlir_tablegen(Ops.cpp.inc -gen-op-defs)
+add_mlir_doc(TritonInstrumentOps TritonInstrumentOps dialects/ -gen-op-doc)
+
+add_public_tablegen_target(TritonInstrumentTableGen)

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

@@ -0,0 +1,12 @@
+#ifndef TRITON_DIALECT_TRITONINSTRUMENT_IR_DIALECT_H_
+#define TRITON_DIALECT_TRITONINSTRUMENT_IR_DIALECT_H_
+
+// TritonInstrument depends on Triton and TritonGPU
+#include "triton/Dialect/Triton/IR/Dialect.h"
+#include "triton/Dialect/TritonGPU/IR/Dialect.h"
+
+#define GET_OP_CLASSES
+#include "triton/Dialect/TritonInstrument/IR/Dialect.h.inc"
+#include "triton/Dialect/TritonInstrument/IR/Ops.h.inc"
+
+#endif // TRITON_DIALECT_TRITONINSTRUMENT_IR_DIALECT_H_

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

@@ -0,0 +1,11 @@
+#ifndef TRITONINSTRUMENT_DIALECT
+#define TRITONINSTRUMENT_DIALECT
+
+include "mlir/IR/OpBase.td"
+
+def TritonInstrument_Dialect : Dialect {
+  let name = "tti";
+  let cppNamespace = "::mlir::triton::instrument";
+}
+
+#endif // TRITONINSTRUMENT_DIALECT

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

@@ -0,0 +1,82 @@
+#ifndef TRITONINSTRUMENT_OPS
+#define TRITONINSTRUMENT_OPS
+
+include "triton/Dialect/TritonInstrument/IR/TritonInstrumentDialect.td"
+include "triton/Dialect/TritonGPU/IR/TritonGPUTypes.td"
+include "triton/Dialect/Triton/IR/TritonTypes.td"
+include "mlir/IR/OpBase.td"
+include "mlir/Interfaces/SideEffectInterfaces.td"
+
+class TTI_Op<string mnemonic, list<Trait> traits = []> :
+    Op<TritonInstrument_Dialect, mnemonic, traits> {
+}
+
+// Define an array of pointers to shared memory buffers
+def TTI_ExperimentalSharedBufferPointersOp : TTI_Op<"experimental_shared_buffer_pointers", [Pure]> {
+  let summary = "definte an array of pointers to shared memory buffers";
+  let description = [{
+    Create a tensor of pointers to shared memory buffers.
+  }];
+  let arguments = (ins DenseI32ArrayAttr:$offsets);
+  let results = (outs TT_Tensor:$result);
+  let assemblyFormat = [{
+    attr-dict `:` type($result)
+  }];
+}
+
+// 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";
+  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.
+  }];
+  let arguments = (ins
+    TTG_MemDescType:$buffer,
+    TTG_MemDescType:$mbar,
+    TT_Tensor:$buffers,
+    TT_Tensor:$states,
+    TT_Tensor:$barriers
+  );
+  let results = (outs
+    TT_Tensor:$outStates,
+    TT_Tensor:$outBarriers
+  );
+  let assemblyFormat = [{
+    $buffer `,` $mbar `{` $buffers `,` $states `,` $barriers `}` attr-dict `:` type($buffer) `,` type($mbar) `,` type($buffers) `,` type($states) `,` type($barriers) `->` type($outStates) `,` type($outBarriers)
+  }];
+  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);
+    }]>
+  ];
+}
+
+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";
+  let description = [{
+    Check if waiting on a mbar is valid and update the barrier state.
+  }];
+  let arguments = (ins
+    TTG_MemDescType:$mbar,
+    TT_Tensor:$barriers,
+    TT_Tensor:$states
+    );
+
+  let results = (outs
+    TT_Tensor:$outStates,
+    TT_Tensor:$outBarriers);
+
+  let assemblyFormat = [{
+    $mbar `{` $states `,` $barriers `}` attr-dict `:` type($mbar) `,` type($states) `,` type($barriers) `->` type($outStates) `,` type($outBarriers)
+  }];
+
+  let builders = [
+    OpBuilder<(ins "Value":$mbar, "Value":$barriers, "Value":$states),
+    [{
+      build($_builder, $_state, {states.getType(), barriers.getType()}, mbar, barriers, states);
+    }]>];
+
+}
+
+#endif // TRITONINSTRUMENT_OPS

include/triton/Dialect/TritonInstrument/Transforms/CMakeLists.txt

@@ -0,0 +1,3 @@
+set(LLVM_TARGET_DEFINITIONS Passes.td)
+mlir_tablegen(Passes.h.inc -gen-pass-decls -name TritonInstrument)
+add_public_tablegen_target(TritonInstrumentTransformsIncGen)

include/triton/Dialect/TritonInstrument/Transforms/Passes.h

@@ -0,0 +1,22 @@
+#ifndef TRITON_DIALECT_TRITONINSTRUMENT_TRANSFORMS_PASSES_H_
+#define TRITON_DIALECT_TRITONINSTRUMENT_TRANSFORMS_PASSES_H_
+
+#include "mlir/Pass/Pass.h"
+#include "triton/Dialect/TritonInstrument/IR/Dialect.h"
+
+namespace mlir {
+namespace triton {
+namespace instrument {
+
+// Generate the pass class declarations.
+#define GEN_PASS_DECL
+#include "triton/Dialect/TritonInstrument/Transforms/Passes.h.inc"
+
+/// Generate the code for registering passes.
+#define GEN_PASS_REGISTRATION
+#include "triton/Dialect/TritonInstrument/Transforms/Passes.h.inc"
+
+} // namespace instrument
+} // namespace triton
+} // namespace mlir
+#endif