Page MenuHomePhabricator

[MLIR] Add tiling validity check to loop tiling pass
ClosedPublic

Authored by kumasento on Wed, Jul 29, 11:52 AM.

Details

Summary

This diff aims to provide a new API, checkTilingLegality, to verify that the loop tiling result still satisifes the dependence constraints of the original loop nest.

Background

Previously, there was no check for the validity of tiling. For instance:

func @diagonal_dependence() {
  %A = alloc() : memref<64x64xf32>

  affine.for %i = 0 to 64 {
    affine.for %j = 0 to 64 {
      %0 = affine.load %A[%j, %i] : memref<64x64xf32>
      %1 = affine.load %A[%i, %j - 1] : memref<64x64xf32>
      %2 = addf %0, %1 : f32
      affine.store %2, %A[%i, %j] : memref<64x64xf32>
    }
  }

  return 
}

You can find more information about this example from the Section 3.11 of [1].

In general, there are three types of dependences here: two flow dependences, one in direction (i, j) = (0, 1) (notation that depicts a vector in the 2D iteration space), one in (i, j) = (1, -1); and one anti dependence in the direction (-1, 1).

Since two of them are along the diagonal in opposite directions, the default tiling method in affine, which tiles the iteration space into rectangles, will violate the legality condition proposed by Irigoin and Triolet [2]. [2] implies two tiles cannot depend on each other, while in the affine tiling case, two rectangles along the same diagonal are indeed dependent, which simply violates the rule.

This diff attempts to put together a validator that checks whether the rule from [2] is violated or not when applying the default tiling method in affine.

Implementation

The canonical way to perform such validation is by examining the effect from adding the constraint from Irigoin and Triolet to the existing dependence constraints.

Since we already have the prior knowlegde that affine tiles in a hyper-rectangular way, and the resulting tiles will be scheduled in the same order as their respective loop indices, we can simplify the solution to just checking whether all dependence components are non-negative along the tiling dimensions.

We put this algorithm into a new API called checkTilingLegality under LoopTiling.cpp. This function iterates every load/store pair, and if there is any dependence between them, we get the dependence component and check whether it has any negative component. This function returns failure if the legality condition is violated.

References

[1]. Bondhugula, Uday. Effective Automatic parallelization and locality optimization using the Polyhedral model. https://dl.acm.org/doi/book/10.5555/1559029
[2]. Irigoin, F. and Triolet, R. Supernode Partitioning. https://dl.acm.org/doi/10.1145/73560.73588

Diff Detail

Event Timeline

kumasento created this revision.Wed, Jul 29, 11:52 AM
Herald added a project: Restricted Project. · View Herald TranscriptWed, Jul 29, 11:52 AM
kumasento requested review of this revision.Wed, Jul 29, 11:52 AM

Hi Uday (@bondhugula)

This diff roughly implements the validation method we've discussed on discourse a while ago.

Sorry that this diff is still very primitive, but I do have some difficulties about the core formulation to make it work well.

I try to create a constraint system with the tiling indices in it. Those indices are simply calculated by i floordiv tileSize, and they will become a local variable in the constraint system. I then add inequality on them to see whether the *opposite of the legality condition (\phi(t) - \phi(s) < 0 based on the notation in the PLUTO paper) is not** satisfied (using isEmpty()).

Theoretically, I think it should work. But in practice, especially the second test case in loop-tile-dependence.mlir, it doesn't work well, i.e., the illegal tiling is not discovered.

I'm wondering whether you could kindly help me pointing out which step is possibly wrong based on your experience. I think it would probably be the floordiv constraint, which I want to be a equality, but by using addLocalFloorDiv I can only get two inequalities. I'm not sure whether it is equivalent to what the original formulation is.

Please let me know if you have any suggestion. Thanks!

kumasento updated this revision to Diff 281746.Wed, Jul 29, 3:06 PM

Fixed problem in constraint system

Please do ignore my previous comment on the current implementation doesn't work well. It does work properly based on the test case I've provided.

