[mlir][affine-loop-fusion] Fix a bug that AffineIfOp prevents fusion of the other loops

[tungld]

Summary

The presence of AffineIfOp inside AffineFor prevents fusion of the other loops to happen. For example:

affine.for %i0 = 0 to 10 {
  affine.store %cf7, %a[%i0] : memref<10xf32>
}
affine.for %i1 = 0 to 10 {
  %v0 = affine.load %a[%i1] : memref<10xf32>
  affine.store %v0, %b[%i1] : memref<10xf32>
}
affine.for %i2 = 0 to 10 {
  affine.if #set(%i2) {
    %v0 = affine.load %b[%i2] : memref<10xf32>
  }
}

The first two loops were not be fused because of affine.if inside the last affine.for.

The issue seems to come from a conservative constraint that does not allow fusion if there are ops whose number of regions != 0 (affine.if is one of them).

This patch just removes such a constraint when`affine.if` is inside affine.for. The existing canFuseLoops method is able to handle affine.if correctly.

[bondhugula]

This would be wrong to do - the checks can't be simply removed. Please see below. Region holding ops would need to be properly handled and they can be in the cases we know which ops those are.

mlir/lib/Transforms/LoopFusion.cpp

747–750

It's not okay to simply remove this check. It would just lead to incorrect fusion. An if op constrains the execution inside and I think the fusion here isn't/may not be currently geared to handle fusion when there are if ops inside.

778–780

We can't simply remove this check without know what kind of op it is! Instead we may just want to allow the ops that we know would not create an issue instead of allowing everything. For eg. this would be wrong to do if there was an scf.execute_region here (you'd have to look inside the op to see what memrefs are accessed) before you do anything at the outer level.

This is also clearly wrong even for affine.if because the if op might be reading/writing to memrefs and you can't fuse other things in general without considering that. Your test cases don't cover that.

[tungld]

Thank you for your comments. I will update the patch to check affine.if only.

This is also clearly wrong even for affine.if because the if op might be reading/writing to memrefs and you can't fuse other things in general without considering that. Your test cases don't cover that.

When I removed these checks, I saw that canFuseLoops method was able to handle affine.if. Actually canFuseLoops checks affine.if in the source and dest ops via calling gatherLoadsAndStores. For example, if destOp or srcOp contains affine.if, there is no fusion for them.
Also, if there is affine.if sandwiched between srcOp and destOp, there is no fusion:

affine.for %arg0 = 0 to 10 {
  affine.for %arg1 = 0 to 10 {
    affine.store %cst, %0[%arg0, %arg1] : memref<10x10xf32>
  }
}
affine.for %arg0 = 0 to 10 {
  affine.for %arg1 = 0 to 10 {
    affine.if #set(%arg1) {
      affine.store %cst, %0[%arg0, %arg1] : memref<10x10xf32>
    }
  }
}
affine.for %arg0 = 0 to 10 {
  affine.for %arg1 = 0 to 10 {
    %2 = affine.load %0[%arg1, %arg0] : memref<10x10xf32>
    affine.store %2, %1[%arg0, %arg1] : memref<10x10xf32>
  }
}

So, it seems that returning false in MemRefDependenceGraph::init is too early for affine.if, which causes all the handling of affine.if inside canFuseLoops unreachable.

[tungld]

Only handle AffineIfOp

[tungld]

@bondhugula I updated the patch to handle only the case that affine.if is inside affine.for. One more test is added for the case where affine.if is sandwiched by the source and dest ops.

[bondhugula]

Some comments. @vinayaka-polymage could you please review this?

mlir/lib/Transforms/LoopFusion.cpp

751

No need of dyn_cast here - instead use isa.

757–758

op here is hiding another variable of the same name. Please rename.

757–762

You don't need to use a flag. Instead, please use interrupt() and wasInterrupted() pattern. Several references for this at other places in the codebase.

[tungld]

Simplify the code according to Uday's comments.

[vinayaka-polymage]

Thanks for this, occurrence an AffineIfOp doesn't need to prevent fusion happening elsewhere, and canFuseLoops already checks for AffineIfOps.

Left a couple of minor suggestions. Also, the commit summary and the example there will have to change now I think, as a top-level AffineIfOp still prevents fusions.

mlir/lib/Transforms/LoopFusion.cpp

756

LoopNestStateCollector is walking the forOp anyway, so you can in fact, set a state inside that itself if AffineIf op is the only reason hasNonForRegion is being set. Only when it is, you can continue with fusion, else return false.

