[mlir][vector] Bring back `maxf`/`minf` reductions

[unterumarmung]

Summary

This patch is part of a larger initiative aimed at fixing floating-point max and min operations in MLIR: https://discourse.llvm.org/t/rfc-fix-floating-point-max-and-min-operations-in-mlir/72671.

In line with the mentioned RFC, this patch tackles tasks 2.3 and 2.4.
It adds LLVM conversions for the maxf/minf reductions to the non-NaN-propagating LLVM intrinsics.

Depends on D158618

[dcaballe]

LGTM, thanks!

@mravishankar, if we convert vector.reduction<minimum/maximum> to vector.reduction<minf/maxf> after this patch, NVTX tests should work.

[mravishankar]

Nice! Thanks @dcaballe and @unterumarmung . Is there a way we can land these patches and the conversion from the NaN propagating version to the non-NaN propagating version in vector dialect. Then we can address downstream issues before we do larger changes.

[unterumarmung]

In D158659#4612105, @mravishankar wrote:

Is there a way we can land these patches and the conversion from the NaN propagating version to the non-NaN propagating version in vector dialect. Then we can address downstream issues before we do larger changes.

Once these patches are applied, we'll end up with two sets of min/max reductions: m**f and m**imumf. The current issue revolves around the m**f reductions, as their lowering is unsupported by certain NVPTX backends. Once these patches are incorporated, the semantics and consequently the lowering of these reductions will shift towards a more conservative approach. This conservative approach has been a part of LLVM for a considerable amount of time. Hence, the older backends shouldn't encounter issues with this altered lowering strategy.

mlir/lib/Conversion/VectorToLLVM/ConvertVectorToLLVM.cpp

 template <>
 struct VectorToScalarMapper<LLVM::vector_reduce_fmaximum> {
   using Type = LLVM::MaximumOp;
 };
 template <>
 struct VectorToScalarMapper<LLVM::vector_reduce_fminimum> {
   using Type = LLVM::MinimumOp;
 };
+template <>
+struct VectorToScalarMapper<LLVM::vector_reduce_fmax> {
+  using Type = LLVM::MaxNumOp;
+};
+template <>
+struct VectorToScalarMapper<LLVM::vector_reduce_fmin> {
+  using Type = LLVM::MinNumOp;
+};
 } // namespace
 
 template <class LLVMRedIntrinOp>
 static Value
 createFPReductionComparisonOpLowering(ConversionPatternRewriter &rewriter,
                                       Location loc, Type llvmType,
                                       Value vectorOperand, Value accumulator) {
   Value result = rewriter.create<LLVMRedIntrinOp>(loc, llvmType, vectorOperand);
matchAndRewrite(vector::ReductionOp reductionOp, OpAdaptor adaptor,
     } else if (kind == vector::CombiningKind::MINIMUMF) {
       result =
           createFPReductionComparisonOpLowering<LLVM::vector_reduce_fminimum>(
               rewriter, loc, llvmType, operand, acc);
     } else if (kind == vector::CombiningKind::MAXIMUMF) {
       result =
           createFPReductionComparisonOpLowering<LLVM::vector_reduce_fmaximum>(
               rewriter, loc, llvmType, operand, acc);
+    } else if (kind == vector::CombiningKind::MINF) {
+      result = createFPReductionComparisonOpLowering<LLVM::vector_reduce_fmin>(
+          rewriter, loc, llvmType, operand, acc);
+    } else if (kind == vector::CombiningKind::MAXF) {
+      result = createFPReductionComparisonOpLowering<LLVM::vector_reduce_fmax>(
+          rewriter, loc, llvmType, operand, acc);
     } else
       return failure();
 
     rewriter.replaceOp(reductionOp, result);
     return success();
   }
 
 private:
case vector::CombiningKind::OR:
           rewriter, loc, llvmType, operand, acc, maskOp.getMask());
       break;
     case vector::CombiningKind::XOR:
       result = lowerReductionWithStartValue<LLVM::VPReduceXorOp,
                                             ReductionNeutralZero>(
           rewriter, loc, llvmType, operand, acc, maskOp.getMask());
       break;
     case vector::CombiningKind::MINF:
-      // FIXME: MLIR's 'minf' and LLVM's 'vector_reduce_fmin' do not handle
-      // NaNs/-0.0/+0.0 in the same way.
       result = lowerReductionWithStartValue<LLVM::VPReduceFMinOp,
                                             ReductionNeutralFPMax>(
           rewriter, loc, llvmType, operand, acc, maskOp.getMask());
       break;
     case vector::CombiningKind::MAXF:
-      // FIXME: MLIR's 'minf' and LLVM's 'vector_reduce_fmin' do not handle
-      // NaNs/-0.0/+0.0 in the same way.
       result = lowerReductionWithStartValue<LLVM::VPReduceFMaxOp,
                                             ReductionNeutralFPMin>(
           rewriter, loc, llvmType, operand, acc, maskOp.getMask());
       break;
     default:
       return rewriter.notifyMatchFailure(
           maskOp,
           "lowering to LLVM is not implemented for this masked operation");

mlir/test/Conversion/VectorToLLVM/vector-to-llvm.mlir

 // CHECK-LABEL: @reduce_fminimum_f32(
 // CHECK-SAME: %[[A:.*]]: vector<16xf32>, %[[B:.*]]: f32)
 //      CHECK: %[[V:.*]] = llvm.intr.vector.reduce.fminimum(%[[A]]) : (vector<16xf32>) -> f32
 //      CHECK: %[[R:.*]] = llvm.intr.minimum(%[[V]], %[[B]]) : (f32, f32) -> f32
 //      CHECK: return %[[R]] : f32
 
 // -----
 
