[mlir][Tensor] Fix ExtractSliceFromReshape transform edge case

[christopherbate]

Summary

The transformation would fail if none of the sliced dimensions were
linearized by the producing tensor.collapse_shape. This is a trivial
edge case but it wasn't correctly tested. Fixes the issue and adds a test.

[nicolasvasilache]

Thanks Christopher!

mlir/include/mlir/Dialect/Utils/ReshapeOpsUtils.h

public:
   /// signature can be described by `(D_0, D_1,.. D_{n-1}) -> (offsets, sizes,
   /// strides)` where `n` the number of "tiled dimensions", which are the
   /// dimensions of the output that are linearized by the collapse shape op and
   /// are also sliced. Each `D_i` is a tuple that must represent a valid
   /// multi-index for the `i-th` tiled dimension. In the example above, there is
   /// only one tiled dimension (D_0) and `arith.delinearize_index` produces the
   /// multi-index (%3) that would be passed to this function to generate the
   /// parameters for the `tensor.extract_slice` op (%4).
-  SmallVector<Range> getExtractSliceParams(ArrayRef<ValueRange> multiIndices);
+  SmallVector<Range> getExtractSliceParams(MLIRContext *ctx,
+                                           ArrayRef<ValueRange> multiIndices);
 
   /// This function takes indices in the index space of the "tiled dimensions"
   /// described above and returns a set of Range variables that describe how the
   /// slice should be inserted into the destination. In the example above, `%iv`
   /// would be passed to this function to generate the parameters for the
   /// `tensor.insert_slice` op producing %6.
-  SmallVector<Range> getInsertSliceParams(ValueRange tileIndices);
+  SmallVector<Range> getInsertSliceParams(MLIRContext *ctx,
+                                          ValueRange tileIndices);
 
 private:
   SmallVector<ReassociationIndices> reassociationIndices;
   SmallVector<OpFoldResult> collapseShapeInputShape;
   SmallVector<OpFoldResult> collapseShapeOutputShape;
   SmallVector<Range> sliceParams;
   llvm::SmallBitVector linearizedDimensions;
   llvm::SmallBitVector slicedDimensions;
 };
 } // namespace mlir
 
 #endif // MLIR_DIALECT_UTILS_RESHAPEOPSUTILS_H

mlir/lib/Dialect/Tensor/Transforms/ExtractSliceFromReshape.cpp

if (linearizedDimensions[i] && slicedDimensions[i]) {
       DimAndIndex tb =
           invertSliceIndexing(builder, loc, sliceParams,
                               std::make_tuple(i, tileInductionVars[loopIdx++]));
       multiIndices.push_back(invertCollapseShapeIndexing(
           builder, loc, reassociationIndices, collapseShapeInputShape, tb));
     }
   }
 
-  auto extractParams = helper.getExtractSliceParams(multiIndices);
+  SmallVector<Range> extractParams =
+      helper.getExtractSliceParams(builder.getContext(), multiIndices);
 
   Value subTileResult = builder.create<tensor::ExtractSliceOp>(
       loc, collapseShapeOp.getSrc(), extractParams);
 
   SmallVector<Range> insertParams =
-      helper.getInsertSliceParams(tileInductionVars);
+      helper.getInsertSliceParams(builder.getContext(), tileInductionVars);
 
   // Collapse the dimensions of the source slice back down.
   Value collapsedResult = builder.create<tensor::CollapseShapeOp>(
       loc, subTileResult, reassociationIndices);
   return std::make_pair(collapsedResult, insertParams);
 }

mlir/lib/Dialect/Utils/ReshapeOpsUtils.cpp

