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..9c2b259a96776 100644 --- a/mlir/lib/Dialect/LLVMIR/IR/NVVMDialect.cpp +++ b/mlir/lib/Dialect/LLVMIR/IR/NVVMDialect.cpp @@ -1563,6 +1563,43 @@ 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::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 " + << 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 " + << reduxType << " type. Only supported type is 'f32'."; + 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..a8a743006fbf8 --- /dev/null +++ b/mlir/test/Target/LLVMIR/nvvm/redux-sync-invalid.mlir @@ -0,0 +1,49 @@ +// 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 {{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 {{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_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 +} + +// ----- + +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 +} + +// ----- + +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 +} + +// ----- + +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 +}