[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()));
+      initSizes.push_back(one);
     } else {
-      initOffsets.push_back(offsets[dim]);
+      llvm_unreachable("Unsupported affine expression type");
----------------
MaheshRavishankar wrote:

I think using an `llvm_unreachable` here is wrong, because this code is very much reachable. I see that I have missed checking for the `AffineMap` being a projected permutation before we get to this point. The options are we check very early on (like here : https://github.com/llvm/llvm-project/blob/125a53ce599dbdbb6eb9f24db1f266fd700561f0/mlir/lib/Dialect/Linalg/Transforms/TilingInterfaceImpl.cpp#L567) for the shape of affine map supported, or we make this method return a `FailureOr<InitSliceInfo>` to gracefully propagate the error. This issue is pre-existing, but is also kind of part of this change. We probably want to fix that.

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