Adding min(f/s/u) and max(f/s/u) cases for vector reduction

[aslepko]

Summary

This PR adds missing AtomicRMWKind::min/max cases which we would like to use for min/max reduction loop vectorizations.

[aartbik]

Thanks for the generalization. Don't you want to put some new tests in test/Dialect/Affine/SuperVectorize just to make sure the new code is covered?

[aslepko]

Hi @aartbik, thank you for your reply. And agree, it would be nice to add testing. I was not sure how to do that for min/max reductions without proper changes in the reduction recognizer (which we do in our local project). But I'd really appreciate any suggestions for what I can do here.
Thanks!

[aartbik]

In D104881#2847109, @aslepko wrote:

proper changes in the reduction recognizer (which we do in our local project).

Do you plan to follow up with these? If so, please make this a parent revision of such a follow up.
I am just trying to avoid checkin in uncovered code....

[aslepko]

Hi @aartbik, my apologies this got dragged out. I was out for a while.
After internal discussion we decided to also add the code to recognize min/max reductions. With that we can add all test cases to cover the initial additions in this current commit.
Since it has been a while, should I do an entirely new commit, or just upload a second patch here?
Thanks.

[aartbik]

In D104881#2971137, @aslepko wrote:

Since it has been a while, should I do an entirely new commit, or just upload a second patch here?

If you rebase with main, let's try to continue on this patch, so we keep full history.
Thanks for adding!

[aslepko]

This update adds code to recognize min/max reductions. This was necessary to test part some of the functionality we had added initially.
Also, this update contains all the tests missing in the original commit.

[aartbik]

Now that I know some testing part through recognition will follow, I am okay breaking up this revision as originally planned, i.e. adding the ops in a first small revision and then adding the recognition later.
That way you get the first revision in a bit quicker, since I think the recognition part may need some discussion.

mlir/lib/Analysis/AffineAnalysis.cpp

105 ↗

