[Mlir-commits] [mlir] [MLIR][NVVM] Add explicit aligned attribute to nvvm.barrier and nvvm.barrier.reduction (PR #200744)
llvmlistbot at llvm.org
llvmlistbot at llvm.org
Mon Jun 1 00:37:39 PDT 2026
https://github.com/xys-syx updated https://github.com/llvm/llvm-project/pull/200744
>From b777b741302c1d322e924cf22107c399f861cca0 Mon Sep 17 00:00:00 2001
From: Yuansui Xu <xuyuansui at outlook.com>
Date: Mon, 1 Jun 2026 02:32:18 -0500
Subject: [PATCH 1/2] add aligned attr for nvvm.barrier and
nvvm.barrier.reduction
---
mlir/include/mlir/Dialect/LLVMIR/NVVMOps.td | 16 ++++--
mlir/lib/Dialect/LLVMIR/IR/NVVMDialect.cpp | 54 +++++++++++++--------
mlir/test/Dialect/LLVMIR/nvvm.mlir | 6 +++
mlir/test/Target/LLVMIR/nvvm/barrier.mlir | 22 +++++++++
4 files changed, 72 insertions(+), 26 deletions(-)
diff --git a/mlir/include/mlir/Dialect/LLVMIR/NVVMOps.td b/mlir/include/mlir/Dialect/LLVMIR/NVVMOps.td
index 01abc7e70f57c..84ca4dcc3ddb1 100644
--- a/mlir/include/mlir/Dialect/LLVMIR/NVVMOps.td
+++ b/mlir/include/mlir/Dialect/LLVMIR/NVVMOps.td
@@ -1131,13 +1131,15 @@ def NVVM_BarrierOp : NVVM_VoidIntrinsicOp<"barrier",
within a CTA (Cooperative Thread Array). It causes executing threads to wait for
all non-exited threads participating in the barrier to arrive.
- The operation takes the following optional operands:
+ The operation takes the following optional operands and attributes:
- `barrierId`: Specifies a logical barrier resource with value 0 through 15.
Each CTA instance has sixteen barriers numbered 0..15. Defaults to 0 if not specified.
- `numberOfThreads`: Specifies the number of threads participating in the barrier.
When specified, the value must be a multiple of the warp size. If not specified,
all threads in the CTA participate in the barrier.
+ - `aligned`: Selects between the `.aligned` and non-`.aligned` forms of the
+ underlying `@llvm.nvvm.barrier.cta.*` intrinsic family. Defaults to true.
Reduction variants of the barrier instruction are modeled by the
`nvvm.barrier.reduction` op.
@@ -1159,7 +1161,8 @@ def NVVM_BarrierOp : NVVM_VoidIntrinsicOp<"barrier",
let arguments = (ins
Optional<I32>:$barrierId,
- Optional<I32>:$numberOfThreads);
+ Optional<I32>:$numberOfThreads,
+ DefaultValuedAttr<BoolAttr, "true">:$aligned);
let assemblyFormat =
"(`id` `=` $barrierId^)? (`number_of_threads` `=` $numberOfThreads^)? "
@@ -1184,10 +1187,12 @@ def NVVM_BarrierReductionOp :
per-thread predicates.
- `reductionPredicate`: The per-thread i32 predicate. It is compared against
zero to form the i1 value fed into the reduction.
+ - `aligned`: Selects between the `.aligned` and non-`.aligned` forms of the
+ underlying `@llvm.nvvm.barrier.cta.red.*` intrinsic family. Defaults to
+ true.
The result is the i32 reduction value computed across all threads
- participating in the barrier. This op always lowers to the aligned form of
- the `@llvm.nvvm.barrier.cta.red.*` intrinsic family.
+ participating in the barrier.
[For more information, see PTX ISA](https://docs.nvidia.com/cuda/parallel-thread-execution/index.html#parallel-synchronization-and-communication-instructions-bar)
}];
@@ -1195,7 +1200,8 @@ def NVVM_BarrierReductionOp :
let arguments = (ins
Optional<I32>:$barrierId,
BarrierReductionAttr:$reductionOp,
- I32:$reductionPredicate);
+ I32:$reductionPredicate,
+ DefaultValuedAttr<BoolAttr, "true">:$aligned);
let results = (outs I32:$res);
let assemblyFormat =
diff --git a/mlir/lib/Dialect/LLVMIR/IR/NVVMDialect.cpp b/mlir/lib/Dialect/LLVMIR/IR/NVVMDialect.cpp
index 00c997ec7a031..8c90f3f43bb49 100644
--- a/mlir/lib/Dialect/LLVMIR/IR/NVVMDialect.cpp
+++ b/mlir/lib/Dialect/LLVMIR/IR/NVVMDialect.cpp
@@ -3449,9 +3449,33 @@ void SubFOp::getCanonicalizationPatterns(RewritePatternSet &patterns,
patterns.add<ConvertFsubToFnegFadd>(context);
}
-//===----------------------------------------------------------------------===//
-// getIntrinsicID/getIntrinsicIDAndArgs methods
-//===----------------------------------------------------------------------===//
+/// Maps the (aligned, hasCount) pair to the `@llvm.nvvm.barrier.cta.sync.*`
+/// intrinsic ID.
+static llvm::Intrinsic::ID getBarrierSyncIntrinsic(bool aligned, bool hasCount) {
+ if (hasCount)
+ return aligned ? llvm::Intrinsic::nvvm_barrier_cta_sync_aligned_count
+ : llvm::Intrinsic::nvvm_barrier_cta_sync_count;
+ return aligned ? llvm::Intrinsic::nvvm_barrier_cta_sync_aligned_all
+ : llvm::Intrinsic::nvvm_barrier_cta_sync_all;
+}
+
+/// Maps the (aligned, kind) pair to the `@llvm.nvvm.barrier.cta.red.*`
+/// intrinsic ID.
+static llvm::Intrinsic::ID
+getBarrierReductionIntrinsic(bool aligned, NVVM::BarrierReduction kind) {
+ switch (kind) {
+ case NVVM::BarrierReduction::AND:
+ return aligned ? llvm::Intrinsic::nvvm_barrier_cta_red_and_aligned_all
+ : llvm::Intrinsic::nvvm_barrier_cta_red_and_all;
+ case NVVM::BarrierReduction::OR:
+ return aligned ? llvm::Intrinsic::nvvm_barrier_cta_red_or_aligned_all
+ : llvm::Intrinsic::nvvm_barrier_cta_red_or_all;
+ case NVVM::BarrierReduction::POPC:
+ return aligned ? llvm::Intrinsic::nvvm_barrier_cta_red_popc_aligned_all
+ : llvm::Intrinsic::nvvm_barrier_cta_red_popc_all;
+ }
+ llvm_unreachable("unknown BarrierReduction kind");
+}
mlir::NVVM::IDArgPair NVVM::BarrierOp::getIntrinsicIDAndArgs(
Operation &op, LLVM::ModuleTranslation &mt, llvm::IRBuilderBase &builder) {
@@ -3459,32 +3483,20 @@ mlir::NVVM::IDArgPair NVVM::BarrierOp::getIntrinsicIDAndArgs(
llvm::Value *barrierId = thisOp.getBarrierId()
? mt.lookupValue(thisOp.getBarrierId())
: builder.getInt32(0);
- llvm::Intrinsic::ID id;
+ bool hasCount = static_cast<bool>(thisOp.getNumberOfThreads());
+ llvm::Intrinsic::ID id =
+ getBarrierSyncIntrinsic(thisOp.getAligned(), hasCount);
llvm::SmallVector<llvm::Value *> args = {barrierId};
- if (thisOp.getNumberOfThreads()) {
- id = llvm::Intrinsic::nvvm_barrier_cta_sync_aligned_count;
+ if (hasCount)
args.push_back(mt.lookupValue(thisOp.getNumberOfThreads()));
- } else {
- id = llvm::Intrinsic::nvvm_barrier_cta_sync_aligned_all;
- }
return {id, std::move(args)};
}
mlir::NVVM::IDArgPair NVVM::BarrierReductionOp::getIntrinsicIDAndArgs(
Operation &op, LLVM::ModuleTranslation &mt, llvm::IRBuilderBase &builder) {
auto thisOp = cast<NVVM::BarrierReductionOp>(op);
- llvm::Intrinsic::ID id;
- switch (thisOp.getReductionOp()) {
- case NVVM::BarrierReduction::AND:
- id = llvm::Intrinsic::nvvm_barrier_cta_red_and_aligned_all;
- break;
- case NVVM::BarrierReduction::OR:
- id = llvm::Intrinsic::nvvm_barrier_cta_red_or_aligned_all;
- break;
- case NVVM::BarrierReduction::POPC:
- id = llvm::Intrinsic::nvvm_barrier_cta_red_popc_aligned_all;
- break;
- }
+ llvm::Intrinsic::ID id =
+ getBarrierReductionIntrinsic(thisOp.getAligned(), thisOp.getReductionOp());
llvm::Value *barrierId = thisOp.getBarrierId()
? mt.lookupValue(thisOp.getBarrierId())
: builder.getInt32(0);
diff --git a/mlir/test/Dialect/LLVMIR/nvvm.mlir b/mlir/test/Dialect/LLVMIR/nvvm.mlir
index 6ba251fe804ad..72935ffd7b7ce 100644
--- a/mlir/test/Dialect/LLVMIR/nvvm.mlir
+++ b/mlir/test/Dialect/LLVMIR/nvvm.mlir
@@ -49,6 +49,10 @@ llvm.func @llvm_nvvm_barrier(%barId : i32, %numberOfThreads : i32) {
nvvm.barrier id = %barId number_of_threads = %numberOfThreads
// CHECK: nvvm.barrier number_of_threads = %[[numberOfThreads]]
nvvm.barrier number_of_threads = %numberOfThreads
+ // CHECK: nvvm.barrier {aligned = false}
+ nvvm.barrier {aligned = false}
+ // CHECK: nvvm.barrier id = %[[barId]] number_of_threads = %[[numberOfThreads]] {aligned = false}
+ nvvm.barrier id = %barId number_of_threads = %numberOfThreads {aligned = false}
llvm.return
}
@@ -63,6 +67,8 @@ llvm.func @llvm_nvvm_barrier_reduction(%barId : i32, %pred : i32) {
%2 = nvvm.barrier.reduction #nvvm.reduction<popc> %pred -> i32
// CHECK: nvvm.barrier.reduction #nvvm.reduction<and> %[[pred]] id = %[[barId]] -> i32
%3 = nvvm.barrier.reduction #nvvm.reduction<and> %pred id = %barId -> i32
+ // CHECK: nvvm.barrier.reduction #nvvm.reduction<and> %[[pred]] -> i32 {aligned = false}
+ %4 = nvvm.barrier.reduction #nvvm.reduction<and> %pred -> i32 {aligned = false}
llvm.return
}
diff --git a/mlir/test/Target/LLVMIR/nvvm/barrier.mlir b/mlir/test/Target/LLVMIR/nvvm/barrier.mlir
index ecfee955904bd..17c086d7737d8 100644
--- a/mlir/test/Target/LLVMIR/nvvm/barrier.mlir
+++ b/mlir/test/Target/LLVMIR/nvvm/barrier.mlir
@@ -33,5 +33,27 @@ llvm.func @llvm_nvvm_barrier(%barID : i32, %numberOfThreads : i32, %redOperand :
// CHECK: %{{.*}} = nvvm.barrier.reduction #nvvm.reduction<and> %{{.*}} id = %{{.*}} -> i32
%3 = nvvm.barrier.reduction #nvvm.reduction<and> %redOperand id = %barID -> i32
+ // Non-aligned sync variants.
+ // LLVM: call void @llvm.nvvm.barrier.cta.sync.all(i32 0)
+ // CHECK: nvvm.barrier {aligned = false}
+ nvvm.barrier {aligned = false}
+ // LLVM: call void @llvm.nvvm.barrier.cta.sync.count(i32 %[[barId]], i32 %[[numThreads]])
+ // CHECK: nvvm.barrier id = %{{.*}} number_of_threads = %{{.*}} {aligned = false}
+ nvvm.barrier id = %barID number_of_threads = %numberOfThreads {aligned = false}
+
+ // Non-aligned reduction variants.
+ // LLVM: %[[redOperandCmp5:.*]] = icmp ne i32 %[[redOperand]], 0
+ // LLVM: %{{.*}} = call i1 @llvm.nvvm.barrier.cta.red.and.all(i32 0, i1 %[[redOperandCmp5]])
+ // CHECK: %{{.*}} = nvvm.barrier.reduction #nvvm.reduction<and> %{{.*}} -> i32 {aligned = false}
+ %4 = nvvm.barrier.reduction #nvvm.reduction<and> %redOperand -> i32 {aligned = false}
+ // LLVM: %[[redOperandCmp6:.*]] = icmp ne i32 %[[redOperand]], 0
+ // LLVM: %{{.*}} = call i1 @llvm.nvvm.barrier.cta.red.or.all(i32 0, i1 %[[redOperandCmp6]])
+ // CHECK: %{{.*}} = nvvm.barrier.reduction #nvvm.reduction<or> %{{.*}} -> i32 {aligned = false}
+ %5 = nvvm.barrier.reduction #nvvm.reduction<or> %redOperand -> i32 {aligned = false}
+ // LLVM: %[[redOperandCmp7:.*]] = icmp ne i32 %[[redOperand]], 0
+ // LLVM: %{{.*}} = call i32 @llvm.nvvm.barrier.cta.red.popc.all(i32 0, i1 %[[redOperandCmp7]])
+ // CHECK: %{{.*}} = nvvm.barrier.reduction #nvvm.reduction<popc> %{{.*}} -> i32 {aligned = false}
+ %6 = nvvm.barrier.reduction #nvvm.reduction<popc> %redOperand -> i32 {aligned = false}
+
llvm.return
}
>From 8a3890dd1694cb4a4f88349a0c0b5b840b558042 Mon Sep 17 00:00:00 2001
From: Yuansui Xu <xuyuansui at outlook.com>
Date: Mon, 1 Jun 2026 02:37:28 -0500
Subject: [PATCH 2/2] fix
---
mlir/lib/Dialect/LLVMIR/IR/NVVMDialect.cpp | 3 +++
1 file changed, 3 insertions(+)
diff --git a/mlir/lib/Dialect/LLVMIR/IR/NVVMDialect.cpp b/mlir/lib/Dialect/LLVMIR/IR/NVVMDialect.cpp
index 8c90f3f43bb49..a37e23c438f43 100644
--- a/mlir/lib/Dialect/LLVMIR/IR/NVVMDialect.cpp
+++ b/mlir/lib/Dialect/LLVMIR/IR/NVVMDialect.cpp
@@ -3448,6 +3448,9 @@ void SubFOp::getCanonicalizationPatterns(RewritePatternSet &patterns,
MLIRContext *context) {
patterns.add<ConvertFsubToFnegFadd>(context);
}
+//===----------------------------------------------------------------------===//
+// getIntrinsicID/getIntrinsicIDAndArgs methods
+//===----------------------------------------------------------------------===//
/// Maps the (aligned, hasCount) pair to the `@llvm.nvvm.barrier.cta.sync.*`
/// intrinsic ID.
More information about the Mlir-commits
mailing list