Page MenuHomePhabricator

[mlir] Normalize dynamic memrefs with a map of tiled-layout.
Needs ReviewPublic

Authored by imaihal on Feb 28 2021, 6:47 PM.

Details

Summary

Steps for normalizing dynamic memrefs for tiled layout map

  1. Check if original map is tiled layout. Only tiled layout is supported.
  2. Create normalized memrefType. Dimensions that include dynamic dimensions in the map output will be dynamic dimensions.
  3. Create new maps to calculate each dimension size of new memref. In tiled layout, the dimension size can be calculated by replacing "floordiv <tile size>" with "ceildiv <tile size>" and "mod <tile size>" with "<tile size>".
  4. Create AffineApplyOp to apply the new maps. The output of AffineApplyOp is dynamicSizes for new AllocOp.
  5. Add the new dynamic sizes in new AllocOp.

This patch also set MemRefsNormalizable trant in CastOp and DimOp since
they used with dynamic memrefs.

Diff Detail

Event Timeline

There are a very large number of changes, so older changes are hidden. Show Older Changes
bondhugula requested changes to this revision.Apr 1 2021, 8:32 PM
bondhugula added inline comments.
mlir/lib/Transforms/Utils/Utils.cpp
394

A tiled layout map has to be clearly defined here so that it's clear what the formal condition to check for is. Please see comment further below.

420

While this handles the common case, this method isn't sound in general: because you'd have to check that for each d_i floordiv t_s, there exists one and only one d_i mod t_s and the former comes before the latter? But even this part isn't clear because a tiled layout map isn't defined in the doc comment of this method - so it'd be ultimately unclear as to what the formal condition to check here is.

427–429

Nit: return tileCount == tileSizes.size();

This revision now requires changes to proceed.Apr 1 2021, 8:32 PM
imaihal updated this revision to Diff 335214.Apr 5 2021, 12:35 AM
imaihal marked 2 inline comments as done.

Rebased and reflected comments.

imaihal updated this revision to Diff 335215.Apr 5 2021, 1:26 AM

Fixed code that accidentally deleted a previous fix, and updated comments.

imaihal updated this revision to Diff 335217.Apr 5 2021, 1:33 AM

Fixed a comment.

@bondhugula Thanks for your review! I reflected your comments. Could you review again?

mlir/lib/Transforms/Utils/Utils.cpp
420

@bondhugula I used AffineExpr.walk here, but it was not correct. I removed it.

This isTiledLayoutMap works as follows in #map_tiled1 = affine_map<(d0, d1, d2, d3) -> (d0, d1, (d2 floordiv 4) floordiv 32, (d3 mod 8) floordiv 64, (d2 floordiv 4) mod 32, (d3 mod 8) mod 64)>,

  1. Finds "<AffineExpr> floordiv <AffineConstantExpr>".

In this example, they are (d2 floordiv 4) floordiv 32 and (d3 mod 8) floordiv 64. So, tileSizes includes {d0 floordiv 4, 32} and {d3 mod 8, 64}.

  1. Finds "<AffineExpr> mod <AffineConstantExpr>", and compares them with the tileSizes

In this example, they are (d2 floordiv 4) mod 32 and (d3 mod 8) mod 64, and the pair of {d2 floordive 4, 32} and {d3 mod 8, 64} exists in the tileSize. So it returns true.

bondhugula requested changes to this revision.Apr 7 2021, 10:16 PM
bondhugula added inline comments.
mlir/include/mlir/Dialect/MemRef/IR/MemRefOps.td
242

Please capture this change in the commit summary as well - thanks.

392

Likewise - please add a brief line on this in the commit summary.

mlir/lib/Transforms/Utils/Utils.cpp
395

But this condition is technically wrong - it's a necessary one but no where close to the sufficient condition. The number of floordiv's being equal to the number of mod's for a dim expression doesn't make the affine map a tiled layout map. It needs to have a specific structure.

This revision now requires changes to proceed.Apr 7 2021, 10:16 PM
imaihal updated this revision to Diff 337347.Tue, Apr 13, 11:47 PM
imaihal marked 4 inline comments as done.

Updated condition of tiled-layout map used in isTiledLayoutMap()

imaihal edited the summary of this revision. (Show Details)Tue, Apr 13, 11:48 PM
imaihal added inline comments.Tue, Apr 13, 11:51 PM
mlir/lib/Transforms/Utils/Utils.cpp
395

@bondhugula Thanks for review again. I updated the condition for tiled-layout map. How do you think of it?

imaihal updated this revision to Diff 337594.Wed, Apr 14, 6:03 PM

Updated to solve warning message in clang-tidy and rebased.

bondhugula requested changes to this revision.Sun, Apr 18, 3:41 AM
bondhugula added inline comments.
mlir/lib/Transforms/Utils/Utils.cpp
406

I think this variable is misnamed - it should be modExprs?

407

Nit: you can use llvm::reverse(map.getResults()).

423–427

Shouldn't all the floordiv exprs come before mod's?

433

Range-based loop?

434–436

This can be written in one line using llvm::is_contained(floorDivs, ...).

This revision now requires changes to proceed.Sun, Apr 18, 3:41 AM
imaihal updated this revision to Diff 338529.Mon, Apr 19, 8:23 AM
imaihal marked 5 inline comments as done.

Reflected comments for isTiledLayoutMap().

imaihal added inline comments.Mon, Apr 19, 8:37 AM
mlir/lib/Transforms/Utils/Utils.cpp
433

All of elements in modExprs need to be found in floorDivs. So, I think we need to check all of the elements.

