[MLIR] Add tiling validity check to loop tiling pass

[kumasento]

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

[kumasento]

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]

Fixed problem in constraint system

[kumasento]

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]

Improved comments a bit.

[bondhugula]

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

[bondhugula]

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.

244

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

[bondhugula]

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?

[kumasento]

In D84882#2184583, @bondhugula wrote:

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 :)

[bondhugula]

In D84882#2184834, @kumasento wrote:

In D84882#2184583, @bondhugula wrote:

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]

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

Updated various other places based on reviews.

[kumasento]

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

[bondhugula]

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?

245–247

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 (+, -).

[bondhugula]

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]

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]

Fixed several issues based on review comments.

[kumasento]

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]

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]

Rephrased the doc comment of checkTilingLegality.

[kumasento]

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]

mlir/lib/Analysis/AffineAnalysis.cpp

868–869

srcOpInst -> srcOp

871–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?

[kumasento]

Fixed OpInst typo and removed unnecessary comments.

mlir/lib/Analysis/AffineAnalysis.cpp

868–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]

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

[kumasento]

In D84882#2199745, @bondhugula wrote:

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

Thank you Uday, yes, that would be great!

mlir/lib/Analysis/AffineAnalysis.cpp

 }
 
 /// Gathers dependence components for dependences between all ops in loop nest
 /// rooted at 'forOp' at loop depths in range [1, maxLoopDepth].
 void mlir::getDependenceComponents(
     AffineForOp forOp, unsigned maxLoopDepth,
     std::vector<SmallVector<DependenceComponent, 2>> *depCompsVec) {
   // Collect all load and store ops in loop nest rooted at 'forOp'.
-  SmallVector<Operation *, 8> loadAndStoreOpInsts;
-  forOp.getOperation()->walk([&](Operation *opInst) {
-    if (isa<AffineReadOpInterface, AffineWriteOpInterface>(opInst))
-      loadAndStoreOpInsts.push_back(opInst);
+  SmallVector<Operation *, 8> loadAndStoreOps;
+  forOp.getOperation()->walk([&](Operation *op) {
+    if (isa<AffineReadOpInterface, AffineWriteOpInterface>(op))
+      loadAndStoreOps.push_back(op);
   });
 
-  unsigned numOps = loadAndStoreOpInsts.size();
+  unsigned numOps = loadAndStoreOps.size();
   for (unsigned d = 1; d <= maxLoopDepth; ++d) {
     for (unsigned i = 0; i < numOps; ++i) {
-      auto *srcOpInst = loadAndStoreOpInsts[i];
-      MemRefAccess srcAccess(srcOpInst);
+      auto *srcOp = loadAndStoreOps[i];
+      MemRefAccess srcAccess(srcOp);
       for (unsigned j = 0; j < numOps; ++j) {
-        auto *dstOpInst = loadAndStoreOpInsts[j];
-        MemRefAccess dstAccess(dstOpInst);
+        auto *dstOp = loadAndStoreOps[j];
+        MemRefAccess dstAccess(dstOp);
 
         FlatAffineConstraints dependenceConstraints;
         SmallVector<DependenceComponent, 2> depComps;
         // TODO: Explore whether it would be profitable to pre-compute and store
         // deps instead of repeatedly checking.
         DependenceResult result = checkMemrefAccessDependence(
             srcAccess, dstAccess, d, &dependenceConstraints, &depComps);
         if (hasDependence(result))
           depCompsVec->push_back(depComps);
       }
     }
   }
 }

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

if (mayBeConstantCount && mayBeConstantCount.getValue() < tileSizes[i]) {
       // No need of the min expression.
       auto dim = b.getAffineDimExpr(0);
       auto ubMap = AffineMap::get(1, 0, dim + tileSizes[i]);
       newLoops[width + i].setUpperBound(newLoops[i].getInductionVar(), ubMap);
     }
   }
 }
 
