Convert towards a proper fix

Author: josel-amd

mlir/lib/Dialect/Linalg/Utils/Utils.cpp

@@ -798,15 +798,8 @@ computeAllSliceParameters(OpBuilder &builder, Location loc, LinalgOp linalgOp,
     // transformations such as padding and bufferization since the
     // extract/insert slice pairs make the accessed iteration argument
     // subdomains explicit.
-    SmallVector<int64_t> domainSizes;
-    // FIXME: tileToPartialReduction adds the new init tensor to the output
-    // but doesn't update the indexing type of the index map causing a crash.
-    // isAllParallelLoops
-    if (linalgOp.getNumParallelLoops() == linalgOp.getNumLoops()) {
-      domainSizes = linalgOp.getStaticLoopRanges();
-    }
     Type operandType = opOperand.get().getType();
-    if (!isTiled(map, tileSizes, domainSizes) &&
+    if (!isTiled(map, tileSizes, linalgOp.getStaticLoopRanges()) &&
         !(isa<RankedTensorType>(operandType) &&
           linalgOp.isDpsInit(&opOperand))) {
       allSliceParams.push_back(std::nullopt);
@@ -818,7 +811,7 @@ computeAllSliceParameters(OpBuilder &builder, Location loc, LinalgOp linalgOp,
 
     allSliceParams.push_back(computeSliceParameters(
         builder, loc, shapedOp, tileSizes, map, lbs, sizeBounds, subShapeSizes,
-        omitPartialTileCheck, domainSizes));
+        omitPartialTileCheck, linalgOp.getStaticLoopRanges()));
   }
 
   return allSliceParams;

mlir/lib/Dialect/SCF/Transforms/TileUsingInterface.cpp

@@ -984,8 +984,7 @@ mlir::scf::tileReductionUsingScf(RewriterBase &b,
 
     // 4a. Clone the operation.
     {
-      auto clonedOp = cast<PartialReductionOpInterface>(
-          cloneOpAndUpdateDestinationArgs(b, op, regionIterArgs));
+      auto clonedOp = cast<PartialReductionOpInterface>(rewriter.clone(*op));
 
       // 4b. Tile the cloned operation.
       FailureOr<TilingResult> partialTilingResult =

mlir/test/Dialect/Linalg/tile-tensors.mlir

@@ -170,16 +170,16 @@ module attributes {transform.with_named_sequence} {
 
 // -----
 
-
-// CHECK-LABEL: func @generic_op_tensors
-//  CHECK-SAME:   %[[ARG0:[a-zA-Z0-9_]+]]: tensor<?x?x?xf32>
-//  CHECK-SAME:   %[[ARG1:[a-zA-Z0-9_]+]]: tensor<?x?x?xf32>
+// CHECK-LABEL: func @non_monotonic_affine_expr
+//  CHECK-SAME:   %[[ARG0:[a-zA-Z0-9_]+]]: tensor<7xf32>
 func.func @non_monotonic_affine_expr(%arg0 : tensor<7xf32>) -> tensor<7xf32> {
   %c0 = arith.constant 0 : index
   %0 = tensor.dim %arg0, %c0 : tensor<7xf32>
   %empty = tensor.empty() : tensor<7xf32>
-  // FIXME: Do proper testing
-  // CHECK: tensor.extract_slice %arg2[0] [7] [1] : tensor<7xf32> to tensor<7xf32>
+
+  // CHECK: %[[OUT:.*]] = tensor.empty() : tensor<7xf32>
+  // CHECK: scf.for {{.*}} to {{.*}} step {{.*}} iter_args(%[[TC0:.*]] = %[[OUT]]) -> (tensor<7xf32>) {
+  // CHECK: tensor.extract_slice %[[TC0]][0] [7] [1] : tensor<7xf32> to tensor<7xf32>
   %generic = linalg.generic
     {indexing_maps = [affine_map<(d0) -> (d0 mod 4)>,
                       affine_map<(d0) -> (d0)>],