[CONSAN] BufferRegion Analysis (#8837)

This change factors out the responsibility for determining what buffer
regions can be independently accessed in the program to a dedicated
analysis. Region is a pair (offset, length), and single shmem/tmem
allocation can be partitioned in many ways using indexing and slicing
operations. Before this change ConSan only did a very superficial IR
tracking to check if an allocation is used in a multi-buferred way.
Generating full list of the possible buffer regions is a pre-requisite
for supporting aliased shmem/tmem in ConSan. In the next step I will add
creating an aliasing matrix and proper checks to ConSan.
This change also introduces some smaller cleanups and refactoring of
FunctionBuilder, like moving the checks for existing auxData to the
`create` methods, and ditching error-prone `operator[]`.

Author: pawelszczerbuk

bin/RegisterTritonDialects.h

@@ -54,6 +54,7 @@ void registerTestAliasPass();
 void registerTestAlignmentPass();
 void registerAMDTestAlignmentPass();
 void registerTestAllocationPass();
+void registerTestBufferRegionPass();
 void registerTestMembarPass();
 void registerTestAMDGPUMembarPass();
 void registerTestTritonAMDGPURangeAnalysis();
@@ -75,6 +76,7 @@ inline void registerTritonDialects(mlir::DialectRegistry &registry) {
   mlir::test::registerTestAlignmentPass();
   mlir::test::registerAMDTestAlignmentPass();
   mlir::test::registerTestAllocationPass();
+  mlir::test::registerTestBufferRegionPass();
   mlir::test::registerTestMembarPass();
   mlir::test::registerTestLoopPeelingPass();
   mlir::test::registerTestAMDGPUMembarPass();

include/triton/Analysis/BufferRegion.h

@@ -0,0 +1,169 @@
+#ifndef TRITON_ANALYSIS_BUFFER_REGION_H
+#define TRITON_ANALYSIS_BUFFER_REGION_H
+
+#include <limits>
+#include <set>
+
+#include "mlir/Analysis/DataFlow/SparseAnalysis.h"
+#include "mlir/IR/Value.h"
+
+namespace mlir::triton {
+
+//===----------------------------------------------------------------------===//
+// BufferRegion: a single logical region derived from an alloc
+//===----------------------------------------------------------------------===//
+struct BufferRegion {
+  uint32_t baseOffset;
+  uint32_t length;
+
+  bool operator==(const BufferRegion &other) const {
+    return baseOffset == other.baseOffset && length == other.length;
+  }
+
+  bool operator<(const BufferRegion &other) const {
+    if (baseOffset != other.baseOffset)
+      return baseOffset < other.baseOffset;
+    return length < other.length;
+  }
+
+  template <typename T> void print(T &os) const {
+    os << "[" << baseOffset << ", " << length << "]";
+  }
+};
+
+} // namespace mlir::triton
+
+namespace llvm {
+
+using namespace mlir::triton;
+
+template <> struct DenseMapInfo<BufferRegion> {
+  static BufferRegion getEmptyKey() {
+    constexpr uint32_t empty = std::numeric_limits<uint32_t>::max();
+    return BufferRegion{empty, empty};
+  }
+  static BufferRegion getTombstoneKey() {
+    constexpr uint32_t tombstone = std::numeric_limits<uint32_t>::max() - 1;
+    return BufferRegion{tombstone, tombstone};
+  }
+  static unsigned getHashValue(const BufferRegion &r) {
+    return llvm::hash_combine(r.baseOffset, r.length);
+  }
+  static bool isEqual(const BufferRegion &a, const BufferRegion &b) {
+    return a == b;
+  }
+};
+
+} // namespace llvm
+
+namespace mlir::triton {
+
+//===----------------------------------------------------------------------===//
+// RegionInfo lattice
+//===----------------------------------------------------------------------===//
+//
+// This wraps a set of BufferRegions and provides lattice semantics
+//
+struct RegionInfo {
+  using RegionList = llvm::DenseSet<BufferRegion>;
+  RegionList regions;
+
+  RegionInfo() = default;
+  RegionInfo(const RegionList &r) : regions(r) {}
+
+  // Lattice join: union of regions
+  static RegionInfo join(const RegionInfo &lhs, const RegionInfo &rhs) {
+    RegionInfo result = lhs;
+    for (const auto &reg : rhs.regions)
+      if (llvm::find(result.regions, reg) == result.regions.end())
+        result.regions.insert(reg);
+    return result;
+  }
+
+  bool operator==(const RegionInfo &other) const {
+    if (regions.size() != other.regions.size())
+      return false;
+    for (auto &r : regions)
+      if (llvm::find(other.regions, r) == other.regions.end())
+        return false;
+    return true;
+  }
+
+  template <typename T> void print(T &os) const {
+    llvm::SmallVector<BufferRegion> sortedRegions(regions.begin(),
+                                                  regions.end());
+    llvm::sort(sortedRegions, [](const BufferRegion &a, const BufferRegion &b) {
+      return a < b;
+    });
+    llvm::interleaveComma(sortedRegions, os,
+                          [&](const BufferRegion &r) { r.print(os); });
+  }
+
+  static RegionInfo getPessimisticValueState(MLIRContext *context = nullptr) {
+    return RegionInfo(); // means "unknown / empty"
+  }
+  static RegionInfo getPessimisticValueState(Value) { return RegionInfo(); }
+};
+
+//===----------------------------------------------------------------------===//
+// BufferRegionAnalysis (Sparse Forward Dataflow)
+//===----------------------------------------------------------------------===//
+//
+// Produces a RegionInfo lattice for each MemDesc/ptr-like SSA value,
+// and also collects a global list of all discovered BufferRegions.
+//
+class BufferRegionAnalysis : public dataflow::SparseForwardDataFlowAnalysis<
+                                 dataflow::Lattice<RegionInfo>> {
+
+public:
+  using Base =
+      dataflow::SparseForwardDataFlowAnalysis<dataflow::Lattice<RegionInfo>>;
+  using Base::getLatticeElement;
+  using Base::SparseForwardDataFlowAnalysis;
+
+  enum RegionType { SHARED_MEMORY, TENSOR_MEMORY, BARRIER, NUM_REGION_TYPES };
+
+  static bool isMemoryAccessOperation(Operation *op);
+
+  // ------------------------------
+  // Public API for ConSan
+  // ------------------------------
+
+  /// Return the list of all unique (alloc,offset,len) buffer regions
+  /// discovered by the analysis.
+  llvm::SmallVector<BufferRegion>
+  getAllUsedBufferRegions(RegionType type) const {
+    return llvm::to_vector(usedBufferRegions[type]);
+  }
+
+  void calculateUsedBufferRegions(Operation *op);
+
+  // ------------------------------
+  // Required overrides
+  // ------------------------------
+
+  void setToEntryState(dataflow::Lattice<RegionInfo> *lat) override {
+    propagateIfChanged(
+        lat, lat->join(RegionInfo::getPessimisticValueState(lat->getAnchor())));
+  }
+
+  LogicalResult visitOperation(
+      Operation *op,
+      llvm::ArrayRef<const dataflow::Lattice<RegionInfo> *> operands,
+      llvm::ArrayRef<dataflow::Lattice<RegionInfo> *> results) override;
+
+  void visitNonControlFlowArguments(
+      Operation *op, const RegionSuccessor &successor,
+      llvm::ArrayRef<dataflow::Lattice<RegionInfo> *> argLattices,
+      unsigned firstIndex) override;
+
+  LogicalResult initialize(Operation *top) override;
+
+private:
+  // Global registry of all regions
+  std::set<BufferRegion> usedBufferRegions[NUM_REGION_TYPES];
+};
+
+} // namespace mlir::triton
+
+#endif // TRITON_ANALYSIS_BUFFER_REGION_H

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

@@ -182,35 +182,35 @@ 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, ValueType buffers,
-                                      ValueType outstandingCommits,
+                                      int thread, Value pred, MemType memType,
+                                      CommitKind::Kind commitKind,
                                       Operation *insertPoint);
   // commitAccesses: convert staged entries to 1 and increment outstanding
   // commits greater than zero for the committing thread.
   void createCommitAccessesCall(ImplicitLocOpBuilder &b, int thread, Value pred,
-                                ValueType outstandingCommits,
+                                CommitKind::Kind commitKind,
                                 Operation *insertPoint);
   // clearOutstandingCommitsTransferWrites: clear entries farther than
   // outstandingNum from the thread and set write visibility for threads in
   // transferThreadMask.
   void createClearOutstandingCommitsTransferWritesCall(
       ImplicitLocOpBuilder &b, int thread, uint64_t transferThreadMask,
-      int outstandingNum, Value pred, ValueType outstandingCommits,
-      ValueType writeVisibility, Operation *insertPoint);
+      int outstandingNum, Value pred, CommitKind::Kind commitKind,
+      MemType memType, Operation *insertPoint);
   // clearOutstandingCommitsTransferReads: clear entries farther than
   // outstandingNum from the thread and set read visibility for threads in
   // transferThreadMask.
   void createClearOutstandingCommitsTransferReadsCall(
       ImplicitLocOpBuilder &b, int thread, uint64_t transferThreadMask,
-      int outstandingNum, Value pred, ValueType outstandingCommits,
-      ValueType readVisibility, Operation *insertPoint);
+      int outstandingNum, Value pred, CommitKind::Kind commitKind,
+      MemType memType, Operation *insertPoint);
   // 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,
                                          StringRef pendingAccessType,
