[tosa] Lower tosa.slice to tensor.slice for dynamic case

[rsuderman]

Summary

Existing slice lowering only supporting static shapes.

mlir/lib/Conversion/TosaToTensor/TosaToTensor.cpp

 //===- TosaToTensor.cpp - Lowering Tosa to Tensor Dialect -------------===//
 //
 // 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
 //
 //===----------------------------------------------------------------------===//
 //
 // These rewriters lower from the Tosa to the Tensor dialect.
 //
 //===----------------------------------------------------------------------===//
 
 #include "mlir/Conversion/TosaToTensor/TosaToTensor.h"
+#include "mlir/Dialect/Arithmetic/IR/Arithmetic.h"
 #include "mlir/Dialect/Tensor/IR/Tensor.h"
 #include "mlir/Dialect/Tosa/IR/TosaOps.h"
 #include "mlir/IR/PatternMatch.h"
 #include "mlir/Transforms/GreedyPatternRewriteDriver.h"
 
 using namespace mlir;
 using namespace tosa;
 
 namespace {
 
 class SliceOpConverter : public OpRewritePattern<tosa::SliceOp> {
 public:
   using OpRewritePattern<tosa::SliceOp>::OpRewritePattern;
 
   LogicalResult matchAndRewrite(tosa::SliceOp sliceOp,
                                 PatternRewriter &rewriter) const final {
+    Location loc = sliceOp.getLoc();
     Value input = sliceOp.input();
     SmallVector<int64_t> strides;
+    auto starts = sliceOp.start();
+    auto sizes = sliceOp.size();
     strides.resize(sliceOp.getType().template cast<ShapedType>().getRank(), 1);
 
-    rewriter.replaceOpWithNewOp<tensor::ExtractSliceOp>(
-        sliceOp, sliceOp.getType(), input, ValueRange({}), ValueRange({}),
-        ValueRange({}), sliceOp.start(), sliceOp.size(),
-        rewriter.getI64ArrayAttr(strides));
+    SmallVector<Value> dynSizes;
+    for (auto i : llvm::enumerate(sizes)) {
+      int64_t size = i.value().cast<IntegerAttr>().getInt();
+      size_t index = i.index();
+      if (size != ShapedType::kDynamicSize)
+        continue;
+
+      auto dim = rewriter.create<tensor::DimOp>(loc, input, index);
+      auto offset = rewriter.create<arith::ConstantOp>(
+          loc,
+          rewriter.getIndexAttr(starts[index].cast<IntegerAttr>().getInt()));
+      dynSizes.push_back(rewriter.create<arith::SubIOp>(loc, dim, offset));
+    }
+
+    auto newSliceOp = rewriter.create<tensor::ExtractSliceOp>(
+        sliceOp.getLoc(), sliceOp.getType(), input, ValueRange({}), dynSizes,
+        ValueRange({}), starts, sizes, rewriter.getI64ArrayAttr(strides));
+
+    rewriter.replaceOp(sliceOp, newSliceOp.getResult());
     return success();
   }
 };
 
 } // namespace
 
 void mlir::tosa::populateTosaToTensorConversionPatterns(
     RewritePatternSet *patterns) {
   patterns->add<SliceOpConverter>(patterns->getContext());
 }

mlir/lib/Conversion/TosaToTensor/TosaToTensorPass.cpp

 //===- TosaToTensorPass.cpp - Lowering Tosa to Tensor Dialect -------------===//
 //
 // 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 transformation pass legalizes Tosa operations to the Tensor dialect.
 //
 //===----------------------------------------------------------------------===//
 
 #include "../PassDetail.h"
 #include "mlir/Conversion/TosaToTensor/TosaToTensor.h"
+#include "mlir/Dialect/Arithmetic/IR/Arithmetic.h"
 #include "mlir/Dialect/Tensor/IR/Tensor.h"
 #include "mlir/Dialect/Tosa/IR/TosaOps.h"
 #include "mlir/Dialect/Tosa/Transforms/PassDetail.h"
 #include "mlir/Dialect/Tosa/Transforms/Passes.h"
 #include "mlir/IR/PatternMatch.h"
 #include "mlir/Pass/PassManager.h"
 #include "mlir/Transforms/DialectConversion.h"
 #include "mlir/Transforms/GreedyPatternRewriteDriver.h"
 
 using namespace mlir;
 using namespace tosa;
 
 namespace {
 struct TosaToTensor : public TosaToTensorBase<TosaToTensor> {
 public:
   void runOnOperation() override {
     RewritePatternSet patterns(&getContext());
     ConversionTarget target(getContext());
     target.addIllegalOp<tosa::SliceOp>();
+    target.addLegalDialect<arith::ArithmeticDialect>();
     target.addLegalDialect<tensor::TensorDialect>();
 
     mlir::tosa::populateTosaToTensorConversionPatterns(&patterns);
 
     if (failed(applyPartialConversion(getOperation(), target,
                                       std::move(patterns))))
       signalPassFailure();
   }
 };
 } // namespace
 
 std::unique_ptr<Pass> mlir::tosa::createTosaToTensor() {
   return std::make_unique<TosaToTensor>();
 }

mlir/test/Conversion/TosaToTensor/tosa-to-tensor.mlir

 // RUN: mlir-opt --split-input-file --tosa-to-tensor %s -o -| FileCheck %s
 
 // CHECK-LABLE: func @slice
 func.func @slice(%arg0: tensor<6xf32>) ->() {
   // CHECK: [[SLICE:%.+]] = tensor.extract_slice %arg0[2] [1] [1]
   %0 = "tosa.slice"(%arg0) {start = [2], size = [1]} : (tensor<6xf32>)  -> (tensor<1xf32>)
   return
 }
+
+// -----
+
+// CHECK-LABLE: func @slice_dyn
+func.func @slice_dyn(%arg0: tensor<?xf32>) -> (tensor<?xf32>) {
+  // CHECK: %[[C0:.+]] = arith.constant 0 : index
+  // CHECK: %[[DIM:.+]] = tensor.dim %arg0, %[[C0]]
+  // CHECK: %[[C2:.+]] = arith.constant 2 : index
+  // CHECK: %[[SUB:.+]] = arith.subi %[[DIM]], %[[C2]]
+  // CHECK: %2 = tensor.extract_slice %arg0[2] [%[[SUB]]] [1]
+  %0 = "tosa.slice"(%arg0) {start = [2], size = [-1]} : (tensor<?xf32>)  -> (tensor<?xf32>)
+  return %0 : tensor<?xf32>
+}

utils/bazel/llvm-project-overlay/mlir/BUILD.bazel

cc_library(
     hdrs = glob([
         "include/mlir/Conversion/TosaToTensor/*.h",
     ]),
     includes = [
         "include",
         "lib/Conversion/TosaToTensor",
     ],
     deps = [
+        ":ArithmeticDialect",
         ":ConversionPassIncGen",
         ":FuncDialect",
         ":IR",
         ":Pass",
         ":TensorDialect",
         ":TosaDialect",
         ":Transforms",
     ],