[mlir][tosa] Add tosa.slice to std.subtensor lowering

[rsuderman]

Summary

Lowering to subtensor is added for tosa.slice operator.

[rsuderman]

Revision was landed but did not automatically close.

mlir/lib/Conversion/TosaToStandard/TosaToStandard.cpp

public:
 
   LogicalResult matchAndRewrite(tosa::ConstOp op,
                                 PatternRewriter &rewriter) const final {
     rewriter.replaceOpWithNewOp<::ConstantOp>(op, op.value());
     return success();
   }
 };
 
+class SliceOpConverter : public OpRewritePattern<tosa::SliceOp> {
+public:
+  using OpRewritePattern<tosa::SliceOp>::OpRewritePattern;
+
+  LogicalResult matchAndRewrite(tosa::SliceOp sliceOp,
+                                PatternRewriter &rewriter) const final {
+    Value input = sliceOp.input();
+    SmallVector<int64_t> strides;
+    strides.resize(sliceOp.getType().template cast<ShapedType>().getRank(), 1);
+
+    rewriter.replaceOpWithNewOp<SubTensorOp>(
+        sliceOp, sliceOp.getType(), input, ValueRange({}), ValueRange({}),
+        ValueRange({}), sliceOp.start(), sliceOp.size(),
+        rewriter.getI64ArrayAttr(strides));
+
+    return success();
+  }
+};
+
 } // namespace
 
 void mlir::tosa::populateTosaToStandardConversionPatterns(
     MLIRContext *context, OwningRewritePatternList *patterns) {
-  patterns->insert<ConstOpConverter>(context);
+  patterns->insert<ConstOpConverter, SliceOpConverter>(context);
 }

mlir/lib/Conversion/TosaToStandard/TosaToStandardPass.cpp

 
 namespace {
 struct TosaToStandard : public TosaToStandardBase<TosaToStandard> {
 public:
   void runOnOperation() override {
     OwningRewritePatternList patterns;
     ConversionTarget target(getContext());
     target.addIllegalOp<tosa::ConstOp>();
-    target.addLegalOp<ConstantOp>();
+    target.addIllegalOp<tosa::SliceOp>();
+    target.addLegalDialect<StandardOpsDialect>();
 
     auto *op = getOperation();
     mlir::tosa::populateTosaToStandardConversionPatterns(op->getContext(),
                                                          &patterns);
     if (failed(applyPartialConversion(op, target, std::move(patterns))))
       signalPassFailure();
   }
 };

mlir/test/Conversion/TosaToStandard/tosa-to-standard.mlir

 // RUN: mlir-opt --split-input-file --tosa-to-standard %s -verify-diagnostics -o -| FileCheck %s
 
 // CHECK-LABEL: func @const_test
 func @const_test() -> (tensor<i32>) {
   // CHECK: [[C3:%.+]] = constant dense<3> : tensor<i32>
   %0 = "tosa.const"() {value = dense<3> : tensor<i32>} : () -> tensor<i32>
 
   // CHECK: return [[C3]]
   return %0 : tensor<i32>
 }
+
+// ----
+
+func @slice(%arg0: tensor<6xf32>) ->() {
+  // CHECK: [[SLICE:%.+]] = subtensor %arg0[2] [1] [1]
+  %0 = "tosa.slice"(%arg0) {start = [2], size = [1]} : (tensor<6xf32>)  -> (tensor<1xf32>)
+  return
+}