[mlir][Linalg] Rewrite SubTensors that take a slice out of a unit-extend dimension.

[mravishankar]

Summary

Subtensor operations that are taking a slice out of a tensor that is
unit-extent along a dimension can be rewritten to drop that
dimension.

[mravishankar]

Fix test error.

[nicolasvasilache]

mlir/lib/Dialect/Linalg/Transforms/DropUnitDims.cpp

488

These days you can significantly simplify all this:

%1 = subtensor %0[0, %o1, 0] [1, %s1, 1] [1, 1, 1] : tensor<1x?x1xf32> to tensor<?xf32>

should just be valid.

You should use this helper to give you the rank-reduced type: https://github.com/llvm/llvm-project/blob/main/mlir/lib/Dialect/StandardOps/IR/Ops.cpp#L1728

[mravishankar]

Actually fix the test

[mravishankar]

THis is doing something different. It is reducing the rank of both the source and the dest, so

%1 = subtensor %0[0, %o1, 0] [1, %s1, 1] [1, 1, 1] : tensor<1x?x1xf32> to tensor<1x?x1xf32>

to

%1 = subtensor %0[%o1] [%s1 [1] : tensor<?xf32> to tensor<?xf32>

with reshapes for the modified source and results to match up with rest of the IR. The reshapes then get canonicalized away if possible. All this is meant for just eliminating the <...x1x...> dimensions that are not really needed (which is what the rest of this pass is for).

[nicolasvasilache]

Ah yes, I see that the reshapes are actually needed for function boundaries.
Removing the blocker then, sorry about the noise.

Still, it seems we aggressively use reshape to drop 1 dimensions even in the absence of function boundaries.

Maybe a solution would be to have extra canonicalizations:

• subtensor + reshape -> rank-reducing subtensor
• reshape + subtensor_insert -> rank-expanding subtensor_insert

The overarching point (that I think I also leaked into other phab reviews I saw fly by) is that subview, subtensor and vector.transfer become better understodd and do well with inplace updates, fusion etc while reshape are still not super-well understood.

Thoughts?

[mravishankar]

Still, it seems we aggressively use reshape to drop 1 dimensions even in the absence of function boundaries.

Maybe a solution would be to have extra canonicalizations:

• subtensor + reshape -> rank-reducing subtensor
• reshape + subtensor_insert -> rank-expanding subtensor_insert

That could work. This solves half the problem (on the target side, but not the source side)

The overarching point (that I think I also leaked into other phab reviews I saw fly by) is that subview, subtensor and vector.transfer become better understodd and do well with inplace updates, fusion etc while reshape are still not super-well understood.

Thoughts?

Totally agree. The goal here is that if the reshapes produced here fold away with other operations. That seems to happen well. Some of the patterns here could well be canonicalizations. But this pass has patterns that are meant to drop unit-trip loop counts within Linalg (which has a couple of good side effects too). It is mostly an experiment, thats why all these patterns are kept within this pass. It is used in IREE, and it works pretty well. Completely on-board though that reshapes should be at the boundaries. As I am bringing up using Linalg on tensors path, the next things i am looking it as propogating reshapes either upto function boundaries, or upto named ops (at which point the reshape stays, and you cant do much with it).

[mravishankar]

@nicolasvasilache did you mean to approve the patch?

In D99226#2650084, @nicolasvasilache wrote:

Ah yes, I see that the reshapes are actually needed for function boundaries.
Removing the blocker then, sorry about the noise.

[mravishankar]

Rebase.

[nicolasvasilache]

or upto named ops (at which point the reshape stays, and you cant do much with it).

For this case, I'd argue the behavior we want is named -> generic -> reshape opt -> named.
This only requires implementing a generic -> named which can be done by comparing a generic against all the named we know of.
There are details and complications but there is some more hope than bailing out on named + reshape.

mlir/lib/Dialect/Linalg/Transforms/DropUnitDims.cpp

LogicalResult replaceBlockArgForUnitDimLoops<IndexedGenericOp>(
   }
   SmallVector<unsigned, 8> unitDimsToErase(unitDims.begin(), unitDims.end());
   entryBlock->eraseArguments(unitDimsToErase);
   return success();
 }
 
 namespace {
 /// Pattern to fold unit-trip count loops in GenericOps.
-// TODO: Generalize this to indexed-generic as well by modifying the region args
-// as well.
 template <typename GenericOpTy>
 struct FoldUnitDimLoops : public OpRewritePattern<GenericOpTy> {
   using OpRewritePattern<GenericOpTy>::OpRewritePattern;
   LogicalResult matchAndRewrite(GenericOpTy op,
                                 PatternRewriter &rewriter) const override {
     SmallVector<AffineMap, 4> indexingMaps = op.getIndexingMaps();
     if (indexingMaps.empty())
       return failure();
for (auto result : llvm::enumerate(replacementOp.getResults())) {
             loc, origResultType, result.value(), reassociationMaps[index]));
       else
         resultReplacements.push_back(result.value());
     }
     rewriter.replaceOp(op, resultReplacements);
     return success();
   }
 };
