[mlir][affine] Normalize constant valued bound loop

[Lewuathe]

Summary

This change aims to resolve the issue reported in https://github.com/llvm/llvm-project/issues/59994.

After calling AffineForOp#setUpperBound and setLowerBound, it makes the original affine map structure inconsistent. It's necessary to keep the original map to build the new induction vars.

[Lewuathe]

Remove unused debug code.

[bondhugula]

Thank you for reporting and attempting to fix https://github.com/llvm/llvm-project/issues/59994. Some comments.

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

568–569

Use the range-based ctor or use llvm::append_range to write this more compactly.

573

Likewise.

579–580

You can move these lines up closer to declaration to keep related things together or move the declaration below.

Group the canonicalize calls below with the code creating origLb/UbOperands.

582–583

origLb/UbMap aren't good names since they aren't really the original maps anymore.

[bondhugula]

mlir/test/Dialect/Affine/affine-loop-normalize.mlir

216

Drop extra line.

[Lewuathe]

Cleanup the code based on the review comments.

[Lewuathe]

Apply clang-format.

[bondhugula]

Comments on test cases.

mlir/test/Dialect/Affine/affine-loop-normalize.mlir

226

#map -> #map{{.*}} or #{{.*}}
Avoid map in CHECK line since its suffix would change when the test case is updated.

239

Both of these test cases really don't exercise all the changes. The fact that the other test cases didn't have to change indicates those aren't impacted by the canonicalization added. Can you include a couple more test cases where the bounds have maps + operands (with symbols) that would get simplified. Here is one.

affine.for %i = 0 to affine_map<(d0, d1) -> (d0)>(%N, %M) step 4 {
   ...
}

(%N, %M) can be func arguments. d1 will be dropped, d0 will become s0.

239

Does this test case exercise anything more than the first one?

240

Capturing IV is unnecessary, you aren't using it.

[Lewuathe]

mlir/test/Dialect/Affine/affine-loop-normalize.mlir

239

The first one is identical to the case in the issue ticket. The second one aimed to just focus on the canonicalization of a constant value in affine.for but we may not need this one.

I'll remove the second one.

[Lewuathe]

Add test cases.

[bondhugula]

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

588–592

How are we ensuring that in newUbExprs, the symbols appear in this order? I couldn't immediately see this.

Overall, I feel adding this canonicalization is unnecessarily complicating this code. From the commit summary, it isn't clear why canonicalization is needed here on the fly. s0 - s0 will simplify to 0 by construction. Something being zero can be readily checked if needed.

598

Please move this line to right above 619.

mlir/test/Dialect/Affine/affine-loop-normalize.mlir

239

Remove numbered values from CHECK lines. Either %arg{{.*}} or capture and match.

295

Drop extra new line.

[Lewuathe]

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

588–592

When we use arith.constant value for the lower and upper bound, it's translated into the affine map ()[s0] -> (s0) irrelevant to the actual value. And it always induces 0 in L.613 ~ L.616 as origUbExprs[i] and origLbExprs[0] are same.

for (unsigned i = 0, e = origUbExprs.size(); i < e; ++i) {
  newUbExprs.push_back(
      (origUbExprs[i] - origLbExprs[0]).ceilDiv(origLoopStep));
}

Without this canonicalization, the first test case here

func.func @constant_lower_bound() {
  %c0 = arith.constant 0 : index
  %c1 = arith.constant 1 : index
  scf.for %j = %c0 to %c1 step %c1 {
    affine.for %i = %c0 to %c1 {
    }
  }
  return
}

went to 0 to 0 as follows.

func.func @constant_lower_bound() {
  %c0 = arith.constant 0 : index
  %c1 = arith.constant 1 : index
  scf.for %j = %c0 to %c1 step %c1 {
    affine.for %i = 0 to 0 {
    }
  }
  return
}

So I thought we did not care about the constant valued bounds properly.

598

We need to use this map in the coming lines L.613 - L.616. So do we need to keep this position?

[bondhugula]

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

588–592

It appears that canonicalization isn't the right fix for the bug you are fixing -- the bug would still be there in other cases I believe when you don't have constants. Also, we shouldn't need to canonicalize bounds, especially if it makes the code complex here. The -canonicalize pass takes care of it.

For eg., does the following work:

affine.for %i = %M to %N

[Lewuathe]

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

588–592

