OnnxToTosa conversion: add option to only convert slice to tosa when all steps are 1

Author: mathmer-amd

src/Conversion/ONNXToTOSA/ConvertONNXToTOSA.cpp

@@ -26,7 +26,7 @@ namespace onnx_mlir {
 
 void populateONNXToTOSAConversionPattern(ConversionTarget &target,
     RewritePatternSet &patterns, TypeConverter &typeConverter, MLIRContext *ctx,
-    int64_t groupedConvThreshold) {
+    int64_t groupedConvThreshold, bool convertSliceOnlyWhenStepOne) {
   // Math
   populateLoweringONNXElementwiseOpToTOSAPattern(
       target, patterns, typeConverter, ctx);
@@ -56,7 +56,7 @@ void populateONNXToTOSAConversionPattern(ConversionTarget &target,
   populateLoweringONNXFlattenOpToTOSAPattern(
       target, patterns, typeConverter, ctx);
   populateLoweringONNXSliceOpToTOSAPattern(
-      target, patterns, typeConverter, ctx);
+      target, patterns, typeConverter, ctx, convertSliceOnlyWhenStepOne);
   populateLoweringONNXSplitOpToTOSAPattern(
       target, patterns, typeConverter, ctx);
   populateLoweringONNXSqueezeOpToTOSAPattern(
@@ -114,6 +114,10 @@ struct FrontendToTosaLoweringPass
                      "into a concatenation of tosa.conv2d operations"),
       llvm::cl::ZeroOrMore,
       llvm::cl::init(std::numeric_limits<int64_t>::max())};
+  Option<bool> convertSliceOnlyWhenStepOne{*this,
+      "convert-slice-only-when-step-one",
+      llvm::cl::desc("If enabled, convert onnx.slice only if all steps are 1"),
+      llvm::cl::ZeroOrMore, llvm::cl::init(false)};
 };
 
 void FrontendToTosaLoweringPass::runOnOperation() {
@@ -144,8 +148,8 @@ void FrontendToTosaLoweringPass::runOnOperation() {
       mlir::arith::ArithDialect, mlir::shape::ShapeDialect>();
 
   // Define patterns
-  populateONNXToTOSAConversionPattern(
-      target, patterns, typeConverter, context, groupedConvThreshold);
+  populateONNXToTOSAConversionPattern(target, patterns, typeConverter, context,
+      groupedConvThreshold, convertSliceOnlyWhenStepOne);
 
   if (failed(applyPartialConversion(module, target, std::move(patterns)))) {
     signalPassFailure();

src/Conversion/ONNXToTOSA/ONNXToTOSACommon.hpp

@@ -163,7 +163,8 @@ void populateLoweringONNXPadOpToTOSAPattern(mlir::ConversionTarget &,
 void populateLoweringONNXFlattenOpToTOSAPattern(mlir::ConversionTarget &,
     mlir::RewritePatternSet &, mlir::TypeConverter &, mlir::MLIRContext *);
 void populateLoweringONNXSliceOpToTOSAPattern(mlir::ConversionTarget &,
-    mlir::RewritePatternSet &, mlir::TypeConverter &, mlir::MLIRContext *);
+    mlir::RewritePatternSet &, mlir::TypeConverter &, mlir::MLIRContext *,
+    bool);
 void populateLoweringONNXSplitOpToTOSAPattern(mlir::ConversionTarget &,
     mlir::RewritePatternSet &, mlir::TypeConverter &, mlir::MLIRContext *);
 void populateLoweringONNXSqueezeOpToTOSAPattern(mlir::ConversionTarget &,

src/Conversion/ONNXToTOSA/Tensor/Slice.cpp

@@ -13,7 +13,6 @@
 //===----------------------------------------------------------------------===//
 
 #include "mlir/Dialect/Tosa/IR/TosaOps.h"
-#include "mlir/Interfaces/InferTypeOpInterface.h"
 #include "src/Conversion/ONNXToTOSA/DialectBuilder.hpp"
 #include "src/Conversion/ONNXToTOSA/ONNXToTOSACommon.hpp"
 #include "src/Conversion/ONNXToTOSA/ONNXToTOSALegalizeUtils.hpp"
@@ -30,6 +29,11 @@ class ONNXSliceLoweringToTOSA : public OpConversionPattern<ONNXSliceOp> {
 public:
   using OpConversionPattern::OpConversionPattern;
   using OpAdaptor = typename ONNXSliceOp::Adaptor;
+
+  ONNXSliceLoweringToTOSA(MLIRContext *ctx, bool convertSliceOnlyWhenStepOne)
+      : OpConversionPattern(ctx),
+        convertSliceOnlyWhenStepOne(convertSliceOnlyWhenStepOne) {}
+
   LogicalResult matchAndRewrite(ONNXSliceOp op, OpAdaptor adaptor,
       ConversionPatternRewriter &rewriter) const override {
 
@@ -74,6 +78,17 @@ class ONNXSliceLoweringToTOSA : public OpConversionPattern<ONNXSliceOp> {
     llvm::SmallVector<int64_t, 4> steps;
     IndexExpr::getLiteral(shapeHelper.steps, steps);
 
+    size_t nbSlicedDims = 0;
+    for (auto [in, out] : llvm::zip(inShape, outShape)) {
+      if (out != in)
+        nbSlicedDims++;
+    }
+    // TODO: remove the check nbSlicedDims == 1 when possible
+    if (convertSliceOnlyWhenStepOne && nbSlicedDims == 1 &&
+        llvm::any_of(steps, [](int64_t step) { return step > 1; })) {
+      return rewriter.notifyMatchFailure(op, "step > 1 are not supported.");
+    }
+
     if (llvm::any_of(steps, [](int64_t step) { return step < 0; })) {
       return rewriter.notifyMatchFailure(op, "negative step not supported.");
     }
@@ -161,14 +176,17 @@ class ONNXSliceLoweringToTOSA : public OpConversionPattern<ONNXSliceOp> {
     rewriter.replaceOp(op, val);
     return success();
   }
+
+private:
+  bool convertSliceOnlyWhenStepOne;
 };
 
 } // namespace
 
-void populateLoweringONNXSliceOpToTOSAPattern(ConversionTarget &target,
-    RewritePatternSet &patterns, TypeConverter &typeConverter,
-    MLIRContext *ctx) {
-  patterns.insert<ONNXSliceLoweringToTOSA>(ctx);
+void populateLoweringONNXSliceOpToTOSAPattern(ConversionTarget & /*target*/,
+    RewritePatternSet &patterns, TypeConverter & /*typeConverter*/,
+    MLIRContext *ctx, bool convertSliceOnlyWhenStepOne) {
+  patterns.insert<ONNXSliceLoweringToTOSA>(ctx, convertSliceOnlyWhenStepOne);
 }
 
 } // namespace onnx_mlir

test/mlir/conversion/onnx_to_tosa/Tensor/Slice.mlir

@@ -1,4 +1,5 @@
 // RUN: onnx-mlir-opt --shape-inference --convert-onnx-to-tosa -cse %s -split-input-file | FileCheck %s
+// RUN: onnx-mlir-opt --shape-inference --convert-onnx-to-tosa="convert-slice-only-when-step-one=true" -cse %s -split-input-file | FileCheck %s --check-prefix=ONLY-STEP1
 
 
 func.func @test_slice_constant_default_steps(%arg0 : tensor<2x4xf32>) -> tensor<1x3xf32> {
@@ -74,6 +75,9 @@ func.func @slice_just_steps(%arg0: tensor<100x200xf32>) -> tensor<20x20xf32> {
 // CHECK: %2 = tosa.reshape %1 {new_shape = array<i64: 20, 20>} : (tensor<20x1x20x1xf32>) -> tensor<20x20xf32>
 // CHECK: return %2 : tensor<20x20xf32>
 
+// ONLY-STEP1-LABEL: func @slice_just_steps
+// ONLY-STEP1: tosa.slice
+
 // -----
 
 func.func @slice_steps_and_edges(%arg0: tensor<100x200xf32>) -> tensor<16x17xf32> {
@@ -91,6 +95,9 @@ func.func @slice_steps_and_edges(%arg0: tensor<100x200xf32>) -> tensor<16x17xf32
 // CHECK: %3 = tosa.reshape %2 {new_shape = array<i64: 16, 17>} : (tensor<16x1x17x1xf32>) -> tensor<16x17xf32>
 // CHECK: return %3 : tensor<16x17xf32>
 
+// ONLY-STEP1-LABEL: func @slice_steps_and_edges
+// ONLY-STEP1: tosa.slice
+
 // -----
 
 func.func @slice_steps_and_edges_with_padding(%arg0: tensor<99x195xf32>) -> tensor<19x19xf32> {
@@ -111,6 +118,9 @@ func.func @slice_steps_and_edges_with_padding(%arg0: tensor<99x195xf32>) -> tens
 // CHECK: %6 = tosa.reshape %5 {new_shape = array<i64: 19, 19>} : (tensor<19x1x19x1xf32>) -> tensor<19x19xf32>
 // CHECK: return %6 : tensor<19x19xf32>
 
+// ONLY-STEP1-LABEL: func @slice_steps_and_edges_with_padding
+// ONLY-STEP1: tosa.slice
+
 // -----
 
 func.func @slice_just_steps_with_padding(%arg0: tensor<99x195xf32>) -> tensor<20x20xf32> {
@@ -130,6 +140,9 @@ func.func @slice_just_steps_with_padding(%arg0: tensor<99x195xf32>) -> tensor<20
 // CHECK: %5 = tosa.reshape %4 {new_shape = array<i64: 20, 20>} : (tensor<20x1x20x1xf32>) -> tensor<20x20xf32>
 // CHECK: return %5 : tensor<20x20xf32>
 
+// ONLY-STEP1-LABEL: func @slice_just_steps_with_padding
+// ONLY-STEP1: tosa.slice
+
 // -----
 
 func.func @slice_negative_steps(%arg0: tensor<100x200xf32>) -> tensor<20x20xf32> {
@@ -195,3 +208,25 @@ func.func @slice_4d(%arg0: tensor<1x56x56x92xf32>) -> tensor<1x28x28x92xf32> {
 // CHECK: %5 = tosa.slice %4 {size = array<i64: 1, 28, 1, 28, 1, 92>, start = array<i64: 0, 0, 0, 0, 0, 0>} : (tensor<1x28x2x28x2x92xf32>) -> tensor<1x28x1x28x1x92xf32>
 // CHECK: %6 = tosa.reshape %5 {new_shape = array<i64: 1, 28, 28, 92>} : (tensor<1x28x1x28x1x92xf32>) -> tensor<1x28x28x92xf32>
 // CHECK: return %6 : tensor<1x28x28x92xf32>
+
+// ONLY-STEP1-LABEL: func @slice_4d
+// ONLY-STEP1: tosa.slice
+
+// -----
+
+func.func @slice_4d_step2_1sliced_dim(%arg0: tensor<1x3x640x640xbf16>) -> tensor<1x3x320x640xbf16> {
+  %starts = "tosa.const"() <{value = dense<0> : tensor<1xi64>}> : () -> tensor<1xi64>
+  %ends = "tosa.const"() <{value = dense<9223372036854775807> : tensor<1xi64>}> : () -> tensor<1xi64>
+  %axes_steps = "tosa.const"() <{value = dense<2> : tensor<1xi64>}> : () -> tensor<1xi64>
+  %3 = "onnx.Slice"(%arg0, %starts, %ends, %axes_steps, %axes_steps) : (tensor<1x3x640x640xbf16>, tensor<1xi64>, tensor<1xi64>, tensor<1xi64>, tensor<1xi64>) -> tensor<1x3x320x640xbf16>
+  return %3 : tensor<1x3x320x640xbf16>
+}
+// CHECK-LABEL: func @slice_4d_step2_1sliced_dim
+// CHECK-SAME:                (%[[ARG_0:.*]]: tensor<1x3x640x640xbf16>) -> tensor<1x3x320x640xbf16>
+// CHECK: %[[VAL_0:.*]] = tosa.reshape %[[ARG_0]] {new_shape = array<i64: 1, 3, 320, 2, 640>} : (tensor<1x3x640x640xbf16>) -> tensor<1x3x320x2x640xbf16>
+// CHECK: %[[VAL_1:.*]] = tosa.slice %[[VAL_0]] {size = array<i64: 1, 3, 320, 1, 640>, start = array<i64: 0, 0, 0, 0, 0>} : (tensor<1x3x320x2x640xbf16>) -> tensor<1x3x320x1x640xbf16>
+// CHECK: %[[VAL_2:.*]] = tosa.reshape %[[VAL_1]] {new_shape = array<i64: 1, 3, 320, 640>} : (tensor<1x3x320x1x640xbf16>) -> tensor<1x3x320x640xbf16>
+// CHECK: return %[[VAL_2]] : tensor<1x3x320x640xbf16>
+
+// ONLY-STEP1-LABEL: func @slice_4d_step2_1sliced_dim
+// ONLY-STEP1-NOT: tosa.slice