This is an archive of the discontinued LLVM Phabricator instance.

[MLIR] Simplify Semi Affine Expressions Using flattening
ClosedPublic

Authored by arnab-oss on Oct 29 2021, 5:06 AM.

Details

Summary
For the semi affine expressions, whenever rhs of a floordiv, ceildiv, mod
or product expression is a symbolic expression, we introduce a local variable
representing the result, and store the floordiv/ceildiv, mod or product
affine expression in LocalExprs. In this way the expression is flattened,
and trivial addition and subtraction related simplifications are performed.

Also rule based matching for detecting and transforming "expr - q * (expr floordiv q)"
to "expr mod q", where q is a symbolic exxpression, in simplifyAdd function.

Diff Detail

Event Timeline

arnab-oss created this revision.Oct 29 2021, 5:06 AM
arnab-oss requested review of this revision.Oct 29 2021, 5:06 AM

Can you split the two commits out into separate revisions (with different titles)? The first one can be set as the parent revision of the second one.

arnab-oss retitled this revision from [MLIR] Simplify Semi Affine Expressions to [MLIR] Simplify Semi Affine Expressions Using flattening.Oct 29 2021, 5:25 AM
arnab-oss edited the summary of this revision. (Show Details)
arnab-oss updated this revision to Diff 383979.EditedNov 1 2021, 11:17 PM

Remove the second commit and keep only the first commit doing simplification of semi affine map in this PR.

Thanks for the contribution, this is really useful. Some first comments, will do a deeper dive in the next days...

mlir/lib/IR/AffineExpr.cpp
932

It seems like you're using -1 in the code below. Comments need updating?

937

d0

937–938

Would be good to have a full example with coefficients, indices and indexToExprMap somewhere here.

967

nit: store in variable to avoid writing it 3 times

980

involving

982

space

984

Use llvm::enumerate, so that you don't have to maintain an idx variable.

986

Is there an idx++ missing here?

999

delete

1003–1004

Is it possible that there is already an entry at {lhsPos, -1}?

1021–1024

This piece of code is a common pattern. I would create a lambda function and reuse it.

Something like

auto addEntry = [&](std::pair<...> index, int64_t coefficient, AffineExpr expr) {
  assert that index is not yet in indices

  indices.push_back(index);
  coefficients.insert(...);
  indexToExprMap.insert(...)
};
utils/arcanist/clang-format.sh
54

I think this is not supposed to be part of the commit.

arnab-oss updated this revision to Diff 385326.Nov 6 2021, 11:53 PM
arnab-oss marked 11 inline comments as done.

Addressed all the comments by springerm.

arnab-oss updated this revision to Diff 385327.Nov 7 2021, 12:10 AM
arnab-oss marked an inline comment as done.

Added an example showing the contents of indices, indexToExprMap and coefficients.

arnab-oss updated this revision to Diff 385328.Nov 7 2021, 12:16 AM

Removed unwanted changes from clang-format.sh

Thanks for this really useful simplification that catches most of the things needed in practice. Mostly minor comments below.

mlir/include/mlir/IR/AffineExpr.h
296–298

Nit: You can rewrap to use the whole width (80 columns). I know it was the existing comment that didn't.

mlir/include/mlir/IR/AffineExprVisitor.h
303

No comma before when.

312

maybe -> may be

312

Please rephrase the first sentence - it isn't really clear.

which represents -> representing ?

314

Spell fix.

315

in proper -> in the proper

mlir/lib/IR/AffineExpr.cpp
954

Function to add entries -> Adds entries

1063–1066

Code comments here.

1094

product, and the -> product; the ...

1133

Using semi-affine (which hyphen) consistently.

1232–1236

But result's size has not been set. If the expectation is for the caller to set it, there has to be an assertion here on the expected size of result.

1263

This should be:
SmallVectorImpl<int64_t> &lhs = ...
?

Thanks a lot for the contribution! LGTM. Just some minor comments.

mlir/lib/IR/AffineExpr.cpp
942

givem -> given

969

j++ -> ++j

mlir/test/Dialect/Affine/simplify-affine-structures.mlir
493

Please, add a CHECK-LABEL for this test

504

Could you please split the test into multiple ones (one for mod, one for floordiv, etc.)? It would help readability and be very helpful to quickly know which simplification is broken when the time comes.

510

Maybe I'm wrong but, unlike affine maps, I think the order of the instructions shouldn't change here so we could remove the -DAG part (?)

arnab-oss updated this revision to Diff 385358.Nov 7 2021, 10:27 AM
arnab-oss marked 15 inline comments as done.

Addressed all the comments by @bondhugula and @dcaballe.

springerm added inline comments.Nov 8 2021, 9:52 PM
mlir/lib/IR/AffineExpr.cpp
969

nit: I would split this in two loops, since you have the j < numDims check below.

1002

I still think there's an idx++ missing here. Can you explain the reasoning for this?

arnab-oss updated this revision to Diff 385776.Nov 9 2021, 4:27 AM
arnab-oss marked an inline comment as done.

Fixed a bug in the loop iterating over localExprs vector. Addressed comments
by @springerm.

Fixed a bug in the loop iterating over localExprs vector. Addressed comments
by @springerm.

Could you add a test case that exercises that part of the code? (if not already)

arnab-oss marked an inline comment as done.

Added an example which is giving correct output now after the previous update bug fix.

Fixed a bug in the loop iterating over localExprs vector. Addressed comments
by @springerm.

Could you add a test case that exercises that part of the code? (if not already)

Added. Thanks.

mlir/lib/IR/AffineExpr.cpp
984

I cannot use llvm::enumerate and replace idx with iterator.index() as idx won't be increasing in every iteration.

986

No

1002

Oh you are correct. In fact this change made me uncover another bug in the code.

1003–1004

No

1263

If I do this, I cannot directly use lhs to initialise dividend later in the code.

mlir/test/Dialect/Affine/simplify-affine-structures.mlir
493

Using CHECK-LABEL for function name and CHECK-SAME for capturing arhguments leads to test failure.

It shows error message "PRODUCT" not defined which I had defined earlier while capturing the map using CHECK-DAG.

dcaballe added inline comments.Nov 11 2021, 2:02 PM
mlir/test/Dialect/Affine/simplify-affine-structures.mlir
493

You have to use $ to capture variables before LABEL. Something like: #[[$MOD:.*]] = affine_map...

arnab-oss updated this revision to Diff 386809.Nov 12 2021, 5:12 AM

Added CHECK-LABEL in unit tests.

bondhugula accepted this revision.Nov 13 2021, 7:31 PM

LGTM - thanks very much!

This revision is now accepted and ready to land.Nov 13 2021, 7:31 PM
arnab-oss marked an inline comment as done.Nov 13 2021, 7:38 PM

ping @springerm and @dcaballe , can you please let me know whether the revision is ok or if any more changes are required to be made? Thanks in advance!

mlir/test/Dialect/Affine/simplify-affine-structures.mlir
493

Done.

ping @springerm and @dcaballe , can you please let me know whether the revision is ok or if any more changes are required to be made? Thanks in advance!

Thanks for addressing the feedback! The tests look good to me. I'll leave the final approval to @springerm.

springerm accepted this revision.Nov 15 2021, 2:56 PM
This revision was automatically updated to reflect the committed changes.

There were three whitespace errors in the test case and the commit title capitalization that I fixed before committing.