[BACKEND] Keren/shared memory barrier (#59)

test/lib/Analysis/CMakeLists.txt

@@ -1,6 +1,7 @@
 add_mlir_library(TritonTestAnalysis
   TestAxisInfo.cpp
   TestAllocation.cpp
+  TestMembar.cpp
 
   LINK_LIBS PUBLIC
   TritonAnalysis

test/lib/Analysis/TestAllocation.cpp

@@ -19,24 +19,29 @@ struct TestAllocationPass
   void runOnOperation() override {
     Operation *operation = getOperation();
     auto &os = llvm::errs();
-    os << "Testing: " << operation->getName() << "\n";
-    AllocationAnalysis analysis(operation);
+    // Convert to std::string can remove quotes from op_name
+    auto op_name = SymbolTable::getSymbolName(operation).getValue().str();
+    os << op_name << "\n";
+    Allocation allocation(operation);
     operation->walk([&](Operation *op) {
+      auto scratchBufferId = allocation.getBufferId(op);
+      if (scratchBufferId != Allocation::InvalidBufferId) {
+        size_t offset = allocation.getOffset(scratchBufferId);
+        size_t size = allocation.getAllocatedSize(scratchBufferId);
+        os << "scratch offset = " << offset << ", size = " << size << "\n";
+      }
       if (op->getNumResults() < 1)
         return;
       for (Value result : op->getResults()) {
-        Type type = result.getType();
-        if (auto tensorType = type.dyn_cast<RankedTensorType>()) {
-          Attribute encoding = tensorType.getEncoding();
-          if (encoding.isa<triton::gpu::TritonGPUSharedEncodingAttr>()) {
-            size_t offset = analysis.getOffset(result);
-            size_t size = analysis.getAllocatedSize(result);
-            os << "offset = " << offset << ", size = " << size << "\n";
-          }
+        auto bufferId = allocation.getBufferId(result);
+        if (bufferId != Allocation::InvalidBufferId) {
+          size_t offset = allocation.getOffset(bufferId);
+          size_t size = allocation.getAllocatedSize(bufferId);
+          os << "offset = " << offset << ", size = " << size << "\n";
         }
       }
     });
-    os << "size = " << analysis.getSharedMemorySize() << "\n";
+    os << "size = " << allocation.getSharedMemorySize() << "\n";
   }
 };
 

test/lib/Analysis/TestMembar.cpp

@@ -0,0 +1,50 @@
+#include "mlir/Dialect/GPU/GPUDialect.h"
+#include "mlir/IR/Dialect.h"
+#include "mlir/Pass/Pass.h"
+#include "triton/Analysis/Allocation.h"
+#include "triton/Analysis/Membar.h"
+
+using namespace mlir;
+
+namespace {
+
+struct TestMembarPass
+    : public PassWrapper<TestMembarPass, OperationPass<FuncOp>> {
+
+  // LLVM15+
+  // MLIR_DEFINE_EXPLICIT_INTERNAL_INLINE_TYPE_ID(TestMembarPass);
+
+  StringRef getArgument() const final { return "test-print-membar"; }
+  StringRef getDescription() const final {
+    return "print the result of the allocation pass";
+  }
+
+  void runOnOperation() override {
+    Operation *operation = getOperation();
+    auto &os = llvm::errs();
+    // Convert to std::string can remove quotes from op_name
+    auto op_name = SymbolTable::getSymbolName(operation).getValue().str();
+    os << op_name << "\n";
+    Allocation allocation(operation);
+    MembarAnalysis analysis(&allocation);
+    size_t operationId = 0;
+    operation->walk<WalkOrder::PreOrder>([&](Operation *op) {
+      if (dyn_cast<gpu::BarrierOp>(op)) {
+        os << "Membar " << operationId << "\n";
+      }
+      if (op->getNumRegions() == 0) {
+        // Don't count parent Operation to simplify the test.
+        operationId++;
+      }
+      return;
+    });
+  }
+};
+
+} // namespace
+
+namespace mlir {
+namespace test {
+void registerTestMembarPass() { PassRegistration<TestMembarPass>(); }
+} // namespace test
+} // namespace mlir