[Mlir-commits] [mlir] [mlir][linalg] Fix tiling with constants in indexing maps (PR #173038)

[llvmlistbot at llvm.org]

================
@@ -436,17 +436,25 @@ static InitSliceInfo getInitSliceInfoForOuterReduction(
     ArrayRef<OpFoldResult> splitReductionIvs, AffineMap partialReductionMap) {
   int64_t initRank = partialReductionMap.getNumResults();
   SmallVector<OpFoldResult> initOffsets, initSizes;
-  Attribute zero = IntegerAttr::get(IndexType::get(context), 0);
-  Attribute one = IntegerAttr::get(IndexType::get(context), 1);
+  Type idxType = IndexType::get(context);
+  Attribute zero = IntegerAttr::get(idxType, 0);
+  Attribute one = IntegerAttr::get(idxType, 1);
   SmallVector<OpFoldResult> initStrides(initRank, one);
-  for (AffineExpr dimExpr : partialReductionMap.getResults()) {
-    unsigned dim = cast<AffineDimExpr>(dimExpr).getPosition();
-    if (reductionDims.contains(dim)) {
-      initOffsets.push_back(zero);
+  for (AffineExpr expr : partialReductionMap.getResults()) {
+    if (auto dimExpr = dyn_cast<AffineDimExpr>(expr)) {
+      unsigned dim = dimExpr.getPosition();
+      if (reductionDims.contains(dim)) {
+        initOffsets.push_back(zero);
+      } else {
+        initOffsets.push_back(offsets[dim]);
+      }
+      initSizes.push_back(sizes[dim]);
+    } else if (auto cstExpr = dyn_cast<AffineConstantExpr>(expr)) {
+      initOffsets.push_back(IntegerAttr::get(idxType, cstExpr.getValue()));
----------------
MaheshRavishankar wrote:

Please add a test here. I am not sure this is correct though. A test will help understand whats happening here better.

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