Methods to work with symbolic tensor shapes.

PiperOrigin-RevId: 190071400

4 files changed, 362 insertions, 0 deletions

diff --git a/tensorflow/core/grappler/optimizers/BUILD b/tensorflow/core/grappler/optimizers/BUILD
index 96ea8f7a83..ac29edd213 100644
--- a/tensorflow/core/grappler/optimizers/BUILD
+++ b/tensorflow/core/grappler/optimizers/BUILD
@@ -5,6 +5,12 @@ load("//tensorflow:tensorflow.bzl", "tf_cc_test_gpu")
 load("//tensorflow:tensorflow.bzl", "tf_kernel_library")
 load("@local_config_cuda//cuda:build_defs.bzl", "if_cuda")
 
+# Platform specific build config
+load(
+    "//tensorflow/core:platform/default/build_config.bzl",
+    "tf_protos_grappler",
+)
+
 filegroup(
     name = "all_files",
     srcs = glob(
@@ -586,3 +592,27 @@ tf_cc_test(
         "//tensorflow/core/grappler/utils:grappler_test",
     ],
 )
+
+cc_library(
+    name = "symbolic_shapes",
+    srcs = ["symbolic_shapes.cc"],
+    hdrs = ["symbolic_shapes.h"],
+    visibility = ["//visibility:public"],
+    deps = [
+        "//tensorflow/core:framework",
+        "//tensorflow/core:protos_all_cc",
+    ] + tf_protos_grappler(),
+)
+
+tf_cc_test(
+    name = "symbolic_shapes_test",
+    srcs = ["symbolic_shapes_test.cc"],
+    deps = [
+        ":symbolic_shapes",
+        "//tensorflow/core:framework",
+        "//tensorflow/core:lib",
+        "//tensorflow/core:protos_all_cc",
+        "//tensorflow/core:test",
+        "//tensorflow/core:test_main",
+    ],
+)
diff --git a/tensorflow/core/grappler/optimizers/symbolic_shapes.cc b/tensorflow/core/grappler/optimizers/symbolic_shapes.cc
new file mode 100644
index 0000000000..cfca2dc0d3
--- /dev/null
+++ b/tensorflow/core/grappler/optimizers/symbolic_shapes.cc
@@ -0,0 +1,177 @@
+/* Copyright 2018 The TensorFlow Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+
+#include "tensorflow/core/grappler/optimizers/symbolic_shapes.h"
+#include "tensorflow/core/util/bcast.h"
+
+namespace tensorflow {
+namespace grappler {
+namespace {
+
+BCast::Vec ShapeDims(const TensorShapeProto& shape) {
+  BCast::Vec dims;
+  dims.reserve(shape.dim_size());
+  for (int i = 0; i < shape.dim_size(); ++i)
+    dims.push_back(shape.dim(i).size());
+  return dims;
+}
+
+}  // namespace
+
+bool IsKnown(const TensorShapeProto::Dim& dim) { return dim.size() >= 0; }
+
+bool IsKnownSymbolically(const TensorShapeProto::Dim& dim) {
+  return dim.size() <= -2;
+}
+
+bool IsUnknown(const TensorShapeProto::Dim& dim) { return dim.size() == -1; }
+
+bool ShapeIsSymbolicallyDefined(const TensorShapeProto& shape) {
+  return !shape.unknown_rank() &&
+         std::all_of(
+             shape.dim().begin(), shape.dim().end(),
+             [](const TensorShapeProto::Dim& dim) { return !IsUnknown(dim); });
+}
+
+bool ShapeIsSymbolicallyDefined(const OpInfo::TensorProperties& properties) {
+  return ShapeIsSymbolicallyDefined(properties.shape());
+}
+
+bool ShapesSymbolicallyEqual(const TensorShapeProto& left,
+                             const TensorShapeProto& right) {
+  if (left.unknown_rank() || right.unknown_rank() ||
+      left.dim_size() != right.dim_size()) {
+    return false;
+  }
+  for (int i = 0; i < left.dim_size(); ++i) {
+    const auto& ldim = left.dim(i);
+    const auto& rdim = right.dim(i);
+    if (IsUnknown(ldim) || IsUnknown(rdim) || ldim.size() != rdim.size()) {
+      return false;
+    }
+  }
+  return true;
+}
+
+bool ShapesSymbolicallyEqual(const OpInfo::TensorProperties& left,
+                             const OpInfo::TensorProperties& right) {
+  return ShapesSymbolicallyEqual(left.shape(), right.shape());
+}
+
+bool ShapesBroadcastable(const TensorShapeProto& left,
+                         const TensorShapeProto& right) {
+  if (!ShapeIsSymbolicallyDefined(left) || !ShapeIsSymbolicallyDefined(right)) {
+    return false;
+  }
+  BCast bcast(ShapeDims(left), ShapeDims(right),
+              /*fewer_dims_optimization*/ false);
+  return bcast.IsValid();
+}
+
+bool ShapesBroadcastable(const OpInfo::TensorProperties& left,
+                         const OpInfo::TensorProperties& right) {
+  return ShapesBroadcastable(left.shape(), right.shape());
+}
+
+bool CompareSymbolicallyShapedTensorSizes(const TensorShapeProto& left,
+                                          const TensorShapeProto& right) {
+  // if one of the ranks is unknown, it's impossible to compare tensor sizes
+  if (left.unknown_rank() || right.unknown_rank()) {
+    return false;
+  }
+
+  // Tensor size, computed as a product of defined dimensions
+  int64 left_defined_size = 1;
+  int64 right_defined_size = 1;
+
+  // Keep how many times each unknown dimension appeared on the left and right
+  std::unordered_map<int64, int64> left_unknown_dims;
+  std::unordered_map<int64, int64> right_unknown_dims;
+
+  // Assign unique id to every unknown dimension (-1). We are going to
+  // assign positive ids, because negative values are already used by
+  // symbolic dimensions.
+  int64 unknown_dim_id = 1;
+
+  // For each shape dimension update "defined tensor size", if shape is defined,
+  // or increment a counter for unknown dim.
+  auto process_dimensions =
+      [&unknown_dim_id](const TensorShapeProto& shape, int64* defined_size,
+                        std::unordered_map<int64, int64>* unknown_dims) {
+        for (int i = 0; i < shape.dim_size(); ++i) {
+          const auto& dim = shape.dim(i);
+          int64 dim_size = dim.size();
+          if (dim_size > 0) {
+            *defined_size *= dim_size;
+          } else if (IsUnknown(dim)) {
+            ++(*unknown_dims)[unknown_dim_id++];
+          } else if (IsKnownSymbolically(dim)) {
+            ++(*unknown_dims)[dim_size];
+          }
+        }
+      };
+
+  process_dimensions(left, &left_defined_size, &left_unknown_dims);
+  process_dimensions(right, &right_defined_size, &right_unknown_dims);
+
+  // Compute a union of unknown dimension ids appeared in both shapes
+  std::set<int64> unknown_dims;
+  for (const auto& el : left_unknown_dims) unknown_dims.insert(el.first);
+  for (const auto& el : right_unknown_dims) unknown_dims.insert(el.first);
+
+  // Cancel unknown dimensions that appeared in both shapes
+  for (int64 unknown_dim : unknown_dims) {
+    int64 co_occurrence = std::min(left_unknown_dims[unknown_dim],
+                                   right_unknown_dims[unknown_dim]);
+    left_unknown_dims[unknown_dim] -= co_occurrence;
+    right_unknown_dims[unknown_dim] -= co_occurrence;
+  }
+
+  // Count unbalanced unknown dimensions
+  int64 left_unbalanced_unknown_dims = 0;
+  int64 right_unbalanced_unknown_dims = 0;
+  for (const auto& el : left_unknown_dims)
+    left_unbalanced_unknown_dims += el.second;
+  for (const auto& el : right_unknown_dims)
+    right_unbalanced_unknown_dims += el.second;
+
+  if (left_unbalanced_unknown_dims == 0 && right_unbalanced_unknown_dims == 0) {
+    // If unknown dimensions cancelled each other, compare tensor sizes
+    // represented by defined dimensions
+    return left_defined_size < right_defined_size;
+  }
+
+  if (left_defined_size <= right_defined_size &&
+      left_unbalanced_unknown_dims == 0 && right_unbalanced_unknown_dims > 0) {
+    // If size of a 'left" tensor computed from defined dimensions less or
+    // equal, and shape on the right has unbalanced unknown dimensions, we can
+    // guarantee that shape on the left is strictly smaller (assuming that
+    // unknown dimension size is larger than 1)
+    return true;
+  }
+
+  // In every other case, assuming that unknown dimensions can be arbitrary
+  // large in size, we can't guarantee any ordering
+  return false;
+}
+
+bool CompareSymbolicallyShapedTensorSizes(
+    const OpInfo::TensorProperties& left,
+    const OpInfo::TensorProperties& right) {
+  return CompareSymbolicallyShapedTensorSizes(left.shape(), right.shape());
+}
+
+}  // end namespace grappler
+}  // end namespace tensorflow
diff --git a/tensorflow/core/grappler/optimizers/symbolic_shapes.h b/tensorflow/core/grappler/optimizers/symbolic_shapes.h
new file mode 100644
index 0000000000..a9dcf44e23
--- /dev/null
+++ b/tensorflow/core/grappler/optimizers/symbolic_shapes.h
@@ -0,0 +1,60 @@
+/* Copyright 2018 The TensorFlow Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+
+#ifndef TENSORFLOW_CORE_GRAPPLER_OPTIMIZERS_SYMBOLIC_SHAPES_H_
+#define TENSORFLOW_CORE_GRAPPLER_OPTIMIZERS_SYMBOLIC_SHAPES_H_
+
+#include "tensorflow/core/framework/tensor_shape.pb.h"
+#include "tensorflow/core/grappler/costs/op_performance_data.pb.h"
+
+namespace tensorflow {
+namespace grappler {
+
+bool IsKnown(const TensorShapeProto::Dim& dim);
+bool IsKnownSymbolically(const TensorShapeProto::Dim& dim);
+bool IsUnknown(const TensorShapeProto::Dim& dim);
+
+// Shape is symbolically defined, if it has a known rank, and each dimension is
+// known (dim_size >= 0), or is a symbolic dimension size (dim_size <= -2).
+bool ShapeIsSymbolicallyDefined(const TensorShapeProto& shape);
+bool ShapeIsSymbolicallyDefined(const OpInfo::TensorProperties& properties);
+
+// Shapes are symbolically equal, if they have the same rank, they are
+// they are known or symbolically defined, and have matching dimensions.
+bool ShapesSymbolicallyEqual(const TensorShapeProto& left,
+                             const TensorShapeProto& right);
+bool ShapesSymbolicallyEqual(const OpInfo::TensorProperties& left,
+                             const OpInfo::TensorProperties& right);
+
+// Check if two shapes can be broadcasted to each other. Both shapes must be at
+// least symbolically defined, and the have valid BCast instance.
+bool ShapesBroadcastable(const TensorShapeProto& left,
+                         const TensorShapeProto& right);
+bool ShapesBroadcastable(const OpInfo::TensorProperties& left,
+                         const OpInfo::TensorProperties& right);
+
+// Return true if can prove, that tensor of size 'left' is smaller than tensor
+// of size 'right'. Return false if it's larger or equal, or it's impossible to
+// compare because of unknown dimensions, or mismatch in symbolic dimensions.
+bool CompareSymbolicallyShapedTensorSizes(const TensorShapeProto& left,
+                                          const TensorShapeProto& right);
+bool CompareSymbolicallyShapedTensorSizes(
+    const OpInfo::TensorProperties& left,
+    const OpInfo::TensorProperties& right);
+
+}  // namespace grappler
+}  // end namespace tensorflow
+
+#endif  // TENSORFLOW_CORE_GRAPPLER_OPTIMIZERS_SYMBOLIC_SHAPES_H_
diff --git a/tensorflow/core/grappler/optimizers/symbolic_shapes_test.cc b/tensorflow/core/grappler/optimizers/symbolic_shapes_test.cc
new file mode 100644
index 0000000000..5ef9f65925
--- /dev/null
+++ b/tensorflow/core/grappler/optimizers/symbolic_shapes_test.cc
@@ -0,0 +1,95 @@
+/* Copyright 2018 The TensorFlow Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+
+#include "tensorflow/core/grappler/optimizers/symbolic_shapes.h"
+#include "tensorflow/core/framework/tensor_shape.pb.h"
+#include "tensorflow/core/platform/test.h"
+
+namespace tensorflow {
+namespace grappler {
+namespace {
+
+class SymbolicShapesTest : public ::testing::Test {
+ protected:
+  TensorShapeProto MakeUnknown() {
+    TensorShapeProto shape;
+    shape.set_unknown_rank(true);
+    return shape;
+  }
+
+  TensorShapeProto MakeShape(std::vector<int> dims) {
+    TensorShapeProto shape;
+    for (int dim_size : dims) {
+      TensorShapeProto::Dim dim;
+      dim.set_size(dim_size);
+      *shape.add_dim() = dim;
+    }
+    return shape;
+  }
+};
+
+bool operator<(const TensorShapeProto& lhs, const TensorShapeProto& rhs) {
+  return CompareSymbolicallyShapedTensorSizes(lhs, rhs);
+}
+
+TEST_F(SymbolicShapesTest, ShapeIsSymbolicallyDefined) {
+  EXPECT_FALSE(ShapeIsSymbolicallyDefined(MakeUnknown()));
+  EXPECT_FALSE(ShapeIsSymbolicallyDefined(MakeShape({-1, 2})));
+
+  EXPECT_TRUE(ShapeIsSymbolicallyDefined(MakeShape({1, 2})));
+  EXPECT_TRUE(ShapeIsSymbolicallyDefined(MakeShape({-2, 2})));
+}
+
+TEST_F(SymbolicShapesTest, ShapesSymbolicallyEqual) {
+  EXPECT_FALSE(ShapesSymbolicallyEqual(MakeUnknown(), MakeUnknown()));
+  EXPECT_FALSE(ShapesSymbolicallyEqual(MakeShape({-1, 2}), MakeShape({-1, 2})));
+  EXPECT_FALSE(ShapesSymbolicallyEqual(MakeShape({-2, 2}), MakeShape({-3, 2})));
+
+  EXPECT_TRUE(ShapesSymbolicallyEqual(MakeShape({1, 2}), MakeShape({1, 2})));
+  EXPECT_TRUE(ShapesSymbolicallyEqual(MakeShape({-2, 2}), MakeShape({-2, 2})));
+}
+
+TEST_F(SymbolicShapesTest, ShapesBroadcastable) {
+  EXPECT_FALSE(ShapesBroadcastable(MakeUnknown(), MakeUnknown()));
+  EXPECT_FALSE(ShapesBroadcastable(MakeShape({-2}), MakeShape({1, -3})));
+  EXPECT_FALSE(ShapesBroadcastable(MakeShape({-1, 2}), MakeShape({-1, 2})));
+  EXPECT_FALSE(ShapesBroadcastable(MakeShape({-2, 2}), MakeShape({-3, 2})));
+  EXPECT_FALSE(ShapesBroadcastable(MakeShape({-2, 4}), MakeShape({-2, 8})));
+
+  EXPECT_TRUE(ShapesBroadcastable(MakeShape({1, 2}), MakeShape({1, 2})));
+  EXPECT_TRUE(ShapesBroadcastable(MakeShape({-2, 2}), MakeShape({-2, 2})));
+  EXPECT_TRUE(ShapesBroadcastable(MakeShape({-2, 32}), MakeShape({-2, 1})));
+  EXPECT_TRUE(ShapesBroadcastable(MakeShape({-2, 1}), MakeShape({1, -2})));
+  EXPECT_TRUE(ShapesBroadcastable(MakeShape({-2, 1}), MakeShape({1, -3})));
+  EXPECT_TRUE(ShapesBroadcastable(MakeShape({-3}), MakeShape({-2, -3})));
+}
+
+TEST_F(SymbolicShapesTest, CompareSymbolicallyShapedTensorSizes) {
+  EXPECT_TRUE(MakeShape({1, 1, 32}) < MakeShape({32, 32}));
+  EXPECT_TRUE(MakeShape({1, 32, 32}) < MakeShape({2048}));
+  EXPECT_TRUE(MakeShape({1, -2, 32}) < MakeShape({-2, 32, 32}));
+  EXPECT_TRUE(MakeShape({1, 32, 32}) < MakeShape({-2, 32, 32}));
+  EXPECT_TRUE(MakeShape({1, 32, 32}) < MakeShape({-1, 32, 32}));
+  EXPECT_TRUE(MakeShape({1, -2, 32}) < MakeShape({-2, -2, 32}));
+
+  EXPECT_FALSE(MakeShape({1, -2, 32}) < MakeShape({-3, 32, 32}));
+  EXPECT_FALSE(MakeShape({1, -1, 32}) < MakeShape({1, -1, 32}));
+  EXPECT_FALSE(MakeShape({1, -1, 32}) < MakeShape({-1, -1, 32}));
+  EXPECT_FALSE(MakeShape({-1, -1, 32}) < MakeShape({1, -1, 32}));
+}
+
+}  // namespace
+}  // namespace grappler
+}  // namespace tensorflow