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

Fix padding in legalization of tf.MaxPool and tf.MaxPoolGrad.

Browse Source 
Padding in reduce window op is expected to be a (N, 2) tensor with each tuple denoting the lower and higher edge padding. For example, if padding for a 4D tensor is ((0,1), (2,3), (4,5), (6,7)) i.e., lower edge padding for 3rd dimension is 4 and higher edge padding is 5, the padding attribute would be [[0,1], [2,3], [4,5], [6,7]].

However, during legalization of MaxPool and MaxPoolGrad to ReduceWindow HLO, we were generating (2, N) tensor with first and second row indicating lower edge and higher edge padding in N dimensions respectively i.e., [[0, 2, 4, 6], [1, 3, 5, 7]]. During export this was interpreted as a (N, 2) tensor and HLO with incorrect padding was emitted.

PiperOrigin-RevId: 288927001
Change-Id: I15fb24b00a7ad43116d352fb4e03b3fe06264eec
This commit is contained in:
Prakalp Srivastava 2020-01-09 10:44:04 -08:00 committed by TensorFlower Gardener
parent a99d0fc812
commit 617297e60a
2 changed files with 5 additions and 5 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

4
tensorflow/compiler/mlir/xla/tests/legalize-tf.mlir
View File

@ -1227,7 +1227,7 @@ func @maxpool_valid_padding(%arg0: tensor<2x12x20x7xi32>) -> tensor<2x3x5x7xi32>
// CHECK-LABEL: maxpool_same_padding
// CHECK-SAME: %[[ARG:.*]]: tensor
func @maxpool_same_padding(%arg0: tensor<2x13x25x7xi32>) -> tensor<2x4x7x7xi32> {
// CHECK: padding = dense<{{\[\[}}0, 0, 1, 0], [0, 1, 1, 0]]> : tensor<2x4xi64>
// CHECK: padding = dense<{{\[\[}}0, 0], [0, 1], [1, 1], [0, 0]]> : tensor<4x2xi64>
%0 = "tf.MaxPool"(%arg0) {data_format = "NHWC", ksize = [1, 2, 3, 1], padding = "SAME", strides = [1, 4, 4, 1]} : (tensor<2x13x25x7xi32>) -> tensor<2x4x7x7xi32>
return %0 : tensor<2x4x7x7xi32>
@ -1263,7 +1263,7 @@ func @max_pool_grad_valid(%orig_input: tensor<10x24x24x64xf32>, %orig_output: te
// CHECK-LABEL: @max_pool_grad_same
func @max_pool_grad_same(%orig_input: tensor<2x13x25x7xf32>, %orig_output: tensor<2x4x7x7xf32>, %grad: tensor<2x4x7x7xf32>) -> tensor<2x13x25x7xf32> {
// CHECK: padding = dense<{{\[\[}}0, 0, 1, 0], [0, 1, 1, 0]]> : tensor<2x4xi64>
// CHECK: padding = dense<{{\[\[}}0, 0], [0, 1], [1, 1], [0, 0]]> : tensor<4x2xi64>
%result = "tf.MaxPoolGrad"(%orig_input, %orig_output, %grad) {
data_format = "NHWC",
ksize = [1, 2, 3, 1],

6
tensorflow/compiler/mlir/xla/transforms/legalize_tf.cc
View File

@ -997,11 +997,11 @@ static DenseIntElementsAttr GetReduceWindowPadding(
int64_t rank = paddings.size();
llvm::SmallVector<int64_t, 8> flatten_paddings(rank * 2);
for (int i = 0; i < rank; i++) {
flatten_paddings[i] = paddings[i].first;
flatten_paddings[rank + i] = paddings[i].second;
flatten_paddings[2 * i] = paddings[i].first;
flatten_paddings[2 * i + 1] = paddings[i].second;
}
return DenseIntElementsAttr::get(
RankedTensorType::get({2, rank}, builder->getIntegerType(64)),
RankedTensorType::get({rank, 2}, builder->getIntegerType(64)),
flatten_paddings);
}