+/// This function checks whether hyper-rectangular loop tiling of the nest
+/// represented by `origLoops` is valid. The validity condition is from Irigoin
+/// and Triolet, which states that two tiles cannot depend on each other. We
+/// simplify such condition to just checking whether there is any negative
+/// dependence direction, since we have the prior knowledge that the tiling
+/// results will be hyper-rectangles, which are scheduled in the
+/// lexicographically increasing order on the vector of loop indices. This
+/// function will return failure when any dependence component is negative along
+/// any of `origLoops`.
+static LogicalResult
+checkTilingLegality(MutableArrayRef<mlir::AffineForOp> origLoops,
+                    ArrayRef<unsigned> tileSizes) {
+  assert(!origLoops.empty() && "no original loops provided");
+
+  // We first find out all dependences we intend to check.
+  SmallVector<Operation *, 8> loadAndStoreOps;
+  origLoops[0].getOperation()->walk([&](Operation *op) {
+    if (isa<AffineReadOpInterface, AffineWriteOpInterface>(op))
+      loadAndStoreOps.push_back(op);
+  });
+
+  unsigned numOps = loadAndStoreOps.size();
+  unsigned numLoops = origLoops.size();
+  FlatAffineConstraints dependenceConstraints;
+  for (unsigned d = 1; d <= numLoops + 1; ++d) {
+    for (unsigned i = 0; i < numOps; ++i) {
+      Operation *srcOp = loadAndStoreOps[i];
+      MemRefAccess srcAccess(srcOp);
+      for (unsigned j = 0; j < numOps; ++j) {
+        Operation *dstOp = loadAndStoreOps[j];
+        MemRefAccess dstAccess(dstOp);
+
+        SmallVector<DependenceComponent, 2> depComps;
+        dependenceConstraints.reset();
+        DependenceResult result = checkMemrefAccessDependence(
+            srcAccess, dstAccess, d, &dependenceConstraints, &depComps);
+
+        // Skip if there is no dependence in this case.
+        if (!hasDependence(result))
+          continue;
+
+        // Check whether there is any negative direction vector in the
+        // dependence components found above, which means that dependence is
+        // violated by the default hyper-rect tiling method.
+        LLVM_DEBUG(llvm::dbgs() << "Checking whether tiling legality violated "
+                                   "for dependence at depth: "
+                                << Twine(d) << " between:\n";);
+        LLVM_DEBUG(srcAccess.opInst->dump(););
+        LLVM_DEBUG(dstAccess.opInst->dump(););
+        for (unsigned k = 0, e = depComps.size(); k < e; k++) {
+          DependenceComponent depComp = depComps[k];
+          if (depComp.lb.hasValue() && depComp.ub.hasValue() &&
+              depComp.lb.getValue() < depComp.ub.getValue() &&
+              depComp.ub.getValue() < 0) {
+            LLVM_DEBUG(llvm::dbgs()
+                       << "Dependence component lb = "
+                       << Twine(depComp.lb.getValue())
+                       << " ub = " << Twine(depComp.ub.getValue())
+                       << " is negative  at depth: " << Twine(d)
+                       << " and thus violates the legality rule.\n");
+            return failure();
+          }
+        }
+      }
+    }
+  }
+
+  return success();
+}
 /// Tiles the specified band of perfectly nested loops creating tile-space loops
 /// and intra-tile loops. A band is a contiguous set of loops.
 //  TODO: handle non hyper-rectangular spaces.
 LogicalResult
 mlir::tilePerfectlyNested(MutableArrayRef<AffineForOp> input,
                           ArrayRef<unsigned> tileSizes,
                           SmallVectorImpl<AffineForOp> *tiledNest) {
   // Check if the supplied for op's are all successively nested.
   assert(!input.empty() && "no loops in input band");
   assert(input.size() == tileSizes.size() && "Too few/many tile sizes");
 
   assert(isPerfectlyNested(input) && "input loops not perfectly nested");
 
   auto origLoops = input;
 
+  // Perform tiling legality test.
+  if (failed(checkTilingLegality(origLoops, tileSizes)))
+    origLoops[0].emitRemark("tiled code is illegal due to dependences");
+
   AffineForOp rootAffineForOp = origLoops[0];
   auto loc = rootAffineForOp.getLoc();
   // Note that width is at least one since band isn't empty.
   unsigned width = input.size();
 
   SmallVector<AffineForOp, 6> tiledLoops(2 * width);
 
   // The outermost among the loops as we add more..

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

+// RUN: mlir-opt %s  -split-input-file -affine-loop-tile="tile-size=32" -verify-diagnostics | FileCheck %s
+
+// -----
+
+// There is no dependence violated in this case. No error should be raised.
+
+// CHECK-DAG: [[$LB:#map[0-9]+]] = affine_map<(d0) -> (d0)>
+// CHECK-DAG: [[$UB:#map[0-9]+]] = affine_map<(d0) -> (d0 + 32)>
+// CHECK-DAG: [[$ID:#map[0-9]+]] = affine_map<() -> (0)>
+// CHECK-DAG: [[$ID_PLUS_21:#map[0-9]+]] = affine_map<() -> (64)>
+
+// CHECK-LABEL: func @legal_loop()
+func @legal_loop() {
+  %0 = alloc() : memref<64xf32>
+
+  affine.for %i = 0 to 64 {
+    %1 = affine.load %0[%i] : memref<64xf32>
+    %2 = addf %1, %1 : f32
+    affine.store %2, %0[%i] : memref<64xf32>
+  }
+
+  return 
+}
+
+// CHECK:   affine.for %{{.*}} = 0 to 64 step 32 {
+// CHECK-NEXT:     affine.for %{{.*}} = [[$LB]](%{{.*}}) to [[$UB]](%{{.*}}) {
+
+// -----
+
+// There are dependences along the diagonal of the 2d iteration space, 
+// specifically, they are of direction (+, -).
+// The default tiling method (hyper-rect) will violate tiling legality.
+// We expect a remark that points that issue out to be emitted.
+
+// CHECK-LABEL: func @illegal_loop_with_diag_dependence
+func @illegal_loop_with_diag_dependence() {
+  %A = alloc() : memref<64x64xf32>
+
+  affine.for %i = 0 to 64 {
+    // expected-remark@above {{tiled code is illegal due to dependences}}
+    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 
+}