Make tl.debug_barrier() a no-op on CPU (#89)

triton-lang · Aug 6, 2024 · 9b7bd67 · 9b7bd67
1 parent ae2509e
commit 9b7bd67
Showing 1 changed file with 20 additions and 0 deletions.
diff --git a/third_party/cpu/lib/TritonCPUToLLVM/DebugOpsToLLVM.cpp b/third_party/cpu/lib/TritonCPUToLLVM/DebugOpsToLLVM.cpp
@@ -3,6 +3,8 @@
 
 #include "cpu/include/TritonCPUToLLVM/Passes.h"
 
+#include "mlir/Dialect/GPU/IR/GPUOps.h.inc"
+
 #include "triton/Conversion/TritonGPUToLLVM/Utility.h"
 #include "triton/Dialect/Triton/IR/Dialect.h"
 #include "triton/Dialect/TritonCPU/IR/Dialect.h"
@@ -164,6 +166,23 @@ struct PrintOpConversion : public ConvertOpToLLVMPattern<triton::PrintOp> {
   }
 };
 
+using BarrierOp = mlir::gpu::BarrierOp;
+
+// This is part of the DebugOps pass because gpu::barrier is generated by
+// tl.debug_barrier.
+struct BarrierOpConversion : public ConvertOpToLLVMPattern<BarrierOp> {
+  explicit BarrierOpConversion(LLVMTypeConverter &typeConverter)
+      : mlir::ConvertOpToLLVMPattern<BarrierOp>(typeConverter) {}
+
+  LogicalResult
+  matchAndRewrite(BarrierOp op, OpAdaptor adaptor,
+                  ConversionPatternRewriter &rewriter) const override {
+    // Just make it a no-op for now
+    rewriter.eraseOp(op);
+    return success();
+  }
+};
+
 struct DebugOpsToLLVM
     : public triton::impl::DebugOpsToLLVMBase<DebugOpsToLLVM> {
   using DebugOpsToLLVMBase::DebugOpsToLLVMBase;
@@ -180,6 +199,7 @@ struct DebugOpsToLLVM
 
     RewritePatternSet patterns(context);
     patterns.add<PrintOpConversion>(typeConverter);
+    patterns.add<BarrierOpConversion>(typeConverter);
     // patterns.add<AssertOpConversion>(typeConverter);
 
     if (failed(applyPartialConversion(mod, convTarget, std::move(patterns)))) {