//===----------------------------------------------------------------------===//
#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"
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();
}
};
#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"
RewritePatternSet patterns(&getContext());
ConversionTarget target(getContext());
target.addIllegalOp<tosa::SliceOp>();
+ target.addLegalDialect<arith::ArithmeticDialect>();
target.addLegalDialect<tensor::TensorDialect>();
mlir::tosa::populateTosaToTensorConversionPatterns(&patterns);
%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>
+}