Handle tiling of rank-reduced insert slice with inner ranks elided. (#7263)

iree/compiler/Dialect/LinalgExt/IR/TiledOpInterface.cpp

diff --git a/iree/compiler/Dialect/LinalgExt/IR/TiledOpInterface.cpp b/iree/compiler/Dialect/LinalgExt/IR/TiledOpInterface.cpp
index 136fcb2..2cd8a14 100644
--- a/iree/compiler/Dialect/LinalgExt/IR/TiledOpInterface.cpp
+++ b/iree/compiler/Dialect/LinalgExt/IR/TiledOpInterface.cpp
@@ -105,7 +105,8 @@
       unsigned opOffsetIndex = opOffset.index();
       OpFoldResult opOffsetVal = opOffset.value();
       Optional<int64_t> opSizeVal = getConstantIntValue(opSizes[opOffsetIndex]);
-      if (opSizeVal && *opSizeVal == 1 && sourceShape[offsetIndex] != 1) {
+      if (offsetIndex >= sourceShape.size() ||
+          (opSizeVal && *opSizeVal == 1 && sourceShape[offsetIndex] != 1)) {
         resultOffsets[opOffsetIndex] = opOffsetVal;
         resultSizes[opOffsetIndex] = oneAttr;
         continue;

iree/compiler/Dialect/LinalgExt/Transforms/test/tiling.mlir

diff --git a/iree/compiler/Dialect/LinalgExt/Transforms/test/tiling.mlir b/iree/compiler/Dialect/LinalgExt/Transforms/test/tiling.mlir
index d439ca4..f632cf6 100644
--- a/iree/compiler/Dialect/LinalgExt/Transforms/test/tiling.mlir
+++ b/iree/compiler/Dialect/LinalgExt/Transforms/test/tiling.mlir
@@ -714,3 +714,34 @@
 //      CHECK:    scf.yield %[[INSERT]]
 //      CHECK:  return %[[RESULT]]
 
+
+// -----
+
+func @insert_slice_rank_reduced_inner(%arg0 : tensor<?xf32>,
+    %arg1 : tensor<?x?x?xf32>, %arg2: index, %arg3 : index, %arg4 : index) -> tensor<?x?x?xf32> {
+  %c0 = constant 0 : index
+  %d0 = tensor.dim %arg0, %c0 : tensor<?xf32>
+  %0 = tensor.insert_slice %arg0 into %arg1[%arg2, %arg3, %arg4] [1, %d0, 1] [1, 1, 1]
+      {__internal_linalg_transform__ = "tiling_input"} : tensor<?xf32> into tensor<?x?x?xf32>
+  return %0 : tensor<?x?x?xf32>
+}
+//  CHECK-DAG: #[[MAP0:.+]] = affine_map<(d0)[s0, s1] -> (10, -d0 + s1)>
+//  CHECK-DAG: #[[MAP1:.+]] = affine_map<(d0)[s0] -> (d0 + s0)>
+//      CHECK: func @insert_slice_rank_reduced_inner(
+// CHECK-SAME:     %[[ARG0:.+]]: tensor<?xf32>
+// CHECK-SAME:     %[[ARG1:.+]]: tensor<?x?x?xf32>
+// CHECK-SAME:     %[[ARG2:[a-zA-Z0-9_]+]]: index
+// CHECK-SAME:     %[[ARG3:[a-zA-Z0-9_]+]]: index
+// CHECK-SAME:     %[[ARG4:[a-zA-Z0-9_]+]]: index
+//  CHECK-DAG:   %[[LB:.+]] = constant 0 : index
+//  CHECK-DAG:   %[[STEP:.+]] = constant 10 : index
+//      CHECK:   %[[UB:.+]] = tensor.dim %[[ARG0]]
+//      CHECK:   %[[RESULT:.+]] = scf.for %[[IV0:[a-zA-Z0-9_]+]] = %[[LB]]
+// CHECK-SAME:       to %[[D0]] step %[[STEP]] iter_args(%[[ARG6:.+]] = %[[ARG1]])
+//      CHECK:     %[[TILESIZE:.+]] = affine.min #[[MAP0]](%[[ARG5]])[%[[STEP]], %[[UB]]]
+//      CHECK:     %[[SLICE:.+]] = tensor.extract_slice %[[ARG0]][%[[IV0]]] [%[[TILESIZE]]]
+//      CHECK:     %[[APPLY:.+]] = affine.apply #[[MAP1]](%[[IV0]])[%[[ARG3]]]
+//      CHECK:     %[[YIELD:.+]] = tensor.insert_slice %[[SLICE]] into %[[ARG6]]
+// CHECK-SAME:         [%[[ARG2]], %[[APPLY]], %[[ARG4]]] [1, %[[TILESIZE]], 1]
+//      CHECK:     scf.yield %[[YIELD]]
+//      CHECK:   return %[[RESULT]]