From d39a26f269385bc02a504f0c634bb3cc75e68321 Mon Sep 17 00:00:00 2001 From: Srinivasa Ravi Date: Fri, 31 Oct 2025 07:43:29 +0000 Subject: [PATCH 1/6] [MLIR][NVVM] Update redux.sync op This change: - Updates the `redux.sync` NVVM Op input and output type constraints - Adds a verifier for the Op to prevent stack dumps and the execution of an `llvm_unreachable` in certain cases of invalid usage, and instead gracefully error out with an informative error message. --- mlir/include/mlir/Dialect/LLVMIR/NVVMOps.td | 8 +++-- mlir/lib/Dialect/LLVMIR/IR/NVVMDialect.cpp | 26 +++++++++++++++ .../Dialect/NVVM/NVVMToLLVMIRTranslation.cpp | 3 -- .../LLVMIR/nvvm/redux-sync-invalid.mlir | 33 +++++++++++++++++++ 4 files changed, 64 insertions(+), 6 deletions(-) create mode 100644 mlir/test/Target/LLVMIR/nvvm/redux-sync-invalid.mlir diff --git a/mlir/include/mlir/Dialect/LLVMIR/NVVMOps.td b/mlir/include/mlir/Dialect/LLVMIR/NVVMOps.td index 4f483859ac18d..3a212f91760f2 100644 --- a/mlir/include/mlir/Dialect/LLVMIR/NVVMOps.td +++ b/mlir/include/mlir/Dialect/LLVMIR/NVVMOps.td @@ -476,9 +476,9 @@ def ReduxKind : I32EnumAttr<"ReduxKind", "NVVM redux kind", def ReduxKindAttr : EnumAttr; def NVVM_ReduxOp : - NVVM_Op<"redux.sync", [NVVMRequiresSM<80>]>, - Results<(outs LLVM_Type:$res)>, - Arguments<(ins LLVM_Type:$val, + NVVM_Op<"redux.sync", [NVVMRequiresSM<80>, AllTypesMatch<["res", "val"]>]>, + Results<(outs AnyTypeOf<[I32, F32]>:$res)>, + Arguments<(ins AnyTypeOf<[I32, F32]>:$val, ReduxKindAttr:$kind, I32:$mask_and_clamp, DefaultValuedAttr:$abs, @@ -496,6 +496,8 @@ def NVVM_ReduxOp : [For more information, see PTX ISA](https://docs.nvidia.com/cuda/parallel-thread-execution/#parallel-synchronization-and-communication-instructions-redux-sync) }]; + let hasVerifier = 1; + string llvmBuilder = [{ auto intId = getReduxIntrinsicId($_resultType, $kind, $abs, $nan); $res = createIntrinsicCall(builder, intId, {$val, $mask_and_clamp}); diff --git a/mlir/lib/Dialect/LLVMIR/IR/NVVMDialect.cpp b/mlir/lib/Dialect/LLVMIR/IR/NVVMDialect.cpp index f0de4dbcc1d4b..c6a5973b4951e 100644 --- a/mlir/lib/Dialect/LLVMIR/IR/NVVMDialect.cpp +++ b/mlir/lib/Dialect/LLVMIR/IR/NVVMDialect.cpp @@ -1563,6 +1563,32 @@ LogicalResult NVVM::ClusterLaunchControlQueryCancelOp::verify() { return success(); } +LogicalResult NVVM::ReduxOp::verify() { + mlir::Type reduxType = getType(); + if (!reduxType.isF32()) { + if (getAbs()) + return emitOpError("abs attribute is supported only for f32 type"); + if (getNan()) + return emitOpError("nan attribute is supported only for f32 type"); + } + + NVVM::ReduxKind kind = getKind(); + switch (kind) { + case NVVM::ReduxKind::FMIN: + case NVVM::ReduxKind::FMAX: + if (!reduxType.isF32()) + return emitOpError("fmin and fmax redux kind must be used with f32 type"); + break; + default: + if (reduxType.isF32()) + return emitOpError( + "only fmin and fmax redux kinds are supported for f32 type"); + break; + } + + return success(); +} + /// Packs the given `field` into the `result`. /// The `result` is 64-bits and each `field` can be 32-bits or narrower. static llvm::Value * diff --git a/mlir/lib/Target/LLVMIR/Dialect/NVVM/NVVMToLLVMIRTranslation.cpp b/mlir/lib/Target/LLVMIR/Dialect/NVVM/NVVMToLLVMIRTranslation.cpp index 3d86b09b32538..0964e1b8c5ef3 100644 --- a/mlir/lib/Target/LLVMIR/Dialect/NVVM/NVVMToLLVMIRTranslation.cpp +++ b/mlir/lib/Target/LLVMIR/Dialect/NVVM/NVVMToLLVMIRTranslation.cpp @@ -36,9 +36,6 @@ using mlir::LLVM::detail::createIntrinsicCall; static llvm::Intrinsic::ID getReduxIntrinsicId(llvm::Type *resultType, NVVM::ReduxKind kind, bool hasAbs, bool hasNaN) { - if (!(resultType->isIntegerTy(32) || resultType->isFloatTy())) - llvm_unreachable("unsupported data type for redux"); - switch (kind) { case NVVM::ReduxKind::ADD: return llvm::Intrinsic::nvvm_redux_sync_add; diff --git a/mlir/test/Target/LLVMIR/nvvm/redux-sync-invalid.mlir b/mlir/test/Target/LLVMIR/nvvm/redux-sync-invalid.mlir new file mode 100644 index 0000000000000..da3d0d11a25e2 --- /dev/null +++ b/mlir/test/Target/LLVMIR/nvvm/redux-sync-invalid.mlir @@ -0,0 +1,33 @@ +// RUN: mlir-translate -verify-diagnostics -split-input-file -mlir-to-llvmir %s + +// ----- + +llvm.func @redux_sync_i32_with_abs(%value: i32, %offset: i32) { + // expected-error@+1 {{'nvvm.redux.sync' op abs attribute is supported only for f32 type}} + %res = nvvm.redux.sync add %value, %offset {abs = true}: i32 -> i32 + llvm.return +} + +// ----- + +llvm.func @redux_sync_i32_with_nan(%value: i32, %offset: i32) { + // expected-error@+1 {{'nvvm.redux.sync' op nan attribute is supported only for f32 type}} + %res = nvvm.redux.sync add %value, %offset {nan = true}: i32 -> i32 + llvm.return +} + +// ----- + +llvm.func @redux_sync_f32_with_invalid_kind(%value: f32, %offset: i32) { + // expected-error@+1 {{'nvvm.redux.sync' op only fmin and fmax redux kinds are supported for f32 type}} + %res = nvvm.redux.sync add %value, %offset: f32 -> f32 + llvm.return +} + +// ----- + +llvm.func @redux_sync_i32_with_invalid_kind(%value: i32, %offset: i32) { + // expected-error@+1 {{'nvvm.redux.sync' op fmin and fmax redux kind must be used with f32 type}} + %res = nvvm.redux.sync fmin %value, %offset: i32 -> i32 + llvm.return +} From dc2092d0f7603e656a88fd68a53dac53ac38b8cb Mon Sep 17 00:00:00 2001 From: Srinivasa Ravi Date: Fri, 31 Oct 2025 10:51:56 +0000 Subject: [PATCH 2/6] address comments --- mlir/lib/Dialect/LLVMIR/IR/NVVMDialect.cpp | 12 ++++++------ 1 file changed, 6 insertions(+), 6 deletions(-) diff --git a/mlir/lib/Dialect/LLVMIR/IR/NVVMDialect.cpp b/mlir/lib/Dialect/LLVMIR/IR/NVVMDialect.cpp index c6a5973b4951e..bd5a759d38faf 100644 --- a/mlir/lib/Dialect/LLVMIR/IR/NVVMDialect.cpp +++ b/mlir/lib/Dialect/LLVMIR/IR/NVVMDialect.cpp @@ -1565,12 +1565,12 @@ LogicalResult NVVM::ClusterLaunchControlQueryCancelOp::verify() { LogicalResult NVVM::ReduxOp::verify() { mlir::Type reduxType = getType(); - if (!reduxType.isF32()) { - if (getAbs()) - return emitOpError("abs attribute is supported only for f32 type"); - if (getNan()) - return emitOpError("nan attribute is supported only for f32 type"); - } + + if (!reduxType.isF32() && getAbs()) + return emitOpError("abs attribute is supported only for f32 type"); + + if (!reduxType.isF32() && getNan()) + return emitOpError("nan attribute is supported only for f32 type"); NVVM::ReduxKind kind = getKind(); switch (kind) { From ed43bae71df6747faa37e958c311a4c22bc9f4cf Mon Sep 17 00:00:00 2001 From: Srinivasa Ravi Date: Mon, 3 Nov 2025 06:43:59 +0000 Subject: [PATCH 3/6] update tests --- .../Target/LLVMIR/nvvm/redux-sync-invalid.mlir | 16 ++++++++++++---- 1 file changed, 12 insertions(+), 4 deletions(-) diff --git a/mlir/test/Target/LLVMIR/nvvm/redux-sync-invalid.mlir b/mlir/test/Target/LLVMIR/nvvm/redux-sync-invalid.mlir index da3d0d11a25e2..26b274d390396 100644 --- a/mlir/test/Target/LLVMIR/nvvm/redux-sync-invalid.mlir +++ b/mlir/test/Target/LLVMIR/nvvm/redux-sync-invalid.mlir @@ -3,7 +3,7 @@ // ----- llvm.func @redux_sync_i32_with_abs(%value: i32, %offset: i32) { - // expected-error@+1 {{'nvvm.redux.sync' op abs attribute is supported only for f32 type}} + // expected-error@+1 {{abs attribute is supported only for f32 type}} %res = nvvm.redux.sync add %value, %offset {abs = true}: i32 -> i32 llvm.return } @@ -11,7 +11,7 @@ llvm.func @redux_sync_i32_with_abs(%value: i32, %offset: i32) { // ----- llvm.func @redux_sync_i32_with_nan(%value: i32, %offset: i32) { - // expected-error@+1 {{'nvvm.redux.sync' op nan attribute is supported only for f32 type}} + // expected-error@+1 {{nan attribute is supported only for f32 type}} %res = nvvm.redux.sync add %value, %offset {nan = true}: i32 -> i32 llvm.return } @@ -19,7 +19,7 @@ llvm.func @redux_sync_i32_with_nan(%value: i32, %offset: i32) { // ----- llvm.func @redux_sync_f32_with_invalid_kind(%value: f32, %offset: i32) { - // expected-error@+1 {{'nvvm.redux.sync' op only fmin and fmax redux kinds are supported for f32 type}} + // expected-error@+1 {{only fmin and fmax redux kinds are supported for f32 type}} %res = nvvm.redux.sync add %value, %offset: f32 -> f32 llvm.return } @@ -27,7 +27,15 @@ llvm.func @redux_sync_f32_with_invalid_kind(%value: f32, %offset: i32) { // ----- llvm.func @redux_sync_i32_with_invalid_kind(%value: i32, %offset: i32) { - // expected-error@+1 {{'nvvm.redux.sync' op fmin and fmax redux kind must be used with f32 type}} + // expected-error@+1 {{fmin and fmax redux kind must be used with f32 type}} %res = nvvm.redux.sync fmin %value, %offset: i32 -> i32 llvm.return } + +// ----- + +llvm.func @redux_sync_non_matching_types(%value: i32, %offset: i32) { + // expected-error@+1 {{failed to verify that all of {res, val} have same type}} + %res = nvvm.redux.sync add %value, %offset: i32 -> f32 + llvm.return +} From 389e86581a38587fc81d9591cdd591518c3599c6 Mon Sep 17 00:00:00 2001 From: Srinivasa Ravi Date: Mon, 3 Nov 2025 11:11:48 +0000 Subject: [PATCH 4/6] address comments --- mlir/lib/Dialect/LLVMIR/IR/NVVMDialect.cpp | 33 ++++++++++++------- .../LLVMIR/nvvm/redux-sync-invalid.mlir | 12 +++++-- 2 files changed, 32 insertions(+), 13 deletions(-) diff --git a/mlir/lib/Dialect/LLVMIR/IR/NVVMDialect.cpp b/mlir/lib/Dialect/LLVMIR/IR/NVVMDialect.cpp index bd5a759d38faf..1a1151808ad1c 100644 --- a/mlir/lib/Dialect/LLVMIR/IR/NVVMDialect.cpp +++ b/mlir/lib/Dialect/LLVMIR/IR/NVVMDialect.cpp @@ -1566,23 +1566,34 @@ LogicalResult NVVM::ClusterLaunchControlQueryCancelOp::verify() { LogicalResult NVVM::ReduxOp::verify() { mlir::Type reduxType = getType(); - if (!reduxType.isF32() && getAbs()) - return emitOpError("abs attribute is supported only for f32 type"); - - if (!reduxType.isF32() && getNan()) - return emitOpError("nan attribute is supported only for f32 type"); + if (!reduxType.isF32()) { + if (getAbs()) + return emitOpError("abs attribute is supported only for f32 type"); + if (getNan()) + return emitOpError("nan attribute is supported only for f32 type"); + } NVVM::ReduxKind kind = getKind(); switch (kind) { + case NVVM::ReduxKind::ADD: + case NVVM::ReduxKind::AND: + case NVVM::ReduxKind::OR: + case NVVM::ReduxKind::XOR: + case NVVM::ReduxKind::MAX: + case NVVM::ReduxKind::MIN: + case NVVM::ReduxKind::UMAX: + case NVVM::ReduxKind::UMIN: + if (!reduxType.isInteger(32)) + return emitOpError("'") + << stringifyEnum(kind) << "' redux kind unsupported with " + << getType() << " type. Only supported type is 'i32'."; + break; case NVVM::ReduxKind::FMIN: case NVVM::ReduxKind::FMAX: if (!reduxType.isF32()) - return emitOpError("fmin and fmax redux kind must be used with f32 type"); - break; - default: - if (reduxType.isF32()) - return emitOpError( - "only fmin and fmax redux kinds are supported for f32 type"); + return emitOpError("'") + << stringifyEnum(kind) << "' redux kind unsupported with " + << getType() << " type. Only supported type is 'f32'."; break; } diff --git a/mlir/test/Target/LLVMIR/nvvm/redux-sync-invalid.mlir b/mlir/test/Target/LLVMIR/nvvm/redux-sync-invalid.mlir index 26b274d390396..988de950d67a6 100644 --- a/mlir/test/Target/LLVMIR/nvvm/redux-sync-invalid.mlir +++ b/mlir/test/Target/LLVMIR/nvvm/redux-sync-invalid.mlir @@ -19,15 +19,23 @@ llvm.func @redux_sync_i32_with_nan(%value: i32, %offset: i32) { // ----- llvm.func @redux_sync_f32_with_invalid_kind(%value: f32, %offset: i32) { - // expected-error@+1 {{only fmin and fmax redux kinds are supported for f32 type}} + // expected-error@+1 {{'add' redux kind unsupported with 'f32' type. Only supported type is 'i32'.}} %res = nvvm.redux.sync add %value, %offset: f32 -> f32 llvm.return } // ----- +llvm.func @redux_sync_f32_with_invalid_kind(%value: f32, %offset: i32) { + // expected-error@+1 {{'and' redux kind unsupported with 'f32' type. Only supported type is 'i32'.}} + %res = nvvm.redux.sync and %value, %offset: f32 -> f32 + llvm.return +} + +// ----- + llvm.func @redux_sync_i32_with_invalid_kind(%value: i32, %offset: i32) { - // expected-error@+1 {{fmin and fmax redux kind must be used with f32 type}} + // expected-error@+1 {{'fmin' redux kind unsupported with 'i32' type. Only supported type is 'f32'.}} %res = nvvm.redux.sync fmin %value, %offset: i32 -> i32 llvm.return } From db57c2a4badf04442adb29a5bb40927eca3f50e0 Mon Sep 17 00:00:00 2001 From: Srinivasa Ravi Date: Mon, 3 Nov 2025 11:24:28 +0000 Subject: [PATCH 5/6] use var instead of calling getType --- mlir/lib/Dialect/LLVMIR/IR/NVVMDialect.cpp | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/mlir/lib/Dialect/LLVMIR/IR/NVVMDialect.cpp b/mlir/lib/Dialect/LLVMIR/IR/NVVMDialect.cpp index 1a1151808ad1c..9c2b259a96776 100644 --- a/mlir/lib/Dialect/LLVMIR/IR/NVVMDialect.cpp +++ b/mlir/lib/Dialect/LLVMIR/IR/NVVMDialect.cpp @@ -1586,14 +1586,14 @@ LogicalResult NVVM::ReduxOp::verify() { if (!reduxType.isInteger(32)) return emitOpError("'") << stringifyEnum(kind) << "' redux kind unsupported with " - << getType() << " type. Only supported type is 'i32'."; + << reduxType << " type. Only supported type is 'i32'."; break; case NVVM::ReduxKind::FMIN: case NVVM::ReduxKind::FMAX: if (!reduxType.isF32()) return emitOpError("'") << stringifyEnum(kind) << "' redux kind unsupported with " - << getType() << " type. Only supported type is 'f32'."; + << reduxType << " type. Only supported type is 'f32'."; break; } From d0c6ddeca421455900d17c65febf66994ecc6ea6 Mon Sep 17 00:00:00 2001 From: Srinivasa Ravi Date: Mon, 3 Nov 2025 13:16:04 +0000 Subject: [PATCH 6/6] fix test names --- mlir/test/Target/LLVMIR/nvvm/redux-sync-invalid.mlir | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/mlir/test/Target/LLVMIR/nvvm/redux-sync-invalid.mlir b/mlir/test/Target/LLVMIR/nvvm/redux-sync-invalid.mlir index 988de950d67a6..a8a743006fbf8 100644 --- a/mlir/test/Target/LLVMIR/nvvm/redux-sync-invalid.mlir +++ b/mlir/test/Target/LLVMIR/nvvm/redux-sync-invalid.mlir @@ -18,7 +18,7 @@ llvm.func @redux_sync_i32_with_nan(%value: i32, %offset: i32) { // ----- -llvm.func @redux_sync_f32_with_invalid_kind(%value: f32, %offset: i32) { +llvm.func @redux_sync_f32_with_invalid_kind_add(%value: f32, %offset: i32) { // expected-error@+1 {{'add' redux kind unsupported with 'f32' type. Only supported type is 'i32'.}} %res = nvvm.redux.sync add %value, %offset: f32 -> f32 llvm.return @@ -26,7 +26,7 @@ llvm.func @redux_sync_f32_with_invalid_kind(%value: f32, %offset: i32) { // ----- -llvm.func @redux_sync_f32_with_invalid_kind(%value: f32, %offset: i32) { +llvm.func @redux_sync_f32_with_invalid_kind_and(%value: f32, %offset: i32) { // expected-error@+1 {{'and' redux kind unsupported with 'f32' type. Only supported type is 'i32'.}} %res = nvvm.redux.sync and %value, %offset: f32 -> f32 llvm.return @@ -34,7 +34,7 @@ llvm.func @redux_sync_f32_with_invalid_kind(%value: f32, %offset: i32) { // ----- -llvm.func @redux_sync_i32_with_invalid_kind(%value: i32, %offset: i32) { +llvm.func @redux_sync_i32_with_invalid_kind_fmin(%value: i32, %offset: i32) { // expected-error@+1 {{'fmin' redux kind unsupported with 'i32' type. Only supported type is 'f32'.}} %res = nvvm.redux.sync fmin %value, %offset: i32 -> i32 llvm.return