c280ebda1b
1. Add missing barriers and revert the previous temporary solution 2. Extract the `run` method from membar analysis because the membar analysis should have two phases, including construction, which doesn't modify any IR, and modification, which adds barrier IRs. Hope this could make the use of membar clear.
53 lines
1.4 KiB
C++
53 lines
1.4 KiB
C++
#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 membarPass(&allocation);
|
|
membarPass.run();
|
|
|
|
size_t operationId = 0;
|
|
operation->walk<WalkOrder::PreOrder>([&](Operation *op) {
|
|
if (isa<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
|