I see. Thinking the case again, we do not need to make sure to canonicalize the bound if we can assume -canonicalize pass. I'm going to apply only the fix for the issue reported originally.

[Lewuathe]

Stop canonicalization in the middle of normalization pass.

[Lewuathe]

Rename variables.

[bondhugula]

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

588–592

Thanks - this would be the right approach instead of combining orthogonal things that are clouding this patch. I'll take another look tomorrow.

[bondhugula]

Can you please update the commit summary? It's no longer reflective of the change.

[Lewuathe]

@bondhugula Thanks. I updated the commit message accordingly.

[bondhugula]

LGTM - thanks.

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

616–619

You can do: opBuilder.getAffineMap(...). (No need of context arg then.)

[Lewuathe]

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

616–619

Thanks for letting me know. But that utility is not prepared yet. OpBuilder seems to have some methods for building purpose specific AffineMaps (e.g. getEmptyAffineMap )

/llvm-project/mlir/lib/Dialect/Affine/Utils/Utils.cpp:613:34: error: no member named 'getAffineMap' in 'mlir::OpBuilder'
  AffineMap newUbMap = opBuilder.getAffineMap(

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

LogicalResult mlir::normalizeAffineFor(AffineForOp op, bool promoteSingleIter) {
   // then skip this op.
   if (op.getLowerBoundMap().getNumResults() != 1)
     return failure();
 
   Location loc = op.getLoc();
   OpBuilder opBuilder(op);
   int64_t origLoopStep = op.getStep();
 
-  // Calculate upperBound for normalized loop.
-  SmallVector<Value, 4> ubOperands;
   AffineBound lb = op.getLowerBound();
+  AffineMap originalLbMap = lb.getMap();
+  SmallVector<Value, 4> origLbOperands;
+  llvm::append_range(origLbOperands, lb.getOperands());
+
   AffineBound ub = op.getUpperBound();
+  AffineMap originalUbMap = ub.getMap();
+  SmallVector<Value, 4> origUbOperands;
+  llvm::append_range(origUbOperands, ub.getOperands());
+
+  // Calculate upperBound for normalized loop.
+  SmallVector<Value, 4> ubOperands;
   ubOperands.reserve(ub.getNumOperands() + lb.getNumOperands());
-  AffineMap origLbMap = lb.getMap();
-  AffineMap origUbMap = ub.getMap();
 
   // Add dimension operands from upper/lower bound.
-  for (unsigned j = 0, e = origUbMap.getNumDims(); j < e; ++j)
+  for (unsigned j = 0, e = originalUbMap.getNumDims(); j < e; ++j)
     ubOperands.push_back(ub.getOperand(j));
-  for (unsigned j = 0, e = origLbMap.getNumDims(); j < e; ++j)
+  for (unsigned j = 0, e = originalLbMap.getNumDims(); j < e; ++j)
     ubOperands.push_back(lb.getOperand(j));
 
   // Add symbol operands from upper/lower bound.
-  for (unsigned j = 0, e = origUbMap.getNumSymbols(); j < e; ++j)
-    ubOperands.push_back(ub.getOperand(origUbMap.getNumDims() + j));
-  for (unsigned j = 0, e = origLbMap.getNumSymbols(); j < e; ++j)
-    ubOperands.push_back(lb.getOperand(origLbMap.getNumDims() + j));
+  for (unsigned j = 0, e = originalUbMap.getNumSymbols(); j < e; ++j)
+    ubOperands.push_back(ub.getOperand(originalUbMap.getNumDims() + j));
+  for (unsigned j = 0, e = originalLbMap.getNumSymbols(); j < e; ++j)
+    ubOperands.push_back(lb.getOperand(originalLbMap.getNumDims() + j));
 
   // Add original result expressions from lower/upper bound map.
-  SmallVector<AffineExpr, 1> origLbExprs(origLbMap.getResults().begin(),
-                                         origLbMap.getResults().end());
-  SmallVector<AffineExpr, 2> origUbExprs(origUbMap.getResults().begin(),
-                                         origUbMap.getResults().end());
+  SmallVector<AffineExpr, 1> origLbExprs(originalLbMap.getResults().begin(),
+                                         originalLbMap.getResults().end());
+  SmallVector<AffineExpr, 2> origUbExprs(originalUbMap.getResults().begin(),
+                                         originalUbMap.getResults().end());
   SmallVector<AffineExpr, 4> newUbExprs;
 
   // The original upperBound can have more than one result. For the new
   // upperBound of this loop, take difference of all possible combinations of
   // the ub results and lb result and ceildiv with the loop step. For e.g.,
   //
   //  affine.for %i1 = 0 to min affine_map<(d0)[] -> (d0 + 32, 1024)>(%i0)
   //  will have an upperBound map as,
   //  affine_map<(d0)[] -> (((d0 + 32) - 0) ceildiv 1, (1024 - 0) ceildiv
   //  1)>(%i0)
   //
   // Insert all combinations of upper/lower bound results.
   for (unsigned i = 0, e = origUbExprs.size(); i < e; ++i) {
     newUbExprs.push_back(
         (origUbExprs[i] - origLbExprs[0]).ceilDiv(origLoopStep));
   }
 
   // Construct newUbMap.
-  AffineMap newUbMap =
-      AffineMap::get(origLbMap.getNumDims() + origUbMap.getNumDims(),
-                     origLbMap.getNumSymbols() + origUbMap.getNumSymbols(),
-                     newUbExprs, opBuilder.getContext());
+  AffineMap newUbMap = AffineMap::get(
+      originalLbMap.getNumDims() + originalUbMap.getNumDims(),
+      originalLbMap.getNumSymbols() + originalUbMap.getNumSymbols(), newUbExprs,
+      opBuilder.getContext());
   canonicalizeMapAndOperands(&newUbMap, &ubOperands);
 
   SmallVector<Value, 4> lbOperands(lb.getOperands().begin(),
                                    lb.getOperands().begin() +
-                                       lb.getMap().getNumDims());
+                                       originalLbMap.getNumDims());
 
   // Normalize the loop.
   op.setUpperBound(ubOperands, newUbMap);
   op.setLowerBound({}, opBuilder.getConstantAffineMap(0));
   op.setStep(1);
 
   // Calculate the Value of new loopIV. Create affine.apply for the value of
   // the loopIV in normalized loop.
   opBuilder.setInsertionPointToStart(op.getBody());
   // Add an extra dim operand for loopIV.
   lbOperands.push_back(op.getInductionVar());
   // Add symbol operands from lower bound.
-  for (unsigned j = 0, e = origLbMap.getNumSymbols(); j < e; ++j)
-    lbOperands.push_back(lb.getOperand(origLbMap.getNumDims() + j));
+  for (unsigned j = 0, e = originalLbMap.getNumSymbols(); j < e; ++j)
+    lbOperands.push_back(origLbOperands[originalLbMap.getNumDims() + j]);
 
-  AffineExpr origIVExpr = opBuilder.getAffineDimExpr(lb.getMap().getNumDims());
-  AffineExpr newIVExpr = origIVExpr * origLoopStep + origLbMap.getResult(0);
-  AffineMap ivMap = AffineMap::get(origLbMap.getNumDims() + 1,
-                                   origLbMap.getNumSymbols(), newIVExpr);
+  AffineExpr origIVExpr =
+      opBuilder.getAffineDimExpr(originalLbMap.getNumDims());
+  AffineExpr newIVExpr = origIVExpr * origLoopStep + originalLbMap.getResult(0);
+  AffineMap ivMap = AffineMap::get(originalLbMap.getNumDims() + 1,
+                                   originalLbMap.getNumSymbols(), newIVExpr);
   canonicalizeMapAndOperands(&ivMap, &lbOperands);
   Operation *newIV = opBuilder.create<AffineApplyOp>(loc, ivMap, lbOperands);
   op.getInductionVar().replaceAllUsesExcept(newIV->getResult(0), newIV);
   return success();
 }
 
 /// Returns true if the memory operation of `destAccess` depends on `srcAccess`
 /// inside of the innermost common surrounding affine loop between the two