Previously I tested the emptiness of the constraint system with the illegal inequalities joined by AND, which means even if tiling on one index violates dependence, the final constraint system still has no solution. I should check each illegal inequality separately. This has been fixed in the latest commit amend.

kumasento updated this revision to Diff 281749.Wed, Jul 29, 3:15 PM

Improved comments a bit.

Can you please rephrase the commit title to avoid the name and instead just have it in natural language?

bondhugula requested changes to this revision.EditedThu, Jul 30, 5:57 AM

Please expand the commit summary a little. The h in h^T . R here are always the canonical directions with the current tiling pass. You really don't need to do anything with FlatAffineConstraints. Instead, the conditions of Irigoin and Triolet would just simplify here to checking whether all dependence components are non-negative along the dimensions you want to tile. And checking this by simply looking at the dependence information is sufficient. You may want to expand on the motivation if you intended something else.

mlir/lib/Dialect/Affine/Transforms/LoopTiling.cpp
178

Avoid auto - spell out the type - here and at other places.

227

Please remove unguarded debug code - if it's debug code, please guard it with LLVM_DEBUG and other information to identify what it is.

247

Please terminate comments with a full stop - here and anywhere else.

mlir/test/Dialect/Affine/loop-tiling-dependence.mlir
25 ↗(On Diff #281749)

Consider rephrasing: illegal in dependence -> invalid due to dependences

This revision now requires changes to proceed.Thu, Jul 30, 5:57 AM
bondhugula added inline comments.Thu, Jul 30, 5:59 AM
mlir/test/Dialect/Affine/loop-tiling-dependence.mlir
9–16 ↗(On Diff #281749)

Missing check lines on whether the tiling actually happened. For eg. it shouldn't have bailed out. Also, can you renaming this file loop-tiling-validity.mlir?

Please expand the commit summary a little. The h in h^T . R here are always the canonical directions with the current tiling pass. You really don't need to do anything with FlatAffineConstraints. Instead, the conditions of Irigoin and Triolet would just simplify here to checking whether all dependence components are non-negative along the dimensions you want to tile. And checking this by simply looking at the dependence information is sufficient. You may want to expand on the motivation if you intended something else.

Thank you for your detailed comments @bondhugula !

Can I ask a quick question about the quoted suggestion: I reckon your point is that, since the loop tiles are always hyperrect in the current implementation in affine, the conditions of Irigoin and Triolet can just be simplified to checking non-negative dependence components?

My implementation is more like a direct translation of the original Irigoin and Triolet conditions, but I guess it is not very necessary here.

And thank you again for your other suggestions, I will fix them later :)

Please expand the commit summary a little. The h in h^T . R here are always the canonical directions with the current tiling pass. You really don't need to do anything with FlatAffineConstraints. Instead, the conditions of Irigoin and Triolet would just simplify here to checking whether all dependence components are non-negative along the dimensions you want to tile. And checking this by simply looking at the dependence information is sufficient. You may want to expand on the motivation if you intended something else.

Thank you for your detailed comments @bondhugula !

Can I ask a quick question about the quoted suggestion: I reckon your point is that, since the loop tiles are always hyperrect in the current implementation in affine, the conditions of Irigoin and Triolet can just be simplified to checking non-negative dependence components?

That's right.

kumasento retitled this revision from [MLIR][Affine] Added checkTilingLegality to [MLIR][Affine] Added legality validator for loop tiling.Fri, Jul 31, 10:09 AM
kumasento edited the summary of this revision. (Show Details)
kumasento marked 4 inline comments as done.
kumasento edited the summary of this revision. (Show Details)

Improved loop tiling legality checking implementation - now based on ensuring non-negative dependence component exists.

Updated various other places based on reviews.

Hi Uday @bondhugula would you mind reviewing this again when you're available? Thanks!

bondhugula accepted this revision.Sat, Aug 1, 12:29 AM

Thanks for revising this - looks good. Some minor comments and this one below.

Previously, there is no validation after loop tiling is performed. For instance:

-> ... there was no check for the validity of tiling.

mlir/lib/Dialect/Affine/Transforms/LoopTiling.cpp
160

This should be marked static and important to document along the lines you have in the commit summary. Need not be that detailed though.

172

OpInsts -> Ops

181

You can hoist this and do dependenceConstraints.reset() to avoid potential repeated allocations.

198

Please use ... k = 0, e = ...; k < e; ... form. Elsewhere as well - prevents repeated evaluation.

200

Nit: Can drop new line.

208

Could we add which dimension/loop depth it is negative along?

248–250

Drop trivial braces.

mlir/test/Dialect/Affine/loop-tiling-validity.mlir
31

Nit: You could mention that the dependence direction vector that makes it invalid to tile here is (+, -).

This revision is now accepted and ready to land.Sat, Aug 1, 12:29 AM

I think the commit summary is more detailed / has additional info that could be in the doc / implementation comment of checkTilingValidity. You could move a few things from the former to the latter.

bondhugula added a comment.EditedSat, Aug 1, 12:32 AM

Consider changing title to something like '[MLIR] Add tiling validity check to loop tiling pass'. You'll have to hit "Edit revision" manually here as arc doesn't automatically update it.

kumasento retitled this revision from [MLIR][Affine] Added legality validator for loop tiling to [MLIR] Add tiling validity check to loop tiling pass.Sat, Aug 1, 3:38 AM
kumasento edited the summary of this revision. (Show Details)
kumasento updated this revision to Diff 282386.Sat, Aug 1, 4:02 AM
kumasento marked 8 inline comments as done.

Fixed several issues based on review comments.

Hi @bondhugula Thank you very much for your detailed comments :) I've improved this diff based on your suggestions. Please do let me know if there is any other I should look at.