+func.func @reduce_fmax_f32(%arg0: vector<16xf32>, %arg1: f32) -> f32 {
+  %0 = vector.reduction <maxf>, %arg0, %arg1 : vector<16xf32> into f32
+  return %0 : f32
+}
+// CHECK-LABEL: @reduce_fmax_f32(
+// CHECK-SAME: %[[A:.*]]: vector<16xf32>, %[[B:.*]]: f32)
+//      CHECK: %[[V:.*]] = llvm.intr.vector.reduce.fmax(%[[A]]) : (vector<16xf32>) -> f32
+//      CHECK: %[[R:.*]] = llvm.intr.maxnum(%[[V]], %[[B]]) : (f32, f32) -> f32
+//      CHECK: return %[[R]] : f32
+
+// -----
+
+func.func @reduce_fmin_f32(%arg0: vector<16xf32>, %arg1: f32) -> f32 {
+  %0 = vector.reduction <minf>, %arg0, %arg1 : vector<16xf32> into f32
+  return %0 : f32
+}
+// CHECK-LABEL: @reduce_fmin_f32(
+// CHECK-SAME: %[[A:.*]]: vector<16xf32>, %[[B:.*]]: f32)
+//      CHECK: %[[V:.*]] = llvm.intr.vector.reduce.fmin(%[[A]]) : (vector<16xf32>) -> f32
+//      CHECK: %[[R:.*]] = llvm.intr.minnum(%[[V]], %[[B]]) : (f32, f32) -> f32
+//      CHECK: return %[[R]] : f32
+
+// -----
+
 func.func @reduce_minui_i32(%arg0: vector<16xi32>) -> i32 {
   %0 = vector.reduction <minui>, %arg0 : vector<16xi32> into i32
   return %0 : i32
 }
 // CHECK-LABEL: @reduce_minui_i32(
 //  CHECK-SAME: %[[A:.*]]: vector<16xi32>)
 //       CHECK: %[[V:.*]] = "llvm.intr.vector.reduce.umin"(%[[A]])
 //       CHECK: return %[[V]] : i32

mlir/test/Integration/Dialect/Vector/CPU/test-reductions-f32-reassoc.mlir

func.func @entry() {
   vector.print %1 : f32
   // CHECK: 6
   %2 = vector.reduction <minimumf>, %v2 : vector<64xf32> into f32
   vector.print %2 : f32
   // CHECK: 1
   %3 = vector.reduction <maximumf>, %v2 : vector<64xf32> into f32
   vector.print %3 : f32
   // CHECK: 3
+  %4 = vector.reduction <minf>, %v2 : vector<64xf32> into f32
+  vector.print %4 : f32
+  // CHECK: 1
+  %5 = vector.reduction <maxf>, %v2 : vector<64xf32> into f32
+  vector.print %5 : f32
+  // CHECK: 3
 
   return
 }

mlir/test/Integration/Dialect/Vector/CPU/test-reductions-f32.mlir

func.func @entry() {
   vector.print %1 : f32
   // CHECK: -5760
   %2 = vector.reduction <minimumf>, %v9 : vector<10xf32> into f32
   vector.print %2 : f32
   // CHECK: -16
   %3 = vector.reduction <maximumf>, %v9 : vector<10xf32> into f32
   vector.print %3 : f32
   // CHECK: 5
+  %4 = vector.reduction <minf>, %v9 : vector<10xf32> into f32
+  vector.print %4 : f32
+  // CHECK: -16
+  %5 = vector.reduction <maxf>, %v9 : vector<10xf32> into f32
+  vector.print %5 : f32
+  // CHECK: 5
 
   return
 }

mlir/test/Integration/Dialect/Vector/CPU/test-reductions-f64-reassoc.mlir

func.func @entry() {
   vector.print %1 : f64
   // CHECK: 6
   %2 = vector.reduction <minimumf>, %v2 : vector<64xf64> into f64
   vector.print %2 : f64
   // CHECK: 1
   %3 = vector.reduction <maximumf>, %v2 : vector<64xf64> into f64
   vector.print %3 : f64
   // CHECK: 3
+  %4 = vector.reduction <minf>, %v2 : vector<64xf64> into f64
+  vector.print %4 : f64
+  // CHECK: 1
+  %5 = vector.reduction <maxf>, %v2 : vector<64xf64> into f64
+  vector.print %5 : f64
+  // CHECK: 3
 
   return
 }

mlir/test/Integration/Dialect/Vector/CPU/test-reductions-f64.mlir

func.func @entry() {
   vector.print %1 : f64
   // CHECK: -5760
   %2 = vector.reduction <minimumf>, %v9 : vector<10xf64> into f64
   vector.print %2 : f64
   // CHECK: -16
   %3 = vector.reduction <maximumf>, %v9 : vector<10xf64> into f64
   vector.print %3 : f64
   // CHECK: 5
+  %4 = vector.reduction <minf>, %v9 : vector<10xf64> into f64
+  vector.print %4 : f64
+  // CHECK: -16
+  %5 = vector.reduction <maxf>, %v9 : vector<10xf64> into f64
+  vector.print %5 : f64
+  // CHECK: 5
 
   return
 }