[Mlir-commits] [mlir] c809c9b - [mlir][spirv] Convert gpu.barrier to spv.ControlBarrier

[Lei Zhang]

Author: Lei Zhang
Date: 2022-03-01T12:04:00-05:00
New Revision: c809c9bd3b75962c9fc7ab90c8e90d3be317b43b

URL: https://github.com/llvm/llvm-project/commit/c809c9bd3b75962c9fc7ab90c8e90d3be317b43b
DIFF: https://github.com/llvm/llvm-project/commit/c809c9bd3b75962c9fc7ab90c8e90d3be317b43b.diff

LOG: [mlir][spirv] Convert gpu.barrier to spv.ControlBarrier

Reviewed By: ThomasRaoux

Differential Revision: https://reviews.llvm.org/D120722

Added: 
    

Modified: 
    mlir/lib/Conversion/GPUToSPIRV/GPUToSPIRV.cpp
    mlir/test/Conversion/GPUToSPIRV/simple.mlir

Removed: 
    


################################################################################
diff  --git a/mlir/lib/Conversion/GPUToSPIRV/GPUToSPIRV.cpp b/mlir/lib/Conversion/GPUToSPIRV/GPUToSPIRV.cpp
index edfc37f867767..8c5627c0aa8a9 100644
--- a/mlir/lib/Conversion/GPUToSPIRV/GPUToSPIRV.cpp
+++ b/mlir/lib/Conversion/GPUToSPIRV/GPUToSPIRV.cpp
@@ -13,12 +13,12 @@
 #include "mlir/Conversion/GPUToSPIRV/GPUToSPIRV.h"
 #include "mlir/Dialect/GPU/GPUDialect.h"
 #include "mlir/Dialect/SPIRV/IR/SPIRVDialect.h"
+#include "mlir/Dialect/SPIRV/IR/SPIRVEnums.h"
 #include "mlir/Dialect/SPIRV/IR/SPIRVOps.h"
 #include "mlir/Dialect/SPIRV/IR/TargetAndABI.h"
 #include "mlir/Dialect/SPIRV/Transforms/SPIRVConversion.h"
 #include "mlir/IR/BuiltinOps.h"
 #include "mlir/Transforms/DialectConversion.h"
-#include "llvm/ADT/StringSwitch.h"
 
 using namespace mlir;
 
@@ -109,6 +109,16 @@ class GPUReturnOpConversion final : public OpConversionPattern<gpu::ReturnOp> {
                   ConversionPatternRewriter &rewriter) const override;
 };
 
+/// Pattern to convert a gpu.barrier op into a spv.ControlBarrier op.
+class GPUBarrierConversion final : public OpConversionPattern<gpu::BarrierOp> {
+public:
+  using OpConversionPattern::OpConversionPattern;
+
+  LogicalResult
+  matchAndRewrite(gpu::BarrierOp barrierOp, OpAdaptor adaptor,
+                  ConversionPatternRewriter &rewriter) const override;
+};
+
 } // namespace
 
 //===----------------------------------------------------------------------===//
@@ -326,6 +336,25 @@ LogicalResult GPUReturnOpConversion::matchAndRewrite(
   return success();
 }
 
+//===----------------------------------------------------------------------===//
+// Barrier.
+//===----------------------------------------------------------------------===//
+
+LogicalResult GPUBarrierConversion::matchAndRewrite(
+    gpu::BarrierOp barrierOp, OpAdaptor adaptor,
+    ConversionPatternRewriter &rewriter) const {
+  MLIRContext *context = getContext();
+  // Both execution and memory scope should be workgroup.
+  auto scope = spirv::ScopeAttr::get(context, spirv::Scope::Workgroup);
+  // Require acquire and release memory semantics for workgroup memory.
+  auto memorySemantics = spirv::MemorySemanticsAttr::get(
+      context, spirv::MemorySemantics::WorkgroupMemory |
+                   spirv::MemorySemantics::AcquireRelease);
+  rewriter.replaceOpWithNewOp<spirv::ControlBarrierOp>(barrierOp, scope, scope,
+                                                       memorySemantics);
+  return success();
+}
+
 //===----------------------------------------------------------------------===//
 // GPU To SPIRV Patterns.
 //===----------------------------------------------------------------------===//
@@ -333,8 +362,8 @@ LogicalResult GPUReturnOpConversion::matchAndRewrite(
 void mlir::populateGPUToSPIRVPatterns(SPIRVTypeConverter &typeConverter,
                                       RewritePatternSet &patterns) {
   patterns.add<
-      GPUFuncOpConversion, GPUModuleConversion, GPUModuleEndConversion,
-      GPUReturnOpConversion,
+      GPUBarrierConversion, GPUFuncOpConversion, GPUModuleConversion,
+      GPUModuleEndConversion, GPUReturnOpConversion,
       LaunchConfigConversion<gpu::BlockIdOp, spirv::BuiltIn::WorkgroupId>,
       LaunchConfigConversion<gpu::GridDimOp, spirv::BuiltIn::NumWorkgroups>,
       LaunchConfigConversion<gpu::ThreadIdOp,

diff  --git a/mlir/test/Conversion/GPUToSPIRV/simple.mlir b/mlir/test/Conversion/GPUToSPIRV/simple.mlir
index d2d983f715bcc..b590d41c358b5 100644
--- a/mlir/test/Conversion/GPUToSPIRV/simple.mlir
+++ b/mlir/test/Conversion/GPUToSPIRV/simple.mlir
@@ -104,3 +104,27 @@ module attributes {gpu.container_module} {
     }
   }
 }
+
+// -----
+
+module attributes {gpu.container_module} {
+  gpu.module @kernels {
+    // CHECK-LABEL: spv.func @barrier
+    gpu.func @barrier(%arg0 : f32, %arg1 : memref<12xf32>) kernel
+      attributes {spv.entry_point_abi = {local_size = dense<[32, 4, 1]>: vector<3xi32>}} {
+      // CHECK: spv.ControlBarrier Workgroup, Workgroup, "AcquireRelease|WorkgroupMemory"
+      gpu.barrier
+      gpu.return
+    }
+  }
+
+  func @main() {
+    %0 = "op"() : () -> (f32)
+    %1 = "op"() : () -> (memref<12xf32>)
+    %cst = arith.constant 1 : index
+    gpu.launch_func @kernels::@barrier
+        blocks in (%cst, %cst, %cst) threads in (%cst, %cst, %cst)
+        args(%0 : f32, %1 : memref<12xf32>)
+    return
+  }
+}