[mlir][affine] fix loop fusion not fusing loops with non-constant loop bounds

[anoopjs]

Summary

Loop fusion fails for loops with non-constant bounds. Addresses the issue mentioned here https://github.com/llvm/llvm-project/issues/61784

[anoopjs]

fix indentation

[bondhugula]

Can you edit the commit summary to include the link to the github issue that this is fixing? https://github.com/llvm/llvm-project/issues/61784

[anoopjs]

In D148559#4286002, @bondhugula wrote:

Can you edit the commit summary to include the link to the github issue that this is fixing? https://github.com/llvm/llvm-project/issues/61784

Done

[bondhugula]

Sorry for the delayed response; some comments/questions below.

mlir/lib/Dialect/Affine/Analysis/Utils.cpp

710

The set difference utility can now handle this? Can you expand the commit summary?

732

Need a code comment.

mlir/test/Transforms/loop-fusion-4.mlir

203–212

A single test case here is insufficient. Can you add one where the bounds don't match, i.e., different symbols, or a symbol which is effectively the other symbol + 1.

mlir/lib/Dialect/Affine/Analysis/Utils.cpp

std::optional<bool> ComputationSliceState::isSliceValid() {
 
   // Create constraints for the source loop nest using which slice is computed.
   FlatAffineValueConstraints srcConstraints;
   // TODO: Store the source's domain to avoid computation at each depth.
   if (failed(getSourceAsConstraints(srcConstraints))) {
     LLVM_DEBUG(llvm::dbgs() << "Unable to compute source's domain\n");
     return std::nullopt;
   }
-  // As the set difference utility currently cannot handle symbols in its
-  // operands, validity of the slice cannot be determined.
-  if (srcConstraints.getNumSymbolVars() > 0) {
-    LLVM_DEBUG(llvm::dbgs() << "Cannot handle symbols in source domain\n");
-    return std::nullopt;
-  }
+
   // TODO: Handle local vars in the source domains while using the 'projectOut'
   // utility below. Currently, aligning is not done assuming that there will be
   // no local vars in the source domain.
   if (srcConstraints.getNumLocalVars() != 0) {
     LLVM_DEBUG(llvm::dbgs() << "Cannot handle locals in source domain\n");
     return std::nullopt;
   }
 
   // Create constraints for the slice loop nest that would be created if the
   // fusion succeeds.
   FlatAffineValueConstraints sliceConstraints;
   if (failed(getAsConstraints(&sliceConstraints))) {
     LLVM_DEBUG(llvm::dbgs() << "Unable to compute slice's domain\n");
     return std::nullopt;
   }
 
   // Projecting out every dimension other than the 'ivs' to express slice's
   // domain completely in terms of source's IVs.
   sliceConstraints.projectOut(ivs.size(),
                               sliceConstraints.getNumVars() - ivs.size());
 
+  srcConstraints.projectOut(ivs.size(),
+                            srcConstraints.getNumVars() - ivs.size());
+
   LLVM_DEBUG(llvm::dbgs() << "Domain of the source of the slice:\n");
   LLVM_DEBUG(srcConstraints.dump());
   LLVM_DEBUG(llvm::dbgs() << "Domain of the slice if this fusion succeeds "
                              "(expressed in terms of its source's IVs):\n");
   LLVM_DEBUG(sliceConstraints.dump());
 
   // TODO: Store 'srcSet' to avoid recalculating for each depth.
   PresburgerSet srcSet(srcConstraints);

mlir/test/Transforms/loop-fusion-4.mlir

 // RUN: mlir-opt -allow-unregistered-dialect %s -pass-pipeline='builtin.module(func.func(affine-loop-fusion{mode=producer}))' -split-input-file | FileCheck %s --check-prefix=PRODUCER-CONSUMER
 // RUN: mlir-opt -allow-unregistered-dialect %s -pass-pipeline='builtin.module(func.func(affine-loop-fusion{fusion-maximal mode=sibling}))' -split-input-file | FileCheck %s --check-prefix=SIBLING-MAXIMAL
+// RUN: mlir-opt -allow-unregistered-dialect %s -pass-pipeline='builtin.module(func.func(affine-loop-fusion{fusion-maximal mode=producer}))' -split-input-file | FileCheck %s --check-prefix=MAXIMAL
 
 // Part I of fusion tests in  mlir/test/Transforms/loop-fusion.mlir.
 // Part II of fusion tests in mlir/test/Transforms/loop-fusion-2.mlir
 // Part III of fusion tests in mlir/test/Transforms/loop-fusion-3.mlir
 
 // Expects fusion of producer into consumer at depth 4 and subsequent removal of
 // source loop.
 // PRODUCER-CONSUMER-LABEL: func @unflatten4d
affine.for %i1 = 0 to 10 {
     %v0 = affine.load %m[%i1] : memref<10xf32>
   }
   // PRODUCER-CONSUMER:      affine.for %{{.*}} = 0 to 10 {
   // PRODUCER-CONSUMER-NEXT:   affine.store %{{.*}}, %{{.*}}[0] : memref<1xf32>
   // PRODUCER-CONSUMER-NEXT:   affine.load %{{.*}}[0] : memref<1xf32>
   // PRODUCER-CONSUMER-NEXT: }
   return
 }
+
+// -----
+
+// MAXIMAL-LABEL func @fusion@nonconstant_bound
+func.func @fusion_nonconstant_bound(%arg0: memref<?xf32>) {
+  %cst = arith.constant 1.000000e+00 : f32
+  %cst_0 = arith.constant 2.000000e+00 : f32
+  %c0 = arith.constant 0 : index
+  %dim = memref.dim %arg0, %c0 : memref<?xf32>
+  affine.for %arg1 = 0 to %dim {
+    %0 = affine.load %arg0[%arg1] : memref<?xf32>
+    %1 = arith.addf %0, %cst : f32
+    affine.store %1, %arg0[%arg1] : memref<?xf32>
+  }
+  affine.for %arg1 = 0 to %dim {
+    %0 = affine.load %arg0[%arg1] : memref<?xf32>
+    %1 = arith.addf %0, %cst_0 : f32
+    affine.store %1, %arg0[%arg1] : memref<?xf32>
+  }
+  // MAXIMAL:      affine.for %[[idx:.*]] = 0 to %{{.*}} {
+  // MAXIMAL-NEXT:   affine.load %[[arr:.*]][%[[idx]]] : memref<?xf32>
+  // MAXIMAL-NEXT:   arith.addf
+  // MAXIMAL-NEXT:   affine.store %{{.*}}, %[[arr]][%[[idx]]] : memref<?xf32>
+  // MAXIMAL-NEXT:   affine.load %[[arr:.*]][%[[idx]]] : memref<?xf32>
+  // MAXIMAL-NEXT:   arith.addf
+  // MAXIMAL-NEXT:   affine.store %{{.*}}, %[[arr]][%[[idx]]] : memref<?xf32>
+  // MAXIMAL-NEXT: }
+  return
+}