mlir/test/Dialect/Affine/affine-loop-normalize.mlir

affine.for %arg1 = 0 to %4 step 32 {
             }
           }
         }
       }
     }
   }
   return
 }
+
+// -----
+
+// CHECK-LABEL: func @constant_lower_bound
+func.func @constant_lower_bound() {
+  %c0 = arith.constant 0 : index
+  %c1 = arith.constant 1 : index
+  scf.for %j = %c0 to %c1 step %c1 {
+    // CHECK: affine.for %[[ARG0:.*]] =
+    affine.for %i = %c0 to %c1 {
+      // CHECK-NEXT: %{{.*}} = affine.apply #map{{.*}}(%[[ARG0]])
+    }
+  }
+  return
+}
+
+// -----
+
+// CHECK-DAG: [[$UB_MAP:#map[0-9]*]] = affine_map<()[s0] -> (s0 ceildiv 4)>
+// CHECK-DAG: [[$IV_MAP:#map[0-9]*]] = affine_map<(d0) -> (d0 * 4)>
+
+// CHECK-LABEL: func @upper_bound_by_symbol
+func.func @upper_bound_by_symbol(%arg0: index, %arg1: index) {
+  // CHECK: affine.for %[[ARG0:.*]] = 0 to [[$UB_MAP]]()[%arg{{.*}}] {
+  affine.for %i = 0 to affine_map<()[s0, s1] -> (s0)>()[%arg0, %arg1] step 4 {
+    // CHECK-NEXT: %[[IV:.*]] = affine.apply [[$IV_MAP]](%[[ARG0]])
+    // CHECK-NEXT: "test.foo"(%[[IV]]) : (index) -> ()
+    "test.foo"(%i) : (index) -> ()
+  }
+  return
+}
+
+// -----
+
+// CHECK-DAG: [[$UB_MAP:#map[0-9]*]] = affine_map<()[s0] -> ((-s0 + 10) ceildiv 4)>
+// CHECK-DAG: [[$IV_MAP:#map[0-9]*]] = affine_map<(d0)[s0] -> (d0 * 4 + s0)>
+
+// CHECK-LABEL: func @lower_bound_by_symbol
+func.func @lower_bound_by_symbol(%arg0: index, %arg1: index) {
+  // CHECK: affine.for %[[ARG0:.*]] = 0 to [[$UB_MAP]]()[%arg{{.*}}] {
+  affine.for %i = affine_map<()[s0, s1] -> (s0)>()[%arg0, %arg1] to 10 step 4 {
+    // CHECK-NEXT: %[[IV:.*]] = affine.apply [[$IV_MAP]](%[[ARG0]])[%arg{{.*}}]
+    // CHECK-NEXT: "test.foo"(%[[IV]]) : (index) -> ()
+    "test.foo"(%i) : (index) -> ()
+  }
+  return
+}
+
+// -----
+
+// CHECK-DAG: [[$UB_MAP:#map[0-9]*]] = affine_map<()[s0] -> (s0 ceildiv 4)>
+// CHECK-DAG: [[$IV_MAP:#map[0-9]*]] = affine_map<(d0) -> (d0 * 4)>
+
+// CHECK-LABEL: func @upper_bound_by_dim
+func.func @upper_bound_by_dim(%arg0: index, %arg1: index) {
+  // CHECK: affine.for %[[ARG0:.*]] = 0 to [[$UB_MAP]]()[%arg{{.*}}] {
+  affine.for %i = 0 to affine_map<(d0, d1) -> (d0)>(%arg0, %arg1) step 4 {
+    // CHECK-NEXT: %[[IV:.*]] = affine.apply [[$IV_MAP]](%[[ARG0]])
+    // CHECK-NEXT: "test.foo"(%[[IV]]) : (index) -> ()
+    "test.foo"(%i) : (index) -> ()
+  }
+  return
+}
+
+// -----
+
+// CHECK-DAG: [[$UB_MAP:#map[0-9]*]] = affine_map<()[s0] -> ((-s0 + 10) ceildiv 4)>
+// CHECK-DAG: [[$IV_MAP:#map[0-9]*]] = affine_map<(d0)[s0] -> (d0 * 4 + s0)>
+
+// CHECK-LABEL: func @upper_bound_by_dim
+func.func @upper_bound_by_dim(%arg0: index, %arg1: index) {
+  // CHECK: affine.for %[[ARG0:.*]] = 0 to [[$UB_MAP]]()[%arg{{.*}}] {
+  affine.for %i = affine_map<(d0, d1) -> (d0)>(%arg0, %arg1) to 10 step 4 {
+    // CHECK-NEXT: %[[IV:.*]] = affine.apply [[$IV_MAP]](%[[ARG0]])[%arg{{.*}}]
+    // CHECK-NEXT: "test.foo"(%[[IV]]) : (index) -> ()
+    "test.foo"(%i) : (index) -> ()
+  }
+  return
+}