[Mlir-commits] [mlir] [MLIR][LLVM][NVVM] ADD LLVM_NToOneIntrOpBase for N:1 intrinsic import (PR #207111)
llvmlistbot at llvm.org
llvmlistbot at llvm.org
Wed Jul 1 19:32:11 PDT 2026
https://github.com/xys-syx created https://github.com/llvm/llvm-project/pull/207111
`LLVM_IntrOpBase` assumes a fixed 1:1 mapping between an `llvm::Intrinsic::ID` and an MLIR op: `emitIntrMLIRBuilders` builds its predicate from a single `llvmEnumName`, and `emitOneMLIRBuilder` relies on a fixed `llvmArgIndices` permutation to operands. This cannot express variants with different operand layouts that all collapse into one op, for example the four PTX `@llvm.nvvm.barrier.cta.sync.{all,count}{,.aligned}` intrinsics that all import into `nvvm.barrier`.
This is an alternative to #202862. Both address the same import gap:
- 202862 keeps the dispatch NVVM-local (no framework changes).
- This PR moves the dispatch into ODS (more declarative and reusable by future N:1 ops).
>From 4096416035a27206ad3f8ac2c5105aef084f7307 Mon Sep 17 00:00:00 2001
From: Yuansui Xu <xuyuansui at outlook.com>
Date: Wed, 1 Jul 2026 21:23:34 -0500
Subject: [PATCH] import llvm barrier intrinics
---
.../include/mlir/Dialect/LLVMIR/LLVMOpBase.td | 14 ++++++
mlir/include/mlir/Dialect/LLVMIR/NVVMOps.td | 21 ++++++++-
.../Dialect/NVVM/LLVMIRToNVVMTranslation.cpp | 29 ++++++++++++
mlir/test/Target/LLVMIR/Import/nvvmir.ll | 44 ++++++++++++++++++-
.../tools/mlir-tblgen/LLVMIRConversionGen.cpp | 37 ++++++++++++++++
5 files changed, 142 insertions(+), 3 deletions(-)
diff --git a/mlir/include/mlir/Dialect/LLVMIR/LLVMOpBase.td b/mlir/include/mlir/Dialect/LLVMIR/LLVMOpBase.td
index a787840663171..7c9fe89535430 100644
--- a/mlir/include/mlir/Dialect/LLVMIR/LLVMOpBase.td
+++ b/mlir/include/mlir/Dialect/LLVMIR/LLVMOpBase.td
@@ -422,6 +422,20 @@ class LLVM_IntrOpBase<Dialect dialect, string opName, string enumName,
}];
}
+// Base class for an MLIR op that maps to multiple LLVM intrinsics on import.
+// The "enumNames" list contains the LLVM intrinsic IDs that import to this op,
+// and "converter" names the function that builds the op from the LLVM call.
+// Unlike `LLVM_IntrOpBase`, this class sets no "mlirBuilder"; the converter
+// handles import. Ops set their own "llvmBuilder" for export.
+class LLVM_NToOneIntrOpBase<Dialect dialect, string opName,
+ list<string> enumNames, string converter,
+ list<Trait> traits = [], int numResults = 0>
+ : LLVM_OpBase<dialect, opName, traits>,
+ Results<!if(!gt(numResults, 0), (outs LLVM_Type:$res), (outs))> {
+ list<string> llvmEnumNames = enumNames;
+ string llvmIRConverter = converter;
+}
+
// Base class for LLVM intrinsic operations, should not be used directly. Places
// the intrinsic into the LLVM dialect and prefixes its name with "intr.".
class LLVM_IntrOp<string mnem, list<int> overloadedResults,
diff --git a/mlir/include/mlir/Dialect/LLVMIR/NVVMOps.td b/mlir/include/mlir/Dialect/LLVMIR/NVVMOps.td
index 40f7f15b694cb..3fb812f78c59a 100644
--- a/mlir/include/mlir/Dialect/LLVMIR/NVVMOps.td
+++ b/mlir/include/mlir/Dialect/LLVMIR/NVVMOps.td
@@ -1123,8 +1123,13 @@ def BarrierReductionAttr
let assemblyFormat = "`<` $value `>`";
}
-def NVVM_BarrierOp : NVVM_VoidIntrinsicOp<"barrier",
- [AttrSizedOperandSegments]> {
+def NVVM_BarrierOp : LLVM_NToOneIntrOpBase<NVVM_Dialect, "barrier",
+ /*enumNames=*/["nvvm_barrier_cta_sync_all",
+ "nvvm_barrier_cta_sync_count",
+ "nvvm_barrier_cta_sync_aligned_all",
+ "nvvm_barrier_cta_sync_aligned_count"],
+ /*converter=*/"convertBarrierSyncIntrinsic",
+ /*traits=*/[AttrSizedOperandSegments]> {
let summary = "CTA Barrier Synchronization Op";
let description = [{
The `nvvm.barrier` operation performs barrier synchronization and communication
@@ -1170,6 +1175,18 @@ def NVVM_BarrierOp : NVVM_VoidIntrinsicOp<"barrier",
let builders = [OpBuilder<(ins), [{
return build($_builder, $_state, Value{}, Value{});
}]>];
+
+ let extraClassDeclaration = [{
+ static NVVM::IDArgPair
+ getIntrinsicIDAndArgs(Operation &op, LLVM::ModuleTranslation &mt,
+ llvm::IRBuilderBase &builder);
+ }];
+
+ string llvmBuilder = [{
+ auto [id, args] = NVVM::BarrierOp::getIntrinsicIDAndArgs(
+ *op, moduleTranslation, builder);
+ createIntrinsicCall(builder, id, builder.getVoidTy(), args);
+ }];
}
def NVVM_BarrierReductionOp :
diff --git a/mlir/lib/Target/LLVMIR/Dialect/NVVM/LLVMIRToNVVMTranslation.cpp b/mlir/lib/Target/LLVMIR/Dialect/NVVM/LLVMIRToNVVMTranslation.cpp
index 55e73e839afcb..82eafedd5b281 100644
--- a/mlir/lib/Target/LLVMIR/Dialect/NVVM/LLVMIRToNVVMTranslation.cpp
+++ b/mlir/lib/Target/LLVMIR/Dialect/NVVM/LLVMIRToNVVMTranslation.cpp
@@ -36,6 +36,35 @@ static ArrayRef<unsigned> getSupportedIntrinsicsImpl() {
return convertibleIntrinsics;
}
+/// 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();
+}
+
/// Converts the LLVM intrinsic to an MLIR NVVM dialect operation if a
/// conversion exits. Returns failure otherwise.
static LogicalResult convertIntrinsicImpl(OpBuilder &odsBuilder,
diff --git a/mlir/test/Target/LLVMIR/Import/nvvmir.ll b/mlir/test/Target/LLVMIR/Import/nvvmir.ll
index 1430f9a44eba1..b2529cc8fac6c 100644
--- a/mlir/test/Target/LLVMIR/Import/nvvmir.ll
+++ b/mlir/test/Target/LLVMIR/Import/nvvmir.ll
@@ -73,11 +73,45 @@ 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]]
+ ; CHECK-NOT: aligned
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 = %{{.*}} {aligned = false}
+ ; 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 = %{{.*}}
+ ; CHECK-NOT: aligned
+ ; CHECK-NOT: number_of_threads
+ 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 = %{{.*}} {aligned = false}
+ 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 = %{{.*}}
+ ; CHECK-NOT: aligned
+ 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 +310,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)
diff --git a/mlir/tools/mlir-tblgen/LLVMIRConversionGen.cpp b/mlir/tools/mlir-tblgen/LLVMIRConversionGen.cpp
index 11a2db4a1cc67..9a2d8bb1d07ca 100644
--- a/mlir/tools/mlir-tblgen/LLVMIRConversionGen.cpp
+++ b/mlir/tools/mlir-tblgen/LLVMIRConversionGen.cpp
@@ -303,6 +303,30 @@ static LogicalResult emitOneMLIRBuilder(const Record &record, raw_ostream &os,
return success();
}
+// Emit a dispatch block for an N-to-1 intrinsic op: an `if` matching any of
+// the listed intrinsic IDs that returns the named converter's result.
+static LogicalResult emitOneNToOneIntrMLIRBuilder(const Record &record,
+ raw_ostream &os) {
+ std::vector<StringRef> enumNames =
+ record.getValueAsListOfStrings("llvmEnumNames");
+ StringRef converter = record.getValueAsString("llvmIRConverter");
+ if (enumNames.empty())
+ return emitError(record, "expected non-empty 'llvmEnumNames'");
+ if (converter.empty())
+ return emitError(record, "expected non-empty 'llvmIRConverter'");
+
+ os << "if (";
+ llvm::interleave(
+ enumNames, os,
+ [&](StringRef id) { os << "intrinsicID == llvm::Intrinsic::" << id; },
+ " ||\n ");
+ os << ") {\n";
+ os << " return " << converter
+ << "(odsBuilder, inst, moduleImport, llvmOperands, llvmOpBundles);\n";
+ os << "}\n";
+ return success();
+}
+
// Emit all intrinsic MLIR builders. Returns false on success because of the
// generator registration requirements.
static bool emitIntrMLIRBuilders(const RecordKeeper &records, raw_ostream &os) {
@@ -316,6 +340,11 @@ static bool emitIntrMLIRBuilders(const RecordKeeper &records, raw_ostream &os) {
if (failed(emitOneMLIRBuilder(*def, os, emitIntrCond)))
return true;
}
+ for (const Record *def :
+ records.getAllDerivedDefinitions("LLVM_NToOneIntrOpBase")) {
+ if (failed(emitOneNToOneIntrMLIRBuilder(*def, os)))
+ return true;
+ }
return false;
}
@@ -554,12 +583,20 @@ static void emitOneIntrinsic(const Record &record, raw_ostream &os) {
os << "llvm::Intrinsic::" << record.getValueAsString("llvmEnumName") << ",\n";
}
+static void emitOneNToOneIntrinsic(const Record &record, raw_ostream &os) {
+ for (StringRef id : record.getValueAsListOfStrings("llvmEnumNames"))
+ os << "llvm::Intrinsic::" << id << ",\n";
+}
+
// Emit the list of LLVM IR intrinsics identifiers that are convertible to a
// matching MLIR LLVM dialect intrinsic operation.
static bool emitConvertibleIntrinsics(const RecordKeeper &records,
raw_ostream &os) {
for (const Record *def : records.getAllDerivedDefinitions("LLVM_IntrOpBase"))
emitOneIntrinsic(*def, os);
+ for (const Record *def :
+ records.getAllDerivedDefinitions("LLVM_NToOneIntrOpBase"))
+ emitOneNToOneIntrinsic(*def, os);
return false;
}
More information about the Mlir-commits
mailing list