Fix affine expr flattener bug introduced by cl/225452174.

- inconsistent local var constraint size when repeatedly using the same
  flattener for all expressions in a map.

PiperOrigin-RevId: 227067836

mlir/lib/Analysis/AffineAnalysis.cpp

@@ -128,7 +128,7 @@ public:
   // In future, consider adding a prepass to determine how big the SmallVector's
   // will be, and linearize this to std::vector<int64_t> to prevent
   // SmallVector moves on re-allocation.
-  std::vector<SmallVector<int64_t, 32>> operandExprStack;
+  std::vector<SmallVector<int64_t, 8>> operandExprStack;
   // Constraints connecting newly introduced local variables (for mod's and
   // div's) to existing (dimensional and symbolic) ones. These are always
   // inequalities.
@@ -426,16 +426,11 @@ static bool getFlattenedAffineExprs(
       return false;
 
     flattener.walkPostOrder(expr);
-
-    SmallVector<int64_t, 8> flattenedExpr;
-    flattenedExpr.reserve(flattener.numDims + flattener.numSymbols +
-                          flattener.numLocals + 1);
-    for (auto v : flattener.operandExprStack.back()) {
-      flattenedExpr.push_back(v);
-    }
-    flattenedExprs->push_back(flattenedExpr);
-    flattener.operandExprStack.pop_back();
   }
+
+  flattenedExprs->insert(flattenedExprs->end(),
+                         flattener.operandExprStack.begin(),
+                         flattener.operandExprStack.end());
   if (localVarCst)
     localVarCst->clearAndCopyFrom(flattener.localVarCst);
 
@@ -861,7 +856,7 @@ addMemRefAccessConstraints(const AffineValueMap &srcAccessMap,
   // Add equality constraints for any dst symbols defined by constant ops.
   addEqForConstOperands(dstOperands);
 
-  // TODO(bondhugula): add srcLocalVarCst, destLocalVarCst to the dependence
+  // TODO(b/122081337): add srcLocalVarCst, destLocalVarCst to the dependence
   // domain.
   return true;
 }

mlir/test/Transforms/memref-dependence-check.mlir

@@ -676,3 +676,26 @@ mlfunc @loop_nest_depth() {
   }
   return
 }
+
+// -----
+// Test case to exercise sanity when flattening multiple expressions involving
+// mod/div's successively.
+// CHECK-LABEL: mlfunc @mod_div_3d() {
+mlfunc @mod_div_3d() {
+  %M = alloc() : memref<2x2x2xi32>
+  %c0 = constant 0 : i32
+  for %i0 = 0 to 8 {
+    for %i1 = 0 to 8 {
+      for %i2 = 0 to 8 {
+        %idx = affine_apply (d0, d1, d2) -> (d0 floordiv 4, d1 mod 2, d2 floordiv 4) (%i0, %i1, %i2)
+        // Dependences below are conservative due to TODO(b/122081337).
+        store %c0, %M[%idx#0, %idx#1, %idx#2] : memref<2 x 2 x 2 x i32>
+        // expected-note@-1 {{dependence from 0 to 0 at depth 1 = [1, 7][-7, 7][-7, 7]}}
+        // expected-note@-2 {{dependence from 0 to 0 at depth 2 = [0, 0][2, 7][-7, 7]}}
+        // expected-note@-3 {{dependence from 0 to 0 at depth 3 = [0, 0][0, 0][1, 7]}}
+        // expected-note@-4 {{dependence from 0 to 0 at depth 4 = false}}
+      }
+    }
+  }
+  return
+}