[mlir][rocdl] add rocdl.barier op.

[whchung]

Summary

• Add rocdl.barrier op.
• Lower gpu.barier to rocdl.barrier in -convert-gpu-to-rocdl.

[rriddle]

mlir/lib/Conversion/GPUToROCDL/LowerGpuOpsToROCDLOps.cpp

50

Just use m.getContext() here.

[whchung]

Address code review comments.

[whchung]

@rriddle Thanks for the review. Wondering may I ask for your help to commit this patch? I am new to contributing to LLVM / MLIR so I don't yet have write access to the repository.

[whchung]

Rebase.

mlir/include/mlir/Dialect/LLVMIR/ROCDLOps.td

def ROCDL_GridDimXOp : ROCDL_DeviceFunctionOp<"grid.dim.x",
                                                "__ockl_get_global_size", 0>;
 
 def ROCDL_GridDimYOp : ROCDL_DeviceFunctionOp<"grid.dim.y",
                                                "__ockl_get_global_size", 1>;
 
 def ROCDL_GridDimZOp : ROCDL_DeviceFunctionOp<"grid.dim.z",
                                                "__ockl_get_global_size", 2>;
 
+//===----------------------------------------------------------------------===//
+// Synchronization primitives
+
+def ROCDL_BarrierOp : ROCDL_Op<"barrier"> {
+  string llvmBuilder = [{
+    llvm::LLVMContext &llvmContext = builder.getContext();
+    builder.CreateFence(llvm::AtomicOrdering::Release,
+                        llvmContext.getOrInsertSyncScopeID("workgroup"));
+    createIntrinsicCall(builder, llvm::Intrinsic::amdgcn_s_barrier);
+    builder.CreateFence(llvm::AtomicOrdering::Acquire,
+                        llvmContext.getOrInsertSyncScopeID("workgroup"));
+  }];
+  let assemblyFormat = "attr-dict";
+}
 
 #endif // ROCDLIR_OPS

mlir/lib/Conversion/GPUToROCDL/CMakeLists.txt

+set(LLVM_TARGET_DEFINITIONS GPUToROCDL.td)
+mlir_tablegen(GPUToROCDL.cpp.inc -gen-rewriters)
+add_public_tablegen_target(MLIRGPUToROCDLIncGen)
+
 add_mlir_conversion_library(MLIRGPUtoROCDLTransforms
   LowerGpuOpsToROCDLOps.cpp
 
   DEPENDS
   MLIRConversionPassIncGen
+  MLIRGPUToROCDLIncGen
   )
+
 target_link_libraries(MLIRGPUtoROCDLTransforms
   PUBLIC
   LLVMSupport
   MLIRGPU
   MLIRLLVMIR
   MLIRROCDLIR
   MLIRPass
   MLIRStandardToLLVM
   )

mlir/lib/Conversion/GPUToROCDL/GPUToROCDL.td

+//==-- GPUToROCDL.td - GPU Ops to ROCDL Patterns -------------*- tablegen -*==//
+//
+// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
+// See https://llvm.org/LICENSE.txt for license information.
+// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+//
+//===----------------------------------------------------------------------===//
+//
+// Defines Patterns to lower GPU ops to ROCDL.
+//
+//===----------------------------------------------------------------------===//
+
+#ifndef MLIR_CONVERSION_GPUTOROCDL_TD
+#define MLIR_CONVERSION_GPUTOROCDL_TD
+
+include "mlir/Dialect/GPU/GPUOps.td"
+include "mlir/Dialect/LLVMIR/ROCDLOps.td"
+
+def : Pat<(GPU_BarrierOp), (ROCDL_BarrierOp)>;
+
+#endif // MLIR_CONVERSION_GPUTOROCDL_TD