-} // namespace
 
-namespace {
 /// Pattern to fold pair of reshape ops where the intermediate has unit-dims for
 /// example:
 ///
 ///  %0 = linalg.tensor_reshape %arg0
 ///    [affine_map<(d0, d1, d2, d3) -> (d0, d1, d2, d3)>]
 ///    : tensor<2048xf32> into tensor<1x4x1x512xf32>
 ///  %1 = linalg.tensor_reshape %0
 ///    [affine_map<(d0, d1, d2, d3) -> (d0, d1, d2)>,
RankedTensorType srcType = reshapeOp.getSrcType(),
                      dstType = reshapeOp.getResultType(),
                      parentSrcType = parentReshapeOp.getSrcType();
     if (!srcType.hasStaticShape() || !dstType.hasStaticShape() ||
         !parentSrcType.hasStaticShape() ||
         srcType.getRank() < dstType.getRank() ||
         parentSrcType.getRank() == dstType.getRank())
       return failure();
 
-    // Check if the result tensor_reshape after folding the reshapeOp and
-    // parentReshapeOp are combined.
-    // If the final tensor_reshape is folding, the parentReshapeOp is
-    // introducing unit-dims, and the reshapeOp does an actual reshape.
-    // If the final tensor_reshape op is expanding, the reshapeOp is
-    // introducing unit-dims, and the parentReshapeOp does an actual reshape.
+    // Check if the result tensor_reshape is folding or expanding after folding
+    // the reshapeOp and parentReshapeOp are combined.  If the final
+    // tensor_reshape is folding, the parentReshapeOp is introducing unit-dims,
+    // and the reshapeOp does an actual reshape.  If the final tensor_reshape op
+    // is expanding, the reshapeOp is introducing unit-dims, and the
+    // parentReshapeOp does an actual reshape.
     bool isFoldingPattern = parentSrcType.getRank() > dstType.getRank();
     ArrayRef<int64_t> expandedShape =
         isFoldingPattern ? parentSrcType.getShape() : dstType.getShape();
     ArrayRef<int64_t> foldedShape =
         isFoldingPattern ? dstType.getShape() : parentSrcType.getShape();
 
     unsigned expandedDim = 0, foldedDim = 0;
     SmallVector<SmallVector<AffineExpr, 4>, 4> reassociationExprs(
SmallVector<AffineMap, 4> reassociationMaps =
                                     rewriter.getContext());
             }));
     rewriter.replaceOpWithNewOp<TensorReshapeOp>(
         reshapeOp, dstType, parentReshapeOp.src(),
         rewriter.getAffineMapArrayAttr(reassociationMaps));
     return success();
   }
 };
