Page MenuHomePhabricator

[mlir][arith] Initial support for fastmath flag attributes in the Arithmetic dialect (v2)
ClosedPublic

Authored by jfurtek on May 24 2022, 8:22 AM.

Details

Summary

This diff adds initial (partial) support for "fastmath" attributes for floating
point operations in the arithmetic dialect. The "fastmath" attributes are
implemented using a default-valued bit enum. The defined flags currently mirror
the fastmath flags in the LLVM dialect (and in LLVM itself). Extending the
set of flags (if necessary) is left as a future task.

In this diff:

  • Definition of FastMathAttr as a custom attribute in the Arithmetic dialect that inherits from the EnumAttr class.
  • Definition of ArithFastMathInterface, which is an interface that is implemented by operations that have an arith::fastmath attribute.
  • Declaration of a default-valued fastmath attribute for unary and (some) binary floating point operations in the Arithmetic dialect.
  • Conversion code to lower arithmetic fastmath flags to LLVM fastmath flags

NOT in this diff (but planned or currently in progress):

  • Documentation of flag meanings
  • Addition of FastMathAttr attributes to other dialects that might lower to the Arithmetic dialect (e.g. Math and Complex)
  • Folding/rewrite implementations that are enabled by fastmath flags
  • Specification of fastmath values from Python bindings (pending other in- progress diffs)

Diff Detail

Event Timeline

There are a very large number of changes, so older changes are hidden. Show Older Changes
jfurtek updated this revision to Diff 431690.May 24 2022, 8:55 AM
  • Fix clang formatting

Adding @mehdi_amini since he had some feedback on a previous iteration of this functionality.

jfurtek updated this revision to Diff 434938.Jun 7 2022, 1:42 PM
  • rebase
mehdi_amini added inline comments.Jun 7 2022, 6:32 PM
mlir/lib/Conversion/LLVMCommon/Pattern.cpp
42

I'd rather make it return just the llvmAttr to avoid rebuilding the StringAttr for the name: either the call-site already have it or they don't and they may not even care about it.
(it also seems more coherent as an API to not return a named attribute when the input is just the attribute value here)