The main reason for suggesting is avoiding multiple walks over forOps.

mlir/test/Transforms/loop-fusion.mlir

526

Can you please change this to contain non-empty regions ? - I am afraid, this AffineIf might get folded away rendering the test case without if.

Also, unless you really want 2D loop nests, even a simpler 1D test case will also simplify checks.

583–601

Check return at the end.

[tungld]

Move the check of AffineIf into LoopNestStateCollector and use 1D tensors in the lit tests.

[tungld]

@vinayaka-polymage I updated the patch according to your comments. Thanks!

[vinayaka-polymage]

Thanks for addressing all the comments. Changes look good!

Really minor suggestion again on test cases: As you are not capturing the values and reusing them, you can get rid of all %{{.*}} and only check affine.for, affine.load etc..

[tungld]

Thanks, @vinayaka-polymage, for your comments! I simplified the lit tests.

[tungld]

@bondhugula @vinayaka-polymage is it the current patch good enough? Thanks!

[vinayaka-polymage]

Thanks @tungld ! It looks good to me.

[tungld]

@bondhugula is this patch ok to be landed? Thanks!

[bondhugula]

mlir/lib/Transforms/LoopFusion.cpp

79

This name has become very convoluted. Isn't the same as:

hasNonAffineRegionOp

?

747–751

Again, really convoluted.

mlir/test/Transforms/loop-fusion.mlir

526

ForInst -> ForOp

[bondhugula]

mlir/test/Transforms/loop-fusion.mlir

566

There's no ForInst.

ForInst -> ForOp everywhere.

[bondhugula]

mlir/lib/Transforms/LoopFusion.cpp

747–751

Isn't this the same as:

Returns false if a region holding op other than affine.for and affine.if ...

[tungld]

Thanks @bondhugula! I changed variable names and comments accoring to your comments.

I used ForInst since the existing tests used it. Now all ForInst are renamed to ForOp, also if_inst to if_op.

[bondhugula]

LGTM - thanks.

mlir/lib/Transforms/LoopFusion.cpp