-                                         Value pred, ValueType buffers,
-                                         ValueType outstandingCommits,
+                                         Value pred, MemType memType,
+                                         CommitKind::Kind commitKind,
                                          Operation *insertPoint);
 
 private:

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

@@ -47,10 +47,16 @@ struct ValueType {
 struct AuxDataMap {
   struct RegionToValueMap {
     DenseMap<Region *, ValueType> values;
-    ValueType &operator[](Region *region) { return values[region]; }
-    ValueType &operator[](Operation *op) {
-      return values[getEnclosingParitionOrFunctionRegion(op)];
+    ValueType at(Region *region) {
+      if (values.find(region) == values.end()) {
+        assert(false && "Region not found in AuxDataMap");
+      }
+      return values[region];
     }
+    ValueType at(Operation *op) {
+      return at(getEnclosingParitionOrFunctionRegion(op));
+    }
+    void insert(Region *region, ValueType value) { values[region] = value; }
     bool empty() const { return values.empty(); }
 
   private:
@@ -75,8 +81,8 @@ struct AuxDataMap {
 
 private:
   void getBuffersAndBarriers(ModuleOp module,
-                             SmallVector<SmallVector<int32_t>, 2> &bufValues,
-                             SmallVector<int32_t> &barrierValues);
+                             SmallVector<SmallVector<uint32_t>, 2> &bufValues,
+                             SmallVector<uint32_t> &barrierValues);
   void passToWarpSpecialize(triton::FuncOp func, ValueType value,
                             RegionToValueMap &map);
   void createInWarpSpecialize(