diff --git a/tensorflow/compiler/mlir/lite/ir/tfl_ops.td b/tensorflow/compiler/mlir/lite/ir/tfl_ops.td index afb82d48468..4c25fe3b545 100644 --- a/tensorflow/compiler/mlir/lite/ir/tfl_ops.td +++ b/tensorflow/compiler/mlir/lite/ir/tfl_ops.td @@ -1990,9 +1990,11 @@ def TFL_TileOp: TFL_Op<"tile", [NoSideEffect, For example, tiling [a b c d] by [2] produces [a b c d a b c d]. }]; - let arguments = (ins AnyTensor:$input, TFL_I32OrI64Tensor:$multiples); + let arguments = (ins + TensorOf<[F32, I1, I32, I64, TFL_Uint8]>:$input, + TFL_I32OrI64Tensor:$multiples); - let results = (outs AnyTensor:$output); + let results = (outs TensorOf<[F32, I1, I32, I64, TFL_Uint8]>:$output); let hasOptions = 0; } diff --git a/tensorflow/compiler/mlir/lite/tests/legalize-tf.mlir b/tensorflow/compiler/mlir/lite/tests/legalize-tf.mlir index 1a6c8113202..e93f7ef5f5e 100644 --- a/tensorflow/compiler/mlir/lite/tests/legalize-tf.mlir +++ b/tensorflow/compiler/mlir/lite/tests/legalize-tf.mlir @@ -651,6 +651,17 @@ func @pad(tensor<2x1x3xf32>, tensor<3x2xi32>) -> tensor { // CHECK: return %0 : tensor } +func @tile(tensor<2x3xf32>, tensor<2xi32>) -> tensor<2x6xf32> { +^bb0(%arg0: tensor<2x3xf32>, %arg1: tensor<2xi32>): + %cst = constant dense<[1, 2]> : tensor<2xi32> + %0 = "tf.Tile"(%arg0, %cst) : (tensor<2x3xf32>, tensor<2xi32>) -> tensor<2x6xf32> + return %0 : tensor<2x6xf32> + + // CHECK-LABEL: tile + // CHECK: %0 = "tfl.tile"(%arg0, %cst) : (tensor<2x3xf32>, tensor<2xi32>) -> tensor<2x6xf32> + // CHECK: return %0 : tensor<2x6xf32> +} + func @padv2(tensor<2x1x3xf32>, tensor<3x2xi32>) -> tensor { ^bb0(%arg0: tensor<2x1x3xf32>, %arg1: tensor<3x2xi32>): %cst = constant dense<2.0> : tensor diff --git a/tensorflow/compiler/mlir/lite/transforms/legalize_patterns.td b/tensorflow/compiler/mlir/lite/transforms/legalize_patterns.td index 410cac51c95..610830e5e1e 100644 --- a/tensorflow/compiler/mlir/lite/transforms/legalize_patterns.td +++ b/tensorflow/compiler/mlir/lite/transforms/legalize_patterns.td @@ -254,6 +254,8 @@ def : Pat<(TF_EqualOp $arg0, $arg1), (TFL_EqualOp $arg0, $arg1)>; def : Pat<(TF_PadOp $arg0, $arg1), (TFL_PadOp $arg0, $arg1)>; +def : Pat<(TF_TileOp $arg0, $arg1), (TFL_TileOp $arg0, $arg1)>; + def : Pat<(TF_PadV2Op $arg0, $arg1, $cst), (TFL_PadV2Op $arg0, $arg1, $cst)>; def : Pat<(TF_MeanOp $arg0, $arg1, BoolAttr:$arg2), (TFL_MeanOp $arg0, $arg1, $arg2)>; diff --git a/tensorflow/compiler/mlir/tensorflow/ir/tf_generated_ops.td b/tensorflow/compiler/mlir/tensorflow/ir/tf_generated_ops.td index 2fd31d193a0..8ac1aaff374 100644 --- a/tensorflow/compiler/mlir/tensorflow/ir/tf_generated_ops.td +++ b/tensorflow/compiler/mlir/tensorflow/ir/tf_generated_ops.td @@ -3518,6 +3518,30 @@ num_elements: optional. If not -1, the number of elements in the list. TF_DerivedResultTypeAttr element_dtype = TF_DerivedResultTypeAttr<0>; } +def TF_TileOp : TF_Op<"Tile", [NoSideEffect]> { + let summary = "Constructs a tensor by tiling a given tensor."; + + let description = [{ +This operation creates a new tensor by replicating `input` `multiples` times. +The output tensor's i'th dimension has `input.dims(i) * multiples[i]` elements, +and the values of `input` are replicated `multiples[i]` times along the 'i'th +dimension. For example, tiling `[a b c d]` by `[2]` produces +`[a b c d a b c d]`. + }]; + + let arguments = (ins + TF_Tensor:$input, + TF_I32OrI64Tensor:$multiples + ); + + let results = (outs + TF_Tensor:$output + ); + + TF_DerivedOperandTypeAttr Tmultiples = TF_DerivedOperandTypeAttr<1>; + TF_DerivedOperandTypeAttr T = TF_DerivedOperandTypeAttr<0>; +} + def TF_TopKV2Op : TF_Op<"TopKV2", [NoSideEffect]> { let summary = [{ Finds values and indices of the `k` largest elements for the last dimension.