[mlir][NVVM] Add nvvm.membar operation

Author: clementval

mlir/include/mlir/Dialect/LLVMIR/NVVMOps.td

@@ -1236,6 +1236,39 @@ def NVVM_FenceProxyAcquireOp : NVVM_Op<"fence.proxy.acquire">,
   let hasVerifier = 1;
 }
 
+// Attrs describing the level of the Memory Operation
+def MemLevelCTA : I32EnumAttrCase<"CTA", 0, "cta">;
+def MemLevelGL : I32EnumAttrCase<"GL", 1, "gl">;
+def MemLevelSys : I32EnumAttrCase<"SYS", 2, "sys">;
+
+def MemLevelKind
+    : I32EnumAttr<
+          "MemLevelKind",
+          "NVVM Memory Level kind", [MemLevelCTA, MemLevelGL, MemLevelSys]> {
+  let genSpecializedAttr = 0;
+  let cppNamespace = "::mlir::NVVM";
+}
+def MemLevelKindAttr : EnumAttr<NVVM_Dialect, MemLevelKind, "mem_level"> {
+  let assemblyFormat = "`<` $value `>`";
+}
+
+def NVVM_MembarOp : NVVM_Op<"membar">,
+                    Arguments<(ins MemLevelKindAttr:$level)> {
+  let summary = "Memory barrier operation";
+  let description = [{
+    `member` operation guarantees that prior memory accesses requested by this
+    thread are performed at the specified `level`, before later memory
+    operations requested by this thread following the membar instruction.
+
+    [For more information, see PTX ISA](https://docs.nvidia.com/cuda/parallel-thread-execution/#parallel-synchronization-and-communication-instructions-membar)
+  }];
+
+  let assemblyFormat = "$level attr-dict";
+  let llvmBuilder = [{
+    createIntrinsicCall(builder, getMembarLevelID($level), {});
+  }];
+}
+
 def NVVM_FenceProxyReleaseOp : NVVM_Op<"fence.proxy.release">,
       Arguments<(ins MemScopeKindAttr:$scope,
                      DefaultValuedAttr<ProxyKindAttr,

mlir/lib/Target/LLVMIR/Dialect/NVVM/NVVMToLLVMIRTranslation.cpp

@@ -291,6 +291,21 @@ static unsigned getUnidirectionalFenceProxyID(NVVM::ProxyKind fromProxy,
   llvm_unreachable("Unsupported proxy kinds");
 }
 
+static unsigned getMembarLevelID(NVVM::MemLevelKind level) {
+  switch (level) {
+  case NVVM::MemLevelKind::CTA: {
+    return llvm::Intrinsic::nvvm_membar_cta;
+  }
+  case NVVM::MemLevelKind::GL: {
+    return llvm::Intrinsic::nvvm_membar_gl;
+  }
+  case NVVM::MemLevelKind::SYS: {
+    return llvm::Intrinsic::nvvm_membar_sys;
+  }
+  }
+  llvm_unreachable("Unknown level for memory barrier");
+}
+
 #define TCGEN05LD(SHAPE, NUM) llvm::Intrinsic::nvvm_tcgen05_ld_##SHAPE##_##NUM
 
 static llvm::Intrinsic::ID

mlir/test/Target/LLVMIR/nvvmir.mlir

@@ -975,3 +975,16 @@ llvm.func @nanosleep() {
   nvvm.nanosleep 4000
   llvm.return
 }
+
+// -----
+
+// CHECK-lABEL: @memorybarrier()
+llvm.func @memorybarrier() {
+  // CHECK: call void @llvm.nvvm.membar.cta()
+  nvvm.membar #nvvm.mem_level<cta>
+  // CHECK: call void @llvm.nvvm.membar.gl()
+  nvvm.membar #nvvm.mem_level<gl>
+  // CHECK: call void @llvm.nvvm.membar.sys()
+  nvvm.membar #nvvm.mem_level<sys>
+  llvm.return
+}