[mlir] Support fast-math friendly constants for identity value

[qcolombet]

Summary

Add an option to the family of getIdentity helper functions so that it is possible to produce fast-math friendly constants.

For instance, for maxf the identity value is -inf, however, if the related operations are lowered with fast-math (noinf in particular), then the value becomes poison and chances are the whole codegen is not going to do what we want.

To avoid this problem, we add an option to getIdentity and friends that specifies whether a finite value needs to be produced or not.

The patch is NFC for all the code but the lowering of linalg::softmax because we know we lower that with fast-math down the line.

I didn't audit the rest of the code to check if it would make sense to set this boolean in more places.

Note: It feels kind of wrong to have to know what the lowering may do, but I don't know what the right (at least short-term) solution is. Long term, we may want a special "neutral element" attribute for the respective ops. I didn't think too much about the implications for that.

[qcolombet]

• Fix Smallest/Largest for maxf/minf respectively.

[rengolin]

Note: It feels kind of wrong to have to know what the lowering may do, but I don't know what the right (at least short-term) solution is. Long term, we may want a special "neutral element" attribute for the respective ops. I didn't think too much about the implications for that.

Fast-math has consequences that go beyond the neutral element, so I think we need a phased approach. First, we introduce patches like this one to "test the waters" and see what we really want in MLIR. After we learn what went right and what went wrong, we can start formulating a strategy.

Just like we did in LLVM, the defaults for each flag/option will be based on least-surprise, most-safe depending on the optimization level. We don't quite have such a thing for MLIR yet, but as we reinforce the dialects and passes upstream, it's bound to happen sooner or later.

My guess is that we'll tune around what GPU toolchains expect, because most MLIR ends up there, but we need to allow the front-ends to control this somehow for the cases it doesn't (ex. Flang, CPU inference).

Overall, the patch looks good with the inline comments, being "safer" than before, but I'll leave time for others to review.

mlir/include/mlir/Dialect/Arith/IR/Arith.h

134

If we eventually add some/all of the fast-math flags from LLVM, we may want this to be a bitfield struct or something.

Not necessarily for this patch, but it's good to start thinking about it.

146

Can we make these getters to set a default value? then we don't need to override on all calls, only when we know we need, for instance, when the front-end / pass knows what it's doing.

mlir/lib/Dialect/Arith/IR/ArithOps.cpp

2397

Later we could discuss if it makes sense to expand fast-math logic into APFloat, so that getSmallest/getLargest would "know what to do" against the fast-math flags.

mlir/test/Dialect/Linalg/transform-op-decompose.mlir

213

With no way to currently pass the flag through mlir-opt, it's hard to test the inf case...

[rengolin]

Fast Math discussion on the Arith dialect: https://discourse.llvm.org/t/rfc-fastmath-flags-support-in-complex-dialect/71981

[qcolombet]

Thanks for the link @rengolin !

mlir/include/mlir/Dialect/Arith/IR/Arith.h

146

Yes we could. I decided against it because I wanted people to think about this problem when they use this API.
Now, since this may be a temporary solution, it may be better to limit the changes.

WDYT?

mlir/test/Dialect/Linalg/transform-op-decompose.mlir

213

Good point, maybe we'll want a sort of index-width option equivalent for fast math handling.

[rengolin]

mlir/include/mlir/Dialect/Arith/IR/Arith.h

146

I wanted people to think about this problem when they use this API.

This works for people adding new code, not the existing code that is being changed with this patch, which the original authors are probably not going to look back.

it may be better to limit the changes.

Also, if the RFC linked goes ahead, the argument will change soon and lead to more mechanical changes.

I'd keep the changes minimal (default value) for now and enter the RFC discussion for the long term strategy.

[dcaballe]

Thanks, Quentin! LGTM. I mostly agree with Renato's feedback. We also have the fact that fmin/fmax semantics in MLIR match LLVM's fmaximum/fminimum. We have started to align MLIR with LLVM, including the reduction flavors, but it's going to take some time...

mlir/include/mlir/Dialect/Arith/IR/Arith.h

125

We use markdown style instead of \p in MLIR

[Hardcode84]

mlir/lib/Dialect/Arith/IR/ArithOps.cpp

2445

While we need to pass flag explicitly to AtomicRMWKind version, for op version we can take fastmath flag from op itself (which already have fastmath flags support I believe).

