Page MenuHomePhabricator

IR changes to add yield semantics for affine.if and affine.parallel, additional affine.parallel canonicalization
ClosedPublic

Authored by jbruestle on Jun 25 2020, 2:09 PM.

Details

Summary

IR and parsing changes to allow yield semantics similar to SCF for affine.if, and to provide a method to perform reductions via affine.parallel. Currently this commit does not modify lowering for affine.if (although lowerings for affine.if that do not yield values are still correct). Additionally, affine.parallel does not currently have a lowering.

Diff Detail

Event Timeline

jbruestle created this revision.Jun 25 2020, 2:09 PM
jbruestle updated this revision to Diff 273514.Jun 25 2020, 2:19 PM
jbruestle updated this revision to Diff 273515.

Attempting to fix patch to add context.

rriddle added inline comments.Jun 25 2020, 2:33 PM
mlir/include/mlir/Dialect/StandardOps/IR/StandardOpsBase.td
2

Please make sure that all lines are wrapped at 80 characters. For this line specifically, you need to shorten it and not wrap.

mlir/lib/Dialect/Affine/IR/AffineOps.cpp
1916

Please drop trivial braces, also why can't you just do result.addTypes(resultTypes)?

1923

nit: Please use builder.createBlock instead.

1929

Same here.

1937

nit: Add a parameter comment(/*...=*/ for the {})

2433

nit: Drop this comment, the code is self-explanatory.

2434

Same comment as above.

2439

nit: Please spell out the type instead of auto, and drop the trivial braces.

2442

nit: The code is self-explanatory, so this comment can be dropped. Also, please use complete sentences in comments with correct grammar/punctuation.

2467

nit: Spell out auto here.

2532

Please drop these trivial braces, here and throughout the commit.

2536

nit: Spell out auto here.

2559

Drop the mlir:: here and many other places.

2561

Please do not use replaceAllUsesWith directly, use the pattern rewriter instead.

2707

Please do not use e as an iteration variable, it is very confusing.

Thanks for reuploading!

jbruestle marked 14 inline comments as done.Jun 25 2020, 3:13 PM
jbruestle updated this revision to Diff 273538.Jun 25 2020, 3:41 PM

Updates from river's review.

flaub added inline comments.Jun 25 2020, 5:34 PM
mlir/include/mlir/Dialect/Affine/IR/AffineOps.td
603–605

To be consistent, use OperationState &result (or whatever is the standard in this file).

mlir/lib/Dialect/Affine/IR/AffineOps.cpp
2396–2397

Parameter comments would be nice here too.

2581

Update comment to match code. Or drop if it's trivial.

mlir/lib/Dialect/Affine/IR/AffineOps.cpp
2545

The canonicalization really feels like it should be in a separate revision?

jbruestle updated this revision to Diff 273744.Jun 26 2020, 8:27 AM
jbruestle edited the summary of this revision. (Show Details)

Code review changes from flaub, and removed additional canonicalizations which are unrelated to the IR change for a separate diff later.

flaub accepted this revision.Jun 26 2020, 12:42 PM
This revision is now accepted and ready to land.Jun 26 2020, 12:42 PM

Thanks Jeremy! Some minor comments.

mlir/include/mlir/Dialect/Affine/IR/AffineOps.td
563

Since this is about SSA value reductions, should we rename "AtomicRMWKind" to something that doesn't imply memory or an atomic implementation? These reductions could be implemented following different approaches, not only using atomic ops, right? Maybe ReductionKind would work?

568

I'm trying to understand if this is intended to achieve the same as scf.reduce.return. Couldn't we do something similar and keep both SCF and Affine aligned? Or is this intentionally different for some reason?

641

I would suggest using 'reduction' terminology instead of aggregates, to align with scf, OpenMp, and others.

mlir/lib/Dialect/Affine/IR/AffineOps.cpp
2644

Same. I would suggest reduce/reduction to align terminology with scf, OpenMp and others.