mlir/lib/Dialect/Arithmetic/IR/ArithmeticOps.cpp
49 ↗(On Diff #434938)

Can we instead print the fast math separately from the other attribute in a more friendly form?

Instead of:

%0 = arith.addf %arg0, %arg1 {fastmath = #arith.fastmath<fast>} : f32

I'd rather read:

%0 = arith.addf %arg0, %arg1 fastmath<fast> : f32
%1 = arith.addf %arg0, %arg1 fastmath<fast> { other_attribute = ... } : f32

Calling arithFMAttr.print() should print it this way directly.

jfurtek updated this revision to Diff 435304.Jun 8 2022, 12:58 PM
  • Change fastmath print/parse format based on review comments.
  • Refactor attribute lowering.
jfurtek marked 2 inline comments as done.Jun 8 2022, 12:59 PM
mehdi_amini accepted this revision.Jun 8 2022, 3:09 PM

Thanks!

This revision is now accepted and ready to land.Jun 8 2022, 3:09 PM
jfurtek planned changes to this revision.Jun 21 2022, 12:05 PM

Subsequent code for python bindings of default-valued attributes has broken the code in this diff - needs to be fixed.

Thank you for working on this @jfurtek!

I guess this may be changed later: in light of extending this to more dialects, does it make sense to make ArithFastMathInterface not-arith specific from the beginning and replace LLVM dialect's FastmathFlagsInterface with the same base interface?

Thank you for working on this

I am hoping that others will find it useful! I have more functionality working locally - I am waiting on some other diffs (regarding Python handling of default-valued attributes) to land before having those reviewed.

I guess this may be changed later: in light of extending this to more dialects, does it make sense to make ArithFastMathInterface not-arith specific from the beginning and replace LLVM dialect's FastmathFlagsInterface with the same base interface?

I tried to minimize changes to the LLVM dialect - at least initially - just to reduce the scope of the diff. It may be possible to have them share the same interface and bit enumeration.

However, I was thinking that at some point MLIR::Arithmetic might want to have flags that are not present in LLVM - for whatever reason. (Hypothetically speaking, such a flag could be derived from the fact that Arithmetic operations can be applied to tensors, or it could be used and removed before lowering to LLVM, or it could be used when lowering to something that is not LLVM.) I don't have a strong opinion, but I would be inclined to keep the Arithmetic and LLVM flags separate. (The code to convert from Arithmetic to LLVM flags is minimal.)

I guess this may be changed later: in light of extending this to more dialects, does it make sense to make ArithFastMathInterface not-arith specific from the beginning and replace LLVM dialect's FastmathFlagsInterface with the same base interface?

I tried to minimize changes to the LLVM dialect - at least initially - just to reduce the scope of the diff. It may be possible to have them share the same interface and bit enumeration.

However, I was thinking that at some point MLIR::Arithmetic might want to have flags that are not present in LLVM - for whatever reason. (Hypothetically speaking, such a flag could be derived from the fact that Arithmetic operations can be applied to tensors, or it could be used and removed before lowering to LLVM, or it could be used when lowering to something that is not LLVM.) I don't have a strong opinion, but I would be inclined to keep the Arithmetic and LLVM flags separate. (The code to convert from Arithmetic to LLVM flags is minimal.)

Thank you for the explanation! It makes sense to keep the two interfaces separate at this point.

Mostly looks good to me.

mlir/include/mlir/Conversion/LLVMCommon/Pattern.h
46

The non-method name should probably sound like convertedAttrs.

mlir/lib/Conversion/LLVMCommon/Pattern.cpp
19

auto should probably be replaces with concrete type.

37

Please replace auto.

Thanks for working on this @jfurtek. Will you be able to land this in the next two weeks? Are there plans for supporting these attributes in the math dialect and func dialect?

mlir/include/mlir/Dialect/Arithmetic/IR/ArithmeticOpsInterfaces.td
41 ↗(On Diff #435304)

Nit: newline

Thanks for working on this @jfurtek. Will you be able to land this in the next two weeks?

I will pick this back up and try to have an updated diff ready for re-review today or tomorrow...

Are there plans for supporting these attributes in the math dialect and func dialect?

I expect to spread fastmath support to math (and some other dialects). I think once this diff lands, it would become easier for anyone else to do it as well.

I'm less clear on how this would work with the func dialect though (since you mentioned it specifically). I seem to recall some discussion in LLVM about whether a function-level attribute makes sense:

https://reviews.llvm.org/D14707

For example, how do you reconcile a function-level fastmath attribute that conflicts with those that might exist on individual ops when inlining? Maybe @mehdi_amini would have some thoughts on this.

https://reviews.llvm.org/D14707
For example, how do you reconcile a function-level fastmath attribute that conflicts with those that might exist on individual ops when inlining? Maybe @mehdi_amini would have some thoughts on this.

This may be tricky to handle with generality, I'd have to think about this and look into scenario.
The first question is what is the semantics of a call with a fast math attribute when the called function does not have it? Maybe conceptually the attribute could be propagate to the entire function? (assuming it is the only call site for the function, or all call sites have the attribute)

https://reviews.llvm.org/D14707
For example, how do you reconcile a function-level fastmath attribute that conflicts with those that might exist on individual ops when inlining? Maybe @mehdi_amini would have some thoughts on this.

This may be tricky to handle with generality, I'd have to think about this and look into scenario.
The first question is what is the semantics of a call with a fast math attribute when the called function does not have it? Maybe conceptually the attribute could be propagate to the entire function? (assuming it is the only call site for the function, or all call sites have the attribute)

Just my two cents. In an example of separate modules compilation and user functions, i.e. function definition is in one translation module and the call site is in another, I think we should give users possibility to use different math optimization settings for the modules. So I think propagating mismatching fast math flags from the call site to the function or even its inlined copy (e.g. with MLIR inter-module inlining) would be incorrect. I believe LLVM inlining actually works this way and it seems correct to me.

Example:
main.c:

#include <math.h>
#include <stdio.h>
extern float foo(float x);

int main(int argc, char *argv[]) {
  printf("%e %e\n", foo(argc), cosf(argc));
  return 0;
}

foo.c:

#include <math.h>
float foo(float x) {
  return sinf(x);
}

clang -O3 -ffast-math main.c -flto -c
clang -O3 foo.c -flto -c
clang -O3 main.o foo.o -flto -lm -Wl,-plugin-opt=-print-after-all 2>log

After inlining we will have:

*** IR Dump After InlinerPass on (main) ***
; Function Attrs: nofree nounwind uwtable
define dso_local i32 @main(i32 noundef %0, ptr nocapture noundef readnone %1) local_unnamed_addr #0 {
...
  %4 = tail call float @sinf(float noundef %3) #4
...
  %6 = tail call fast float @llvm.cos.f32(float %3)
...
}

As you may see the inliner preserved the function definition's behavior instead of propagating the call-site attributes to the inlined copy.

What is the intended meaning of the fast math attribute on the call to foo in main?

The LLVM Langref might help you with these questions.

What is the intended meaning of the fast math attribute on the call to foo in main?

This is a good question :)
Fast math flags can be attached to LLVM call instructions, so I guess clang does this unconditionally based on the driver option. I am not sure if it is intentional, but it seems that nnan, ninf and nsz can potentially be applied even to user function call. Whether these flags have any effect on user function calls - I do not know.

The LLVM Langref might help you with these questions.

Can you elaborate maybe? Did you found something more than I did?

The optional fast-math flags marker indicates that the call has one or more fast-math flags, which are optimization hints to enable otherwise unsafe floating-point optimizations. Fast-math flags are only valid for calls that return a floating-point scalar or vector type, or an array (nested to any depth) of floating-point scalar or vector types.

This isn't helpful here...

Will you be able to land this in the next two weeks?

Apologies for the delay - I am still finishing this up. The code has been rebased, and I am just clearing up new test failures.

Are there plans for supporting these attributes in the ... func dialect?

We can try to be compatible with LLVM, but I was thinking that we might not need the fastmath attribute at the function and module levels. (There is some brief discussion from 2016 about LLVM moving away from function-level attributes for fastmath in this video.) I suspect that for most practical module and function scenarios (setting flags globally, propagating attributes through inlining, ...) MLIR passes and the Transform dialect could accomplish the same thing.

Will you be able to land this in the next two weeks?

Apologies for the delay - I am still finishing this up. The code has been rebased, and I am just clearing up new test failures.

Are there plans for supporting these attributes in the ... func dialect?

We can try to be compatible with LLVM, but I was thinking that we might not need the fastmath attribute at the function and module levels. (There is some brief discussion from 2016 about LLVM moving away from function-level attributes for fastmath in this video.) I suspect that for most practical module and function scenarios (setting flags globally, propagating attributes through inlining, ...) MLIR passes and the Transform dialect could accomplish the same thing.

Hi @jfurtek, so Hal was talking about making math "settings" more exact by attaching the flags to instructions rather than keeping them on the function/module level. I believe this implies that the call instructions may have math flags attached, and I think this is what we probably want to have in MLIR as well (e.g. for func.call operations produced by ToLibm convertors). There indeed seem to be no reason to have math flags attached to func.func operations.

Will you be able to land this in the next two weeks?

Apologies for the delay - I am still finishing this up. The code has been rebased, and I am just clearing up new test failures.

Are there plans for supporting these attributes in the ... func dialect?

We can try to be compatible with LLVM, but I was thinking that we might not need the fastmath attribute at the function and module levels. (There is some brief discussion from 2016 about LLVM moving away from function-level attributes for fastmath in this video.) I suspect that for most practical module and function scenarios (setting flags globally, propagating attributes through inlining, ...) MLIR passes and the Transform dialect could accomplish the same thing.

Hi @jfurtek, so Hal was talking about making math "settings" more exact by attaching the flags to instructions rather than keeping them on the function/module level. I believe this implies that the call instructions may have math flags attached, and I think this is what we probably want to have in MLIR as well (e.g. for func.call operations produced by ToLibm convertors). There indeed seem to be no reason to have math flags attached to func.func operations.

OK. I was just asking because clang-generated code seemed to have these fast function attributes. If clang/llvm's plan is to move away then let us not support it for now in the func dialect. We can revisit if necessary later.

We need the fast attribute support in Flang. We have started working on the driver to support this. I believe @vzakhari plans to extend this to the math dialect. We would like to generate some numbers for Flang with the Ofast flag. It will be great to see this submitted soon.

tblah added a subscriber: tblah.Oct 17 2022, 2:49 AM
jfurtek updated this revision to Diff 468632.Oct 18 2022, 11:28 AM
  • Fix clang formatting
  • Change fastmath print/parse format based on review comments. Refactor attribute lowering.
  • Update MLIRArithmeticDialect CMake reference
  • Add fastmath attribute binding to arith canonicalizations
  • Fixes for compile failures after rebase
  • Modify LLVM fastmath interface to be attribute-based instead of value-based.
This revision is now accepted and ready to land.Oct 18 2022, 11:28 AM
jfurtek planned changes to this revision.Oct 18 2022, 11:30 AM
jfurtek updated this revision to Diff 469067.Oct 19 2022, 3:59 PM
  • Update LLVM lowering logic to drop fastmath attributes for targets that do not support the interface.
This revision is now accepted and ready to land.Oct 19 2022, 3:59 PM
jfurtek marked 3 inline comments as done.Oct 19 2022, 5:22 PM
vzakhari accepted this revision.Oct 19 2022, 6:19 PM

LGTM. Thank you!

tpopp added a subscriber: tpopp.Oct 20 2022, 6:45 AM
tpopp added inline comments.
mlir/lib/Conversion/LLVMCommon/CMakeLists.txt
18 ↗(On Diff #469067)

It seems weird to add a dependency on the Arith Dialect here.

jfurtek planned changes to this revision.Oct 20 2022, 8:29 AM

It seems weird to add a dependency on the Arith Dialect here.

This is a good point - let me see if I can avoid that...

tpopp added a comment.Oct 21 2022, 6:59 AM

It seems weird to add a dependency on the Arith Dialect here.

This is a good point - let me see if I can avoid that...

I might be wrong, and I was just leaving a passing comment as I saw it, so please don't wait on my approval, whenever you and your reviewers are happy with the change :)

jfurtek updated this revision to Diff 469770.Oct 21 2022, 2:31 PM
  • Fix clang formatting
  • Change fastmath print/parse format based on review comments. Refactor attribute lowering.
  • Update MLIRArithmeticDialect CMake reference
  • Add fastmath attribute binding to arith canonicalizations
  • Fixes for compile failures after rebase
  • Modify LLVM fastmath interface to be attribute-based instead of value-based.
  • Update LLVM lowering logic to drop fastmath attributes for targets that do not support the interface.
  • Remove arith dependency from LLVMCommon. Move attribute conversion to the specific conversion module (in this case ArithToLLVM).
This revision is now accepted and ready to land.Oct 21 2022, 2:31 PM
jfurtek edited the summary of this revision. (Show Details)Oct 21 2022, 3:47 PM
jfurtek updated this revision to Diff 469812.Oct 21 2022, 3:53 PM
  • Fix clang formatting
  • Change fastmath print/parse format based on review comments. Refactor attribute lowering.
  • Update MLIRArithmeticDialect CMake reference
  • Add fastmath attribute binding to arith canonicalizations
  • Fixes for compile failures after rebase
  • Modify LLVM fastmath interface to be attribute-based instead of value-based.
  • Update LLVM lowering logic to drop fastmath attributes for targets that do not support the interface.
  • Remove arith dependency from LLVMCommon. Move attribute conversion to the specific conversion module (in this case ArithToLLVM).
  • clang-formatting fix
  • clang formatting fix

I might be wrong, and I was just leaving a passing comment as I saw it

I think the comment was spot on - thanks for the suggestion. I fixed it in a subsequent commit.

vzakhari added inline comments.Oct 24 2022, 10:48 AM
mlir/include/mlir/Conversion/LLVMCommon/VectorPattern.h
99

Can we avoid specifying SourceOp by templatizing the constructors of AttrConvertPassThrough, AttrConvertFastMath and AttrDropFastMath by their argument type? It seems to be redundant to specify SourceOp as long as it is deducible from the constructor argument.

vzakhari added inline comments.Oct 24 2022, 11:26 AM
mlir/include/mlir/Conversion/LLVMCommon/VectorPattern.h
99

Here is an example of where this will help to avoid some redundancy: https://reviews.llvm.org/D136312#change-SuMljwYyXLk3 (ConvertFastMathHelper wrapper currently has 6 uses in math::ExpM1Op, math::Log1pOp and math::RsqrtOp converters).

jfurtek added inline comments.Oct 24 2022, 3:08 PM
mlir/include/mlir/Conversion/LLVMCommon/VectorPattern.h
99

I disagree that templatizing the constructor is preferable. For the trivial case here that works, but if I look ahead to (for example) an attribute converter that chains together conversions for more than one attribute, I think that being able to specialize on the converter class will be more flexible and clean than specializing or overloading the constructor.

Your linked example already specifies the target type in the class template and the matchAndRewrite() function anyway, so it doesn't seem to be to be a huge amount of additional redundancy.

This is just my opinion - feel free to make that change as part of your child diff.

vzakhari added inline comments.Oct 24 2022, 5:04 PM
mlir/include/mlir/Conversion/LLVMCommon/VectorPattern.h
99

I was referring to places like this:

ConvertFastMathHelper<LLVM::ExpOp> expAttrs(op);
ConvertFastMathHelper<LLVM::FSubOp> subAttrs(op);

If I were to use AttrConvert interface directly it would look like this:

ConvertFastMath<math::ExpM1Op, LLVM::ExpOp> expAttrs(op);
ConvertFastMath<math::ExpM1Op, LLVM::FSubOp> subAttrs(op);

I agree, this is not a huge amount of redundancy, so I am okay with it as it is.

jfurtek updated this revision to Diff 470626.Oct 25 2022, 3:25 PM
  • Fix clang formatting
  • Change fastmath print/parse format based on review comments. Refactor attribute lowering.
  • Update MLIRArithmeticDialect CMake reference
  • Add fastmath attribute binding to arith canonicalizations
  • Fixes for compile failures after rebase
  • Modify LLVM fastmath interface to be attribute-based instead of value-based.
  • Update LLVM lowering logic to drop fastmath attributes for targets that do not support the interface.
  • Remove arith dependency from LLVMCommon. Move attribute conversion to the specific conversion module (in this case ArithToLLVM).
  • clang-format fix
  • clang formatting fix
  • More clang-formatting
rkayaith added inline comments.
mlir/include/mlir/Dialect/Arith/IR/ArithOps.td
86

you should be able to avoid the custom formatting now:

You have got a notification about this breakage https://lab.llvm.org/buildbot/#/builders/61/builds/34355 ?

I tried to fix in https://github.com/llvm/llvm-project/commit/8f4320cf1d5f115fa98652ce8765157557aeaaea ; hopefully it's enough.

Thank you, Mehdi! I wish the committers were also notified about the breakages.

jfurtek added inline comments.Oct 27 2022, 3:42 PM
mlir/include/mlir/Dialect/Arith/IR/ArithOps.td
86

The documentation for Optional Groups says that "only optional attributes can be marked as the anchor." In this case, the attribute is not optional, it is default-valued. Will this still work?

rkayaith added inline comments.Oct 27 2022, 3:53 PM
mlir/include/mlir/Dialect/Arith/IR/ArithOps.td
86

It should work since D134993. I need to update that line in the docs though, thanks for pointing it out.