[qcolombet]

mlir/lib/Dialect/Arith/IR/ArithOps.cpp

2445

I think the fastmath flags are only on arith operations, not on the generic ones.
I.e., I'm not sure this is generally possible.

[Hardcode84]

mlir/lib/Dialect/Arith/IR/ArithOps.cpp

2445

Current getNeutralElement(op) implementation supports only arith ops now, so I don't see problem here.

[qcolombet]

• Rebase
• Set default value for getIdentity
• Use markdown instead of \p.

[qcolombet]

mlir/lib/Dialect/Arith/IR/ArithOps.cpp

2445

Good point. Let me give it a try then.

[qcolombet]

• Use the fastmath flags attached to arith op when using getIdentityElement
• Add tests catching that behavior

[qcolombet]

mlir/include/mlir/Dialect/Arith/IR/Arith.h

146

Used default value that keeps the old behavior.

mlir/lib/Dialect/Arith/IR/ArithOps.cpp

2445

Please take a look @Hardcode84

mlir/test/Dialect/Linalg/transform-op-decompose.mlir

213

Added a test through the reduction split test following @Hardcode84 recommendation of using the fastmath flag attached to the operation when using getIdentityElement.

[Hardcode84]

mlir/lib/Dialect/Arith/IR/ArithOps.cpp

2469

There is ArithFastMathInterface, I think you can get flags from it instead of iteration through all attrs.

[qcolombet]

• Use ArithFastMathInterface instead of iterating through the attributes

[qcolombet]

mlir/lib/Dialect/Arith/IR/ArithOps.cpp

2469

Ah thanks!

[Hardcode84]

Thanks!

[qcolombet]

Hi @rengolin,

Does this look good to you?

Cheers,
-Quentin

mlir/include/mlir/Dialect/Arith/IR/Arith.h

bool applyCmpPredicate(arith::CmpIPredicate predicate, const APInt &lhs,
                        const APInt &rhs);
 
 /// Compute `lhs` `pred` `rhs`, where `pred` is one of the known floating point
 /// comparison predicates.
 bool applyCmpPredicate(arith::CmpFPredicate predicate, const APFloat &lhs,
                        const APFloat &rhs);
 
 /// Returns the identity value attribute associated with an AtomicRMWKind op.
+/// `useOnlyFiniteValue` defines whether the identity value should steer away
+/// from infinity representations or anything that is not a proper finite
+/// number.
+/// E.g., The identity value for maxf is in theory `-Inf`, but if we want to
+/// stay in the finite range, it would be `BiggestRepresentableNegativeFloat`.
+/// The purpose of this boolean is to offer constants that will play nice
+/// with fast math related optimizations.
 TypedAttr getIdentityValueAttr(AtomicRMWKind kind, Type resultType,
-                               OpBuilder &builder, Location loc);
+                               OpBuilder &builder, Location loc,
+                               bool useOnlyFiniteValue = false);
 
 /// Return the identity numeric value associated to the give op. Return
 /// std::nullopt if there is no known neutral element.
+/// If `op` has `FastMathFlags::ninf`, only finite values will be used
+/// as neutral element.
 std::optional<TypedAttr> getNeutralElement(Operation *op);
 
 /// Returns the identity value associated with an AtomicRMWKind op.
+/// \see getIdentityValueAttr for a description of what `useOnlyFiniteValue`
+/// does.
 Value getIdentityValue(AtomicRMWKind op, Type resultType, OpBuilder &builder,
-                       Location loc);
+                       Location loc, bool useOnlyFiniteValue = false);
 
 /// Returns the value obtained by applying the reduction operation kind
 /// associated with a binary AtomicRMWKind op to `lhs` and `rhs`.
 Value getReductionOp(AtomicRMWKind op, OpBuilder &builder, Location loc,
                      Value lhs, Value rhs);
 
 arith::CmpIPredicate invertPredicate(arith::CmpIPredicate pred);
 } // namespace arith
 } // namespace mlir
 
 #endif // MLIR_DIALECT_ARITH_IR_ARITH_H_

