[mlir] bubble up tensor.extract_slice above linalg op

[okkwon]

Summary

Bubble up extract_slice above Linalg operation.

A sequence of operations

%0 = linalg.<op> ... arg0, arg1, ...
%1 = tensor.extract_slice %0 ...

can be replaced with

%0 = tensor.extract_slice %arg0
%1 = tensor.extract_slice %arg1
%2 = linalg.<op> ... %0, %1, ...

This results in the reduce computation of the linalg operation.

[mravishankar]

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

29

Instead of doing this, I'd just convert OpFoldResult to Value by insert arith.constant <val> : index operations for cases where it is holding an integer attribute and then use applyMapToValues. This will canonicalize back anyway.

82

Actually this might be too conservative. The only condition that is needed I think is that the output indexing map must be a "permutation". The use of applyMapToValues should handle all other valid indexing maps.

97

This is not enough. The output might not be accessed as an identity, but a permutation. So you need to "invert" the indexing map used for the output, compose it with the indexing map for the inputs and them apply the map.

You can use [inversePermutation](https://github.com/llvm/llvm-project/blob/14e4d2e5643e505039588d50b31e29e3042009f5/mlir/include/mlir/IR/AffineMap.h#L430) for inverting the output map, and compose it with the indexing map for the operand as follows

AffineMap outputIndexMap = ...
AffineMap inputIndexingMap = ...
AffineMap inverseOutputIndexingMap = inversePermutation(outputIndexingMap);
AffineMap composedMap = inputIndexingMap.compose(inverseOutputIndexingMap);
.... applyMapToValues(composedMap, sliceOffsets)
.... applyMapToValues(composedMap, sliceSizes)

99

Strides are tricky... I'd rather just check that strides for the insert slice are all 1s.

mlir/test/Dialect/Linalg/bubble-up-extract-slice-op.mlir

92

Would be interesting to try some other LinalgOps as well. Like linalg.matmul linalg.batch_matmul linalg.conv_2d_nhwc_hwcf . See more named ops list here : https://github.com/llvm/llvm-project/blob/main/mlir/python/mlir/dialects/linalg/opdsl/ops/core_named_ops.py. The python file is the source of truth for these operations. They are used to generate the ODS definition for these ops automatically (there is also an intermediate YAML serialization step thrown in for good measure).

[okkwon]

Address Mahesh's comments

[okkwon]

Minor updates

[okkwon]

Update the commit message

[okkwon]

mlir/test/Dialect/Linalg/bubble-up-extract-slice-op.mlir

92

Matmul and other non-trivial operations will be covered later since they show more complex index mapping than the simple permutation. I left a TODO comment for that.

[okkwon]

Update the commit message

[okkwon]

check projected permutation for inputs

This seems to be needed because ApplyMapToValues() only handles AffineDimExpr
and AffineSymExpr; More complicated affine expresssions are not handled. To
simplify the eligible cases and the transformation, let's handle the inputs
only when their indexing map is a projected permutation.

[okkwon]

merge and rebase

[mravishankar]

mlir/include/mlir/Dialect/Linalg/Passes.h

81 ↗

(On Diff #415552)

I'd drop the pass from here. These passses have been hard to maintain. I've found passes here (especially in Linalg) to be not very useful. It is best to expose the patterns using populate*Pattern method. I'd do the same here.

mlir/include/mlir/Dialect/Linalg/Passes.td

383 ↗

(On Diff #415552)

Same here, drop the pass. You anyway have the test pass to run the pattern.

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

60

Probably also need to check for the linalg op to have a single result. We can generalize it later on.

74

Nit : replace with

if (llvm::any_of(linalgOp.getIndexingMaps(), [](AffineMap map) { return !map.isPorjectedPermutation(); })) {
  return false;
}

124

This can run into type mismatch issues if the tensor.insert_slice return type and the newOp return type dont match. Its always just easier to check the types match and insert a tensor.cast if the types dont match, with the expectation that the cast folds away during canonicalizations.

[okkwon]

Address Mahesh's comments

[okkwon]

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

124

Would you mind giving me an example? I am using the same offsets, sizes, and slides, so the types are supposed to be the same.

[okkwon]

The plan is changed because

1. I made a big code change to support matmul cases, and
2. Mahesh pointed me to the tiling implementation, which works for any Linalg operations. I need to take a look at the code and understand how it works for all without any restrictions on the linalg operation, and
3. Apply the same or similar approach to my implementation.

Thanks Mahesh for the inputs.

mlir/include/mlir/Dialect/Linalg/Transforms/Transforms.h

 
 /// Patterns to fold unit-extent dimensions in operands/results of linalg ops on
 /// tensors.
 void populateFoldUnitExtentDimsPatterns(RewritePatternSet &patterns);
 
 /// Patterns that are used to inline constant operands into linalg generic ops.
 void populateInlineConstantOperandsPatterns(RewritePatternSet &patterns);
 
+/// Patterns taht are used to bubble up extract slice op above linalg op.
+void populateBubbleUpExtractSliceOpPatterns(RewritePatternSet &patterns);
+
 /// Options that control fusion of elementwise operations.
 struct LinalgElementwiseFusionOptions {
   /// Enable fusion of reshapes into the shape with elementwise operations. By
   /// default it is disabled for unit dimensions reshape.
   ControlElementwiseOpsFusionFn controlFoldingReshapesFn = skipUnitDimReshape;
 
   LinalgElementwiseFusionOptions &
   setControlFoldingReshapes(ControlElementwiseOpsFusionFn fun) {

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

+//===- BubbleUpExtractSlice.cpp - bubble up tensor.extract_slice ----------===//
+//
+// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
+// See https://llvm.org/LICENSE.txt for license information.
+// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+//
+//===----------------------------------------------------------------------===//
+//
+// This file implements patterns that transforms linalg.<op> +
+// tensor.extract_slice into tensor.extract_slice + linalg.<op> to reduce
+// the computation for the linalg op.
+//
+//===----------------------------------------------------------------------===//
+
+#include "PassDetail.h"
+#include "mlir/Dialect/Affine/IR/AffineOps.h"
+#include "mlir/Dialect/Arithmetic/Utils/Utils.h"
+#include "mlir/Dialect/Linalg/IR/Linalg.h"
+#include "mlir/Dialect/Linalg/Passes.h"
+#include "mlir/Dialect/Linalg/Transforms/Transforms.h"
+#include "mlir/Transforms/GreedyPatternRewriteDriver.h"
+
+using namespace mlir;
+using namespace mlir::linalg;
+
+namespace {
+/// Bubble up extract_slice above Linalg operation.
+///
+/// A sequence of operations
+///
+/// ```mlir
+/// %0 = linalg.<op> ... arg0, arg1, ...
+/// %1 = tensor.extract_slice %0 ...
+/// ```
+///
+/// can be replaced with
+///
+/// ```mlir
+/// %0 = tensor.extract_slice %arg0
+/// %1 = tensor.extract_slice %arg1
+/// %2 = linalg.<op> ... %0, %1, ...
+/// ```
+///
+/// This results in the reduce computation of the linalg operation.
+///
+struct BubbleUpExtractSliceOpPattern
+    : OpRewritePattern<tensor::ExtractSliceOp> {
+  using OpRewritePattern<tensor::ExtractSliceOp>::OpRewritePattern;
+
+  LogicalResult matchAndRewrite(tensor::ExtractSliceOp sliceOp,
+                                PatternRewriter &rewriter) const final {
+    Value source = sliceOp.source();
+    auto linalgOp = source.getDefiningOp<LinalgOp>();
+    if (!linalgOp) {
+      return rewriter.notifyMatchFailure(sliceOp,
+                                         "expected source to be linalg op");
+    }
+
+    if (!linalgOp->hasOneUse()) {
+      return rewriter.notifyMatchFailure(sliceOp,
+                                         "expected single use of linalg op");
+    }
+
+    if (linalgOp.getNumOutputs() != 1) {
+      return rewriter.notifyMatchFailure(sliceOp,
+                                         "expected single output of linalg op");
+    }
+
+    if (!linalgOp.hasTensorSemantics()) {
+      return rewriter.notifyMatchFailure(sliceOp,
+                                         "expected tensor of linalg op");
+    }
+
+    if (!sliceOp.hasUnitStride())
+      return rewriter.notifyMatchFailure(sliceOp, "expected unit stride");
+
+    // Check all input indexing maps are a projected permutation.
+    if (llvm::any_of(linalgOp.getIndexingMaps(), [](AffineMap map) {
+          return !map.isProjectedPermutation();
+        })) {
+      return rewriter.notifyMatchFailure(sliceOp, "expected a projected "
+                                                  "permutation for input");
+    }
+
+    auto resultNumber = source.cast<OpResult>().getResultNumber();
+    AffineMap outputMap =
+        linalgOp.getTiedIndexingMap(linalgOp.getOutputOperand(resultNumber));
+
+    // To apply the slice to the arguments of the linalg operation, we need a
+    // permutation relationship between the iteration space and the output.
+    // TODO: Relax the condition and add partial dimension slicing to support
+    // more complicated cases such as matmul.
+    if (!outputMap.isPermutation()) {
+      return rewriter.notifyMatchFailure(sliceOp,
+                                         "expected a permutation for output");
+    }
+
+    AffineMap inverseMap = inversePermutation(outputMap);
+
+    // Bubble up extract slice for each operand.
+    auto sliceLoc = sliceOp.getLoc();
+    auto sliceOffsets = getValueOrCreateConstantIndexOp(
+        rewriter, sliceLoc, sliceOp.getMixedOffsets());
+    auto sliceSizes = getValueOrCreateConstantIndexOp(rewriter, sliceLoc,
+                                                      sliceOp.getMixedSizes());
+    auto sliceStrides = getValueOrCreateConstantIndexOp(
+        rewriter, sliceLoc, sliceOp.getMixedStrides());
+
+    SmallVector<Value> newOperands;
+    for (OpOperand *operand : linalgOp.getInputAndOutputOperands()) {
+      AffineMap inputMap = linalgOp.getTiedIndexingMap(operand);
+      AffineMap composedMap = inputMap.compose(inverseMap);
+
+      auto newOffsets = getAsOpFoldResult(
+          applyMapToValues(rewriter, sliceLoc, composedMap, sliceOffsets));
+      auto newSizes = getAsOpFoldResult(
+          applyMapToValues(rewriter, sliceLoc, composedMap, sliceSizes));
+      auto newStrides = getAsOpFoldResult(
+          applyMapToValues(rewriter, sliceLoc, composedMap, sliceStrides));
+      Value newSlice = rewriter.create<tensor::ExtractSliceOp>(
+          sliceOp.getLoc(), operand->get(), newOffsets, newSizes, newStrides);
+      newOperands.push_back(newSlice);
+    }
+
+    Operation *newOp = linalgOp.clone(rewriter, linalgOp.getLoc(),
+                                      sliceOp.getType(), newOperands);
+    rewriter.replaceOp(sliceOp, newOp->getResults());
+    return success();
+  }
+};
+} // namespace
+
+void mlir::linalg::populateBubbleUpExtractSliceOpPatterns(
+    RewritePatternSet &patterns) {
+  auto *context = patterns.getContext();
+  patterns.add<BubbleUpExtractSliceOpPattern>(context);
+}

mlir/lib/Dialect/Linalg/Transforms/CMakeLists.txt

 add_mlir_dialect_library(MLIRLinalgTransforms
+  BubbleUpExtractSlice.cpp
   BufferizableOpInterfaceImpl.cpp
   Bufferize.cpp
   CodegenStrategy.cpp
   ComprehensiveBufferizePass.cpp
   Detensorize.cpp
   DropUnitDims.cpp
   ElementwiseOpFusion.cpp
   ElementwiseToLinalg.cpp

mlir/test/Dialect/Linalg/bubble-up-extract-slice-op.mlir

+// RUN: mlir-opt -test-linalg-transform-patterns=test-bubble-up-extract-slice-op-pattern -split-input-file %s | FileCheck %s
+
+#map0 = affine_map<(d0, d1) -> (d0, d1)>
+#map1 = affine_map<(d0, d1) -> (d1)>
+
+func @dynamic(%arg0: tensor<?x?xf32>, %arg1: tensor<?xf32>, %arg2: index, %arg3: index, %arg4: index, %arg5:index) -> tensor<?x?xf32> {
+  %0 = linalg.generic {
+    indexing_maps = [#map0, #map1, #map0],
+    iterator_types = ["parallel", "parallel"]
+  } ins(%arg0, %arg1 : tensor<?x?xf32>, tensor<?xf32>)
+    outs(%arg0 : tensor<?x?xf32>) {
+    ^bb0(%b0 : f32, %b1 : f32, %b2 : f32):
+      %add = arith.addf %b0, %b1 : f32
+      linalg.yield %add : f32
+  } -> tensor<?x?xf32>
+  %1 = tensor.extract_slice %0 [%arg2, %arg3] [%arg4, %arg5] [1, 1]
+    : tensor<?x?xf32> to tensor<?x?xf32>
+  return %1 : tensor<?x?xf32>
+}
+
+//       CHECK: func @dynamic
+//       CHECK: %[[SLICE0:.+]] = tensor.extract_slice %arg0[%arg2, %arg3] [%arg4, %arg5] [1, 1] : tensor<?x?xf32> to tensor<?x?xf32>
+//       CHECK: %[[SLICE1:.+]] = tensor.extract_slice %arg1[%arg3] [%arg5] [1] : tensor<?xf32> to tensor<?xf32>
+//       CHECK: %[[SLICE2:.+]] = tensor.extract_slice %arg0[%arg2, %arg3] [%arg4, %arg5] [1, 1] : tensor<?x?xf32> to tensor<?x?xf32>
+//       CHECK: %[[GENERIC:.+]] = linalg.generic {indexing_maps = [#map0, #map1, #map0], iterator_types = ["parallel", "parallel"]}
+// CHECK-SAME: ins(%[[SLICE0]], %[[SLICE1]] : tensor<?x?xf32>, tensor<?xf32>) outs(%[[SLICE2]] : tensor<?x?xf32>)
+//      CHECK: return %[[GENERIC]] : tensor<?x?xf32>
+
+// -----
+
+#map0 = affine_map<(d0, d1) -> (d0, d1)>
+#map1 = affine_map<(d0, d1) -> (d1)>
+
+func @static(%arg0: tensor<16x8xf32>, %arg1: tensor<8xf32>) -> tensor<4x2xf32> {
+  %0 = linalg.generic {
+    indexing_maps = [#map0, #map1, #map0],
+    iterator_types = ["parallel", "parallel"]
+  } ins(%arg0, %arg1 : tensor<16x8xf32>, tensor<8xf32>)
+    outs(%arg0 : tensor<16x8xf32>) {
+    ^bb0(%b0 : f32, %b1 : f32, %b2 : f32):
+      %add = arith.addf %b0, %b1 : f32
+      linalg.yield %add : f32
+  } -> tensor<16x8xf32>
+  %1 = tensor.extract_slice %0 [8, 4] [4, 2] [1, 1]
+    : tensor<16x8xf32> to tensor<4x2xf32>
+  return %1 : tensor<4x2xf32>
+}
+
+//       CHECK: func @static
+//       CHECK: %[[SLICE0:.+]] = tensor.extract_slice %arg0[8, 4] [4, 2] [1, 1] : tensor<16x8xf32> to tensor<4x2xf32>
+//       CHECK: %[[SLICE1:.+]] = tensor.extract_slice %arg1[4] [2] [1] : tensor<8xf32> to tensor<2xf32>
+//       CHECK: %[[SLICE2:.+]] = tensor.extract_slice %arg0[8, 4] [4, 2] [1, 1] : tensor<16x8xf32> to tensor<4x2xf32>
+//       CHECK: %[[GENERIC:.+]] = linalg.generic {indexing_maps = [#map0, #map1, #map0], iterator_types = ["parallel", "parallel"]}
+// CHECK-SAME: ins(%[[SLICE0]], %[[SLICE1]] : tensor<4x2xf32>, tensor<2xf32>) outs(%[[SLICE2]] : tensor<4x2xf32>)
+//      CHECK: return %[[GENERIC]] : tensor<4x2xf32>
+
+// -----
+
+#map0 = affine_map<(d0, d1) -> (d0, d1)>
+#map1 = affine_map<(d0, d1) -> (d1)>
+
+func @mixed(%arg0: tensor<?x8xf32>, %arg1: tensor<8xf32>, %arg2: index, %arg3: index) -> tensor<?x2xf32> {
+  %0 = linalg.generic {
+    indexing_maps = [#map0, #map1, #map0],
+    iterator_types = ["parallel", "parallel"]
+  } ins(%arg0, %arg1 : tensor<?x8xf32>, tensor<8xf32>)
+    outs(%arg0 : tensor<?x8xf32>) {
+    ^bb0(%b0 : f32, %b1 : f32, %b2 : f32):
+      %add = arith.addf %b0, %b1 : f32
+      linalg.yield %add : f32
+  } -> tensor<?x8xf32>
+  %1 = tensor.extract_slice %0 [8, %arg2] [%arg3, 2] [1, 1]
+    : tensor<?x8xf32> to tensor<?x2xf32>
+  return %1 : tensor<?x2xf32>
+}
+
+//       CHECK: func @mixed
+//       CHECK: %[[SLICE0:.+]] = tensor.extract_slice %arg0[8, %arg2] [%arg3, 2] [1, 1] : tensor<?x8xf32> to tensor<?x2xf32>
+//       CHECK: %[[SLICE1:.+]] = tensor.extract_slice %arg1[%arg2] [2] [1] : tensor<8xf32> to tensor<2xf32>
+//       CHECK: %[[SLICE2:.+]] = tensor.extract_slice %arg0[8, %arg2] [%arg3, 2] [1, 1] : tensor<?x8xf32> to tensor<?x2xf32>
+//       CHECK: %[[GENERIC:.+]] = linalg.generic {indexing_maps = [#map0, #map1, #map0], iterator_types = ["parallel", "parallel"]}
+// CHECK-SAME: ins(%[[SLICE0]], %[[SLICE1]] : tensor<?x2xf32>, tensor<2xf32>) outs(%[[SLICE2]] : tensor<?x2xf32>)
+//      CHECK: return %[[GENERIC]] : tensor<?x2xf32>
+
+// -----
+
+#map0 = affine_map<(d0, d1) -> (d0, d1)>
+#map1 = affine_map<(d0, d1) -> (d1)>
+
+func @dynamic_to_static(%arg0: tensor<?x?xf32>, %arg1: tensor<?xf32>) -> tensor<4x2xf32> {
+  %0 = linalg.generic {
+    indexing_maps = [#map0, #map1, #map0],
+    iterator_types = ["parallel", "parallel"]
+  } ins(%arg0, %arg1 : tensor<?x?xf32>, tensor<?xf32>)
+    outs(%arg0 : tensor<?x?xf32>) {
+    ^bb0(%b0 : f32, %b1 : f32, %b2 : f32):
+      %add = arith.addf %b0, %b1 : f32
+      linalg.yield %add : f32
+  } -> tensor<?x?xf32>
+  %1 = tensor.extract_slice %0 [8, 4] [4, 2] [1, 1]
+    : tensor<?x?xf32> to tensor<4x2xf32>
+  return %1 : tensor<4x2xf32>
+}
+
+//       CHECK: func @dynamic_to_static
+//       CHECK: %[[SLICE0:.+]] = tensor.extract_slice %arg0[8, 4] [4, 2] [1, 1] : tensor<?x?xf32> to tensor<4x2xf32>
+//       CHECK: %[[SLICE1:.+]] = tensor.extract_slice %arg1[4] [2] [1] : tensor<?xf32> to tensor<2xf32>
+//       CHECK: %[[SLICE2:.+]] = tensor.extract_slice %arg0[8, 4] [4, 2] [1, 1] : tensor<?x?xf32> to tensor<4x2xf32>
+//       CHECK: %[[GENERIC:.+]] = linalg.generic {indexing_maps = [#map0, #map1, #map0], iterator_types = ["parallel", "parallel"]}
+// CHECK-SAME: ins(%[[SLICE0]], %[[SLICE1]] : tensor<4x2xf32>, tensor<2xf32>) outs(%[[SLICE2]] : tensor<4x2xf32>)
+//      CHECK: return %[[GENERIC]] : tensor<4x2xf32>

mlir/test/lib/Dialect/Linalg/TestLinalgTransforms.cpp

Option<bool> skipPartial{
       *this, "skip-partial",
       llvm::cl::desc("Skip loops inside partial iterations during peeling"),
       llvm::cl::init(false)};
   Option<std::string> loopType{
       *this, "loop-type",
       llvm::cl::desc("Specify the type of loops to generate: for, parallel or "
                      "tiled_loop"),
       llvm::cl::init("for")};
+  Option<bool> testBubbleUpExtractSliceOpPattern{
+      *this, "test-bubble-up-extract-slice-op-pattern",
+      llvm::cl::desc("Test rewrite of linalgOp + extract_slice into "
+                     "extract_slice + linalgOp"),
+      llvm::cl::init(false)};
 };
 } // namespace
 
 static void applyPatterns(FuncOp funcOp) {
   MLIRContext *ctx = funcOp.getContext();
   RewritePatternSet patterns(ctx);
 
   //===--------------------------------------------------------------------===//
if (scalarizeDynamicDims) {
     linalgTilingOptions.setTileSizes(tileSizes);
   }
   linalg::LinalgTransformationFilter f(StringAttr::get(context, "tile"));
   TilingPatterns<linalg::MatmulOp, linalg::GenericOp>::insert(
       tilingPattern, linalgTilingOptions, f);
   (void)applyPatternsAndFoldGreedily(funcOp, std::move(tilingPattern));
 }
 
+static void applyBubbleUpExtractSliceOpPattern(FuncOp funcOp) {
+  RewritePatternSet patterns(funcOp.getContext());
+  populateBubbleUpExtractSliceOpPatterns(patterns);
+  (void)applyPatternsAndFoldGreedily(funcOp, std::move(patterns));
+}
+
 /// Apply transformations specified as patterns.
 void TestLinalgTransforms::runOnOperation() {
   auto lambda = [&](void *) {
     getOperation().walk([](LinalgOp op) {
       op->removeAttr(LinalgTransforms::kLinalgTransformMarker);
     });
   };
   std::unique_ptr<void, decltype(lambda)> cleanupGuard{(void *)1, lambda};
void TestLinalgTransforms::runOnOperation() {
   if (testSwapSubTensorPadTensor)
     return applyExtractSliceOfPadTensorSwapPattern(getOperation());
   if (testTilePattern)
     return applyTilePattern(getOperation(), loopType, tileSizes, peeledLoops,
                             /*scalarizeDynamicDims=*/false);
   if (testTileScalarizeDynamicDims)
     return applyTilePattern(getOperation(), loopType, tileSizes,
                             /*peeledLoops=*/{}, /*scalarizeDynamicDims=*/true);
+  if (testBubbleUpExtractSliceOpPattern)
+    return applyBubbleUpExtractSliceOpPattern(getOperation());
 }
 
 namespace mlir {
 namespace test {
 void registerTestLinalgTransforms() {
   PassRegistration<TestLinalgTransforms>();
 }
 } // namespace test
 } // namespace mlir