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

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

Browse Source 
PiperOrigin-RevId: 337998263
Change-Id: I0b58817043c190af183b67a5ae45a0f3607a5b80
This commit is contained in:
Eugene Brevdo 2020-10-19 22:14:16 -07:00 committed by TensorFlower Gardener
parent d345c40688
commit 9fea4acc26
3 changed files with 203 additions and 0 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

61
tensorflow/python/tools/BUILD
View File

@ -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",
]),
)

60
tensorflow/python/tools/aot_compiled_test.cc
View File

@ -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:

82
tensorflow/python/tools/make_aot_compile_models.py Normal file
View File

@ -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)