This is an archive of the discontinued LLVM Phabricator instance.

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

Authored by qcolombet on Jul 27 2023, 11:30 AM.

Details

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.

Diff Detail

Event Timeline

qcolombet created this revision.Jul 27 2023, 11:30 AM
Herald added a reviewer: kuhar. · View Herald Transcript
Herald added a project: Restricted Project. · View Herald Transcript
qcolombet requested review of this revision.Jul 27 2023, 11:30 AM
  • Fix Smallest/Largest for maxf/minf respectively.

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

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 added inline comments.Jul 28 2023, 9:00 AM
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 accepted this revision.Jul 28 2023, 4:42 PM

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

This revision is now accepted and ready to land.Jul 28 2023, 4:42 PM
Hardcode84 added inline comments.
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 added inline comments.Aug 7 2023, 4:43 AM
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 added inline comments.Aug 7 2023, 4:54 AM
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 updated this revision to Diff 547731.Aug 7 2023, 5:13 AM
  • Rebase
  • Set default value for getIdentity
  • Use markdown instead of \p.
qcolombet added inline comments.Aug 7 2023, 5:33 AM
mlir/lib/Dialect/Arith/IR/ArithOps.cpp
2445

Good point. Let me give it a try then.

qcolombet updated this revision to Diff 547775.Aug 7 2023, 7:11 AM
  • Use the fastmath flags attached to arith op when using getIdentityElement
  • Add tests catching that behavior
qcolombet marked 5 inline comments as done.Aug 7 2023, 7:13 AM
qcolombet added inline comments.
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 added inline comments.Aug 7 2023, 7:20 AM
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 updated this revision to Diff 547784.Aug 7 2023, 7:35 AM
qcolombet marked an inline comment as done.
  • Use ArithFastMathInterface instead of iterating through the attributes
qcolombet marked an inline comment as done.Aug 7 2023, 7:36 AM
qcolombet added inline comments.
mlir/lib/Dialect/Arith/IR/ArithOps.cpp
2469

Ah thanks!

Hardcode84 accepted this revision.Aug 7 2023, 7:41 AM

Thanks!

qcolombet marked an inline comment as done.Aug 8 2023, 12:36 AM

Hi @rengolin,

Does this look good to you?

Cheers,
-Quentin

rengolin accepted this revision.Aug 8 2023, 6:19 AM