[Mlir-commits] [mlir] [MLIR][NVVM]Import llvm intrinsics for nvvm.barrier (PR #202862)

[llvmlistbot at llvm.org]

llvmorg-github-actions[bot] wrote:


<!--LLVM PR SUMMARY COMMENT-->

@llvm/pr-subscribers-mlir-llvm

Author: xys-syx

<details>
<summary>Changes</summary>

This PR adds LLVM IR import support for the `bar.sync` intrinsic family, lifting the four variants into `nvvm.barrier` with the appropriate `aligned` attribute and optional thread-count operand.

`LLVM_IntrOpBase` and its auto-generated import is one to one, but four `bar.sync` variants are many-to-one, where they all lift into `nvvm.barrier`, differing only in `aligned` and the optional `numberOfThreads` operand.

An alternative implementation would extend mlir-tblgen with many-to-one intrinsic-to-op mapping support so the conversion could live in `NVVMOps.td`. But the change is bigger than manually import `bar.sync` intrinsic family.

---
Full diff: https://github.com/llvm/llvm-project/pull/202862.diff


2 Files Affected:

- (modified) mlir/lib/Target/LLVMIR/Dialect/NVVM/LLVMIRToNVVMTranslation.cpp (+59-7) 
- (modified) mlir/test/Target/LLVMIR/Import/nvvmir.ll (+41-1) 


``````````diff
diff --git a/mlir/lib/Target/LLVMIR/Dialect/NVVM/LLVMIRToNVVMTranslation.cpp b/mlir/lib/Target/LLVMIR/Dialect/NVVM/LLVMIRToNVVMTranslation.cpp
index 55e73e839afcb..f15965e01761a 100644
--- a/mlir/lib/Target/LLVMIR/Dialect/NVVM/LLVMIRToNVVMTranslation.cpp
+++ b/mlir/lib/Target/LLVMIR/Dialect/NVVM/LLVMIRToNVVMTranslation.cpp
@@ -18,21 +18,68 @@
 using namespace mlir;
 using namespace mlir::NVVM;
 
-/// Returns true if the LLVM IR intrinsic is convertible to an MLIR NVVM dialect
-/// intrinsic. Returns false otherwise.
+/// LLVM intrinsic IDs of the four `bar.sync` variants imported into
+/// `nvvm.barrier`.
+static constexpr llvm::Intrinsic::ID kBarrierSyncIntrinsics[] = {
+    llvm::Intrinsic::nvvm_barrier_cta_sync_all,
+    llvm::Intrinsic::nvvm_barrier_cta_sync_aligned_all,
+    llvm::Intrinsic::nvvm_barrier_cta_sync_count,
+    llvm::Intrinsic::nvvm_barrier_cta_sync_aligned_count,
+};
+
+/// Returns true if `id` is one of the `bar.sync` intrinsic IDs imported into
+/// `nvvm.barrier`.
+static bool isBarrierSyncIntrinsic(llvm::Intrinsic::ID id) {
+  return llvm::is_contained(kBarrierSyncIntrinsics, id);
+}
+
+/// Imports one of the four `bar.sync` LLVM intrinsic variants into a single
+/// `nvvm.barrier` op, deriving the `aligned` attribute and the optional
+/// `numberOfThreads` operand from the specific intrinsic ID.
+static LogicalResult
+convertBarrierSyncIntrinsic(OpBuilder &odsBuilder, llvm::CallInst *inst,
+                            LLVM::ModuleImport &moduleImport,
+                            ArrayRef<llvm::Value *> llvmOperands,
+                            ArrayRef<llvm::OperandBundleUse> llvmOpBundles) {
+  llvm::Intrinsic::ID id = inst->getIntrinsicID();
+  bool aligned = (id == llvm::Intrinsic::nvvm_barrier_cta_sync_aligned_all ||
+                  id == llvm::Intrinsic::nvvm_barrier_cta_sync_aligned_count);
+  bool hasCount = (id == llvm::Intrinsic::nvvm_barrier_cta_sync_count ||
+                   id == llvm::Intrinsic::nvvm_barrier_cta_sync_aligned_count);
+
+  SmallVector<Value> mlirOperands;
+  SmallVector<NamedAttribute> mlirAttrs;
+  if (failed(moduleImport.convertIntrinsicArguments(
+          llvmOperands, llvmOpBundles, /*requiresOpBundles=*/false, {}, {},
+          mlirOperands, mlirAttrs)))
+    return failure();
+
+  auto op = NVVM::BarrierOp::create(
+      odsBuilder, moduleImport.translateLoc(inst->getDebugLoc()),
+      mlirOperands.front(), hasCount ? mlirOperands.back() : Value{},
+      odsBuilder.getBoolAttr(aligned));
+  moduleImport.mapNoResultOp(inst, op);
+  return success();
+}
+
+/// Returns true if the LLVM intrinsic `id` has a corresponding MLIR NVVM op.
 static bool isConvertibleIntrinsic(llvm::Intrinsic::ID id) {
   static const DenseSet<unsigned> convertibleIntrinsics = {
 #include "mlir/Dialect/LLVMIR/NVVMConvertibleLLVMIRIntrinsics.inc"
   };
-  return convertibleIntrinsics.contains(id);
+  return convertibleIntrinsics.contains(id) || isBarrierSyncIntrinsic(id);
 }
 
-/// Returns the list of LLVM IR intrinsic identifiers that are convertible to
-/// MLIR NVVM dialect intrinsics.
+/// Returns the LLVM intrinsic IDs that have a corresponding MLIR NVVM op.
 static ArrayRef<unsigned> getSupportedIntrinsicsImpl() {
-  static const SmallVector<unsigned> convertibleIntrinsics = {
+  static const SmallVector<unsigned> convertibleIntrinsics = []() {
+    SmallVector<unsigned> ids = {
 #include "mlir/Dialect/LLVMIR/NVVMConvertibleLLVMIRIntrinsics.inc"
-  };
+    };
+    ids.append(std::begin(kBarrierSyncIntrinsics),
+               std::end(kBarrierSyncIntrinsics));
+    return ids;
+  }();
   return convertibleIntrinsics;
 }
 
@@ -54,6 +101,11 @@ static LogicalResult convertIntrinsicImpl(OpBuilder &odsBuilder,
     for (unsigned i = 0; i < inst->getNumOperandBundles(); ++i)
       llvmOpBundles.push_back(inst->getOperandBundleAt(i));
 
+    // Route `bar.sync` intrinsics to `nvvm.barrier`.
+    if (isBarrierSyncIntrinsic(intrinsicID))
+      return convertBarrierSyncIntrinsic(odsBuilder, inst, moduleImport,
+                                         llvmOperands, llvmOpBundles);
+
 #include "mlir/Dialect/LLVMIR/NVVMFromLLVMIRConversions.inc"
   }
 
diff --git a/mlir/test/Target/LLVMIR/Import/nvvmir.ll b/mlir/test/Target/LLVMIR/Import/nvvmir.ll
index 1430f9a44eba1..1fe3451cd0dd0 100644
--- a/mlir/test/Target/LLVMIR/Import/nvvmir.ll
+++ b/mlir/test/Target/LLVMIR/Import/nvvmir.ll
@@ -73,11 +73,43 @@ define float @nvvm_rcp(float %0) {
 
 ; CHECK-LABEL: @llvm_nvvm_barrier0()
 define void @llvm_nvvm_barrier0() {
-  ; CHECK: llvm.nvvm.barrier.cta.sync.aligned.all
+  ; CHECK: %[[c0:.*]] = llvm.mlir.constant(0 : i32) : i32
+  ; CHECK: nvvm.barrier id = %[[c0]] {aligned = true}
   call void @llvm.nvvm.barrier0()
   ret void
 }
 
+; CHECK-LABEL: @llvm_nvvm_barrier_sync_all
+define void @llvm_nvvm_barrier_sync_all(i32 %bar) {
+  ; CHECK: nvvm.barrier id = %{{.*}}
+  ; CHECK-NOT: aligned
+  ; CHECK-NOT: number_of_threads
+  call void @llvm.nvvm.barrier.cta.sync.all(i32 %bar)
+  ret void
+}
+
+; CHECK-LABEL: @llvm_nvvm_barrier_sync_aligned_all
+define void @llvm_nvvm_barrier_sync_aligned_all(i32 %bar) {
+  ; CHECK: nvvm.barrier id = %{{.*}} {aligned = true}
+  call void @llvm.nvvm.barrier.cta.sync.aligned.all(i32 %bar)
+  ret void
+}
+
+; CHECK-LABEL: @llvm_nvvm_barrier_sync_count
+define void @llvm_nvvm_barrier_sync_count(i32 %bar, i32 %n) {
+  ; CHECK: nvvm.barrier id = %{{.*}} number_of_threads = %{{.*}}
+  ; CHECK-NOT: aligned
+  call void @llvm.nvvm.barrier.cta.sync.count(i32 %bar, i32 %n)
+  ret void
+}
+
+; CHECK-LABEL: @llvm_nvvm_barrier_sync_aligned_count
+define void @llvm_nvvm_barrier_sync_aligned_count(i32 %bar, i32 %n) {
+  ; CHECK: nvvm.barrier id = %{{.*}} number_of_threads = %{{.*}} {aligned = true}
+  call void @llvm.nvvm.barrier.cta.sync.aligned.count(i32 %bar, i32 %n)
+  ret void
+}
+
 ; CHECK-LABEL: @llvm_nvvm_bar_warp_sync
 define void @llvm_nvvm_bar_warp_sync(i32 %mask) {
   ; CHECK: nvvm.bar.warp.sync %{{.*}} : i32
@@ -276,6 +308,14 @@ declare float @llvm.nvvm.rcp.approx.ftz.f(float)
 
 declare void @llvm.nvvm.barrier0()
 
+declare void @llvm.nvvm.barrier.cta.sync.all(i32)
+
+declare void @llvm.nvvm.barrier.cta.sync.aligned.all(i32)
+
+declare void @llvm.nvvm.barrier.cta.sync.count(i32, i32)
+
+declare void @llvm.nvvm.barrier.cta.sync.aligned.count(i32, i32)
+
 declare void @llvm.nvvm.bar.warp.sync(i32)
 
 declare i32 @llvm.nvvm.shfl.sync.bfly.i32(i32, i32, i32, i32)

``````````

</details>


https://github.com/llvm/llvm-project/pull/202862