jbruestle marked 3 inline comments as done.Jun 26 2020, 4:01 PM
jbruestle added inline comments.
mlir/include/mlir/Dialect/Affine/IR/AffineOps.td
563

I do think ReductionKind would be a better name, but I wanted to minimize unrelated changes such as in atomicrmw, so I was considering do a NFC commit afterwards for both atomicrmw and affine.parallel. Actually, I was considering if it might be good to rename atomicrmw to atomic_reduce or something similar since there is already another more general atomic rmw that has an arbitrary interior and is typically lowered though atomic compare and exchange, and the enum based one is all reductions.

568

SCF uses a somewhat more general mechanism for reductions, but that adds a fair amount of complexity, in particular regarding initialization. I'd rather keep reductions for affine.parallel to be a closed set with known identity elements until we have a reason not to, since it's easier to pattern match to specialized hardware, and makes a number of the transforms easier as well. In practice this set of reduction operations seems to cover all of the ML use cases we've run into.

641

Totally agree, I will rename.

jbruestle updated this revision to Diff 273857.Jun 26 2020, 4:24 PM

Naming change for better consistency as suggested by @dcaballe.

flaub added inline comments.Jun 26 2020, 4:57 PM
mlir/include/mlir/Dialect/Affine/IR/AffineOps.td
562

Typo: acccumulated

594

Multiple plurals, seems a bit confusing.

jbruestle updated this revision to Diff 274160.Jun 29 2020, 9:50 AM

Typo fixes

bondhugula requested changes to this revision.Jun 29 2020, 7:13 PM

It's nice to see this. Sorry about the delay - I missed this patch. I think this is an important design update to be mentioned on discourse. Could you please post a note there - just the op doc descriptions + the example snippets you have here should be sufficient to start with.

mlir/include/mlir/Dialect/Affine/IR/AffineOps.td
305–315

Missing doc description here on the meaning on the yield operands. Note that the semantics of the terminator depend on the parent op holding it, and so this has to be described here. Please also update the example to include affine.yield.

This revision now requires changes to proceed.Jun 29 2020, 7:13 PM
jbruestle updated this revision to Diff 276543.Jul 8 2020, 1:38 PM

Improve docs, rebase off of most recent commit.

Could you please add a single line guard in hoistAffineIfOp with a TODO to bail out on any affine.if ops that have results? The loop unswitching will no longer be correct with such affine.if ops.

mlir/include/mlir/Dialect/Affine/IR/AffineOps.td
794–795

This has to be marked ReturnLike.

mlir/lib/Conversion/AffineToStandard/AffineToStandard.cpp
329–330

Nit: Yield ops (terminators for affine regions) are removed.

bondhugula requested changes to this revision.Jul 9 2020, 12:53 AM

Could you please update the commit summary to reflect this revision? "reduce" isn't mentioned there. You can drop the "... I do not intend .." part from the summary. I think the last line is also outdated.

mlir/lib/Dialect/Affine/IR/AffineOps.cpp
2568

Please avoid auto here.

This revision now requires changes to proceed.Jul 9 2020, 12:53 AM
bondhugula added inline comments.Jul 9 2020, 6:03 AM
mlir/lib/Dialect/Affine/IR/AffineOps.cpp
2643–2666

Some comments here for this whole block code.

2647

Comment here.

2654

Avoid auto here please.

2699–2700

You can use variadic isa here.

jbruestle updated this revision to Diff 276766.Jul 9 2020, 9:54 AM
jbruestle marked 14 inline comments as done.
jbruestle edited the summary of this revision. (Show Details)

Changes requested by @bondhugula.

bondhugula accepted this revision.Jul 9 2020, 11:26 AM

Update commit summary to say you are adding reduction support for affine.parallel?

mlir/include/mlir/Dialect/Affine/IR/AffineOps.td
342

Nit: load -> affine.load (won't parse otherwise).

This revision is now accepted and ready to land.Jul 9 2020, 11:26 AM

Nit: I think the commit title could be shortened to start with "Add yield semantics ..." ('IR changes to' is obvious).

This revision was automatically updated to reflect the committed changes.