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

[Krzysztof Drewniak]

================
@@ -75,18 +74,70 @@ 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);
+
+  // If all strides and sizes are constant, we can compute the result
+  // directly without creating the AffineMaxOp.
+  int64_t constResult = 0;
+  int64_t constStride = 0;
+  int64_t constSize = 0;
+  bool isAllConstant = true;
+  for (unsigned i = 0; i < sourceRank; ++i) {
+    if (auto constantStride = getConstantIntValue(strides[i])) {
+      constStride = *constantStride;
+    } else {
+      isAllConstant = false;
+      break;
+    }
+    if (auto constantSize = getConstantIntValue(sizes[i])) {
+      constSize = *constantSize;
+    } else {
+      isAllConstant = false;
----------------
krzysz00 wrote:

Quick question - do we need this logic? `affine.max` could be `createOrFold<>`ed and thus folded to a constant if all its arguments are constants?

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