+
+/// Pattern to fold subtensors that are just taking a slice of unit-dimension
+/// tensor. For example
+///
+/// %1 = subtensor %0[0, %o1, 0] [1, %s1, 1] [1, 1, 1]
+///     : tensor<1x?x1xf32> to tensor<1x?x1xf32>
+///
+/// can be replaced with
+///
+/// %0 = linalg.tensor_reshape %0 [affine_map<(d0, d1, d2) -> (d0, d1, d2)>]
+///     : tensor<1x?x1xf32> into tensor<?xf32>
+/// %1 = subtensor %0[%o1] [%s1] [1] : tensor<?xf32> to tensor<?xf32>
+/// %2 = linalg.tensor_reshape %1 [affine_map<(d0, d1, d2) -> (d0, d1, d2)>]
+///     : tensor<?xf32> into tensor<1x?x1xf32>
+///
+/// The additional tensor_reshapes will hopefully get canonicalized away with
+/// other reshapes that drop unit dimensions. Three condiitions to fold a
+/// dimension
+/// - The offset must be 0
+/// - The size must be 1
+/// - The dimension of the source type must be 1.
+struct FoldUnitDimSubTensorOp : public OpRewritePattern<SubTensorOp> {
+  using OpRewritePattern<SubTensorOp>::OpRewritePattern;
+
+  LogicalResult matchAndRewrite(SubTensorOp subTensorOp,
+                                PatternRewriter &rewriter) const override {
+    SmallVector<OpFoldResult> mixedOffsets = subTensorOp.getMixedOffsets();
+    SmallVector<OpFoldResult> mixedSizes = subTensorOp.getMixedSizes();
+    SmallVector<OpFoldResult> mixedStrides = subTensorOp.getMixedStrides();
+    auto hasValue = [](OpFoldResult valueOrAttr, int64_t val) {
+      auto attr = valueOrAttr.dyn_cast<Attribute>();
+      return attr && attr.cast<IntegerAttr>().getInt() == val;
+    };
+
+    if (llvm::any_of(mixedStrides, [&](OpFoldResult valueOrAttr) {
+          return !hasValue(valueOrAttr, 1);
+        }))
+      return failure();
+
+    // Find the expanded unit dimensions.
+    SmallVector<ReassociationIndices> reassociation;
+    SmallVector<OpFoldResult> newOffsets, newSizes;
+    ArrayRef<int64_t> sourceShape = subTensorOp.getSourceType().getShape();
+    ReassociationIndices curr;
+    for (int64_t dim : llvm::seq<int64_t>(0, mixedOffsets.size())) {
+      curr.push_back(dim);
+      if (sourceShape[dim] == 1 && hasValue(mixedOffsets[dim], 0) &&
+          hasValue(mixedSizes[dim], 1)) {
+        continue;
+      }
+      newOffsets.push_back(mixedOffsets[dim]);
+      newSizes.push_back(mixedSizes[dim]);
+      reassociation.emplace_back(ReassociationIndices{});
+      std::swap(reassociation.back(), curr);
+    }
+    if (newOffsets.size() == mixedOffsets.size())
+      return failure();
+    reassociation.back().append(curr.begin(), curr.end());
+    SmallVector<OpFoldResult> newStrides(newOffsets.size(),
+                                         rewriter.getI64IntegerAttr(1));
+    Location loc = subTensorOp->getLoc();
+    auto srcReshape = rewriter.create<TensorReshapeOp>(
+        loc, subTensorOp.source(), reassociation);
+    auto newSubTensorOp = rewriter.create<SubTensorOp>(
+        loc, srcReshape, newOffsets, newSizes, newStrides);
+    rewriter.replaceOpWithNewOp<TensorReshapeOp>(
+        subTensorOp, subTensorOp.getType(), newSubTensorOp, reassociation);
+    return success();
+  }
+};
+
 } // namespace
 
 /// Patterns that are used to canonicalize the use of unit-extent dims for
 /// broadcasting.
 void mlir::populateLinalgFoldUnitExtentDimsPatterns(
     MLIRContext *context, OwningRewritePatternList &patterns) {
   patterns
       .insert<FoldUnitDimLoops<GenericOp>, FoldUnitDimLoops<IndexedGenericOp>,
               ReplaceUnitExtentTensors<GenericOp>,
-              ReplaceUnitExtentTensors<IndexedGenericOp>>(context);
+              ReplaceUnitExtentTensors<IndexedGenericOp>,
+              FoldUnitDimSubTensorOp>(context);
   TensorReshapeOp::getCanonicalizationPatterns(patterns, context);
   patterns.insert<FoldReshapeOpWithUnitExtent>(context);
   populateFoldUnitDimsReshapeOpsByLinearizationPatterns(context, patterns);
 }
 
 namespace {
 /// Pass that removes unit-extent dims within generic ops.
 struct LinalgFoldUnitExtentDimsPass

mlir/test/Dialect/Linalg/drop-unit-extent-dims.mlir

 func @drop_one_trip_loops_indexed_generic
   (%arg0 : tensor<?x1x?xi32>, %shape: tensor<?x1x?x1x?xi32>) -> tensor<?x1x?x1x?xi32>
 {
   %0 = linalg.indexed_generic #trait
      ins(%arg0 : tensor<?x1x?xi32>)
     outs(%shape: tensor<?x1x?x1x?xi32>) {
        ^bb0(%arg1 : index, %arg2 : index, %arg3 : index, %arg4 : index,
             %arg5 : index, %arg6 : i32, %arg7 : i32) :
-   %1 = addi %arg1, %arg2 : index
-   %2 = addi %1, %arg3 : index
-   %3 = addi %2, %arg4 : index
-   %4 = addi %3, %arg5 : index
-   %5 = index_cast %4 : index to i32
-   %6 = addi %5, %arg6 : i32
+         %1 = addi %arg1, %arg2 : index
+         %2 = addi %1, %arg3 : index
+         %3 = addi %2, %arg4 : index
+         %4 = addi %3, %arg5 : index
+         %5 = index_cast %4 : index to i32
+         %6 = addi %5, %arg6 : i32
          linalg.yield %6 : i32
        } -> tensor<?x1x?x1x?xi32>
   return %0 : tensor<?x1x?x1x?xi32>
 }
 // CHECK-LABEL: func @drop_one_trip_loops_indexed_generic
 //       CHECK:   linalg.indexed_generic
 //       CHECK:   ^{{.+}}(
 //  CHECK-SAME:     %[[ARG1:[a-zA-Z0-9]+]]: index, %[[ARG2:[a-zA-Z0-9]+]]: index

 
 func @drop_all_loops_indexed_generic
   (%arg0 : tensor<1x1xi32>) -> tensor<1x1xi32>{
   %0 = linalg.indexed_generic #trait
      ins(%arg0 : tensor<1x1xi32>)
     outs(%arg0 : tensor<1x1xi32>) {
        ^bb0(%arg1 : index, %arg2 : index, %arg3: i32, %arg4: i32) :
          %1 = addi %arg1, %arg2 : index
-   %2 = index_cast %1 : index to i32
-   %3 = addi %2, %arg3 : i32
+         %2 = index_cast %1 : index to i32
+         %3 = addi %2, %arg3 : i32
          linalg.yield %3 : i32
        } -> tensor<1x1xi32>
   return %0 : tensor<1x1xi32>
 }
 
 // CHECK-LABEL: func @drop_all_loops_indexed_generic
 //       CHECK:   linalg.indexed_generic
 //       CHECK:   ^{{.+}}(%[[ARG1:.+]]: i32, %[[ARG2:.+]]: i32)
%2 = linalg.generic {i64, indexing_maps = [#map1, #map0],
     } -> tensor<1x2x5xf32>
   %3 = linalg.tensor_reshape %2 [#map3, #map4]
     : tensor<1x2x5xf32> into tensor<2x5xf32>
   return %3 : tensor<2x5xf32>
 }
 // CHECK-LABEL: func @fold_unit_dim_tensor_reshape_op
 //       CHECK:   %[[RESULT:.+]] = linalg.generic
 //       CHECK:   return %[[RESULT]]
+
+// -----
+
+func @fold_subtensor(
+    %arg0 : tensor<1x?x?x1x?x1x1xf32>, %arg1 : index, %arg2 : index,
+    %arg3 : index, %arg4 : index, %arg5 : index, %arg6 : index)
+    -> tensor<1x?x?x1x?x1x1xf32> {
+  %0 = subtensor %arg0[0, %arg1, %arg2, 0, %arg3, 0, 0]
+                      [1, %arg4, %arg5, 1, %arg6, 1, 1] [1, 1, 1, 1, 1, 1, 1] :
+      tensor<1x?x?x1x?x1x1xf32> to tensor<1x?x?x1x?x1x1xf32>
+  return %0 : tensor<1x?x?x1x?x1x1xf32>
+}
+//  CHECK-DAG: #[[MAP0:.+]] = affine_map<(d0, d1, d2, d3, d4, d5, d6) -> (d0, d1)>
+//  CHECK-DAG: #[[MAP1:.+]] = affine_map<(d0, d1, d2, d3, d4, d5, d6) -> (d2)>
+//  CHECK-DAG: #[[MAP2:.+]] = affine_map<(d0, d1, d2, d3, d4, d5, d6) -> (d3, d4, d5, d6)>
+//      CHECK: func @fold_subtensor
+// CHECK-SAME:   %[[ARG0:.+]]: tensor<1x?x?x1x?x1x1xf32>
+// CHECK-SAME:   %[[ARG1:[a-z0-9]+]]: index
+// CHECK-SAME:   %[[ARG2:[a-z0-9]+]]: index
+// CHECK-SAME:   %[[ARG3:[a-z0-9]+]]: index
+// CHECK-SAME:   %[[ARG4:[a-z0-9]+]]: index
+// CHECK-SAME:   %[[ARG5:[a-z0-9]+]]: index
+// CHECK-SAME:   %[[ARG6:[a-z0-9]+]]: index
+//      CHECK:   %[[SRC_RESHAPE:.+]] = linalg.tensor_reshape %[[ARG0]]
+// CHECK-SAME:       [#[[MAP0]], #[[MAP1]], #[[MAP2]]]
+//      CHECK:   %[[SUBTENSOR:.+]] = subtensor %[[SRC_RESHAPE]]
+// CHECK-SAME:       [%[[ARG1]], %[[ARG2]], %[[ARG3]]]
+// CHECK-SAME:       [%[[ARG4]], %[[ARG5]], %[[ARG6]]]
+//      CHECK:   %[[RESULT_RESHAPE:.+]] = linalg.tensor_reshape %[[SUBTENSOR]]
+// CHECK-SAME:       [#[[MAP0]], #[[MAP1]], #[[MAP2]]]
+//      CHECK:   return %[[RESULT_RESHAPE]]
+
+// -----
+
+func @no_fold_subtensor(
+    %arg0 : tensor<1x?x?x?x?x1x1xf32>, %arg1 : index, %arg2 : index,
+    %arg3 : index, %arg4 : index, %arg5 : index, %arg6 : index)
+    -> tensor<1x?x?x1x?x1x1xf32> {
+  %0 = subtensor %arg0[%arg1, 0, %arg2, 0, 0, %arg3, 0]
+                      [1, %arg4, %arg5, 1, %arg6, 1, 1] [1, 1, 1, 1, 1, 1, 1] :
+      tensor<1x?x?x?x?x1x1xf32> to tensor<1x?x?x1x?x1x1xf32>
+  return %0 : tensor<1x?x?x1x?x1x1xf32>
+}
+//  CHECK-DAG: #[[MAP0:.+]] = affine_map<(d0, d1, d2, d3, d4, d5, d6) -> (d0)>
+//  CHECK-DAG: #[[MAP1:.+]] = affine_map<(d0, d1, d2, d3, d4, d5, d6) -> (d1)>
+//  CHECK-DAG: #[[MAP2:.+]] = affine_map<(d0, d1, d2, d3, d4, d5, d6) -> (d2)>
+//  CHECK-DAG: #[[MAP3:.+]] = affine_map<(d0, d1, d2, d3, d4, d5, d6) -> (d3)>
+//  CHECK-DAG: #[[MAP4:.+]] = affine_map<(d0, d1, d2, d3, d4, d5, d6) -> (d4)>
+//  CHECK-DAG: #[[MAP5:.+]] = affine_map<(d0, d1, d2, d3, d4, d5, d6) -> (d5, d6)>
+//      CHECK: func @no_fold_subtensor
+// CHECK-SAME:   %[[ARG0:.+]]: tensor<1x?x?x?x?x1x1xf32>
+// CHECK-SAME:   %[[ARG1:[a-z0-9]+]]: index
+// CHECK-SAME:   %[[ARG2:[a-z0-9]+]]: index
+// CHECK-SAME:   %[[ARG3:[a-z0-9]+]]: index
+// CHECK-SAME:   %[[ARG4:[a-z0-9]+]]: index
+// CHECK-SAME:   %[[ARG5:[a-z0-9]+]]: index
+// CHECK-SAME:   %[[ARG6:[a-z0-9]+]]: index
+//      CHECK:   %[[SRC_RESHAPE:.+]] = linalg.tensor_reshape %[[ARG0]]
+// CHECK-SAME:       [#[[MAP0]], #[[MAP1]], #[[MAP2]], #[[MAP3]], #[[MAP4]], #[[MAP5]]]
+//      CHECK:   %[[SUBTENSOR:.+]] = subtensor %[[SRC_RESHAPE]]
+// CHECK-SAME:       [%[[ARG1]], 0, %[[ARG2]], 0, 0, %[[ARG3]]]
+// CHECK-SAME:       [1, %[[ARG4]], %[[ARG5]], 1, %[[ARG6]], 1]
+//      CHECK:   %[[RESULT_RESHAPE:.+]] = linalg.tensor_reshape %[[SUBTENSOR]]
+// CHECK-SAME:       [#[[MAP0]], #[[MAP1]], #[[MAP2]], #[[MAP3]], #[[MAP4]], #[[MAP5]]]
+//      CHECK:   return %[[RESULT_RESHAPE]]