[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.

Author: Wolfram70

mlir/include/mlir/Dialect/LLVMIR/NVVMOps.td

@@ -476,9 +476,9 @@ def ReduxKind : I32EnumAttr<"ReduxKind", "NVVM redux kind",
 def ReduxKindAttr : EnumAttr<NVVM_Dialect, ReduxKind, "redux_kind">;
 
 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<BoolAttr, "false">:$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});

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 *

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;

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
+}