experiments/STT-tensorflow
1
0
Fork 0
Code Issues Pull Requests Projects Releases Wiki Activity

Add 'Slice' in TF lite MLIR dialect.

Browse Source 
PiperOrigin-RevId: 256467939
This commit is contained in:
Haoliang Zhang 2019-07-03 17:40:53 -07:00 committed by TensorFlower Gardener
parent 9097812f34
commit 02f4e49d60
3 changed files with 31 additions and 0 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

23
tensorflow/compiler/mlir/lite/ir/tfl_ops.td
View File

@ -1078,6 +1078,29 @@ def TFL_MeanOp : TFL_Op<"mean", [NoSideEffect]> {
let customOption = "ReducerOptions";
}
def TFL_SliceOp : TFL_Op<"slice", [NoSideEffect]> {
let summary = "Return a slice from 'input'.";
let description = [{
The output tensor is a tensor with dimensions described by 'size'
whose values are extracted from 'input' starting at the offsets in
'begin'.
*Requirements*:
0 <= begin[i] <= begin[i] + size[i] <= Di for i in [0, n)
}];
let arguments = (ins
AnyTensor:$input,
TFL_I32OrI64Tensor:$begin,
TFL_I32OrI64Tensor:$size
);
let results = (outs
AnyTensor:$output
);
}
def TFL_SumOp: TFL_Op<"sum", [NoSideEffect]> {
let summary = "Sum operator";

7
tensorflow/compiler/mlir/lite/tests/legalize-tf.mlir
View File

@ -822,3 +822,10 @@ func @strided_slice(%arg0: tensor<12x2x2x5xf32>, %arg1: tensor<1xi32>, %arg2: te
// CHECK-LABEL: strided_slice
// CHECK: "tfl.strided_slice"(%arg0, %arg1, %arg2, %arg3) {begin_mask = 0 : i32, ellipsis_mask = 0 : i32, end_mask = 0 : i32, new_axis_mask = 0 : i32, shrink_axis_mask = 0 : i32} : (tensor<12x2x2x5xf32>, tensor<1xi32>, tensor<1xi32>, tensor<1xi32>) -> tensor<1x2x2x5xf32>
}
func @slice1Tensor(%arg0: tensor<2x3x5xf32>, %arg1: tensor<3xi32>, %arg2: tensor<3xi32>) -> tensor<?x3x5xf32> {
%0 = "tf.Slice"(%arg0, %arg1, %arg2) : (tensor<2x3x5xf32>, tensor<3xi32>, tensor<3xi32>) -> tensor<?x3x5xf32>
return %0 : tensor<?x3x5xf32>
// CHECK-LABEL: slice1Tensor
// CHECK: "tfl.slice"(%arg0, %arg1, %arg2) : (tensor<2x3x5xf32>, tensor<3xi32>, tensor<3xi32>) -> tensor<?x3x5xf32>
}

1
tensorflow/compiler/mlir/lite/transforms/legalize_patterns.td
View File

@ -128,6 +128,7 @@ def : Pat<(TF_SelectOp $cond, $x, $y), (TFL_SelectOp $cond, $x, $y)>;
def : Pat<(TF_ShapeOp $arg), (TFL_ShapeOp $arg)>;
def : Pat<(TF_SigmoidOp $arg), (TFL_LogisticOp $arg)>;
def : Pat<(TF_SinOp F32Tensor:$arg), (TFL_SinOp $arg)>;
def : Pat<(TF_SliceOp $input, $begin, $size), (TFL_SliceOp $input, $begin, $size)>;
def : Pat<(TF_SoftmaxOp $arg), (TFL_SoftmaxOp $arg, ConstF32Attr<"1.0">)>;
def : Pat<(TF_SqueezeOp $arg, $squeeze_dims), (TFL_SqueezeOp $arg, $squeeze_dims)>;
def : Pat<(TF_TransposeOp $arg, $perm), (TFL_TransposeOp $arg, $perm)>;