[TF] [saved_model_cli] Add unit test for single threaded AOT compiled large matmul.

PiperOrigin-RevId: 337998263
Change-Id: I0b58817043c190af183b67a5ae45a0f3607a5b80

tensorflow/python/tools/BUILD

@@ -368,6 +368,64 @@ py_test(
     ],
 )
 
+py_binary(
+    name = "make_aot_compile_models",
+    srcs = ["make_aot_compile_models.py"],
+    python_version = "PY3",
+    deps = [
+        "//tensorflow/core:protos_all_py",
+        "//tensorflow/python:array_ops",
+        "//tensorflow/python:framework",
+        "//tensorflow/python:math_ops",
+        "//tensorflow/python:platform",
+        "//tensorflow/python/eager:def_function",
+        "//tensorflow/python/saved_model",
+        "@absl_py//absl:app",
+        "@absl_py//absl/flags",
+        "@six_archive//:six",
+    ],
+)
+
+EMITTED_AOT_SAVE_MODEL_OBJECTS = [
+    "x_matmul_y_large/saved_model.pb",
+    "x_matmul_y_large/variables/variables.index",
+    "x_matmul_y_small/saved_model.pb",
+    "x_matmul_y_small/variables/variables.index",
+]
+
+genrule(
+    name = "create_models_for_aot_compile",
+    outs = EMITTED_AOT_SAVE_MODEL_OBJECTS,
+    cmd = (
+        "$(location :make_aot_compile_models) --out_dir $(@D)"
+    ),
+    exec_tools = [":make_aot_compile_models"],
+    tags = ["no_rocm"],
+)
+
+filegroup(
+    name = "aot_saved_models",
+    srcs = EMITTED_AOT_SAVE_MODEL_OBJECTS,
+)
+
+saved_model_compile_aot(
+    name = "aot_compiled_x_matmul_y_large",
+    cpp_class = "XMatmulYLarge",
+    directory = "//tensorflow/python/tools:x_matmul_y_large",
+    filegroups = [":aot_saved_models"],
+    force_without_xla_support_flag = False,
+    tags = ["no_rocm"],
+)
+
+saved_model_compile_aot(
+    name = "aot_compiled_x_matmul_y_small",
+    cpp_class = "XMatmulYSmall",
+    directory = "//tensorflow/python/tools:x_matmul_y_small",
+    filegroups = [":aot_saved_models"],
+    force_without_xla_support_flag = False,
+    tags = ["no_rocm"],
+)
+
 saved_model_compile_aot(
     name = "aot_compiled_x_plus_y",
     cpp_class = "XPlusY",
@@ -413,8 +471,11 @@ tf_cc_test(
     ] + if_xla_available([
         ":aot_compiled_vars_and_arithmetic",
         ":aot_compiled_vars_and_arithmetic_frozen",
+        ":aot_compiled_x_matmul_y_large",
+        ":aot_compiled_x_matmul_y_small",
         ":aot_compiled_x_plus_y",
         "//tensorflow/core:test",
+        "//third_party/eigen3",
         "//tensorflow/core/platform:logging",
     ]),
 )

tensorflow/python/tools/aot_compiled_test.cc

@@ -13,10 +13,13 @@ See the License for the specific language governing permissions and
 limitations under the License.
 ==============================================================================*/
 
+#include "third_party/eigen3/unsupported/Eigen/CXX11/Tensor"
 #include "tensorflow/core/platform/logging.h"
 #include "tensorflow/core/platform/test.h"
 #include "tensorflow/python/tools/aot_compiled_vars_and_arithmetic.h"
 #include "tensorflow/python/tools/aot_compiled_vars_and_arithmetic_frozen.h"
