Non hyper-rectangular loop tiling in the Affine dialect

Hi there

My current project may require the affine dialect to support loop tiling on non hyper-rectangular iteration spaces. I’m not very sure but I assume that common polyhedral tools may already support this feature, so maybe it is possibly doable :slight_smile:

If I’m thinking of extending the LoopTiling pass to support non hyper-rectangular space, to accomplish my project, I’m wondering if there is any recommended steps that I need to take, e.g., papers, algorithms, etc.

Please let me know if there is any suggestion, thanks!

Best regards

Mainstream polyhedral works tackle code generation in a generic way not distinguishing between say hyper-rectangular and non hyperrectangular or even based on the kind of transformation: they just work on schedules encoded as affine functions/relations. But for your project, it appears you are just looking for tiling a band of successively nested affine.for ops, which is simpler / more focused. You could look at some of the earlier works including those outside of the mainstream polyhedral domain/schedule ones. For eg:

You could also look at how the restricted case you are dealing with simplifies with state-of-the-art polyhedral code generators (ISL codegen).

Polyhedral approaches for code generation will typically rely on successive projection (eliminating dimensions from a constraint system). The new GBR-based approach committed to MLIR could help with integer exactness (getting rid of redundancy) in conjunction with a simple projection algorithm like Fourier-Motzkin, but it may be expensive unless carefully used – this hasn’t been evaluated but could be great for an exploratory project. Note that generating nests ensuring zero redundancy in bounds is very different than not caring about redundant bounds – in terms of costs and techniques. Both approaches will lead to correct code, but redundancy in bounds may only get harder to get rid of later, and impact other transformations negatively.

Hi Uday,

Thank you very much for your very detailed reply! I’m indeed grateful for your insight on this matter, regarding polyhedral analysis and automatic loop tiling.

I’ve been thinking through your comments once I got it, and here are my intuitive thoughts, hope you won’t mind:

First of all, about my project -

It is partially true. I’m currently working on transforming polyhedral models to hardware descriptions described by synchronous data-flow (part of my PhD thesis), which should take into account all the tiling opportunities that can be represented by mainstream polyhedral works. Non hyper-rectangular successive loop nests is one case that I’m thinking of at the moment.

I could definitely establish my work on well-developed polyhedral projects, e.g., PolyMage, but since MLIR is a much bigger infrastructure that covers more applications, and it is an objective to rebuild existing polyhedral methods in MLIR (based on your IMPACT '20 keynote), I feel that it is a greater opportunity to contirbute to the polyhedral community and get wider potential audience group of my work.

Now back to this specific question on non hyper-rectangular tiling -

This is a great suggestion. I just fininished reading that paper (by Goumas et al.) and it is very nicely written. One of its contributions that uses non-unimodular transformation to transform the iteration space to hyper-rectangular can be useful to address this problem.

But it might not be very straightforward to put this into MLIR, as far as what I’ve discovered. The current tiling implementation in affine, please correct me if I’m wrong, is ad-hoc and doesn’t explicitly formulates the affine transformation (using matrices) for tiling. For instance, in LoopTiling.cpp, we basically insert loops for tiles and update existing loop nests’ bounds, by ad-hoc rules, and then derive a concise set of constraints by Fourier-Motzkin elimination. To support non hyper-rectangular tiling, we might need to add an additional step to encode the loop transformation in a generic form (e.g., matrix) before altering the AST (please do let me know if there is one already!), and in this encoding step we can support various kinds of tiling cases, including non hyper-rectangular. This design might benefit future polyhedral method implementation as well.

I definitely agree with this. I’ve only tried out to reproduce some examples in your PLUTO paper by CLooG though, but I assume using ISL and other established polyhedral codegen tools will be very easy for this task :slight_smile:

Actually, regarding the polyhedral approach, I do have an additional question on the current affine implementation of tiling. As far as I know, it doesn’t take care of dependence and it may produce illegal code.

For example, I’ve created a short testcast based on the example in the Supernode Partitioning paper:

func @diagonal(%A: memref<?x?xf32>) -> () {
  %c0 = constant 0 : index
  %c1 = constant 1 : index

  %M = dim %A, %c0: memref<?x?xf32>
  %N = dim %A, %c1: memref<?x?xf32>

  affine.for %i = 0 to %M {
    affine.for %j = 0 to %N {
      %0 = affine.load %A[%i, %j]: memref<?x?xf32>
      %1 = affine.load %A[%i + %c1, %j]: memref<?x?xf32>
      %2 = affine.load %A[%i, %j + %c1]: memref<?x?xf32>
      %3 = affine.load %A[%i - %c1, %j]: memref<?x?xf32>
      %4 = affine.load %A[%i, %j - %c1]: memref<?x?xf32>

      %5 = addf %0, %1 : f32
      %6 = addf %5, %2 : f32
      %7 = addf %6, %3 : f32
      %8 = addf %7, %4 : f32 %8, %A[%i, %j] : memref<?x?xf32>



It’s ideal tiling scheme would be partitioning the itertaion space into diamond shape, but if I run mlir-opt with -affine-loop-tile="tile-size=32" I simply get a two additional outer loops over tiles, which obviously means the tiles are rectangular and the dependences are violated.

One possible solution is to use the PLUTO algorithm, which ensures the tiling legality is satisfied. But I’m not sure whether this is somewhere in MLIR, or it will be implemented somehow in the future.

I will be happy to spend some time to extend the current affine tiling method to support non hyper-rectangular iteration space and try to take into account dependence when tiling. Also, I’m interested to know more about the exploratory project you’ve mentioned about getting rid of redundancy :slight_smile:

Please let me know if you have any further suggestions or comments!

Many thanks

This isn’t really accurate. FM is not used and not needed for deriving bounds to tile hyper-rectangular nests. The approach in LoopTiling.cpp is syntactic with some checks to avoid redundant bounds.

As I mentioned in my previous response, there isn’t one like this.

Yes, it’s really a test pass - it has the mechanics of tiling, but doesn’t check for validity. The infrastructure to check for validity is already there, and it’s pretty straightforward to use the memref dependence analysis utilities to obtain the necessary information for the validity of tiling.

This would be a good start especially if you’d like to start by addressing the validity part of it. I’ll be happy to review any patches.

Thank you very much Uday! And so sorry for my late reply. I’ve been viewing through the codebase recently.

Sure, I’m working on that track at the moment. I may constantly create some patches to resolve some TODOs in the affine codebase first to understand better how it works :slight_smile:
Here is the latest one I just created:

1 Like