llvm::SmallBitVector mlir::getLinearizedDimensions(
     ArrayRef<ReassociationIndices> reassociationIndices) {
   llvm::SmallBitVector result(reassociationIndices.size());
   for (const auto &it : llvm::enumerate(reassociationIndices))
     result[it.index()] = it.value().size() > 1;
   return result;
 }
 
 SmallVector<Range> SliceFromCollapseHelper::getExtractSliceParams(
-    ArrayRef<ValueRange> multiIndices) {
-  assert(!multiIndices.empty() && !multiIndices[0].empty() &&
-         "multiIndices should not be empty");
+    MLIRContext *ctx, ArrayRef<ValueRange> multiIndices) {
   unsigned loopIdx = 0;
-  MLIRContext *ctx = multiIndices[0][0].getContext();
   auto oneAttr = IntegerAttr::get(IndexType::get(ctx), 1);
   auto zeroAttr = IntegerAttr::get(IndexType::get(ctx), 0);
   SmallVector<Range> offsetsSizesAndStrides;
   offsetsSizesAndStrides.reserve(collapseShapeInputShape.size());
   for (const auto &it : llvm::enumerate(reassociationIndices)) {
     // Case 1: Linearized dimensions that have also been sliced. These
     // are size of 1 because we are iterating over these dimensions. The
     // offsets are exactly the de-linearized multi-indices.
for (const auto &it : llvm::enumerate(reassociationIndices)) {
 
     // Case 3: A single index, but it may be sliced.
     offsetsSizesAndStrides.push_back(sliceParams[it.index()]);
   }
   return offsetsSizesAndStrides;
 }
 
 SmallVector<Range>
-SliceFromCollapseHelper::getInsertSliceParams(ValueRange tileIndices) {
-  MLIRContext *ctx = tileIndices[0].getContext();
+SliceFromCollapseHelper::getInsertSliceParams(MLIRContext *ctx,
+                                              ValueRange tileIndices) {
   auto one = IntegerAttr::get(IndexType::get(ctx), 1);
   auto zero = IntegerAttr::get(IndexType::get(ctx), 0);
   SmallVector<Range> insertParams;
   insertParams.reserve(linearizedDimensions.size());
   unsigned loopIdx = 0;
   for (unsigned i = 0; i < linearizedDimensions.size(); i++) {
     if (linearizedDimensions[i] && slicedDimensions[i]) {
       insertParams.push_back(Range{tileIndices[loopIdx++], one, one});
       continue;
     }
     insertParams.push_back(Range{zero, sliceParams[i].size, one});
   }
   return insertParams;
 }

mlir/test/Dialect/Tensor/extract-slice-from-collapse-shape.mlir

func.func @extract_slice_non_sliced_linearized_dim(%input: tensor<3x?x?x11x2xf32>, %offt: index, %size: index) -> tensor<?x22xf32> {
   %collapsed = tensor.collapse_shape %input [[0, 1, 2], [3, 4]] : tensor<3x?x?x11x2xf32> into tensor<?x22xf32>  
   %slice = tensor.extract_slice %collapsed [%offt, 0] [%size, 22] [1, 1] : tensor<?x22xf32> to tensor<?x22xf32>
   // CHECK: scf.for
   // CHECK-NOT: scf.for
   // CHECK: %[[multiIndex:.+]]:3 = affine.delinearize_index
   // CHECK: tensor.extract_slice %[[arg0]][%[[multiIndex]]#0, %[[multiIndex]]#1, %[[multiIndex]]#2, 0, 0] [1, 1, 1, 11, 2] [1, 1, 1, 1, 1]
   return %slice : tensor<?x22xf32>
 }
+
+// -----
+
+// CHECK: @no_sliced_linearized_dims(%[[arg0:.+]]: tensor<{{.*}}>, %[[arg1:.+]]: index, %[[arg2:.+]]: index
+func.func @no_sliced_linearized_dims(%input: tensor<30x11x100xf32>, %offt: index, %size: index) -> tensor<330x?xf32> {
+  %collapsed = tensor.collapse_shape %input [[0, 1], [2]] : tensor<30x11x100xf32> into tensor<330x100xf32>
+  %slice = tensor.extract_slice %collapsed [0, %offt] [330, %size] [1, 1] : tensor<330x100xf32> to tensor<330x?xf32>
+  // CHECK-NOT: scf.for  
+  // CHECK: %[[init:.+]] = linalg.init_tensor [330, %[[arg2]]]
+  // CHECK: %[[e:.+]] = tensor.extract_slice %[[arg0]][0, 0, %[[arg1]]] [30, 11, %[[arg2]]] [1, 1, 1]
+  // CHECK: %[[c:.+]] = tensor.collapse_shape %[[e]] {{\[}}[0, 1], [2]]
+  // CHECK: %[[res:.+]] = tensor.insert_slice %[[c]] into %[[init]]
+  // CHECK: return %[[res]]
+  return %slice : tensor<330x?xf32>
+}

mlir/test/lib/Dialect/Tensor/TestTensorTransforms.cpp

LogicalResult emitReplacement(tensor::ExtractSliceOp op, Value dest,
                                 tensor::ExtractSliceFromCollapseHelper &helper,
                                 PatternRewriter &rewriter) const override {
     Location loc = op.getLoc();
     const unsigned numTiledDims = helper.getIterationSpaceSizes().size();
     auto zero = rewriter.create<arith::ConstantIndexOp>(loc, 0);
     auto one = rewriter.create<arith::ConstantIndexOp>(loc, 1);
     SmallVector<Value> lbs(numTiledDims, zero);
     SmallVector<Value> steps(numTiledDims, one);
+
+    // Below, we pass out the result of the loop body builder lambda via the
+    // `insertResult` variable. In certain cases, no loops will be created, but
+    // the body builder will still execute. In this case, the results will not
+    // be passed to the LoopNest object.
+    // TODO: remove this workaround if `scf::buildLoopNest` behavior is updated.
+    Value insertResult = nullptr;
     scf::LoopNest nest = scf::buildLoopNest(
         rewriter, loc, lbs, helper.getIterationSpaceSizes(), steps, dest,
         [&](OpBuilder &nestedBuilder, Location loc, ValueRange outputIvs,
             ValueRange iterArgs) -> scf::ValueVector {
           auto [tile, insertParams] =
               helper.emitLoopNestBody(nestedBuilder, loc, outputIvs);
 
           // Insert the slice into the destination.
-          Value result = nestedBuilder.create<tensor::InsertSliceOp>(
+          insertResult = nestedBuilder.create<tensor::InsertSliceOp>(
               loc, tile, iterArgs[0], insertParams);
-          return {result};
+          return {insertResult};
         });
-    rewriter.replaceOp(op, nest.getResults()[0]);
+
+    if (!nest.loops.empty())
+      rewriter.replaceOp(op, nest.getResults());
+    else
+      rewriter.replaceOp(op, insertResult);
+
     return success();
   }
 };
 
 struct RewriteExtractSliceFromCollapseShapeUsingScfForeach
     : public RewriteExtractSliceFromCollapseShapeBase {
   RewriteExtractSliceFromCollapseShapeUsingScfForeach(MLIRContext *context)
       : RewriteExtractSliceFromCollapseShapeBase(context) {}