+#include "tensorflow/python/tools/aot_compiled_x_matmul_y_large.h"
+#include "tensorflow/python/tools/aot_compiled_x_matmul_y_small.h"
 #include "tensorflow/python/tools/aot_compiled_x_plus_y.h"
 
 namespace tensorflow {
@@ -30,6 +33,63 @@ TEST(AOTCompiledSavedModelTest, XPlusY) {
   ASSERT_NEAR(model.result_fetch_output_0(), 7.0f, /*abs_error=*/1e-6f);
 }
 
+TEST(AOTCompiledSavedModelTest, XMatmulYLarge) {
+  XMatmulYLarge model;
+  // Calculation is: output_0 = x @ y.
+  EXPECT_EQ(model.arg0_size(), sizeof(float) * 3000 * 5000);
+  EXPECT_EQ(model.arg1_size(), sizeof(float) * 5000 * 4000);
+  EXPECT_EQ(model.result0_size(), sizeof(float) * 3000 * 4000);
+
+  Eigen::Tensor<float, 2, Eigen::RowMajor> arg0(3000, 5000);
+  Eigen::Tensor<float, 2, Eigen::RowMajor> arg1(5000, 4000);
+  arg0.setRandom();
+  arg1.setRandom();
+
+  // Set up dimensions for standard matmul.
+  const Eigen::array<Eigen::IndexPair<int>, 1> product_dims = {
+      Eigen::IndexPair<int>(1, 0)};
+  // Ground truth matmul.
+  const Eigen::Tensor<float, 2, Eigen::RowMajor> expected_output0 =
+      arg0.contract(arg1, product_dims);
+
+  model.set_arg_feed_x_data(arg0.data());
+  model.set_arg_feed_y_data(arg1.data());
+  CHECK(model.Run());
+  EXPECT_NEAR(model.result_fetch_output_0(0, 0), expected_output0(0, 0),
+              /*abs_error=*/1e-6f);
+  EXPECT_NEAR(model.result_fetch_output_0(2999, 3999),
+              expected_output0(2999, 3999),
+              /*abs_error=*/1e-6f);
+}
+
+TEST(AOTCompiledSavedModelTest, XMatmulYSmall) {
+  XMatmulYSmall model;
+  // Calculation is: output_0 = x @ y.
+  EXPECT_EQ(model.arg0_size(), sizeof(float) * 3 * 5);
+  EXPECT_EQ(model.arg1_size(), sizeof(float) * 5 * 4);
+  EXPECT_EQ(model.result0_size(), sizeof(float) * 3 * 4);
+
+  Eigen::Tensor<float, 2, Eigen::RowMajor> arg0(3, 5);
+  Eigen::Tensor<float, 2, Eigen::RowMajor> arg1(5, 4);
+  arg0.setRandom();
+  arg1.setRandom();
+
+  // Set up dimensions for standard matmul.
+  const Eigen::array<Eigen::IndexPair<int>, 1> product_dims = {
+      Eigen::IndexPair<int>(1, 0)};
+  // Ground truth matmul.
+  const Eigen::Tensor<float, 2, Eigen::RowMajor> expected_output0 =
+      arg0.contract(arg1, product_dims);
+
+  model.set_arg_feed_x_data(arg0.data());
+  model.set_arg_feed_y_data(arg1.data());
+  CHECK(model.Run());
+  EXPECT_NEAR(model.result_fetch_output_0(0, 0), expected_output0(0, 0),
+              /*abs_error=*/1e-6f);
+  EXPECT_NEAR(model.result_fetch_output_0(2, 3), expected_output0(2, 3),
+              /*abs_error=*/1e-6f);
+}
+
 TEST(AOTCompiledSavedModelTest, VarsAndArithmetic) {
   VarsAndArithmeticFrozen frozen_model;
   // Calculation is:

tensorflow/python/tools/make_aot_compile_models.py

@@ -0,0 +1,82 @@
+# Copyright 2020 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.
+# ==============================================================================
+"""Generate some SavedModels for use by AOT compilation tests."""
+
+from __future__ import absolute_import
+from __future__ import division
+from __future__ import print_function
+
+import os
+
+from absl import flags
+
+from tensorflow.python.eager import def_function
+from tensorflow.python.framework import dtypes
+from tensorflow.python.framework import tensor_spec
+from tensorflow.python.ops import array_ops
+from tensorflow.python.ops import math_ops
+from tensorflow.python.platform import app
+from tensorflow.python.saved_model import save
+from tensorflow.python.training.tracking import tracking
+
+
+flags.DEFINE_string('out_dir', None,
+                    'Directory to output saved models to.')
+
+FLAGS = flags.FLAGS
+
+
+def create_large_matmul_savedmodel(out_dir):
+  """Create a SavedModel that performs a large matmul."""
+  root = tracking.AutoTrackable()
+  root.f = def_function.function(
+      lambda x, y: math_ops.matmul(x, y),  # pylint: disable=unnecessary-lambda
+      input_signature=[tensor_spec.TensorSpec([3000, 5000], dtypes.float32),
+                       tensor_spec.TensorSpec([5000, 4000], dtypes.float32),])
+  root.f(x=array_ops.zeros((3000, 5000)),
+         y=array_ops.zeros((5000, 4000)))
+  save_dir = os.path.join(out_dir, 'x_matmul_y_large')
+  save.save(root, save_dir, root.f)
+  # This simple SavedModel lacks any variables, but we need to create a
+  # variables.index file to make bazel genrule happy.
+  with open(os.path.join(save_dir, 'variables', 'variables.index'), 'w'):
+    pass
+
+
+def create_small_matmul_savedmodel(out_dir):
+  """Create a SavedModel that performs a small matmul."""
+  root = tracking.AutoTrackable()
+  root.f = def_function.function(
+      lambda x, y: math_ops.matmul(x, y),  # pylint: disable=unnecessary-lambda
+      input_signature=[tensor_spec.TensorSpec([3, 5], dtypes.float32),
+                       tensor_spec.TensorSpec([5, 4], dtypes.float32),])
+  root.f(x=array_ops.zeros((3, 5)),
+         y=array_ops.zeros((5, 4)))
+  save_dir = os.path.join(out_dir, 'x_matmul_y_small')
+  save.save(root, save_dir, root.f)
+  # This simple SavedModel lacks any variables, but we need to create a
+  # variables.index file to make bazel genrule happy.
+  with open(os.path.join(save_dir, 'variables', 'variables.index'), 'w'):
+    pass
+
+
+def main(unused_args):
+  create_small_matmul_savedmodel(FLAGS.out_dir)
+  create_large_matmul_savedmodel(FLAGS.out_dir)
+
+
+if __name__ == '__main__':
+  flags.mark_flag_as_required('out_dir')
+  app.run(main)