mlir/lib/Conversion/GPUToROCDL/LowerGpuOpsToROCDLOps.cpp

 #include "../GPUCommon/IndexIntrinsicsOpLowering.h"
 #include "../GPUCommon/OpToFuncCallLowering.h"
 #include "../PassDetail.h"
 
 using namespace mlir;
 
 namespace {
 
+/// Import the GPU Ops to ROCDL Patterns.
+#include "GPUToROCDL.cpp.inc"
+
 // A pass that replaces all occurrences of GPU device operations with their
 // corresponding ROCDL equivalent.
 //
 // This pass only handles device code and is not meant to be run on GPU host
 // code.
 class LowerGpuOpsToROCDLOpsPass
     : public ConvertGpuOpsToROCDLOpsBase<LowerGpuOpsToROCDLOpsPass> {
 public:
   void runOnOperation() override {
     gpu::GPUModuleOp m = getOperation();
 
     LLVMTypeConverter converter(m.getContext());
 
     OwningRewritePatternList patterns;
 
     populateGpuRewritePatterns(m.getContext(), patterns);
     applyPatternsAndFoldGreedily(m, patterns);
     patterns.clear();
 
     populateVectorToLLVMConversionPatterns(converter, patterns);
     populateStdToLLVMConversionPatterns(converter, patterns);
if (failed(applyPartialConversion(m, target, patterns, &converter)))
       signalPassFailure();
   }
 };
 
 } // anonymous namespace
 
 void mlir::populateGpuToROCDLConversionPatterns(
     LLVMTypeConverter &converter, OwningRewritePatternList &patterns) {
+  populateWithGenerated(converter.getDialect()->getContext(), &patterns);
   patterns.insert<
       GPUIndexIntrinsicOpLowering<gpu::ThreadIdOp, ROCDL::ThreadIdXOp,
                                   ROCDL::ThreadIdYOp, ROCDL::ThreadIdZOp>,
       GPUIndexIntrinsicOpLowering<gpu::BlockDimOp, ROCDL::BlockDimXOp,
                                   ROCDL::BlockDimYOp, ROCDL::BlockDimZOp>,
       GPUIndexIntrinsicOpLowering<gpu::BlockIdOp, ROCDL::BlockIdXOp,
                                   ROCDL::BlockIdYOp, ROCDL::BlockIdZOp>,
       GPUIndexIntrinsicOpLowering<gpu::GridDimOp, ROCDL::GridDimXOp,

mlir/test/Conversion/GPUToROCDL/gpu-to-rocdl.mlir

std.return %tIdX, %tIdY, %tIdZ, %bDimX, %bDimY, %bDimZ,
         : index, index, index, index, index, index,
           index, index, index, index, index, index
   }
 }
 
 // -----
 
 gpu.module @test_module {
+  // CHECK-LABEL: func @gpu_sync()
+  func @gpu_sync() {
+    // CHECK: rocdl.barrier
+    gpu.barrier
+    std.return
+  }
+}
+
+// -----
+
+gpu.module @test_module {
   // CHECK: llvm.func @__ocml_fabs_f32(!llvm.float) -> !llvm.float
   // CHECK: llvm.func @__ocml_fabs_f64(!llvm.double) -> !llvm.double
   // CHECK-LABEL: func @gpu_fabs
   func @gpu_fabs(%arg_f32 : f32, %arg_f64 : f64) -> (f32, f64) {
     %result32 = std.absf %arg_f32 : f32
     // CHECK: llvm.call @__ocml_fabs_f32(%{{.*}}) : (!llvm.float) -> !llvm.float
     %result64 = std.absf %arg_f64 : f64
     // CHECK: llvm.call @__ocml_fabs_f64(%{{.*}}) : (!llvm.double) -> !llvm.double

mlir/test/Dialect/LLVMIR/rocdl.mlir

func @rocdl_special_regs() -> !llvm.i32 {
   // CHECK: rocdl.grid.dim.x : !llvm.i32
   %9 = rocdl.grid.dim.x : !llvm.i32
   // CHECK: rocdl.grid.dim.y : !llvm.i32
   %10 = rocdl.grid.dim.y : !llvm.i32
   // CHECK: rocdl.grid.dim.z : !llvm.i32
   %11 = rocdl.grid.dim.z : !llvm.i32
   llvm.return %0 : !llvm.i32
 }
+
+func @rocdl.barrier() {
+  // CHECK: rocdl.barrier
+  rocdl.barrier
+  llvm.return
+}

mlir/test/Target/rocdl.mlir

llvm.func @rocdl_special_regs() -> !llvm.i32 {
   %12 = rocdl.grid.dim.z : !llvm.i64
   llvm.return %1 : !llvm.i32
 }
 
 llvm.func @kernel_func() attributes {gpu.kernel} {
   // CHECK-LABEL: amdgpu_kernel void @kernel_func
   llvm.return
 }
+
+llvm.func @rocdl.barrier() {
+  // CHECK:      fence syncscope("workgroup") release
+  // CHECK-NEXT: call void @llvm.amdgcn.s.barrier()
+  // CHECK-NEXT: fence syncscope("workgroup") acquire
+  rocdl.barrier
+  llvm.return
+}