BTW I've changed the naming of opInst to op in AffineAnalysis.cpp as well. Haven't fixed others though.

bondhugula added inline comments.Sat, Aug 1, 7:36 AM
mlir/lib/Dialect/Affine/Transforms/LoopTiling.cpp
160

I would rephrase / suggest changes to this.

Checks whether hyper-rectangular loop tiling of the nest represented by `origLoops` ... is valid.

?

validation rule -> validity condition

Please use the whole line - no need to break for different sentences.

will be in hyper-rectangles <- please rephrase.

-> in the lexicographically increasing order on the vector of loop indices.

will raise failure result -> will return failure

when any dependence component is negative -> when any dependence component is negative along any of origLoops
?

kumasento updated this revision to Diff 282399.Sat, Aug 1, 8:24 AM

Rephrased the doc comment of checkTilingLegality.

kumasento added inline comments.Sat, Aug 1, 8:26 AM
mlir/lib/Dialect/Affine/Transforms/LoopTiling.cpp
160

Thank you very much for your suggestions on the doc comment! The previous version was indeed hard to read. Would you mind checking again this version? Thanks! :)

bondhugula accepted this revision.Sat, Aug 1, 8:47 AM
bondhugula requested changes to this revision.
bondhugula added inline comments.
mlir/lib/Analysis/AffineAnalysis.cpp
869

srcOpInst -> srcOp

872

dstOpInst -> dstOp

mlir/lib/Dialect/Affine/Transforms/LoopTiling.cpp
172

This isn't accurate - you are just checking whether they can be tiled. They aren't already tiled, right?

This revision now requires changes to proceed.Sat, Aug 1, 8:48 AM
kumasento updated this revision to Diff 282400.Sat, Aug 1, 8:57 AM
kumasento marked 2 inline comments as done.

Fixed OpInst typo and removed unnecessary comments.

mlir/lib/Analysis/AffineAnalysis.cpp
869

Ah sorry I just missed them again :(

mlir/lib/Dialect/Affine/Transforms/LoopTiling.cpp
172

You are absolutely right! It was remained from my previous implementation, which is placed after the tiling is done. I forgot to remove this comment afterwards.

bondhugula accepted this revision.Thu, Aug 6, 7:19 AM
This revision is now accepted and ready to land.Thu, Aug 6, 7:19 AM

Looks great - let me know if you'd like me to commit.

Looks great - let me know if you'd like me to commit.

Thank you Uday, yes, that would be great!

This revision was automatically updated to reflect the committed changes.