[mlir][vector] Implement Workaround Lowerings for Masked `fm**imum` 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.

Within LLVM, there are no masked reduction counterparts for vector reductions such as fmaximum and fminimum.
More information can be found here: https://github.com/llvm/llvm-project/issues/64940#issuecomment-1690694156.

To address this issue in MLIR, where we need to generate appropriate lowerings for these cases, we employ regular non-masked intrinsics.
However, we modify the input vector using the arith.select operation to effectively deactivate undesired elements using a "neutral mask value".
The neutral mask value is the smallest possible value for the fmaximum reduction and the largest possible value for the fminimum reduction.

Depends on D158618

[unterumarmung]

Rewrite to use only LLVM Dialect operations, without arith.constant, arith.select and vector.broadcast

[dcaballe]

LGTM! Just minor comments. You can address them before commiting. Thanks!

mlir/lib/Conversion/VectorToLLVM/ConvertVectorToLLVM.cpp

599

nit: Finish comment with '.' per coding standards. Here and a couple of times below.

622–623

nit: value -> neutralVal or something more specific.
nit: denseValue -> vecNeutral or something more specific.

mlir/lib/Conversion/VectorToLLVM/ConvertVectorToLLVM.cpp

 
 #include "mlir/Conversion/LLVMCommon/TypeConverter.h"
 #include "mlir/Conversion/LLVMCommon/VectorPattern.h"
 #include "mlir/Dialect/Arith/IR/Arith.h"
 #include "mlir/Dialect/Arith/Utils/Utils.h"
 #include "mlir/Dialect/LLVMIR/FunctionCallUtils.h"
 #include "mlir/Dialect/LLVMIR/LLVMDialect.h"
 #include "mlir/Dialect/MemRef/IR/MemRef.h"
+#include "mlir/Dialect/Vector/IR/VectorOps.h"
 #include "mlir/Dialect/Vector/Interfaces/MaskableOpInterface.h"
 #include "mlir/Dialect/Vector/Transforms/LoweringPatterns.h"
 #include "mlir/Dialect/Vector/Transforms/VectorTransforms.h"
 #include "mlir/IR/BuiltinTypes.h"
 #include "mlir/IR/TypeUtilities.h"
 #include "mlir/Target/LLVMIR/TypeToLLVM.h"
 #include "mlir/Transforms/DialectConversion.h"
 #include "llvm/Support/Casting.h"
if (accumulator) {
     result =
         rewriter.create<typename VectorToScalarMapper<LLVMRedIntrinOp>::Type>(
             loc, result, accumulator);
   }
 
   return result;
 }
 
+/// Reduction neutral classes for overloading
+class MaskNeutralFMaximum {};
+class MaskNeutralFMinimum {};
+
+/// Create the mask neutral floating point maximum value
+static Value createMaskNeutralValue(MaskNeutralFMaximum,
+                                    ConversionPatternRewriter &rewriter,
+                                    Location loc, Type llvmType) {
+  auto floatType = cast<FloatType>(llvmType);
+  return rewriter.create<LLVM::ConstantOp>(
+      loc, llvmType,
+      rewriter.getFloatAttr(
+          llvmType, llvm::APFloat::getSmallest(floatType.getFloatSemantics(),
+                                               /*Negative=*/true)));
+}
+
+/// Create the mask neutral floating point minimum value
+static Value createMaskNeutralValue(MaskNeutralFMinimum,
+                                    ConversionPatternRewriter &rewriter,
+                                    Location loc, Type llvmType) {
+  auto floatType = cast<FloatType>(llvmType);
+  return rewriter.create<LLVM::ConstantOp>(
+      loc, llvmType,
+      rewriter.getFloatAttr(
+          llvmType, llvm::APFloat::getLargest(floatType.getFloatSemantics(),
+                                              /*Negative=*/false)));
+}
+
+/// Lowers masked `fmaximum` and `fminimum` reductions using the non-masked
+/// intrinsics. It is a workaround to overcome the lack of masked intrinsics for
+/// `fmaximum`/`fminimum`.
+/// More information: https://github.com/llvm/llvm-project/issues/64940
+template <class LLVMRedIntrinOp, class MaskNeutral>
+static Value lowerMaskedReductionWithRegular(
+    ConversionPatternRewriter &rewriter, Location loc, Type llvmType,
+    Value vectorOperand, Value accumulator, Value mask) {
+  const Value singleMaskNeutral =
+      createMaskNeutralValue(MaskNeutral{}, rewriter, loc, llvmType);
+  const Value vectorMaskNeutral = rewriter.create<BroadcastOp>(
+      loc, vectorOperand.getType(), singleMaskNeutral);
+  const Value selectedVectorByMask = rewriter.create<arith::SelectOp>(
+      loc, mask, vectorOperand, vectorMaskNeutral);
+  return createFPReductionComparisonOpLowering<LLVMRedIntrinOp>(
+      rewriter, loc, llvmType, selectedVectorByMask, accumulator);
+}
+
 /// Overloaded methods to lower a reduction to an llvm instrinsic that requires
 /// a start value. This start value format spans across fp reductions without
 /// mask and all the masked reduction intrinsics.
 template <class LLVMVPRedIntrinOp, class ReductionNeutral>
 static Value lowerReductionWithStartValue(ConversionPatternRewriter &rewriter,
                                           Location loc, Type llvmType,
                                           Value vectorOperand,
                                           Value accumulator) {
case vector::CombiningKind::MINF:
       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");
+    case CombiningKind::MAXIMUMF:
+      result = lowerMaskedReductionWithRegular<LLVM::vector_reduce_fmaximum,
+                                               MaskNeutralFMaximum>(
+          rewriter, loc, llvmType, operand, acc, maskOp.getMask());
+      break;
+    case CombiningKind::MINIMUMF:
+      result = lowerMaskedReductionWithRegular<LLVM::vector_reduce_fminimum,
+                                               MaskNeutralFMinimum>(
+          rewriter, loc, llvmType, operand, acc, maskOp.getMask());
+      break;
     }
 
     // Replace `vector.mask` operation altogether.
     rewriter.replaceOp(maskOp, result);
     return success();
   }
 };
 

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

 // CHECK-SAME:                                      %[[INPUT:.*]]: vector<16xf32>,
 // CHECK-SAME:                                      %[[MASK:.*]]: vector<16xi1>) -> f32 {
 // CHECK:           %[[NEUTRAL:.*]] = llvm.mlir.constant(0x7FC00000 : f32) : f32
 // CHECK:           %[[VL:.*]] = llvm.mlir.constant(16 : i32) : i32
 // CHECK:           "llvm.intr.vp.reduce.fmax"(%[[NEUTRAL]], %[[INPUT]], %[[MASK]], %[[VL]]) : (f32, vector<16xf32>, vector<16xi1>, i32) -> f32
 
 // -----
 