bondhugula added inline comments.Mon, Apr 19, 1:47 PM
mlir/lib/Transforms/Utils/Utils.cpp
432

Nit: avoid auto.

433

Yes, I was only asking for a switch to a range-based for loop from an iterator one. You could check all elements.

bondhugula requested changes to this revision.Mon, Apr 19, 2:05 PM
bondhugula added inline comments.
mlir/lib/Transforms/Utils/Utils.cpp
396

I'm sorry but what you have here again is a necessary condition but not sufficient. For eg. the following is not the kind of tiled layout maps you want to detect:
d0 / 32, d1 / 32, d0 / 64, d1 mod 32)
The sufficient condition is actually much simpler - you need a specific k dimensions being floordiv'ed and then those same k dimensions appearing with a mod w.r.t the same tile sizes and nothing else (no other expressions) involving those dimensions. Also, this function should provide those tile sizes back to the call site (as an extra output argument) - since a user would in nearly all cases want those for further action.

517–519

But looks like you are assuming here that every floordiv has a corresponding mod which isn't the case given the check you have. Consider
d0 / 32, d1 / 32, d0 / 64, d0 mod 32, d1 mod 32 - will this work correctly?
I think isTiledLayoutMap should just return the tile sizes and the corresponding positions (since the tile-space ordering can be different from the intra-tile one), and you could use it here. This way you would unify the logic for createDimSizeExprForTiledLayout with isTiledLayoutMap?

641

Nit: is -> is a

This revision now requires changes to proceed.Mon, Apr 19, 2:05 PM
imaihal added inline comments.Mon, Apr 19, 7:56 PM
mlir/lib/Transforms/Utils/Utils.cpp
396

@bondhugula Thanks for your review. Sorry, I may be confused about tiled layout.

  • d0 / 32, d1 / 32, d0 / 64, d1 mod 32)

I implemented the function returns true for this map because d1 dimension meets tiled layout pattern (tiled in d1 dimension using tile size of 32, and not tiled in d0 dimension), but you suggested this is not true.
In this map, d0 dimension is floordived, but no mod for d0. Is this the reason why we can not say this is the tiled layout?

  • d0 / 32, d1 / 32, d0 / 64, d0 mod 32, d1 mod 32

In another comment, you showed this map. Can we say this is the tiled layout ? (tiled in d0 dimension using tile size of 32, and tiled in d1 dimension using tile size of 32?)

imaihal added inline comments.Mon, Apr 19, 8:14 PM
mlir/lib/Transforms/Utils/Utils.cpp
517–519

@bondhugula

d0 / 32, d1 / 32, d0 / 64, d0 mod 32, d1 mod 32 - will this work correctly?

This does not work correctly in current implementation. I will update using the tile sizes and the positions.

imaihal updated this revision to Diff 339647.Thu, Apr 22, 8:13 AM
imaihal marked 4 inline comments as done.

Updated check routine for a tiled layout. It checks dimensions being floordiv'ed and then those same dimensions appearing with a mod and no other expressions involving those dimensions. Also, the routine provides those tile sizes back to the call site.

imaihal updated this revision to Diff 339843.Thu, Apr 22, 7:16 PM

Fixed errors in clang-tidy.

@bondhugula I updated the patch.

The sufficient condition is actually much simpler - you need a specific k dimensions being floordiv'ed and then those same k dimensions appearing with a mod w.r.t the same tile sizes and nothing else (no other expressions) involving those dimensions. Also, this function should provide those tile sizes back to the call site (as an extra output argument) - since a user would in nearly all cases want those for further action.

I implemented this and changed the function name from isTiledLayoutMap() to getTileSizePos(). I hope this meets your request.

imaihal updated this revision to Diff 342725.Tue, May 4, 7:25 AM

Rebased.

Hi, @bondhugula , I updated this patch to reflect your suggestions. How do you think of this?

Hi, @bondhugula , I updated this patch to reflect your suggestions. How do you think of this?

Sorry about the delay. I should be able to review this (will hopefully be the final round) in the next one day.

mlir/lib/Transforms/Utils/Utils.cpp
399–400

Can you use an output parameter instead of returning? Sometimes RVO may not happen and will lead to an extra copy.

496

Nit: Can you add the output parameter as the last argument?

imaihal updated this revision to Diff 344051.Mon, May 10, 7:16 AM

Used output parameters to reflect comments.

imaihal updated this revision to Diff 344307.Tue, May 11, 1:04 AM

Rebased and removed unnecessary include file.

imaihal marked 2 inline comments as done.Tue, May 11, 1:07 AM
bondhugula added inline comments.Tue, May 11, 4:33 AM
mlir/lib/Transforms/NormalizeMemRefs.cpp
42 ↗(On Diff #342725)

This can be instead specified in a single line in Passes.td.

mlir/lib/Transforms/Utils/Utils.cpp
395

Many typos here:
no other expression doesn't involve -> no other expression involves
are appeared -> appear
by a tile size <-- this is inaccurate since tile sizes can be different

-> by respective tile sizes
398

an empty vector is returned.

422

It's more readable to return an empty vector - does {} or SmallVector<...>{}/() work?

478

Why do you need a separate vector just to return empty? Does {} or SmallVector<...>{}/() work?

513

This isn't clear - which dimension size is being referred to here?

641

single -> single level ?

703

tileSizePos.empty() instead of size() == 0.

713–719

if shape[d] < 0, you are adding an invalid constraint here: 0 <= s <= -2

imaihal updated this revision to Diff 344820.Wed, May 12, 7:23 AM

Reflected comments.