[MLIR] Simplify Semi Affine Expressions Using flattening

[arnab-oss]

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.

[bondhugula]

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]

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

[springerm]

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 ↗

(On Diff #383979)

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

[arnab-oss]

Addressed all the comments by springerm.

[arnab-oss]

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

[arnab-oss]

Removed unwanted changes from clang-format.sh

[bondhugula]

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

295–296

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.

1099

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

1138

Using semi-affine (which hyphen) consistently.

1237–1241

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.

1271

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

[dcaballe]

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]

Addressed all the comments by @bondhugula and @dcaballe.

[springerm]

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]

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

[bondhugula]

In D112808#3118161, @arnab-oss wrote:

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]

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

[arnab-oss]

In D112808#3122209, @bondhugula wrote:

In D112808#3118161, @arnab-oss wrote:

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

1271

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]

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]

Added CHECK-LABEL in unit tests.

[bondhugula]

LGTM - thanks very much!

[arnab-oss]

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.

[dcaballe]

In D112808#3129562, @arnab-oss wrote:

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.

[bondhugula]

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