+func.func @masked_reduce_maximumf_f32(%arg0: vector<16xf32>, %mask : vector<16xi1>) -> f32 {
+  %0 = vector.mask %mask { vector.reduction <maximumf>, %arg0 : vector<16xf32> into f32 } : vector<16xi1> -> f32
+  return %0 : f32
+}
+
+// CHECK-LABEL:   func.func @masked_reduce_maximumf_f32(
+// CHECK-SAME:                                      %[[INPUT:.*]]: vector<16xf32>,
+// CHECK-SAME:                                      %[[MASK:.*]]: vector<16xi1>) -> f32 {
+// CHECK:           %[[MASK_NEUTRAL:.*]] = arith.constant dense<-1.401300e-45> : vector<16xf32>
+// CHECK:           %[[MASKED:.*]] = arith.select %[[MASK]], %[[INPUT]], %[[MASK_NEUTRAL]] : vector<16xi1>, vector<16xf32>
+// CHECK:           %[[RESULT:.*]] = llvm.intr.vector.reduce.fmaximum(%[[MASKED]])  : (vector<16xf32>) -> f32
+// CHECK:           return %[[RESULT]]
+
+// -----
+
+func.func @masked_reduce_minimumf_f32(%arg0: vector<16xf32>, %mask : vector<16xi1>) -> f32 {
+  %0 = vector.mask %mask { vector.reduction <minimumf>, %arg0 : vector<16xf32> into f32 } : vector<16xi1> -> f32
+  return %0 : f32
+}
+
+// CHECK-LABEL:   func.func @masked_reduce_minimumf_f32(
+// CHECK-SAME:                                      %[[INPUT:.*]]: vector<16xf32>,
+// CHECK-SAME:                                      %[[MASK:.*]]: vector<16xi1>) -> f32 {
+// CHECK:           %[[MASK_NEUTRAL:.*]] = arith.constant dense<3.40282347E+38> : vector<16xf32>
+// CHECK:           %[[MASKED:.*]] = arith.select %[[MASK]], %[[INPUT]], %[[MASK_NEUTRAL]] : vector<16xi1>, vector<16xf32>
+// CHECK:           %[[RESULT:.*]] = llvm.intr.vector.reduce.fminimum(%[[MASKED]])  : (vector<16xf32>) -> f32
+// CHECK:           return %[[RESULT]]
+
+// -----
+
 func.func @masked_reduce_add_i8(%arg0: vector<32xi8>, %mask : vector<32xi1>) -> i8 {
   %0 = vector.mask %mask { vector.reduction <add>, %arg0 : vector<32xi8> into i8 } : vector<32xi1> -> i8
   return %0 : i8
 }
 
 // CHECK-LABEL:   func.func @masked_reduce_add_i8(
 // CHECK-SAME:                             %[[INPUT:.*]]: vector<32xi8>,
 // CHECK-SAME:                             %[[MASK:.*]]: vector<32xi1>) -> i8 {