[Mlir-commits] [mlir] [MLIR] Fixing the memref linearization size computation (PR #138922)

[Zhuoran Yin]

================
@@ -75,18 +74,48 @@ std::pair<LinearizedMemRefInfo, OpFoldResult> getLinearizedMemRefOffsetAndSize(
     addMulMap = addMulMap + symbols[offsetIdx] * symbols[offsetIdx + 1];
     offsetValues[offsetIdx] = indicesVec[i];
     offsetValues[offsetIdx + 1] = strides[i];
-
-    mulMap = mulMap * symbols[i];
   }
 
   // Adjust linearizedIndices and size by the scale factor (dstBits / srcBits).
   int64_t scaler = dstBits / srcBits;
-  mulMap = mulMap.floorDiv(scaler);
+  size_t symbolIndex = 0;
+  SmallVector<Value> values;
+  SmallVector<AffineExpr> productExpressions;
+  for (unsigned i = 0; i < sourceRank; ++i) {
+    AffineExpr strideExpr, sizeExpr;
+    OpFoldResult stride = strides[i];
+    OpFoldResult size = sizes[i];
+    if (auto constantStride = getConstantIntValue(stride)) {
+      strideExpr = builder.getAffineConstantExpr(*constantStride);
+    } else {
+      strideExpr = symbols[symbolIndex++];
+      values.push_back(getValueOrCreateConstantIndexOp(builder, loc, stride));
+    }
+
+    if (auto constantSize = getConstantIntValue(size)) {
+      sizeExpr = builder.getAffineConstantExpr(*constantSize);
+    } else {
+      sizeExpr = symbols[symbolIndex++];
+      values.push_back(getValueOrCreateConstantIndexOp(builder, loc, size));
+    }
+
+    productExpressions.push_back((strideExpr * sizeExpr).floorDiv(scaler));
+  }
+  AffineMap maxMap = AffineMap::get(
+      /*dimCount=*/0, /*symbolCount=*/symbolIndex, productExpressions,
+      builder.getContext());
+
+  OpFoldResult linearizedSize;
+  Value totalSize =
+      builder.createOrFold<affine::AffineMaxOp>(loc, maxMap, values);
+  if (auto constantSize = getConstantIntValue(totalSize)) {
+    linearizedSize = builder.getIndexAttr(*constantSize);
+  } else {
+    linearizedSize = totalSize;
+  }
----------------
jerryyin wrote:

This suggestion has seriously make my life easier! Having seen how convenient `makeComposedFoldedAffineApply` was, I wanted to find this but somehow didn't find the corresponding one for affine::max. So it ended in the cumbersome logic you saw with the initial version of the PR. I'm very happy to get rid of all those thanks to both your and Krzysztof's code reviews!

https://github.com/llvm/llvm-project/pull/138922