mlir::createLoopFusionPass(unsigned fastMemorySpace,
                            uint64_t localBufSizeThreshold, bool maximalFusion) {
   return std::make_unique<LoopFusion>(fastMemorySpace, localBufSizeThreshold,
                                       maximalFusion);
 }
 
 namespace {
 
 // LoopNestStateCollector walks loop nests and collects load and store
-// operations, and whether or not an IfInst was encountered in the loop nest.
+// operations, and whether or not a region holding op other than ForOp and IfOp
+// was encountered in the loop nest.
 struct LoopNestStateCollector {
   SmallVector<AffineForOp, 4> forOps;
   SmallVector<Operation *, 4> loadOpInsts;
   SmallVector<Operation *, 4> storeOpInsts;
-  bool hasNonForRegion = false;
+  bool hasNonAffineRegionOp = false;
 
   void collect(Operation *opToWalk) {
     opToWalk->walk([&](Operation *op) {
       if (isa<AffineForOp>(op))
         forOps.push_back(cast<AffineForOp>(op));
-      else if (op->getNumRegions() != 0)
-        hasNonForRegion = true;
+      else if (op->getNumRegions() != 0 && !isa<AffineIfOp>(op))
+        hasNonAffineRegionOp = true;
       else if (isa<AffineReadOpInterface>(op))
         loadOpInsts.push_back(op);
       else if (isa<AffineWriteOpInterface>(op))
         storeOpInsts.push_back(op);
     });
   }
 };
 
if (!llvm::hasSingleElement(f))
     return false;
 
   DenseMap<Operation *, unsigned> forToNodeMap;
   for (auto &op : f.front()) {
     if (auto forOp = dyn_cast<AffineForOp>(op)) {
       // Create graph node 'id' to represent top-level 'forOp' and record
       // all loads and store accesses it contains.
       LoopNestStateCollector collector;
       collector.collect(&op);
-      // Return false if a non 'affine.for' region was found (not currently
-      // supported).
-      if (collector.hasNonForRegion)
+      // Return false if a region holding op other than 'affine.for' and
+      // 'affine.if' was found (not currently supported).
+      if (collector.hasNonAffineRegionOp)
         return false;
       Node node(nextNodeId++, &op);
       for (auto *opInst : collector.loadOpInsts) {
         node.loads.push_back(opInst);
         auto memref = cast<AffineReadOpInterface>(opInst).getMemRef();
         memrefAccesses[memref].insert(node.id);
       }
       for (auto *opInst : collector.storeOpInsts) {
         node.stores.push_back(opInst);
         auto memref = cast<AffineWriteOpInterface>(opInst).getMemRef();
         memrefAccesses[memref].insert(node.id);
       }
       forToNodeMap[&op] = node.id;
       nodes.insert({node.id, node});
     } else if (auto loadOp = dyn_cast<AffineReadOpInterface>(op)) {
       // Create graph node for top-level load op.
       Node node(nextNodeId++, &op);
       node.loads.push_back(&op);
       auto memref = cast<AffineReadOpInterface>(op).getMemRef();
       memrefAccesses[memref].insert(node.id);
       nodes.insert({node.id, node});
     } else if (auto storeOp = dyn_cast<AffineWriteOpInterface>(op)) {
       // Create graph node for top-level store op.
       Node node(nextNodeId++, &op);
       node.stores.push_back(&op);
       auto memref = cast<AffineWriteOpInterface>(op).getMemRef();
       memrefAccesses[memref].insert(node.id);
       nodes.insert({node.id, node});
     } else if (op.getNumRegions() != 0) {
       // Return false if another region is found (not currently supported).
       return false;
     } else if (op.getNumResults() > 0 && !op.use_empty()) {
       // Create graph node for top-level producer of SSA values, which
       // could be used by loop nest nodes.
       Node node(nextNodeId++, &op);
       nodes.insert({node.id, node});
     } else if (isa<CallOpInterface>(op)) {
       // Create graph node for top-level Call Op that takes any argument of
       // memref type. Call Op that returns one or more memref type results

mlir/test/Transforms/loop-fusion.mlir

func @should_fuse_no_top_level_access() {
   // CHECK-NEXT: return
   return
 }
 
 // -----
 
 #set0 = affine_set<(d0) : (1 == 0)>
 
-// CHECK-LABEL: func @should_not_fuse_if_inst_at_top_level() {
-func @should_not_fuse_if_inst_at_top_level() {
+// CHECK-LABEL: func @should_not_fuse_if_op_at_top_level() {
+func @should_not_fuse_if_op_at_top_level() {
   %m = memref.alloc() : memref<10xf32>
   %cf7 = constant 7.0 : f32
 
   affine.for %i0 = 0 to 10 {
     affine.store %cf7, %m[%i0] : memref<10xf32>
   }
   affine.for %i1 = 0 to 10 {
     %v0 = affine.load %m[%i1] : memref<10xf32>
func @should_not_fuse_if_op_at_top_level() {
   // CHECK-NEXT: }
   return
 }
 
 // -----
 
 #set0 = affine_set<(d0) : (1 == 0)>
 
-// CHECK-LABEL: func @should_not_fuse_if_inst_in_loop_nest() {
-func @should_not_fuse_if_inst_in_loop_nest() {
+// CHECK-LABEL: func @should_not_fuse_if_op_in_loop_nest() {
+func @should_not_fuse_if_op_in_loop_nest() {
   %m = memref.alloc() : memref<10xf32>
   %cf7 = constant 7.0 : f32
   %c4 = constant 4 : index
 
   affine.for %i0 = 0 to 10 {
     affine.store %cf7, %m[%i0] : memref<10xf32>
   }
   affine.for %i1 = 0 to 10 {
     affine.if #set0(%c4) {
     }
     %v0 = affine.load %m[%i1] : memref<10xf32>
   }
 
-  // IfOp in ForInst should prevent fusion.
+  // IfOp in ForOp should prevent fusion.
   // CHECK:      affine.for %{{.*}} = 0 to 10 {
   // CHECK-NEXT:   affine.store %{{.*}}, %{{.*}}[%{{.*}}] : memref<10xf32>
   // CHECK-NEXT: }
   // CHECK:      affine.for %{{.*}} = 0 to 10 {
   // CHECK-NEXT:   affine.if #set(%{{.*}}) {
   // CHECK-NEXT:   }
   // CHECK-NEXT:   affine.load %{{.*}}[%{{.*}}] : memref<10xf32>
   // CHECK-NEXT: }
   return
 }
 
 // -----
 
+#set = affine_set<(d0) : (d0 - 1 >= 0)>
+
+// CHECK-LABEL: func @should_fuse_if_op_in_loop_nest_not_sandwiched() -> memref<10xf32> {
+func @should_fuse_if_op_in_loop_nest_not_sandwiched() -> memref<10xf32> {
+  %a = memref.alloc() : memref<10xf32>
+  %b = memref.alloc() : memref<10xf32>
+  %cf7 = constant 7.0 : f32
+
+  affine.for %i0 = 0 to 10 {
+    affine.store %cf7, %a[%i0] : memref<10xf32>
+  }
+  affine.for %i1 = 0 to 10 {
+    %v0 = affine.load %a[%i1] : memref<10xf32>
+    affine.store %v0, %b[%i1] : memref<10xf32>
+  }
+  affine.for %i2 = 0 to 10 {
+    affine.if #set(%i2) {
+      %v0 = affine.load %b[%i2] : memref<10xf32>
+    }
+  }
+
+  // IfOp in ForOp should not prevent fusion if it does not in between the
+  // source and dest ForOp ops.
+
+  // CHECK:      affine.for
+  // CHECK-NEXT:   affine.store
+  // CHECK-NEXT:   affine.load
+  // CHECK-NEXT:   affine.store
+  // CHECK:      affine.for
+  // CHECK-NEXT:   affine.if
+  // CHECK-NEXT:     affine.load
+  // CHECK-NOT:  affine.for
+  // CHECK:      return
+
+  return %a : memref<10xf32>
+}
+
+// -----
+
+#set = affine_set<(d0) : (d0 - 1 >= 0)>
+
+// CHECK-LABEL: func @should_not_fuse_if_op_in_loop_nest_between_src_and_dest() -> memref<10xf32> {
+func @should_not_fuse_if_op_in_loop_nest_between_src_and_dest() -> memref<10xf32> {
+  %a = memref.alloc() : memref<10xf32>
+  %b = memref.alloc() : memref<10xf32>
+  %cf7 = constant 7.0 : f32
+
+  affine.for %i0 = 0 to 10 {
+    affine.store %cf7, %a[%i0] : memref<10xf32>
+  }
+  affine.for %i1 = 0 to 10 {
+    affine.if #set(%i1) {
+      affine.store %cf7, %a[%i1] : memref<10xf32>
+    }
+  }
+  affine.for %i3 = 0 to 10 {
+    %v0 = affine.load %a[%i3] : memref<10xf32>
+    affine.store %v0, %b[%i3] : memref<10xf32>
+  }
+  return %b : memref<10xf32>
+
+  // IfOp in ForOp which modifies the memref should prevent fusion if it is in
+  // between the source and dest ForOp.
+
+  // CHECK:      affine.for
+  // CHECK-NEXT:   affine.store
+  // CHECK:      affine.for
+  // CHECK-NEXT:   affine.if
+  // CHECK-NEXT:     affine.store
+  // CHECK:      affine.for
+  // CHECK-NEXT:   affine.load
+  // CHECK-NEXT:   affine.store
+  // CHECK:      return
+}
+
+// -----
+
 // CHECK-LABEL: func @permute_and_fuse() {
 func @permute_and_fuse() {
   %m = memref.alloc() : memref<10x20x30xf32>
 
   %cf7 = constant 7.0 : f32
   affine.for %i0 = 0 to 10 {
     affine.for %i1 = 0 to 20 {
       affine.for %i2 = 0 to 30 {
         affine.store %cf7, %m[%i0, %i1, %i2] : memref<10x20x30xf32>
       }
     }
   }
   affine.for %i3 = 0 to 30 {
     affine.for %i4 = 0 to 10 {
       affine.for %i5 = 0 to 20 {
         %v0 = affine.load %m[%i4, %i5, %i3] : memref<10x20x30xf32>
         "foo"(%v0) : (f32) -> ()
       }
     }
   }
 // CHECK:       affine.for %{{.*}} = 0 to 30 {
 // CHECK-NEXT:    affine.for %{{.*}} = 0 to 10 {
 // CHECK-NEXT:      affine.for %{{.*}} = 0 to 20 {
 // CHECK-NEXT:        affine.store %{{.*}}, %{{.*}}[0, 0, 0] : memref<1x1x1xf32>
 // CHECK-NEXT:        affine.load %{{.*}}[0, 0, 0] : memref<1x1x1xf32>
 // CHECK-NEXT:        "foo"(%{{.*}}) : (f32) -> ()
 // CHECK-NEXT:      }
 // CHECK-NEXT:    }