mlir/lib/Dialect/Arith/IR/ArithOps.cpp

 }
 
 //===----------------------------------------------------------------------===//
 // Atomic Enum
 //===----------------------------------------------------------------------===//
 
 /// Returns the identity value attribute associated with an AtomicRMWKind op.
 TypedAttr mlir::arith::getIdentityValueAttr(AtomicRMWKind kind, Type resultType,
-                                            OpBuilder &builder, Location loc) {
+                                            OpBuilder &builder, Location loc,
+                                            bool useOnlyFiniteValue) {
   switch (kind) {
-  case AtomicRMWKind::maxf:
-    return builder.getFloatAttr(
-        resultType,
-        APFloat::getInf(llvm::cast<FloatType>(resultType).getFloatSemantics(),
-                        /*Negative=*/true));
+  case AtomicRMWKind::maxf: {
+    const llvm::fltSemantics &semantic =
+        llvm::cast<FloatType>(resultType).getFloatSemantics();
+    APFloat identity = useOnlyFiniteValue
+                           ? APFloat::getSmallest(semantic, /*Negative=*/true)
+                           : APFloat::getInf(semantic, /*Negative=*/true);
+    return builder.getFloatAttr(resultType, identity);
+  }
   case AtomicRMWKind::addf:
   case AtomicRMWKind::addi:
   case AtomicRMWKind::maxu:
   case AtomicRMWKind::ori:
     return builder.getZeroAttr(resultType);
   case AtomicRMWKind::andi:
     return builder.getIntegerAttr(
         resultType,
         APInt::getAllOnes(llvm::cast<IntegerType>(resultType).getWidth()));
   case AtomicRMWKind::maxs:
     return builder.getIntegerAttr(
         resultType, APInt::getSignedMinValue(
                         llvm::cast<IntegerType>(resultType).getWidth()));
-  case AtomicRMWKind::minf:
-    return builder.getFloatAttr(
-        resultType,
-        APFloat::getInf(llvm::cast<FloatType>(resultType).getFloatSemantics(),
-                        /*Negative=*/false));
+  case AtomicRMWKind::minf: {
+    const llvm::fltSemantics &semantic =
+        llvm::cast<FloatType>(resultType).getFloatSemantics();
+    APFloat identity = useOnlyFiniteValue
+                           ? APFloat::getLargest(semantic, /*Negative=*/false)
+                           : APFloat::getInf(semantic, /*Negative=*/false);
+
+    return builder.getFloatAttr(resultType, identity);
+  }
   case AtomicRMWKind::mins:
     return builder.getIntegerAttr(
         resultType, APInt::getSignedMaxValue(
                         llvm::cast<IntegerType>(resultType).getWidth()));
   case AtomicRMWKind::minu:
     return builder.getIntegerAttr(
         resultType,
         APInt::getMaxValue(llvm::cast<IntegerType>(resultType).getWidth()));
   case AtomicRMWKind::muli:
     return builder.getIntegerAttr(resultType, 1);
   case AtomicRMWKind::mulf:
     return builder.getFloatAttr(resultType, 1);
   // TODO: Add remaining reduction operations.
   default:
     (void)emitOptionalError(loc, "Reduction operation type not supported");
     break;
   }
   return nullptr;
 }
 
 /// Return the identity numeric value associated to the give op.
 std::optional<TypedAttr> mlir::arith::getNeutralElement(Operation *op) {
   std::optional<AtomicRMWKind> maybeKind =
       llvm::TypeSwitch<Operation *, std::optional<AtomicRMWKind>>(op)
           // Floating-point operations.
           .Case([](arith::AddFOp op) { return AtomicRMWKind::addf; })
           .Case([](arith::MulFOp op) { return AtomicRMWKind::mulf; })
           .Case([](arith::MaxFOp op) { return AtomicRMWKind::maxf; })
           .Case([](arith::MinFOp op) { return AtomicRMWKind::minf; })
           // Integer operations.
           .Case([](arith::AddIOp op) { return AtomicRMWKind::addi; })
           .Case([](arith::OrIOp op) { return AtomicRMWKind::ori; })
           .Case([](arith::XOrIOp op) { return AtomicRMWKind::ori; })
           .Case([](arith::AndIOp op) { return AtomicRMWKind::andi; })
           .Case([](arith::MaxUIOp op) { return AtomicRMWKind::maxu; })
           .Case([](arith::MinUIOp op) { return AtomicRMWKind::minu; })
           .Case([](arith::MaxSIOp op) { return AtomicRMWKind::maxs; })
           .Case([](arith::MinSIOp op) { return AtomicRMWKind::mins; })
           .Case([](arith::MulIOp op) { return AtomicRMWKind::muli; })
           .Default([](Operation *op) { return std::nullopt; });
   if (!maybeKind) {
     op->emitError() << "Unknown neutral element for: " << *op;
     return std::nullopt;
   }
 
+  bool useOnlyFiniteValue = false;
+  for (const auto &attr : op->getAttrs()) {
+    auto fastAttr =
+        llvm::dyn_cast<mlir::arith::FastMathFlagsAttr>(attr.getValue());
+    if (fastAttr && bitEnumContainsAny(fastAttr.getValue(),
+                                       mlir::arith::FastMathFlags::ninf)) {
+      useOnlyFiniteValue = true;
+      break;
+    }
+  }
+
   // Builder only used as helper for attribute creation.
   OpBuilder b(op->getContext());
   Type resultType = op->getResult(0).getType();
 
-  return getIdentityValueAttr(*maybeKind, resultType, b, op->getLoc());
+  return getIdentityValueAttr(*maybeKind, resultType, b, op->getLoc(),
+                              useOnlyFiniteValue);
 }
 
 /// Returns the identity value associated with an AtomicRMWKind op.
 Value mlir::arith::getIdentityValue(AtomicRMWKind op, Type resultType,
-                                    OpBuilder &builder, Location loc) {
-  auto attr = getIdentityValueAttr(op, resultType, builder, loc);
+                                    OpBuilder &builder, Location loc,
+                                    bool useOnlyFiniteValue) {
+  auto attr =
+      getIdentityValueAttr(op, resultType, builder, loc, useOnlyFiniteValue);
   return builder.create<arith::ConstantOp>(loc, attr);
 }
 
 /// Return the value obtained by applying the reduction operation kind
 /// associated with a binary AtomicRMWKind op to `lhs` and `rhs`.
 Value mlir::arith::getReductionOp(AtomicRMWKind op, OpBuilder &builder,
                                   Location loc, Value lhs, Value rhs) {
   switch (op) {

mlir/lib/Dialect/Linalg/IR/LinalgOps.cpp

FailureOr<SmallVector<Value>> SoftmaxOp::decomposeOperation(OpBuilder &b) {
   Type elementType = inputType.getElementType();
   int64_t reductionDim = getDimension();
   SmallVector<OpFoldResult> dims = tensor::getMixedSizes(b, loc, input);
   Value output = getOutput();
   dims.erase(dims.begin() + reductionDim);
   // Step 1: Compute max along dim.
   Value outputReduce = b.create<tensor::EmptyOp>(loc, dims, elementType);
   Value neutralForMaxF =
-      arith::getIdentityValue(arith::AtomicRMWKind::maxf, elementType, b, loc);
+      arith::getIdentityValue(arith::AtomicRMWKind::maxf, elementType, b, loc,
+                              /*useOnlyFiniteValue=*/true);
   Value neutralForMaxFInit =
       b.create<linalg::FillOp>(loc, Value{neutralForMaxF}, outputReduce)
           .result();
   Value max =
       reduce<arith::MaxFOp>(b, loc, input, neutralForMaxFInit, reductionDim);
 
   // Step 2: Subtract max from input and exponentiate.
   Value numerator = buildSubAndExpOp(b, loc, input, max, output, reductionDim);
 
   // Step 3: Compute sum along dim.
-  Value zero =
-      arith::getIdentityValue(arith::AtomicRMWKind::addf, elementType, b, loc);
+  Value zero = arith::getIdentityValue(arith::AtomicRMWKind::addf, elementType,
+                                       b, loc, /*useOnlyFiniteValue=*/true);
   Value zeroInit =
       b.create<linalg::FillOp>(loc, Value{zero}, outputReduce).result();
   Value denominator =
       reduce<arith::AddFOp>(b, loc, numerator, zeroInit, reductionDim);
 
   // Step 4: Compute softmax.
   Value result =
       buildDivOp(b, loc, numerator, denominator, output, reductionDim);

mlir/test/Dialect/Linalg/transform-op-decompose.mlir

 func.func @softmax(%arg0: tensor<2x16x32xf32>, %dst: tensor<2x16x32xf32>) -> tensor<2x16x32xf32> {
   %1 = linalg.softmax dimension(2) ins(%arg0 : tensor<2x16x32xf32>) outs(%dst: tensor<2x16x32xf32>) -> tensor<2x16x32xf32>
   return %1 : tensor<2x16x32xf32>
 }
 
 // CHECK-LABEL:      func.func @softmax(
 // CHECK-SAME:           %[[ARG0:[a-zA-Z0-9_]+]]: tensor<2x16x32xf32>, %[[DST:[a-zA-Z0-9_]+]]: tensor<2x16x32xf32>) -> tensor<2x16x32xf32> {
 // CHECK-DAG:        %[[D1:.+]] = tensor.empty() : tensor<2x16xf32>
-// CHECK-DAG:        %[[CST:.+]] = arith.constant 0xFF800000 : f32
+// CHECK-DAG:        %[[CST:.+]] = arith.constant -1.401300e-45 : f32
 // CHECK:        %[[D2:.+]] = linalg.fill ins(%[[CST]] : f32) outs(%[[D1]] : tensor<2x16xf32>) -> tensor<2x16xf32>
 // CHECK:        %[[D3:.+]] = linalg.generic {indexing_maps = [#[[$MAP]], #[[$MAP1]]], iterator_types = ["parallel",
 // CHECK-SAME:     "parallel", "reduction"]} ins(%[[ARG0]] : tensor<2x16x32xf32>) outs(%[[D2]] : tensor<2x16xf32>) {
 // CHECK:        ^bb0(%[[IN:.+]]: f32, %[[OUT:.+]]: f32):
 // CHECK:          %[[D8:.+]] = arith.maxf %[[IN]], %[[OUT]] : f32
 // CHECK:          linalg.yield %[[D8]] : f32
 // CHECK:        } -> tensor<2x16xf32>
 // CHECK:        %[[D4:.+]] = linalg.generic {indexing_maps = [#[[$MAP]], #[[$MAP1]], #[[$MAP]]], iterator_types =

mlir/test/Dialect/Linalg/transform-op-split-reduction.mlir

 ^bb1(%arg1: !transform.any_op):
   %0 = transform.structured.match ops{["linalg.generic"]} in %arg1 : (!transform.any_op) -> !transform.any_op
   %1:4 = transform.structured.split_reduction %0 { split_factor = 4, insert_split_dimension = 2}
     : (!transform.any_op) -> (!transform.any_op, !transform.any_op, !transform.any_op, !transform.any_op)
 }
 
 // -----
 
+// Check that we don't use -inf as the neutral element for maxf when maxf has
+// ninf. Instead check that we use the smallest finite floating point value.
+// Also check that the fastmath flags are set on the created maxf
+// instructions.
+func.func @generic_split_3d_ninf(%input: tensor<32x2xf32>, %input_2: tensor<5x32xf32>, %output: tensor<5x2xf32>)
+  -> tensor<5x2xf32>
+{
+  %0 = linalg.generic {
+      indexing_maps = [
+        affine_map<(d0, d1, d2) -> (d1, d0)>,
+        affine_map<(d0, d1, d2) -> (d2, d1)>,
+        affine_map<(d0, d1, d2) -> (d2, d0)>
+      ],
+      iterator_types = ["parallel", "reduction", "parallel"]
+    } ins(%input, %input_2 : tensor<32x2xf32>, tensor<5x32xf32>) outs(%output : tensor<5x2xf32>) {
+    ^bb0(%arg0: f32, %arg1: f32, %arg2: f32):
+      %3 = arith.addf %arg0, %arg1 : f32
+      %4 = arith.maxf %3, %arg2 fastmath<nnan,ninf> : f32
+      linalg.yield %4 : f32
+    } -> tensor<5x2xf32>
+  return %0 : tensor<5x2xf32>
+}
+
+//  CHECK-DAG: #[[$MAP0:.*]] = affine_map<(d0, d1, d2, d3) -> (d2, d1, d0)>
+//  CHECK-DAG: #[[$MAP1:.*]] = affine_map<(d0, d1, d2, d3) -> (d3, d2, d1)>
+//  CHECK-DAG: #[[$MAP2:.*]] = affine_map<(d0, d1, d2, d3) -> (d3, d0, d2)>
+//  CHECK-DAG: #[[$MAP3:.*]] = affine_map<(d0, d1, d2) -> (d0, d1, d2)>
+//  CHECK-DAG: #[[$MAP4:.*]] = affine_map<(d0, d1, d2) -> (d0, d1)>
+// CHECK-LABEL:  func @generic_split_3d_ninf
+//  CHECK-DAG: %[[ID:.*]] = arith.constant -1.401300e-45 : f32
+//  CHECK-DAG: %[[I1:.*]] = tensor.expand_shape %{{.*}}[0, 1], [2]] : tensor<32x2xf32> into tensor<4x8x2xf32>
+//  CHECK-DAG: %[[I2:.*]] = tensor.expand_shape %{{.*}}[0], [1, 2]] : tensor<5x32xf32> into tensor<5x4x8xf32>
+//  CHECK-DAG: %[[INI:.*]] = tensor.empty() : tensor<5x2x4xf32>
+//      CHECK: %[[F:.*]] = linalg.fill ins(%[[ID]] : f32) outs(%[[INI]] : tensor<5x2x4xf32>) -> tensor<5x2x4xf32>
+//      CHECK: %[[G:.*]] = linalg.generic {indexing_maps = [#[[$MAP0]], #[[$MAP1]], #[[$MAP2]]], iterator_types = ["parallel", "reduction", "parallel", "parallel"]}
+// CHECK-SAME:   ins(%[[I1]], %[[I2]] : tensor<4x8x2xf32>, tensor<5x4x8xf32>) outs(%[[F]] : tensor<5x2x4xf32>) {
+//      CHECK:   arith.addf
+//      CHECK:   arith.maxf {{.*}} fastmath<nnan,ninf>
+//      CHECK:   linalg.yield
+//      CHECK: } -> tensor<5x2x4xf32>
+//      CHECK: %[[R:.*]] = linalg.generic {indexing_maps = [#[[$MAP3]], #[[$MAP4]]], iterator_types = ["parallel", "parallel", "reduction"]}
+// CHECK-SAME:   ins(%[[G]] : tensor<5x2x4xf32>) outs(%{{.*}} : tensor<5x2xf32>) {
+//      CHECK:   arith.maxf {{.*}} fastmath<nnan,ninf>
+//      CHECK:   linalg.yield
+//      CHECK:  } -> tensor<5x2xf32>
+//      CHECK: return %[[R]] : tensor<5x2xf32>
+
+transform.sequence failures(propagate) {
+^bb1(%arg1: !transform.any_op):
+  %0 = transform.structured.match ops{["linalg.generic"]} in %arg1 : (!transform.any_op) -> !transform.any_op
+  %1:4 = transform.structured.split_reduction %0 { split_factor = 4, insert_split_dimension = 2}
+    : (!transform.any_op) -> (!transform.any_op, !transform.any_op, !transform.any_op, !transform.any_op)
+}
+
+// -----
+
 func.func @matmul_split(%A : tensor<16x256xf32>, %B: tensor<256x32xf32>, %C: tensor<16x32xf32>) -> tensor<16x32xf32> {
   %0 = linalg.matmul ins(%A, %B: tensor<16x256xf32>, tensor<256x32xf32>)
                     outs(%C: tensor<16x32xf32>) -> tensor<16x32xf32>
   return %0: tensor<16x32xf32>
 }
 
 //  CHECK-DAG: #[[$MAP0:.*]] = affine_map<(d0, d1, d2, d3) -> (d0, d2, d3)>
 //  CHECK-DAG: #[[$MAP1:.*]] = affine_map<(d0, d1, d2, d3) -> (d2, d3, d1)>

 //      CHECK: return %[[R]] : tensor<5x2xf32>
 
 transform.sequence failures(propagate) {
 ^bb1(%arg1: !transform.any_op):
   %0 = transform.structured.match ops{["linalg.generic"]} in %arg1 : (!transform.any_op) -> !transform.any_op
   %1:4 = transform.structured.split_reduction %0 { split_factor = 4, insert_split_dimension = 2, inner_parallel}
     : (!transform.any_op) -> (!transform.any_op, !transform.any_op, !transform.any_op, !transform.any_op)
 }
+
+// -----
+
+// Check that we don't use +inf as the neutral element for minf when minf has
+// ninf. Instead check that we use the largest finite floating point value.
+// Also check that the fastmath flags are set on the created minf
+// instructions.
+func.func @generic_split_3d(%input: tensor<32x2xf32>, %input_2: tensor<5x32xf32>, %output: tensor<5x2xf32>)
+  -> tensor<5x2xf32>
+{
+  %0 = linalg.generic {
+      indexing_maps = [
+        affine_map<(d0, d1, d2) -> (d1, d0)>,
+        affine_map<(d0, d1, d2) -> (d2, d1)>,
+        affine_map<(d0, d1, d2) -> (d2, d0)>
+      ],
+      iterator_types = ["parallel", "reduction", "parallel"]
+    } ins(%input, %input_2 : tensor<32x2xf32>, tensor<5x32xf32>) outs(%output : tensor<5x2xf32>) {
+    ^bb0(%arg0: f32, %arg1: f32, %arg2: f32):
+      %3 = arith.addf %arg0, %arg1 : f32
+      %4 = arith.minf %3, %arg2 fastmath<ninf> : f32
+      linalg.yield %4 : f32
+    } -> tensor<5x2xf32>
+  return %0 : tensor<5x2xf32>
+}
+
+//  CHECK-DAG: #[[$MAP0:.*]] = affine_map<(d0, d1, d2, d3) -> (d1, d2, d0)>
+//  CHECK-DAG: #[[$MAP1:.*]] = affine_map<(d0, d1, d2, d3) -> (d3, d1, d2)>
+//  CHECK-DAG: #[[$MAP2:.*]] = affine_map<(d0, d1, d2, d3) -> (d3, d0, d2)>
+//  CHECK-DAG: #[[$MAP3:.*]] = affine_map<(d0, d1, d2) -> (d0, d1, d2)>
+//  CHECK-DAG: #[[$MAP4:.*]] = affine_map<(d0, d1, d2) -> (d0, d1)>
+// CHECK-LABEL:  func @generic_split_3d
+//  CHECK-DAG: %[[ID:.*]] = arith.constant 3.40282347E+38 : f32
+//  CHECK-DAG: %[[I1:.*]] = tensor.expand_shape %{{.*}}[0, 1], [2]] : tensor<32x2xf32> into tensor<8x4x2xf32>
+//  CHECK-DAG: %[[I2:.*]] = tensor.expand_shape %{{.*}}[0], [1, 2]] : tensor<5x32xf32> into tensor<5x8x4xf32>
+//  CHECK-DAG: %[[INI:.*]] = tensor.empty() : tensor<5x2x4xf32>
+//      CHECK: %[[F:.*]] = linalg.fill ins(%[[ID]] : f32) outs(%[[INI]] : tensor<5x2x4xf32>) -> tensor<5x2x4xf32>
+//      CHECK: %[[G:.*]] = linalg.generic {indexing_maps = [#[[$MAP0]], #[[$MAP1]], #[[$MAP2]]], iterator_types = ["parallel", "reduction", "parallel", "parallel"]}
+// CHECK-SAME:   ins(%[[I1]], %[[I2]] : tensor<8x4x2xf32>, tensor<5x8x4xf32>) outs(%[[F]] : tensor<5x2x4xf32>) {
+//      CHECK:   arith.addf
+//      CHECK:   arith.minf {{.*}} fastmath<ninf>
+//      CHECK:   linalg.yield
+//      CHECK: } -> tensor<5x2x4xf32>
+//      CHECK: %[[R:.*]] = linalg.generic {indexing_maps = [#[[$MAP3]], #[[$MAP4]]], iterator_types = ["parallel", "parallel", "reduction"]}
+// CHECK-SAME:   ins(%[[G]] : tensor<5x2x4xf32>) outs(%{{.*}} : tensor<5x2xf32>) {
+//      CHECK:   arith.minf {{.*}} fastmath<ninf>
+//      CHECK:   linalg.yield
+//      CHECK:  } -> tensor<5x2xf32>
+//      CHECK: return %[[R]] : tensor<5x2xf32>
+
+transform.sequence failures(propagate) {
+^bb1(%arg1: !transform.any_op):
+  %0 = transform.structured.match ops{["linalg.generic"]} in %arg1 : (!transform.any_op) -> !transform.any_op
+  %1:4 = transform.structured.split_reduction %0 { split_factor = 4, insert_split_dimension = 2, inner_parallel}
+    : (!transform.any_op) -> (!transform.any_op, !transform.any_op, !transform.any_op, !transform.any_op)
+}