From ac40eeeec83d96f335ba59ec30c4f0cae9505ff4 Mon Sep 17 00:00:00 2001 From: Tzu-Wei Sung Date: Wed, 12 Aug 2020 00:47:34 -0700 Subject: [PATCH 1/2] Verify and fold tf.Tile Fold when both input and multiples are constant Fix tests Do not fold when input is constant Empty line Fix typo --- .../mlir/tensorflow/ir/tf_generated_ops.td | 6 +- .../compiler/mlir/tensorflow/ir/tf_ops_n_z.cc | 86 +++++++++++++++++++ .../mlir/tensorflow/tests/canonicalize.mlir | 8 ++ .../mlir/tensorflow/tests/tf-ops.mlir | 59 +++++++++++++ 4 files changed, 156 insertions(+), 3 deletions(-) diff --git a/tensorflow/compiler/mlir/tensorflow/ir/tf_generated_ops.td b/tensorflow/compiler/mlir/tensorflow/ir/tf_generated_ops.td index 081903d13cf..8561303809e 100644 --- a/tensorflow/compiler/mlir/tensorflow/ir/tf_generated_ops.td +++ b/tensorflow/compiler/mlir/tensorflow/ir/tf_generated_ops.td @@ -10925,9 +10925,9 @@ array([[1, 2, 3, 1, 2, 3], TF_DerivedOperandTypeAttr Tmultiples = TF_DerivedOperandTypeAttr<1>; TF_DerivedOperandTypeAttr T = TF_DerivedOperandTypeAttr<0>; - // TODO(parkers): Add folds for multiples = [1,...]. - // TODO(parkers): Add errors for negative multiples and multiples.size() != - // input.rank() + let verifier = [{ return Verify(*this); }]; + + let hasFolder = 1; } def TF_ToBoolOp : TF_Op<"ToBool", [NoSideEffect]> { diff --git a/tensorflow/compiler/mlir/tensorflow/ir/tf_ops_n_z.cc b/tensorflow/compiler/mlir/tensorflow/ir/tf_ops_n_z.cc index ffedcb47f7e..b45c310e08a 100644 --- a/tensorflow/compiler/mlir/tensorflow/ir/tf_ops_n_z.cc +++ b/tensorflow/compiler/mlir/tensorflow/ir/tf_ops_n_z.cc @@ -1712,6 +1712,92 @@ static LogicalResult Verify(TensorScatterUpdateOp op) { return success(); } +//===----------------------------------------------------------------------===// +// TileOp +//===----------------------------------------------------------------------===// + +// Verifies that, +// +// - input has at least rank 1 +// - multiples is rank 1 +// - multiples.size() == input.rank() +// - input.rank() == output.rank() +// - Elements in multiples are non-negative +// - input.shape[i] * multiples[i] == output.shape[i] +// for i in [0, input.rank() - 1] + +static LogicalResult Verify(TileOp op) { + auto input_type = op.input().getType().dyn_cast(); + auto multiples_type = op.multiples().getType().dyn_cast(); + auto output_type = op.output().getType().dyn_cast(); + + if (input_type && input_type.getRank() < 1) { + return op.emitOpError() + << "expected input to be rank 1 or higher, got rank = " + << input_type.getRank(); + } + + if (multiples_type && multiples_type.getRank() != 1) { + return op.emitOpError() << "expected multiples to be rank 1, got rank = " + << multiples_type.getRank(); + } + + if (input_type && multiples_type && multiples_type.hasStaticShape() && + input_type.getRank() != multiples_type.getNumElements()) { + return op.emitOpError() + << "expected size of multiples equal to rank of input" + << ", got multiples of size " << multiples_type.getNumElements() + << ", and input of rank " << input_type.getRank(); + } + + if (input_type && output_type) { + if (input_type.getRank() != output_type.getRank()) { + return op.emitOpError() + << "expected rank of input to equal to rank of output" + << ", got input of rank " << input_type.getRank() + << ", and output of rank " << output_type.getRank(); + } + + DenseIntElementsAttr multiples_attr; + if (matchPattern(op.multiples(), m_Constant(&multiples_attr))) { + for (int32_t i = 0; i < input_type.getRank(); ++i) { + const int64_t input_dim = input_type.getDimSize(i); + const int64_t output_dim = output_type.getDimSize(i); + const int64_t m = multiples_attr.getValue(i).getSExtValue(); + + if (m < 0) { + return op.emitOpError() + << "expected multiples to be non-negative, got " + << "multiples[" << i << "] = " << m; + } + + if (!ShapedType::isDynamic(input_dim) && + !ShapedType::isDynamic(output_dim) && output_dim != input_dim * m) { + return op.emitOpError() + << "requires input.shape[" << i << "] (" << input_dim << ")" + << " * " << m << " to be equal to " + << "output.shape[" << i << "] (" << output_dim << ")"; + } + } + } + } + + return success(); +} + +OpFoldResult TileOp::fold(ArrayRef operands) { + DenseIntElementsAttr multiples_attr; + if (matchPattern(multiples(), m_Constant(&multiples_attr))) { + // Return input directly when multiples are all ones, + // regardless what input is. + if (multiples_attr.isSplat() && + multiples_attr.getSplatValue().getSExtValue() == 1) { + return input(); + } + } + return {}; +} + //===----------------------------------------------------------------------===// // TopKV2Op //===----------------------------------------------------------------------===// diff --git a/tensorflow/compiler/mlir/tensorflow/tests/canonicalize.mlir b/tensorflow/compiler/mlir/tensorflow/tests/canonicalize.mlir index 595bdce5be4..f014f46775a 100644 --- a/tensorflow/compiler/mlir/tensorflow/tests/canonicalize.mlir +++ b/tensorflow/compiler/mlir/tensorflow/tests/canonicalize.mlir @@ -560,6 +560,14 @@ func @testSelectElseUnranked(%arg0: tensor<3xi1>, %arg1: tensor<3x2xf16>, %arg2: return %0: tensor<*xf16> } +// CHECK-LABEL: testTileMultiplesAllOnes +func @testTileMultiplesAllOnes(%arg0: tensor<2x3xf32>) -> tensor<2x3xf32> { + %cst = constant dense <[1, 1]> : tensor<2xi32> + // CHECK: return %arg0 + %0 = "tf.Tile"(%arg0, %cst) : (tensor<2x3xf32>, tensor<2xi32>) -> tensor<2x3xf32> + return %0: tensor<2x3xf32> +} + // CHECK-LABEL: testLogicalNotOfEqual func @testLogicalNotOfEqual(%arg0: tensor<8x16xf32>, %arg1: tensor<8x16xf32>) -> tensor<8x16xi1> { %0 = "tf.Equal"(%arg0, %arg1) : (tensor<8x16xf32>, tensor<8x16xf32>) -> tensor<8x16xi1> diff --git a/tensorflow/compiler/mlir/tensorflow/tests/tf-ops.mlir b/tensorflow/compiler/mlir/tensorflow/tests/tf-ops.mlir index 20a0e22c48e..e49233a55f2 100644 --- a/tensorflow/compiler/mlir/tensorflow/tests/tf-ops.mlir +++ b/tensorflow/compiler/mlir/tensorflow/tests/tf-ops.mlir @@ -3313,3 +3313,62 @@ func @testBatchToSpaceInvalidOutputDepth(%arg0: tensor<16x8x8x3xf32>, %arg1: ten %0 = "tf.BatchToSpace"(%arg0, %arg1) {block_size = 2 : i64} : (tensor<16x8x8x3xf32>, tensor<*xi32>) -> tensor<4x8x8x8xf32> return } + + +// ----- + +func @testTile(%arg0: tensor<2x3x?xf32>) { + %cst = constant dense <[2, 3, 4]> : tensor<3xi32> + %0 = "tf.Tile"(%arg0, %cst) : (tensor<2x3x?xf32>, tensor<3xi32>) -> tensor<4x9x?xf32> + return +} + +// ----- + +func @testTileRank0Input(%arg0: tensor, %arg1: tensor) { + // expected-error @+1 {{expected input to be rank 1 or higher, got rank = 0}} + %0 = "tf.Tile"(%arg0, %arg1) : (tensor, tensor) -> tensor + return +} + +// ----- + +func @testTileMultipleNotRank1(%arg0: tensor<2x3xf32>, %arg1: tensor<1x1xi32>) { + // expected-error @+1 {{expected multiples to be rank 1, got rank = 2}} + %0 = "tf.Tile"(%arg0, %arg1) : (tensor<2x3xf32>, tensor<1x1xi32>) -> tensor<2x3xf32> + return +} + +// ----- + +func @testTileInputRankNotEqualToMultiplesSize(%arg0: tensor<2x3xf32>, %arg1: tensor<3xi32>) { + // expected-error @+1 {{expected size of multiples equal to rank of input, got multiples of size 3, and input of rank 2}} + %0 = "tf.Tile"(%arg0, %arg1) : (tensor<2x3xf32>, tensor<3xi32>) -> tensor<2x3xf32> + return +} + +// ----- + +func @testTileInputRankNotEqualToOutputRank(%arg0: tensor<2x3xf32>, %arg1: tensor<2xi32>) { + // expected-error @+1 {{expected rank of input to equal to rank of output, got input of rank 2, and output of rank 3}} + %0 = "tf.Tile"(%arg0, %arg1) : (tensor<2x3xf32>, tensor<2xi32>) -> tensor<2x3x1xf32> + return +} + +// ----- + +func @testTileNegativeMultiples(%arg0: tensor<2x3xf32>) { + %cst = constant dense <[-1, 1]> : tensor<2xi32> + // expected-error @+1 {{expected multiples to be non-negative, got multiples[0] = -1}} + %0 = "tf.Tile"(%arg0, %cst) : (tensor<2x3xf32>, tensor<2xi32>) -> tensor<2x3xf32> + return +} + +// ----- + +func @testTileInvalidOutputShape(%arg0: tensor<2x3xf32>) { + %cst = constant dense <[2, 3]> : tensor<2xi32> + // expected-error @+1 {{requires input.shape[1] (3) * 3 to be equal to output.shape[1] (6)}} + %0 = "tf.Tile"(%arg0, %cst) : (tensor<2x3xf32>, tensor<2xi32>) -> tensor<4x6xf32> + return +} From 1685f52428698d31239d8e668183c24b2ea42f27 Mon Sep 17 00:00:00 2001 From: Tzu-Wei Sung Date: Fri, 14 Aug 2020 23:14:05 -0700 Subject: [PATCH 2/2] Solve conflicts --- tensorflow/compiler/mlir/tensorflow/tests/tf-ops.mlir | 1 + 1 file changed, 1 insertion(+) diff --git a/tensorflow/compiler/mlir/tensorflow/tests/tf-ops.mlir b/tensorflow/compiler/mlir/tensorflow/tests/tf-ops.mlir index 660420c0344..a5afe76280c 100644 --- a/tensorflow/compiler/mlir/tensorflow/tests/tf-ops.mlir +++ b/tensorflow/compiler/mlir/tensorflow/tests/tf-ops.mlir @@ -3376,6 +3376,7 @@ func @testCaseRegionMismatchedResultTypes(%arg0: tensor, %arg1: tensor // ----- +// Test valid tf.Cumsum func @testCumsum(%arg: tensor<8x16xf32>, %axis: tensor) -> tensor<8x16xf32> { %0 = "tf.Cumsum"(%arg, %axis) : (tensor<8x16xf32>, tensor) -> tensor<8x16xf32> return %0 : tensor<8x16xf32>