(On Diff #369596)

please align parameters (clang-format will fix all those issues for you)

129 ↗

(On Diff #369596)

L108 to here could use some comments, in particular some visual representation of what you are looking for

154 ↗

(On Diff #369596)

IEEE floating-point comparisons as very tricky and we need to make sure we preserve the semantics when rewriting them into min/max . Have you made sure these rewriting preserve behaviors for e.g. NaN operands?

[aslepko]

@aartbik, that sounds reasonable. Should I then basically go ahead and remove the tests+recognizer for a new/third patch and upload that patch. Then, make a new PR once this is merged for the recognition part? Thanks.

[aartbik]

In D104881#2987247, @aslepko wrote:

@aartbik, that sounds reasonable. Should I then basically go ahead and remove the tests+recognizer for a new/third patch and upload that patch. Then, make a new PR once this is merged for the recognition part? Thanks.

Yes please. That way we can get your first patch in quickly, and discuss some details on the second.

[aslepko]

As discussed, this update removes the min/max recognition and testing part from the prior update (which will be added through a separate, new PR).

[aartbik]

make sure you have green bot builds prior to submitting but lgtm

mlir/lib/Dialect/StandardOps/IR/Ops.cpp

static LogicalResult verify(AtomicRMWOp op) {
   }
   return success();
 }
 
 /// Returns the identity value attribute associated with an AtomicRMWKind op.
 Attribute mlir::getIdentityValueAttr(AtomicRMWKind kind, Type resultType,
                                      OpBuilder &builder, Location loc) {
   switch (kind) {
+  case AtomicRMWKind::maxf:
+    return builder.getFloatAttr(
+        resultType,
+        APFloat::getInf(resultType.cast<FloatType>().getFloatSemantics(),
+                        /*Negative=*/true));
   case AtomicRMWKind::addf:
   case AtomicRMWKind::addi:
+  case AtomicRMWKind::maxu:
     return builder.getZeroAttr(resultType);
+  case AtomicRMWKind::maxs:
+    return builder.getIntegerAttr(
+        resultType,
+        APInt::getSignedMinValue(resultType.cast<IntegerType>().getWidth()));
+  case AtomicRMWKind::minf:
+    return builder.getFloatAttr(
+        resultType,
+        APFloat::getInf(resultType.cast<FloatType>().getFloatSemantics(),
+                        /*Negative=*/false));
+  case AtomicRMWKind::mins:
+    return builder.getIntegerAttr(
+        resultType,
+        APInt::getSignedMaxValue(resultType.cast<IntegerType>().getWidth()));
+  case AtomicRMWKind::minu:
+    return builder.getIntegerAttr(
+        resultType,
+        APInt::getMaxValue(resultType.cast<IntegerType>().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;
Value mlir::getReductionOp(AtomicRMWKind op, OpBuilder &builder, Location loc,
   case AtomicRMWKind::addf:
     return builder.create<AddFOp>(loc, lhs, rhs);
   case AtomicRMWKind::addi:
     return builder.create<AddIOp>(loc, lhs, rhs);
   case AtomicRMWKind::mulf:
     return builder.create<MulFOp>(loc, lhs, rhs);
   case AtomicRMWKind::muli:
     return builder.create<MulIOp>(loc, lhs, rhs);
+  case AtomicRMWKind::maxf:
+    return builder.create<SelectOp>(
+        loc, builder.create<CmpFOp>(loc, CmpFPredicate::OGT, lhs, rhs), lhs,
+        rhs);
+  case AtomicRMWKind::minf:
+    return builder.create<SelectOp>(
+        loc, builder.create<CmpFOp>(loc, CmpFPredicate::OLT, lhs, rhs), lhs,
+        rhs);
+  case AtomicRMWKind::maxs:
+    return builder.create<SelectOp>(
+        loc, builder.create<CmpIOp>(loc, CmpIPredicate::sgt, lhs, rhs), lhs,
+        rhs);
+  case AtomicRMWKind::mins:
+    return builder.create<SelectOp>(
+        loc, builder.create<CmpIOp>(loc, CmpIPredicate::slt, lhs, rhs), lhs,
+        rhs);
+  case AtomicRMWKind::maxu:
+    return builder.create<SelectOp>(
+        loc, builder.create<CmpIOp>(loc, CmpIPredicate::ugt, lhs, rhs), lhs,
+        rhs);
+  case AtomicRMWKind::minu:
+    return builder.create<SelectOp>(
+        loc, builder.create<CmpIOp>(loc, CmpIPredicate::ult, lhs, rhs), lhs,
+        rhs);
   // TODO: Add remaining reduction operations.
   default:
     (void)emitOptionalError(loc, "Reduction operation type not supported");
     break;
   }
   return nullptr;
 }
 

mlir/lib/Dialect/Vector/VectorOps.cpp

case AtomicRMWKind::addi:
     return builder.create<vector::ReductionOp>(vector.getLoc(), scalarType,
                                                builder.getStringAttr("add"),
                                                vector, ValueRange{});
   case AtomicRMWKind::mulf:
   case AtomicRMWKind::muli:
     return builder.create<vector::ReductionOp>(vector.getLoc(), scalarType,
                                                builder.getStringAttr("mul"),
                                                vector, ValueRange{});
+  case AtomicRMWKind::minf:
+  case AtomicRMWKind::mins:
+  case AtomicRMWKind::minu:
+    return builder.create<vector::ReductionOp>(vector.getLoc(), scalarType,
+                                               builder.getStringAttr("min"),
+                                               vector, ValueRange{});
+  case AtomicRMWKind::maxf:
+  case AtomicRMWKind::maxs:
+  case AtomicRMWKind::maxu:
+    return builder.create<vector::ReductionOp>(vector.getLoc(), scalarType,
+                                               builder.getStringAttr("max"),
+                                               vector, ValueRange{});
   // TODO: Add remaining reduction operations.
   default:
     (void)emitOptionalError(loc, "Reduction operation type not supported");
     break;
   }
   return nullptr;
 }