wehub/tensorflow--tensorflow
1
0
Fork 0
Code Issues Pull Requests Actions 3 Deployments Models Agent Notes Packages Projects Releases Wiki Activity

chore: import upstream snapshot with attribution
cffconvert / validate (push) Has been skipped
Details
License Check / license-check (push) Failing after 2s
Details

Browse Source
This commit is contained in:
wehub-resource-sync
2026-07-13 12:14:16 +08:00
commit 8a852e4b4e
36502 changed files with 9277225 additions and 0 deletions
Download Patch File Download Diff File Expand all files Collapse all files
tensorflow/compiler/OWNERS_FORCE_SUBMIT_APPROVERS
+19
View File
@@ -0,0 +1,19 @@
# This file contains the list of senior ICs in the XLA team who can
# approve force submits in addition to manager approval. See
# go/foundations-force-submit-policy for full policy.
// go/keep-sorted start
andydavis
asabne
berkin
chenyin
jikimjikim
jsreeram
ksmurthy
majnemer
masonchang
mkuper
pifon
pizzud
// go/keep-sorted end
tensorflow/compiler/aot/BUILD
+423
View File
@@ -0,0 +1,423 @@
load("//tensorflow:tensorflow.bzl", "if_google", "if_oss", "tf_cc_binary", "tf_cc_test")
load("//tensorflow/compiler/aot:tfcompile.bzl", "tf_library")
load(
"//tensorflow/core/platform:build_config_root.bzl",
"if_llvm_aarch32_available",
"if_llvm_aarch64_available",
"if_llvm_hexagon_available",
"if_llvm_powerpc_available",
"if_llvm_system_z_available",
"if_llvm_x86_available",
)
load("//tensorflow/core/platform:rules_cc.bzl", "cc_library")
package(
# copybara:uncomment default_applicable_licenses = ["//tensorflow:license"],
default_visibility = ["//visibility:private"],
licenses = ["notice"],
)
# Don't depend on this directly; this is only used for the benchmark test
# generated by tf_library.
cc_library(
name = "tf_library_test_main",
testonly = 1,
visibility = ["//visibility:public"],
deps = if_oss([
"//tensorflow/core:test_main",
]) + if_google([
"//tensorflow/core/platform/default/build_config:test_main",
]),
)
filegroup(
name = "quantize_header",
srcs = ["quantize.h"],
visibility = ["//visibility:public"],
)
cc_library(
name = "thunk_proto_execution_deserializer",
srcs = ["thunk_proto_execution_deserializer.cc"],
hdrs = ["thunk_proto_execution_deserializer.h"],
visibility = ["//visibility:public"],
deps = [
"@com_google_absl//absl/numeric:int128",
"@com_google_absl//absl/status",
"@com_google_absl//absl/status:statusor",
"@com_google_absl//absl/strings",
"@xla//xla:debug_options_flags",
"@xla//xla:shape_util",
"@xla//xla:status_macros",
"@xla//xla:util",
"@xla//xla:xla_data_proto_cc",
"@xla//xla/backends/cpu/runtime:convolution_dims",
"@xla//xla/backends/cpu/runtime:dot_dims",
"@xla//xla/backends/cpu/runtime:kernel_thunk",
"@xla//xla/backends/cpu/runtime:thunk_proto_cc",
"@xla//xla/service/cpu:executable_proto_cc",
"@xla//xla/tsl/platform:statusor",
],
)
cc_library(
name = "tfcompile_lib",
srcs = [
"codegen.cc",
"compile.cc",
"flags.cc",
],
hdrs = [
"codegen.h",
"compile.h",
"flags.h",
"quantize.h",
],
compatible_with = [],
defines = if_llvm_aarch32_available(["TF_LLVM_AARCH32_AVAILABLE=1"]) + if_llvm_aarch64_available([
"TF_LLVM_AARCH64_AVAILABLE=1",
]) + if_llvm_hexagon_available([
"TF_LLVM_HEXAGON_AVAILABLE=1",
]) + if_llvm_powerpc_available([
"TF_LLVM_POWERPC_AVAILABLE=1",
]) + if_llvm_system_z_available([
"TF_LLVM_S390X_AVAILABLE=1",
]) + if_llvm_x86_available([
"TF_LLVM_X86_AVAILABLE=1",
]),
visibility = [
"//tensorflow:__pkg__",
"//tensorflow/python:__pkg__",
],
deps = [
":aot_only_var_handle_op",
":embedded_protocol_buffers",
":llvm_targets", # fixdeps: keep
":thunk_proto_execution_deserializer",
"//tensorflow/compiler/tf2xla",
"//tensorflow/compiler/tf2xla:allocator",
"//tensorflow/compiler/tf2xla:encoded_buffer_allocation_info",
"//tensorflow/compiler/tf2xla:mlir_tf2xla", # fixdeps: keep
"//tensorflow/compiler/tf2xla:tf2xla_proto_cc",
"//tensorflow/compiler/tf2xla:tf2xla_util",
"//tensorflow/compiler/tf2xla:xla_compiler",
"//tensorflow/compiler/tf2xla/kernels:xla_dummy_ops",
"//tensorflow/compiler/tf2xla/kernels:xla_ops",
"//tensorflow/core:core_cpu_internal",
"//tensorflow/core:framework_internal",
"//tensorflow/core:lib",
"//tensorflow/core:lib_internal",
"//tensorflow/core:protos_all_cc",
"//tensorflow/core/platform:regexp",
"//tensorflow/core/platform:status",
"@com_google_absl//absl/algorithm:container",
"@com_google_absl//absl/base",
"@com_google_absl//absl/log:check",
"@com_google_absl//absl/status",
"@com_google_absl//absl/status:statusor",
"@com_google_absl//absl/strings",
"@com_google_absl//absl/strings:str_format",
"@com_google_absl//absl/types:span",
"@llvm-project//llvm:Support",
"@llvm-project//llvm:Target",
"@xla//xla:debug_options_flags",
"@xla//xla:shape_util",
"@xla//xla:status_macros",
"@xla//xla:util",
"@xla//xla:xla_data_proto_cc",
"@xla//xla/backends/cpu:buffer_allocation_info",
"@xla//xla/backends/cpu:buffer_allocation_info_util",
"@xla//xla/backends/cpu/codegen:symbol_name_util",
"@xla//xla/backends/cpu/runtime:thunk_proto_cc",
"@xla//xla/backends/cpu/runtime:thunk_proto_serdes",
"@xla//xla/client:client_library",
"@xla//xla/client:compile_only_client",
"@xla//xla/hlo/builder:xla_computation",
"@xla//xla/service:compiler",
"@xla//xla/service/cpu:cpu_aot_compilation_result",
"@xla//xla/service/cpu:cpu_compiler",
"@xla//xla/service/cpu:cpu_executable",
"@xla//xla/stream_executor:platform",
"@xla//xla/stream_executor:platform_manager",
"@xla//xla/util:embedded_constant_buffers",
],
)
tf_cc_test(
name = "codegen_test",
srcs = ["codegen_test.cc"],
defines = if_llvm_x86_available(["TF_LLVM_X86_AVAILABLE=1"]),
deps = [
":tfcompile_lib",
"//tensorflow/core:lib",
"//tensorflow/core:protos_all_cc",
"//tensorflow/core:test",
"//tensorflow/core:test_main",
"//tensorflow/core/platform:resource_loader",
"@com_google_absl//absl/memory",
"@com_google_absl//absl/strings",
"@llvm-project//llvm:Support", # fixdeps: keep
"@xla//xla:shape_util",
"@xla//xla/service/cpu:cpu_aot_compilation_result",
] + if_llvm_x86_available([
"@llvm-project//llvm:X86CodeGen", # fixdeps: keep
]),
)
tf_cc_binary(
name = "tfcompile",
visibility = ["//visibility:public"],
deps = [":tfcompile_main"],
)
cc_library(
name = "llvm_targets",
visibility = ["//tensorflow/python:__pkg__"],
deps = [
] + if_llvm_aarch32_available([
"@llvm-project//llvm:ARMAsmParser", # fixdeps: keep
"@llvm-project//llvm:ARMCodeGen", # fixdeps: keep
]) + if_llvm_aarch64_available([
"@llvm-project//llvm:AArch64AsmParser", # fixdeps: keep
"@llvm-project//llvm:AArch64CodeGen", # fixdeps: keep
]) + if_llvm_hexagon_available([
"@llvm-project//llvm:HexagonAsmParser", # fixdeps: keep
"@llvm-project//llvm:HexagonCodeGen", # fixdeps: keep
]) + if_llvm_powerpc_available([
"@llvm-project//llvm:PowerPCAsmParser", # fixdeps: keep
"@llvm-project//llvm:PowerPCCodeGen", # fixdeps: keep
]) + if_llvm_system_z_available([
"@llvm-project//llvm:SystemZAsmParser", # fixdeps: keep
"@llvm-project//llvm:SystemZCodeGen", # fixdeps: keep
]) + if_llvm_x86_available([
"@llvm-project//llvm:X86AsmParser", # fixdeps: keep
"@llvm-project//llvm:X86CodeGen", # fixdeps: keep
]),
)
cc_library(
name = "tfcompile_main",
srcs = ["tfcompile_main.cc"],
compatible_with = [],
visibility = ["//visibility:public"],
deps = [
":tfcompile_lib",
"//tensorflow/compiler/tf2xla:tf2xla_proto_cc",
"//tensorflow/core:core_cpu_internal",
"//tensorflow/core:framework_internal",
"//tensorflow/core:lib",
"//tensorflow/core:protos_all_cc",
"@com_google_absl//absl/log:check",
"@com_google_absl//absl/status",
"@com_google_absl//absl/strings",
"@xla//xla:debug_options_flags",
"@xla//xla/tsl/platform:status",
],
)
# NOTE: Most end-to-end tests are in the "tests" subdirectory, to ensure that
# tfcompile.bzl correctly handles usage from outside of the package that it is
# defined in.
# A simple test of tf_library from a text protobuf, to enable benchmark_test.
# This test uses an incompleted graph with a node that is not defined. The
# compilation works because the undefined node is a feed node.
tf_library(
name = "test_graph_tfadd",
testonly = 1,
config = "test_graph_tfadd.config.pbtxt",
cpp_class = "AddComp",
graph = "test_graph_tfadd.pbtxt",
mlir_components = "None",
tags = [
"manual",
],
)
tf_library(
name = "test_graph_tfadd_mlir_bridge",
testonly = 1,
config = "test_graph_tfadd.config.pbtxt",
cpp_class = "AddComp",
graph = "test_graph_tfadd.pbtxt",
mlir_components = "Bridge",
tags = [
"manual",
],
)
# A test of tf_library that includes a graph with an unknown op, but where
# the compilation works because the node with the unknown op is not needed
# for the fetches.
tf_library(
name = "test_graph_tfunknownop",
testonly = 1,
config = "test_graph_tfunknownop.config.pbtxt",
cpp_class = "UnknownOpAddComp",
graph = "test_graph_tfunknownop.pbtxt",
mlir_components = "None",
tags = [
"manual",
],
)
tf_library(
name = "test_graph_tfunknownop_mlir_bridge",
testonly = 1,
config = "test_graph_tfunknownop.config.pbtxt",
cpp_class = "UnknownOpAddComp",
graph = "test_graph_tfunknownop.pbtxt",
mlir_components = "Bridge",
tags = [
"manual",
],
)
# A test of tf_library that includes a graph with an unknown op, but where
# the compilation works because the node with the unknown op is only used as
# an input of a feed node.
tf_library(
name = "test_graph_tfunknownop2",
testonly = 1,
config = "test_graph_tfunknownop2.config.pbtxt",
cpp_class = "UnknownOpAddComp",
graph = "test_graph_tfunknownop.pbtxt",
mlir_components = "None",
tags = [
"manual",
],
)
tf_library(
name = "test_graph_tfunknownop2_mlir_bridge",
testonly = 1,
config = "test_graph_tfunknownop2.config.pbtxt",
cpp_class = "UnknownOpAddComp",
graph = "test_graph_tfunknownop.pbtxt",
mlir_components = "Bridge",
tags = [
"manual",
],
)
# A test of tf_library that includes a graph with an unknown op, but where
# the compilation works because the node with the unknown op is a feed node.
tf_library(
name = "test_graph_tfunknownop3",
testonly = 1,
config = "test_graph_tfunknownop3.config.pbtxt",
cpp_class = "UnknownOpAddComp",
graph = "test_graph_tfunknownop.pbtxt",
mlir_components = "None",
tags = [
"manual",
],
)
tf_library(
name = "test_graph_tfunknownop3_mlir_bridge",
testonly = 1,
config = "test_graph_tfunknownop3.config.pbtxt",
cpp_class = "UnknownOpAddComp",
graph = "test_graph_tfunknownop.pbtxt",
mlir_components = "Bridge",
tags = [
"manual",
],
)
# Utility library for benchmark binaries, used by the *_benchmark rules that are
# added by the tfcompile bazel macro.
cc_library(
name = "benchmark",
srcs = ["benchmark.cc"],
hdrs = ["benchmark.h"],
visibility = ["//visibility:public"],
deps = [
# The purpose of the benchmark library is to support building an aot
# binary with minimal dependencies, to demonstrate small binary sizes.
#
# KEEP THE DEPENDENCIES MINIMAL.
"//tensorflow/core:framework_lite",
],
)
cc_library(
name = "benchmark_extra_android",
tags = [
"manual",
],
visibility = ["//visibility:public"],
)
cc_library(
name = "embedded_protocol_buffers",
srcs = ["embedded_protocol_buffers.cc"],
hdrs = ["embedded_protocol_buffers.h"],
visibility = ["//visibility:public"],
deps = [
"//tensorflow/core:lib",
"@com_google_absl//absl/log:check",
"@com_google_absl//absl/memory",
"@com_google_absl//absl/status:statusor",
"@com_google_absl//absl/strings",
"@com_google_absl//absl/types:span",
"@llvm-project//llvm:Core",
"@llvm-project//llvm:MC",
"@llvm-project//llvm:Target",
"@llvm-project//llvm:TargetParser",
"@xla//xla:util",
"@xla//xla/service/llvm_ir:llvm_type_conversion_util",
],
)
cc_library(
name = "aot_only_var_handle_op",
srcs = ["aot_only_var_handle_op.cc"],
hdrs = ["aot_only_var_handle_op.h"],
visibility = [
"//tensorflow/compiler/tf2xla:__pkg__",
],
deps = [
"//tensorflow/compiler/tf2xla:xla_compiler",
"//tensorflow/compiler/tf2xla:xla_context",
"//tensorflow/compiler/tf2xla:xla_op_registry",
"//tensorflow/core:framework",
],
alwayslink = 1,
)
tf_cc_test(
name = "benchmark_test",
srcs = ["benchmark_test.cc"],
tags = ["manual"],
deps = [
":benchmark",
":test_graph_tfadd",
"//tensorflow/core:test",
"//tensorflow/core:test_main",
],
)
test_suite(
name = "all_tests",
tags = ["manual"],
tests = [
":benchmark_test",
":codegen_test",
":test_graph_tfadd_mlir_bridge_test",
":test_graph_tfadd_test",
":test_graph_tfunknownop2_mlir_bridge_test",
":test_graph_tfunknownop2_test",
":test_graph_tfunknownop3_mlir_bridge_test",
":test_graph_tfunknownop3_test",
":test_graph_tfunknownop_mlir_bridge_test",
":test_graph_tfunknownop_test",
"//tensorflow/compiler/aot/tests:all_tests",
],
)
exports_files([
"benchmark_main.template", # used by tf_library(...,gen_benchmark=True)
if_oss("test.cc", "test_google.cc"), # used by tf_library(...,gen_test=True)
])
tensorflow/compiler/aot/aot_only_var_handle_op.cc
+85
View File
@@ -0,0 +1,85 @@
/* Copyright 2019 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/compiler/aot/aot_only_var_handle_op.h"
#include "tensorflow/compiler/tf2xla/xla_context.h"
#include "tensorflow/compiler/tf2xla/xla_op_kernel.h"
#include "tensorflow/compiler/tf2xla/xla_op_registry.h"
#include "tensorflow/core/framework/shape_inference.h"
namespace tensorflow {
namespace {
// Implementation of varhandle that binds a VarHandleOp to an XlaResource of the
// same name. It is not safe to use this op in a JIT context.
class XlaAotOnlyVarHandleOp : public XlaOpKernel {
public:
explicit XlaAotOnlyVarHandleOp(OpKernelConstruction* c);
void Compile(XlaOpKernelContext* context) override;
private:
std::string name_;
};
XlaAotOnlyVarHandleOp::XlaAotOnlyVarHandleOp(OpKernelConstruction* c)
: XlaOpKernel(c) {
OP_REQUIRES_OK(c, c->GetAttr("shared_name", &name_));
}
void XlaAotOnlyVarHandleOp::Compile(XlaOpKernelContext* context) {
// Look for a resource of the same name. TF also keys that on the container
// and type attributes, but that doesn't seem necessary.
for (const auto& resource : context->xla_context()->resources()) {
if (resource->kind() == XlaResource::kVariable &&
resource->name() == name_) {
context->SetResourceOutput(0, resource.get());
return;
}
}
context->SetStatus(absl::InvalidArgumentError(
absl::StrCat("Variable: ", name_, " not configured")));
}
} // namespace
REGISTER_OP(tfcompile::kXlaAotOnlyVarHandleOp)
.Doc(R"doc(
Internal VarHandleOp registration used for XLA AOT compilation.
)doc")
.Attr("container: string = ''")
.Attr("shared_name: string = ''")
.Attr("debug_name: string = ''")
.Attr("dtype: type")
.Attr("shape: shape")
.Output("resource: resource")
.SetIsStateful()
.SetShapeFn([](shape_inference::InferenceContext* c) {
c->set_output(0, c->Scalar());
DataType t;
TF_RETURN_IF_ERROR(c->GetAttr("dtype", &t));
PartialTensorShape p;
TF_RETURN_IF_ERROR(c->GetAttr("shape", &p));
shape_inference::ShapeHandle s;
TF_RETURN_IF_ERROR(c->MakeShapeFromPartialTensorShape(p, &s));
c->set_output_handle_shapes_and_types(
0, std::vector<shape_inference::ShapeAndType>{{s, t}});
return absl::OkStatus();
});
REGISTER_XLA_OP(Name(tfcompile::kXlaAotOnlyVarHandleOp).CompilationOnly(),
XlaAotOnlyVarHandleOp);
} // namespace tensorflow
tensorflow/compiler/aot/aot_only_var_handle_op.h
+27
View File
@@ -0,0 +1,27 @@
/* 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.
==============================================================================*/
#ifndef TENSORFLOW_COMPILER_AOT_AOT_ONLY_VAR_HANDLE_OP_H_
#define TENSORFLOW_COMPILER_AOT_AOT_ONLY_VAR_HANDLE_OP_H_
namespace tensorflow {
namespace tfcompile {
static constexpr const char* const kXlaAotOnlyVarHandleOp =
"_XlaAotOnlyVarHandleOp";
} // namespace tfcompile
} // namespace tensorflow
#endif // TENSORFLOW_COMPILER_AOT_AOT_ONLY_VAR_HANDLE_OP_H_
tensorflow/compiler/aot/benchmark.cc
+138
View File
@@ -0,0 +1,138 @@
/* Copyright 2017 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.
==============================================================================*/
// The purpose of the benchmark library is to support building an aot binary
// with minimal dependencies, to demonstrate small binary sizes.
//
// KEEP THE DEPENDENCIES MINIMAL.
#include "tensorflow/compiler/aot/benchmark.h"
#include <sys/time.h>
#include <algorithm>
#include <string>
#include <utility>
#include <vector>
#include "tensorflow/core/platform/types.h"
namespace tensorflow {
namespace tfcompile {
namespace benchmark {
// Returns current wall time in micros.
//
// TODO(b/33546473): Refactor tensorflow::Env::NowMicros() so that we can re-use
// the implementation without pulling in all of the Env dependencies.
static uint64_t NowMicros() {
struct timeval tv;
gettimeofday(&tv, nullptr);
return static_cast<uint64_t>(tv.tv_sec) * 1000000 + tv.tv_usec;
}
void DumpStatsToStdout(const Stats& stats) {
// Compute stats.
std::vector<int64_t> sorted_us(stats.per_iter_us);
std::sort(sorted_us.begin(), sorted_us.end());
const size_t count_us = sorted_us.size();
double sum_us = 0;
size_t count_us_trimmed = 0;
double sum_us_trimmed = 0;
size_t count_us_best = 0;
double sum_us_best = 0;
static constexpr float trim_ratio = 0.25;
static constexpr float best_ratio = 0.1;
const size_t count_trimmed = count_us * trim_ratio;
const size_t count_best = count_us * best_ratio;
for (size_t i = 0; i < sorted_us.size(); ++i) {
const int64_t us = sorted_us[i];
sum_us += us;
if (i >= count_trimmed && i < count_us - count_trimmed) {
sum_us_trimmed += us;
++count_us_trimmed;
}
if (i < count_best) {
sum_us_best += us;
++count_us_best;
}
}
// Prepare nicely-formatted data.
const int kBufSize = 1000;
char buf[kBufSize];
snprintf(buf, kBufSize, "Mean with %2.0f%% trimmed:", trim_ratio * 100);
std::string label_trimmed(buf);
snprintf(buf, kBufSize, "Mean of %2.0f%% best:", best_ratio * 100);
std::string label_best(buf);
std::vector<std::pair<std::string, double>> groups = {
{"Best:", sorted_us.front()},
{"Worst:", sorted_us.back()},
{"Median:", sorted_us[count_us / 2]},
{"Mean:", sum_us / count_us},
{std::move(label_trimmed), sum_us_trimmed / count_us_trimmed},
{std::move(label_best), sum_us_best / count_us_best},
};
int max_label_size = 0;
double max_us = 0;
for (const auto& g : groups) {
if (g.first.size() > max_label_size) {
max_label_size = g.first.size();
}
if (g.second > max_us) {
max_us = g.second;
}
}
int max_digits = 1;
while (max_us >= 10.0) {
max_us /= 10.0;
++max_digits;
}
// Dump stats out.
printf("Benchmark ran %zu iterations over %lld us\n", count_us,
static_cast<long long>(stats.total_us)); // NOLINT
for (const auto& g : groups) {
printf(" %-*s %*.3f us\n", max_label_size, g.first.c_str(), max_digits + 4,
g.second);
}
}
void Benchmark(const Options& options, const BenchmarkFn& fn, Stats* stats) {
// If neither max_seconds or max_iters is set, stop at kDefaultMicros.
const int64_t max_us = (options.max_micros <= 0 && options.max_iters <= 0)
? Options::kDefaultMicros
: options.max_micros;
// NOLINTNEXTLINE
printf("Running benchmark for %lld us\n", static_cast<long long>(max_us));
const int64_t start_us = NowMicros();
int64_t iters = 0;
while (true) {
const int64_t iter_start_us = NowMicros();
fn();
const int64_t end_us = NowMicros();
// Collect stats and decide whether to stop.
stats->per_iter_us.push_back(end_us - iter_start_us);
const int64_t total_us = end_us - start_us;
++iters;
if ((max_us > 0 && total_us >= max_us) ||
(options.max_iters > 0 && iters >= options.max_iters)) {
stats->total_us = total_us;
break;
}
}
}
} // namespace benchmark
} // namespace tfcompile
} // namespace tensorflow
tensorflow/compiler/aot/benchmark.h
+70
View File
@@ -0,0 +1,70 @@
/* Copyright 2017 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.
==============================================================================*/
// Contains benchmark functions used with the code-generated benchmarks that can
// be used to test a model on android. See also code generation rules in
// tfcompile.bzl.
//
// This is separate from the built-in micro-benchmarks, because we want to:
// 1. show a binary with minimal dependencies, to show a close-to-lower-bound
// binary size.
// 2. compile on Android.
#ifndef TENSORFLOW_COMPILER_AOT_BENCHMARK_H_
#define TENSORFLOW_COMPILER_AOT_BENCHMARK_H_
#include <functional>
#include <string>
#include <vector>
#include "tensorflow/core/platform/types.h"
namespace tensorflow {
namespace tfcompile {
namespace benchmark {
// Options specifies options for benchmarks of functions generated by tfcompile.
struct Options {
// kDefaultMicros specifies the default time to run the benchmark, and is used
// if neither max_iters nor max_micros is set.
static constexpr int64_t kDefaultMicros = 3000000;
int64_t max_iters = 0; // Maximum iterations to run, ignored if <= 0.
int64_t max_micros = 0; // Maximum microseconds to run, ignored if <= 0.
};
// Stats holds statistics collected during benchmarking.
struct Stats {
std::vector<int64_t> per_iter_us; // Per-iteration deltas in us.
int64_t total_us; // Total time in us.
Stats() : total_us(0) { per_iter_us.reserve(5000); }
};
// DumpStatsToStdout printfs to stdout stats in a multi-line human-friendly
// form.
void DumpStatsToStdout(const Stats& stats);
// BenchmarkFn is the signature of the function generated by tfcompile.
typedef std::function<void()> BenchmarkFn;
// Benchmark runs a benchmark of the function `fn`, collecting stats in `stats`.
// Use `options` to configure benchmarking options.
void Benchmark(const Options& options, const BenchmarkFn& fn, Stats* stats);
} // namespace benchmark
} // namespace tfcompile
} // namespace tensorflow
#endif // TENSORFLOW_COMPILER_AOT_BENCHMARK_H_
tensorflow/compiler/aot/benchmark_main.template
+51
View File
@@ -0,0 +1,51 @@
// Generated by the tf_library build rule. DO NOT EDIT!
//
// This file contains the main function and logic for benchmarking code
// generated by tfcompile. All tokens of the form `{{TFCOMPILE_*}}` must be
// rewritten to real values before this file can be compiled.
//
// TFCOMPILE_HEADER : Path to the header file generated by tfcompile.
// TFCOMPILE_CPP_CLASS : Name of the C++ class generated by tfcompile.
//
// The tf_library bazel macro in tfcompile.bzl performs the token rewriting, and
// generates a cc_binary rule for you.
// These macros must be defined before eigen files are included.
#define EIGEN_USE_THREADS
#define EIGEN_USE_CUSTOM_THREAD_POOL
// clang-format off
#include "{{TFCOMPILE_HEADER}}" // NOLINT(whitespace/braces)
// clang-format on
#include "tensorflow/compiler/aot/benchmark.h"
#include "unsupported/Eigen/CXX11/Tensor"
// Macros that expand to tokens based on the entry point name.
// clang-format off
#define CPP_CLASS {{TFCOMPILE_CPP_CLASS}} // NOLINT(whitespace/braces)
// clang-format on
namespace tensorflow {
namespace tfcompile {
int Main(int argc, char** argv) {
Eigen::ThreadPool pool(1 /* num_threads */);
Eigen::ThreadPoolDevice device(&pool, pool.NumThreads());
CPP_CLASS computation;
computation.set_thread_pool(&device);
benchmark::Options options;
benchmark::Stats stats;
benchmark::Benchmark(options, [&] { computation.Run(); }, &stats);
benchmark::DumpStatsToStdout(stats);
return 0;
}
} // namespace tfcompile
} // namespace tensorflow
int main(int argc, char** argv) {
return tensorflow::tfcompile::Main(argc, argv);
}
tensorflow/compiler/aot/benchmark_test.cc
+46
View File
@@ -0,0 +1,46 @@
/* Copyright 2017 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/compiler/aot/benchmark.h"
#include "tensorflow/compiler/aot/test_graph_tfadd.h"
#include "tensorflow/core/platform/test.h"
namespace tensorflow {
namespace tfcompile {
namespace benchmark {
namespace {
// There isn't much we can verify in a stable fashion, so we just run the
// benchmark with max_iters, and ensure we end up with that many iter stats.
TEST(Benchmark, Benchmark) {
AddComp add;
Options options;
options.max_iters = 1;
Stats stats1;
Benchmark(options, [&] { add.Run(); }, &stats1);
EXPECT_EQ(stats1.per_iter_us.size(), 1);
options.max_iters = 5;
Stats stats5;
Benchmark(options, [&] { add.Run(); }, &stats5);
EXPECT_EQ(stats5.per_iter_us.size(), 5);
}
} // namespace
} // namespace benchmark
} // namespace tfcompile
} // namespace tensorflow
tensorflow/compiler/aot/codegen.cc
+1355
View File
File diff suppressed because it is too large Load Diff
tensorflow/compiler/aot/codegen.h
+126
View File
@@ -0,0 +1,126 @@
/* Copyright 2017 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_COMPILER_AOT_CODEGEN_H_
#define TENSORFLOW_COMPILER_AOT_CODEGEN_H_
#include <vector>
#include "absl/status/status.h"
#include "absl/status/statusor.h"
#include "absl/strings/string_view.h"
#include "tensorflow/compiler/aot/compile.h"
#include "tensorflow/compiler/tf2xla/tf2xla.pb.h"
#include "xla/util/embedded_constant_buffers.h"
#include "tensorflow/core/platform/types.h"
namespace tensorflow {
namespace tfcompile {
// CodegenOpts specifies code generation options for the generated header file
// and the generated metadata object file.
struct CodegenOpts {
// The name of the generated C++ class, wrapping the generated function.
std::string class_name;
// Target triple for the architecture we're targeting.
std::string target_triple;
// Namespaces specifies a list of C++ namespaces to add to the generated
// header. If empty, all symbols will be in the global namespace.
std::vector<std::string> namespaces;
// If true, generate name-to-index data for Lookup{Arg,Result}Index methods.
bool gen_name_to_index = false;
// If true, generate program shape data for the ProgramShape method.
bool gen_program_shape = false;
// If true, emit a serialized HloProfilePrinterData protobuf that can be used
// to pretty print HLO profile counters.
bool gen_hlo_profile_printer_data = false;
// If true, sets this executable as an XLA Runtime one.
bool use_xla_runtime = false;
// If true, use xla cpu nanort runtime.
bool use_xla_nanort_runtime = false;
};
// Describes a generated metadata object file.
struct MetadataResult {
// These are top level "extern C" declarations that are expected to be visible
// wherever program_shape_access_shim is emitted.
std::vector<std::string> header_variable_decls;
// program_shape_access_shim is a C++ expression that constructs the
// xla::ProgramShapeProto instance for the CompileResult passed to
// GenerateMetadata.
std::string program_shape_access_shim;
// hlo_profile_printer_data_access_shim is a C++ expression that constructs
// the xla::HloProfilePrinterData instance for the CompileResult passed to
// GenerateMetadata. If the xla::HloProfilePrinterData is null then this is a
// C++ expression that evaluates to nullptr at runtime.
// This is set only for AOT legacy.
std::string hlo_profile_printer_data_access_shim;
// cpu_executable_access_shim is a C++ expression that constructs
// a protobuf required to construct a CpuExecutable.
// This is set only for AOT thunks.
std::string cpu_executable_access_shim;
// The contents of the object (".o") file.
std::string object_file_data;
};
// Generates a set of constant buffers embedded into an object file.
absl::StatusOr<xla::EmbeddedConstantBuffers> GenerateConstantBuffersData(
const CodegenOpts& opts, const CompileResult& compile_result);
// Generates a metadata object file according to `opts` and
// `compile_result`. The generated object file is returned via
// `metadata_result`.
absl::Status GenerateMetadata(const CodegenOpts& opts,
const CompileResult& compile_result,
MetadataResult* metadata_result);
// GenerateHeader uses the meta-information from compile_result to generate a
// C++ header giving access to the function in the generated object file. The
// header includes API usage documentation.
//
// metadata_result is an instance of MetadataResult obtained by a previous
// invocation to GenerateMetadata.
absl::Status GenerateHeader(
const CodegenOpts& opts, const tf2xla::Config& config,
const CompileResult& compile_result, const MetadataResult& metadata_result,
const xla::EmbeddedConstantBuffers& embedded_constant_buffers,
std::string* header);
// ParseCppClass parses `cpp_class` into its `class_name` and `namespaces`
// components. The syntax is [[<optional_namespace>::],...]<class_name>. This
// mirrors the C++ syntax for referring to a class, where multiple namespaces
// may precede the class name, separated by double-colons.
absl::Status ParseCppClass(const std::string& cpp_class,
std::string* class_name,
std::vector<std::string>* namespaces);
// ValidateCppIdent returns OK iff ident is a valid C++ identifier. The msg is
// appended to error messages.
absl::Status ValidateCppIdent(absl::string_view ident, absl::string_view msg);
} // namespace tfcompile
} // namespace tensorflow
#endif // TENSORFLOW_COMPILER_AOT_CODEGEN_H_
tensorflow/compiler/aot/codegen_test.cc
+184
View File
@@ -0,0 +1,184 @@
/* Copyright 2017 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/compiler/aot/codegen.h"
#include <algorithm>
#include <string>
#include <vector>
#include "absl/memory/memory.h"
#include "absl/strings/match.h"
#include "absl/strings/string_view.h"
#include "llvm/Support/TargetSelect.h"
#include "tensorflow/compiler/aot/compile.h"
#include "xla/service/cpu/cpu_aot_compilation_result.h"
#include "xla/shape_util.h"
#include "tensorflow/core/framework/tensor_shape.pb.h"
#include "tensorflow/core/framework/types.pb.h"
#include "tensorflow/core/lib/core/status.h"
#include "tensorflow/core/lib/core/status_test_util.h"
#include "tensorflow/core/lib/io/path.h"
#include "tensorflow/core/platform/env.h"
#include "tensorflow/core/platform/resource_loader.h"
#include "tensorflow/core/platform/test.h"
namespace tensorflow {
namespace tfcompile {
namespace {
using ::xla::cpu::BufferAllocationInfo;
void ExpectErrorContains(const absl::Status& status, absl::string_view str) {
EXPECT_NE(absl::OkStatus(), status);
EXPECT_TRUE(absl::StrContains(status.message(), str))
<< "expected error: " << status.message() << " to contain: " << str;
}
TEST(ValidateCppIdent, Simple) {
TF_EXPECT_OK(ValidateCppIdent("a", ""));
TF_EXPECT_OK(ValidateCppIdent("abc", ""));
TF_EXPECT_OK(ValidateCppIdent("_abc", ""));
TF_EXPECT_OK(ValidateCppIdent("_abc123", ""));
// Make sure we didn't skip a valid letter or digit
std::string ident;
for (char c = 'a'; c <= 'z'; c++) {
ident.append(1, c);
}
for (char c = 'A'; c <= 'Z'; c++) {
ident.append(1, c);
}
for (char c = '0'; c <= '9'; c++) {
ident.append(1, c);
}
ident += "_";
TF_EXPECT_OK(ValidateCppIdent(ident, ""));
ExpectErrorContains(ValidateCppIdent("", ""), "empty identifier");
ExpectErrorContains(ValidateCppIdent(" ", ""), "illegal leading char");
ExpectErrorContains(ValidateCppIdent("0", ""), "illegal leading char");
ExpectErrorContains(ValidateCppIdent(".", ""), "illegal leading char");
ExpectErrorContains(ValidateCppIdent(":", ""), "illegal leading char");
ExpectErrorContains(ValidateCppIdent("a.", ""), "illegal char");
ExpectErrorContains(ValidateCppIdent("a:", ""), "illegal char");
ExpectErrorContains(ValidateCppIdent("a:", ""), "illegal char");
}
class ParseCppClassTest : public ::testing::Test {
protected:
void ExpectOK(const std::string& cpp_class,
const std::string& want_class_name,
const std::vector<std::string>& want_namespaces) {
std::string class_name;
std::vector<std::string> namespaces;
TF_EXPECT_OK(ParseCppClass(cpp_class, &class_name, &namespaces));
EXPECT_EQ(class_name, want_class_name);
EXPECT_EQ(namespaces, want_namespaces);
}
void ExpectFail(const std::string& cpp_class) {
std::string class_name;
std::vector<std::string> namespaces;
EXPECT_NE(ParseCppClass(cpp_class, &class_name, &namespaces),
absl::OkStatus())
<< cpp_class;
}
};
TEST_F(ParseCppClassTest, ParseOK) {
ExpectOK("MyClass", "MyClass", {});
ExpectOK("_MyClass", "_MyClass", {});
ExpectOK("a::MyClass", "MyClass", {"a"});
ExpectOK("a::foo::MyClass", "MyClass", {"a", "foo"});
ExpectOK("a::foo::b::MyClass", "MyClass", {"a", "foo", "b"});
ExpectOK("a::foo::b::bar::MyClass", "MyClass", {"a", "foo", "b", "bar"});
ExpectOK("foo::MyClass", "MyClass", {"foo"});
ExpectOK("_foo::MyClass", "MyClass", {"_foo"});
ExpectOK("_foo::_MyClass", "_MyClass", {"_foo"});
ExpectOK("::foo::bar::MyClass", "MyClass", {"foo", "bar"});
ExpectOK("::_foo::MyClass", "MyClass", {"_foo"});
ExpectOK("::_foo::_MyClass", "_MyClass", {"_foo"});
// Make sure we didn't skip a valid letter or digit
std::string ident;
for (char c = 'a'; c <= 'z'; c++) {
ident.append(1, c);
}
for (char c = 'A'; c <= 'Z'; c++) {
ident.append(1, c);
}
for (char c = '0'; c <= '9'; c++) {
ident.append(1, c);
}
ident += "_";
ExpectOK(ident, ident, {});
ExpectOK(ident + "::" + ident, ident, {ident});
ExpectOK(ident + "::" + ident + "::" + ident, ident, {ident, ident});
}
TEST_F(ParseCppClassTest, ParseFail) {
ExpectFail("");
ExpectFail("::");
ExpectFail("0");
ExpectFail("a.b");
ExpectFail("a:b");
ExpectFail(":foo::bar");
ExpectFail("good::.bad");
ExpectFail("good:::bad");
ExpectFail("good::bad::");
ExpectFail("good::::bad");
ExpectFail("::::bad");
ExpectFail("good:: bad");
ExpectFail("good::0bad");
}
static void CompareWithGoldenFile(
const std::string& tensorflow_relative_golden_file_name,
const std::string& expected_contents, bool ignore_cr) {
// Get rid of all CR characters, we may be running under windows.
std::string sanitized_expected_contents(expected_contents);
if (ignore_cr) {
sanitized_expected_contents.erase(
std::remove(sanitized_expected_contents.begin(),
sanitized_expected_contents.end(), '\r'),
sanitized_expected_contents.end());
}
// To update the golden file, flip update_golden to true and run the
// following:
// blaz test --test_strategy=local \
// "third_party/tensorflow/compiler/aot:codegen_test"
const bool update_golden = false;
std::string golden_file_name =
GetDataDependencyFilepath(tensorflow_relative_golden_file_name);
if (update_golden) {
TF_EXPECT_OK(
WriteStringToFile(Env::Default(), golden_file_name, expected_contents));
}
std::string golden_file_contents;
TF_ASSERT_OK(ReadFileToString(Env::Default(), golden_file_name,
&golden_file_contents));
if (ignore_cr) {
golden_file_contents.erase(std::remove(golden_file_contents.begin(),
golden_file_contents.end(), '\r'),
golden_file_contents.end());
}
EXPECT_EQ(golden_file_contents, expected_contents);
}
} // namespace
} // namespace tfcompile
} // namespace tensorflow
tensorflow/compiler/aot/compile.cc
+381
View File
@@ -0,0 +1,381 @@
/* Copyright 2017 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/compiler/aot/compile.h"
#include <cstddef>
#include <memory>
#include <set>
#include <string>
#include <utility>
#include <vector>
#include "absl/base/call_once.h"
#include "absl/log/check.h"
#include "absl/status/status.h"
#include "absl/status/statusor.h"
#include "absl/strings/match.h"
#include "absl/strings/str_cat.h"
#include "absl/strings/str_join.h"
#include "absl/strings/str_split.h"
#include "absl/strings/string_view.h"
#include "llvm-c/Target.h"
#include "llvm/Support/ManagedStatic.h"
#include "tensorflow/compiler/aot/codegen.h"
#include "tensorflow/compiler/aot/flags.h"
#include "tensorflow/compiler/aot/quantize.h"
#include "tensorflow/compiler/tf2xla/tf2xla.h"
#include "tensorflow/compiler/tf2xla/tf2xla_util.h"
#include "xla/backends/cpu/codegen/symbol_name_util.h"
#include "xla/client/client_library.h"
#include "xla/client/compile_only_client.h"
#include "xla/hlo/builder/xla_computation.h"
#include "xla/service/compiler.h"
#include "xla/service/cpu/cpu_aot_compilation_result.h"
#include "xla/shape.h"
#include "xla/status_macros.h"
#include "xla/stream_executor/platform.h"
#include "xla/stream_executor/platform_manager.h"
#include "xla/tsl/platform/errors.h"
#include "xla/tsl/platform/statusor.h"
#include "xla/util.h"
#include "xla/util/embedded_constant_buffers.h"
#include "xla/xla_data.pb.h"
#include "tensorflow/core/framework/graph.pb.h"
#include "tensorflow/core/lib/core/errors.h"
#include "tensorflow/core/lib/strings/proto_serialization.h"
#include "tensorflow/core/platform/env.h"
#include "tensorflow/core/platform/logging.h"
#include "tensorflow/core/platform/regexp.h" // IWYU pragma: keep
#include "tensorflow/core/platform/types.h"
namespace tensorflow {
namespace tfcompile {
static llvm::ManagedStatic<QuantizeXlaFn> quantize_xla;
bool RegisterQuantizeFn(const QuantizeXlaFn& fn) {
if (*quantize_xla) return false;
*quantize_xla = fn;
return true;
}
namespace {
// Compiles the XLA computation into executable code.
absl::Status CompileXla(xla::CompileOnlyClient* client,
const xla::XlaComputation& computation,
const xla::cpu::CpuAotCompilationOptions& aot_opts,
CompileResult* compile_result) {
// Retrieves arg and result layouts from the computation.
// TODO(toddw): Should we let the user choose the major/minor ordering?
absl::StatusOr<std::unique_ptr<xla::ProgramShape>> pshape_or =
client->GetComputationShape(computation);
if (!pshape_or.ok()) {
return absl::UnknownError(absl::StrCat("Couldn't get XLA program shape: ",
pshape_or.status().message()));
}
compile_result->program_shape = pshape_or.value()->ToProto();
xla::ProgramShapeProto* pshape = &compile_result->program_shape;
// AotXlaComputationInstance::argument_layouts is a vector of Shape
// pointers. Accumulate the Shape objects themselves in a separate vector
// while building the vector of pointers.
std::vector<const xla::Shape*> arg_layout_ptrs(pshape->parameters_size());
std::vector<xla::Shape> arg_layouts(pshape->parameters_size());
for (int i = 0; i < pshape->parameters_size(); ++i) {
TF_ASSIGN_OR_RETURN(arg_layouts[i],
xla::Shape::FromProto(pshape->parameters(i)));
arg_layout_ptrs[i] = &arg_layouts[i];
}
xla::CompileOnlyClient::AotXlaComputationInstance instance;
instance.computation = &computation;
instance.argument_layouts = std::move(arg_layout_ptrs);
TF_ASSIGN_OR_RETURN(xla::Shape result_shape,
xla::Shape::FromProto(pshape->result()));
instance.result_layout = &result_shape;
absl::StatusOr<std::vector<std::unique_ptr<xla::CompiledModule>>> aot_or =
client->CompileAheadOfTime(instance, aot_opts);
if (!aot_or.ok()) {
return absl::UnknownError(
absl::StrCat("XLA compilation failed: ", aot_or.status().message()));
}
compile_result->aot =
xla::unique_ptr_down_cast<xla::cpu::CpuAotCompilationResult>(
std::move(aot_or.value().back()));
compile_result->entry_point = aot_opts.entry_point_name();
compile_result->pointer_size =
xla::CompileOnlyClient::PointerSizeForTriple(aot_opts.triple());
return absl::OkStatus();
}
// Renames the computation proto to ensure unique symbol names to avoid linking
// errors when linking multiple tf_library targets.
absl::Status ConfigureKernelNamingConvention(
xla::cpu::CpuAotCompilationOptions& aot_opts,
xla::XlaComputation& computation, const std::string& cpp_class) {
aot_opts.mutable_debug_options()
->set_xla_cpu_generate_unique_c_style_kernel_entry_points(true);
TF_ASSIGN_OR_RETURN(std::string class_name_as_valid_c_name,
xla::cpu::ConvertToCName(cpp_class));
// Prefix the computation name. We use this to blacklist the generated symbols
// from dfsan instrumentation.
constexpr absl::string_view kModuleNameGeneratorPrefix =
"tfcompile_xla_generated";
// Rename proto to ensure unique symbol names.
*computation.mutable_proto()->mutable_name() =
absl::StrCat(kModuleNameGeneratorPrefix, "_", computation.proto().name(),
"_", class_name_as_valid_c_name);
return absl::OkStatus();
}
} // namespace
absl::Status CompileGraph(GraphDef graph_def, const tf2xla::Config& config,
const MainFlags& flags,
CompileResult* compile_result) {
// Converts the graph into an XLA computation, and compiles the
// computation.
// TODO(toddw): Should we let the user pick the XLA cpu vs. gpu client?
se::Platform* cpu_platform =
se::PlatformManager::PlatformWithName("Host").value();
xla::CompileOnlyClient* client =
xla::ClientLibrary::GetOrCreateCompileOnlyClient(cpu_platform).value();
xla::XlaComputation computation;
bool use_mlir_bridge = false;
if (!flags.mlir_components.empty() && flags.mlir_components != "None") {
for (auto component : absl::StrSplit(flags.mlir_components, ',')) {
if (component == "Bridge") {
use_mlir_bridge = true;
} else {
return absl::UnknownError(
absl::StrCat("Unknown mlir_component ", component));
}
}
}
if (use_mlir_bridge) {
TF_RETURN_IF_ERROR(ConvertGraphDefToXlaViaMlir(
graph_def, config, &computation, flags.debug_info,
flags.debug_info_path_begin_marker));
} else {
TF_RETURN_IF_ERROR(ConvertGraphDefToXla(std::move(graph_def), config,
client, &computation));
}
if (flags.experimental_quantize && *quantize_xla) {
TF_RETURN_IF_ERROR((*quantize_xla)(config, &computation));
}
if (!flags.out_session_module.empty()) {
TF_ASSIGN_OR_RETURN(std::unique_ptr<xla::HloSnapshot> module,
computation.Snapshot());
// Serialize the HloSnapshot deterministically so that all the outputs of a
// tf_library genrule are deterministic.
const size_t size = module->ByteSizeLong();
auto serialized = std::make_unique<char[]>(size);
TF_RET_CHECK(
SerializeToBufferDeterministic(*module, serialized.get(), size));
TF_RETURN_IF_ERROR(
WriteStringToFile(Env::Default(), flags.out_session_module,
absl::string_view(serialized.get(), size)));
}
xla::cpu::CpuAotCompilationOptions aot_opts(
flags.target_triple, flags.target_cpu, flags.target_features,
flags.entry_point,
xla::cpu::CpuAotCompilationOptions::RelocationModel::BigPic,
/*compile_copy_as_llvm_kernel=*/true);
if (flags.sanitize_dataflow) {
aot_opts.set_sanitize_dataflow(flags.sanitize_dataflow);
aot_opts.set_sanitize_abilists_dataflow(absl::StrSplit(
flags.sanitize_abilists_dataflow, ',', absl::SkipEmpty()));
}
TF_RETURN_IF_ERROR(
ConfigureKernelNamingConvention(aot_opts, computation, flags.cpp_class));
return CompileXla(client, computation, aot_opts, compile_result);
}
static absl::Status ReadProtoFile(const std::string& fname,
protobuf::Message* proto) {
if (absl::EndsWith(fname, ".pbtxt")) {
return ReadTextProto(Env::Default(), fname, proto);
} else {
return ReadBinaryProto(Env::Default(), fname, proto);
}
}
static absl::once_flag targets_init;
static void InitializeTargets() {
// Initialize all LLVM targets so we can cross compile.
#if TF_LLVM_AARCH32_AVAILABLE
LLVMInitializeARMTarget();
LLVMInitializeARMTargetInfo();
LLVMInitializeARMTargetMC();
LLVMInitializeARMAsmParser();
LLVMInitializeARMAsmPrinter();
#endif
#if TF_LLVM_AARCH64_AVAILABLE
LLVMInitializeAArch64Target();
LLVMInitializeAArch64TargetInfo();
LLVMInitializeAArch64TargetMC();
LLVMInitializeAArch64AsmParser();
LLVMInitializeAArch64AsmPrinter();
#endif
#if TF_LLVM_HEXAGON_AVAILABLE
LLVMInitializeHexagonTarget();
LLVMInitializeHexagonTargetInfo();
LLVMInitializeHexagonTargetMC();
LLVMInitializeHexagonAsmParser();
LLVMInitializeHexagonAsmPrinter();
#endif
#if TF_LLVM_POWERPC_AVAILABLE
LLVMInitializePowerPCTarget();
LLVMInitializePowerPCTargetInfo();
LLVMInitializePowerPCTargetMC();
LLVMInitializePowerPCAsmParser();
LLVMInitializePowerPCAsmPrinter();
#endif
#if TF_LLVM_S390X_AVAILABLE
LLVMInitializeSystemZTarget();
LLVMInitializeSystemZTargetInfo();
LLVMInitializeSystemZTargetMC();
LLVMInitializeSystemZAsmParser();
LLVMInitializeSystemZAsmPrinter();
#endif
#if TF_LLVM_X86_AVAILABLE
LLVMInitializeX86Target();
LLVMInitializeX86TargetInfo();
LLVMInitializeX86TargetMC();
LLVMInitializeX86AsmParser();
LLVMInitializeX86AsmPrinter();
#endif
}
// Replaces {{tag.type tag.name}} in the error message with tag_name.
// TODO(bixia): We currently only handlge tag.type == "node".
//
// In the error message, a graph node is represented as {{tag.type, tag.name}},
// to allow a Python debugger to insert source information about the graph node.
// For example, a Python add expression may be represented as
// {{node, x_y_sum}} = Add(x, y) in the error message. See routine interpolate
// in tensorflow/python/framework/error_interpolation.py for more detail.
static std::string InterpolateErrorMessage(std::string message) {
// See _NAME_REGEX in tensorflow/python/framework/error_interpolation.py
// Change "prefix {{node tag.name}} suffix" to "prefix tag.name suffix".
static LazyRE2 pattern{"(.*){{node (.*)}}(.*)"};
RE2::GlobalReplace(&message, *pattern, "\\1\\2\\3");
return message;
}
absl::Status Main(const MainFlags& flags) {
absl::call_once(targets_init, &InitializeTargets);
// Process config.
tf2xla::Config config;
if (flags.config.empty()) {
return absl::InvalidArgumentError("Must specify --config");
}
TF_RETURN_IF_ERROR(ReadProtoFile(flags.config, &config));
TF_RETURN_IF_ERROR(ValidateConfig(config));
if (flags.dump_fetch_nodes) {
std::set<std::string> nodes;
for (const tf2xla::Fetch& fetch : config.fetch()) {
nodes.insert(fetch.id().node_name());
}
std::cout << absl::StrJoin(nodes, ",");
return absl::OkStatus();
}
// Read and initialize the graph.
if (flags.graph.empty()) {
return absl::InvalidArgumentError("Must specify --graph");
}
GraphDef graph_def;
TF_RETURN_IF_ERROR(ReadProtoFile(flags.graph, &graph_def));
CompileResult compile_result;
absl::Status status =
CompileGraph(std::move(graph_def), config, flags, &compile_result);
if (!status.ok()) {
return errors::CreateWithUpdatedMessage(
status, InterpolateErrorMessage(std::string(status.message())));
}
// Write output files.
Env* env = Env::Default();
const auto obj_files = compile_result.aot->obj_files();
DCHECK_EQ(obj_files.size(), 1);
const absl::string_view obj_file = obj_files[0];
TF_RETURN_IF_ERROR(
WriteStringToFile(env, flags.out_function_object, obj_file));
CodegenOpts codegen_opts;
codegen_opts.gen_name_to_index = flags.gen_name_to_index;
codegen_opts.gen_program_shape = flags.gen_program_shape;
codegen_opts.target_triple = flags.target_triple;
codegen_opts.use_xla_nanort_runtime = flags.use_xla_nanort_runtime;
// Set the XLA Runtime bit if this is an HloLowering.
if (!flags.mlir_components.empty() && flags.mlir_components != "None") {
for (auto component : absl::StrSplit(flags.mlir_components, ',')) {
if (component == "HloLowering") {
codegen_opts.use_xla_runtime = true;
}
}
}
if (flags.cpp_class.empty()) {
return absl::InvalidArgumentError("Must specify --cpp_class");
}
codegen_opts.gen_hlo_profile_printer_data =
xla::GetDebugOptionsFromFlags().xla_hlo_profile();
TF_RETURN_IF_ERROR(ParseCppClass(flags.cpp_class, &codegen_opts.class_name,
&codegen_opts.namespaces));
xla::EmbeddedConstantBuffers embedded_constant_buffers;
if (flags.out_constant_buffers_object.empty()) {
return absl::InvalidArgumentError(
"Must specify --out_constant_buffers_object when using AOT thunks");
}
TF_ASSIGN_OR_RETURN(
embedded_constant_buffers,
GenerateConstantBuffersData(codegen_opts, compile_result));
TF_RETURN_IF_ERROR(
WriteStringToFile(env, flags.out_constant_buffers_object,
embedded_constant_buffers.object_file_data));
MetadataResult metadata_result;
TF_RETURN_IF_ERROR(
GenerateMetadata(codegen_opts, compile_result, &metadata_result));
TF_RETURN_IF_ERROR(WriteStringToFile(env, flags.out_metadata_object,
metadata_result.object_file_data));
std::string header;
TF_RETURN_IF_ERROR(GenerateHeader(codegen_opts, config, compile_result,
metadata_result, embedded_constant_buffers,
&header));
TF_RETURN_IF_ERROR(WriteStringToFile(env, flags.out_header, header));
return absl::OkStatus();
}
} // namespace tfcompile
} // namespace tensorflow
tensorflow/compiler/aot/compile.h
+59
View File
@@ -0,0 +1,59 @@
/* Copyright 2017 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_COMPILER_AOT_COMPILE_H_
#define TENSORFLOW_COMPILER_AOT_COMPILE_H_
#include <memory>
#include <utility>
#include "absl/status/status.h"
#include "absl/status/statusor.h"
#include "tensorflow/compiler/aot/flags.h"
#include "tensorflow/compiler/tf2xla/tf2xla.pb.h"
#include "xla/service/cpu/cpu_aot_compilation_result.h"
#include "xla/xla_data.pb.h"
#include "tensorflow/core/framework/graph.pb.h"
#include "tensorflow/core/platform/types.h"
#include "tsl/platform/casts.h"
namespace tensorflow {
namespace tfcompile {
// CompileResult describes the output of CompileGraph, where the object file
// data and meta-information is available in aot.
struct CompileResult {
// Contains object file and meta-info.
std::unique_ptr<xla::cpu::CpuAotCompilationResult> aot;
xla::ProgramShapeProto program_shape; // Static shape of args and results.
std::string entry_point; // Name of generated function.
int pointer_size = 0; // Size of a pointer in bytes.
};
// CompileGraph compiles the graph_def into an object file containing a function
// that performs the graph operations.
//
// The XLA compilation options are specified in the flags.
absl::Status CompileGraph(GraphDef graph_def, const tf2xla::Config& config,
const MainFlags& flags,
CompileResult* compile_result);
// The full compilation method, for reuse in a library setting.
absl::Status Main(const MainFlags& flags);
} // namespace tfcompile
} // namespace tensorflow
#endif // TENSORFLOW_COMPILER_AOT_COMPILE_H_
tensorflow/compiler/aot/embedded_protocol_buffers.cc
+159
View File
@@ -0,0 +1,159 @@
/* Copyright 2017 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/compiler/aot/embedded_protocol_buffers.h"
#include <memory>
#include <optional>
#include <string>
#include "absl/log/check.h"
#include "absl/memory/memory.h"
#include "absl/strings/str_replace.h"
#include "llvm/IR/Constants.h"
#include "llvm/IR/GlobalVariable.h"
#include "llvm/IR/LLVMContext.h"
#include "llvm/IR/LegacyPassManager.h"
#include "llvm/IR/Module.h"
#include "llvm/MC/TargetRegistry.h"
#include "llvm/Target/TargetMachine.h"
#include "llvm/Target/TargetOptions.h"
#include "llvm/TargetParser/Triple.h"
#include "xla/service/llvm_ir/llvm_type_conversion_util.h"
#include "xla/tsl/platform/statusor.h"
#include "xla/util.h"
#include "tensorflow/core/lib/strings/proto_serialization.h"
namespace tensorflow {
namespace tfcompile {
using xla::llvm_ir::AsStringRef;
static void AddEmbeddedProtocolBufferToLlvmModule(
llvm::Module* module, const ::tensorflow::protobuf::MessageLite& proto,
absl::string_view unique_identifier,
std::string* protobuf_array_symbol_name, int64_t* protobuf_array_size) {
std::string protobuf_array_contents;
CHECK(::tensorflow::SerializeToStringDeterministic(proto,
&protobuf_array_contents));
*protobuf_array_symbol_name =
absl::StrCat(unique_identifier, "_protobuf_array_contents");
*protobuf_array_size = protobuf_array_contents.size();
llvm::Constant* protobuf_array_initializer =
llvm::ConstantDataArray::getString(module->getContext(),
AsStringRef(protobuf_array_contents),
/*AddNull=*/false);
new llvm::GlobalVariable(
*module, protobuf_array_initializer->getType(),
/*isConstant=*/true, llvm::GlobalValue::ExternalLinkage,
protobuf_array_initializer, AsStringRef(*protobuf_array_symbol_name));
}
static std::string CreateCPPShimExpression(
absl::string_view qualified_cpp_protobuf_name,
absl::string_view protobuf_array_symbol_name, int64_t protobuf_array_size) {
std::string code =
"[]() {\n"
" {{PROTOBUF_NAME}}* proto = new {{PROTOBUF_NAME}};\n"
" proto->ParseFromArray(&{{ARRAY_SYMBOL}}[0], {{ARRAY_SIZE}});\n"
" return proto;\n"
" }()";
return absl::StrReplaceAll(
code,
{
{"{{ARRAY_SYMBOL}}", absl::StrCat(protobuf_array_symbol_name)},
{"{{ARRAY_SIZE}}", absl::StrCat(protobuf_array_size)},
{"{{PROTOBUF_NAME}}", absl::StrCat(qualified_cpp_protobuf_name)},
});
}
static absl::StatusOr<std::string> CodegenModule(
llvm::TargetMachine* target_machine, std::unique_ptr<llvm::Module> module) {
llvm::SmallVector<char, 0> stream_buffer;
llvm::raw_svector_ostream ostream(stream_buffer);
llvm::legacy::PassManager codegen_passes;
if (target_machine->addPassesToEmitFile(codegen_passes, ostream, nullptr,
llvm::CodeGenFileType::ObjectFile)) {
return xla::Internal(
"Could not create pass pipeline to generate object file");
}
codegen_passes.run(*module);
return std::string(stream_buffer.begin(), stream_buffer.end());
}
static absl::StatusOr<std::unique_ptr<llvm::TargetMachine>>
GetTargetMachineFromTriple(absl::string_view target_triple) {
std::string error;
llvm::Triple normalized_triple(
llvm::Triple::normalize(AsStringRef(absl::string_view(target_triple))));
const llvm::Target* target =
llvm::TargetRegistry::lookupTarget(normalized_triple, error);
if (target == nullptr) {
return xla::Internal("TargetRegistry::lookupTarget failed: %s",
error.c_str());
}
return absl::WrapUnique(target->createTargetMachine(
normalized_triple, /*CPU=*/"",
/*Features=*/"", llvm::TargetOptions(), std::nullopt));
}
absl::StatusOr<EmbeddedProtocolBuffers> CreateEmbeddedProtocolBuffers(
absl::string_view target_triple,
absl::Span<const ProtobufToEmbed> protobufs_to_embed) {
TF_ASSIGN_OR_RETURN(std::unique_ptr<llvm::TargetMachine> target_machine,
GetTargetMachineFromTriple(target_triple));
llvm::LLVMContext llvm_context;
auto module_with_serialized_proto =
std::make_unique<llvm::Module>("embedded_data_module", llvm_context);
EmbeddedProtocolBuffers result;
for (const ProtobufToEmbed& protobuf_to_embed : protobufs_to_embed) {
std::string cpp_shim, cpp_variable_decl;
if (protobuf_to_embed.message) {
std::string protobuf_array_symbol_name;
int64_t protobuf_array_size;
AddEmbeddedProtocolBufferToLlvmModule(
module_with_serialized_proto.get(), *protobuf_to_embed.message,
protobuf_to_embed.symbol_prefix, &protobuf_array_symbol_name,
&protobuf_array_size);
cpp_shim = CreateCPPShimExpression(
protobuf_to_embed.qualified_cpp_protobuf_name,
protobuf_array_symbol_name, protobuf_array_size);
cpp_variable_decl =
absl::StrCat("extern \"C\" char ", protobuf_array_symbol_name, "[];");
} else {
cpp_shim = "nullptr";
}
result.cpp_shims.push_back({cpp_shim, cpp_variable_decl});
}
TF_ASSIGN_OR_RETURN(result.object_file_data,
CodegenModule(target_machine.get(),
std::move(module_with_serialized_proto)));
return result;
}
} // namespace tfcompile
} // namespace tensorflow
tensorflow/compiler/aot/embedded_protocol_buffers.h
+92
View File
@@ -0,0 +1,92 @@
/* 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.
==============================================================================*/
// This file defines utilities to help "embed" protocol buffers into object
// (".o") files. These C++ binaries and shared objects can link in these .o to
// get access to said protocol buffers at runtime.
#ifndef TENSORFLOW_COMPILER_AOT_EMBEDDED_PROTOCOL_BUFFERS_H_
#define TENSORFLOW_COMPILER_AOT_EMBEDDED_PROTOCOL_BUFFERS_H_
#include "absl/status/statusor.h"
#include "absl/types/span.h"
#include "tensorflow/core/platform/protobuf.h"
namespace tensorflow {
namespace tfcompile {
using absl::StatusOr;
// Represents a set of protocol buffers embedded into an object file and
// describes how to access them at runtime.
struct EmbeddedProtocolBuffers {
// Each instance CPPShim describes how to generate C++ code to instantiate a
// protobuf instance from the corresponding static data emitted into the
// object file.
struct CPPShim {
// `expression` is a C++ expression that creates an instance of said
// protocol buffer when executed.
std::string expression;
// `variable_decl` is an "extern C" array declaration that is used in
// `expression`. It must be visible wherever `expression` is emitted.
std::string variable_decl;
};
// Each cpp_shim corresponds to one embedded protocol buffer.
std::vector<CPPShim> cpp_shims;
// The contents of the object (".o") file the protocol buffers are embbed in.
// This needs to be linked in to any program that wants to execute any of the
// expressions in `cpp_shims`.
std::string object_file_data;
};
// Describes a protocol buffer to embed into an object file.
struct ProtobufToEmbed {
// `symbol_prefix` is prefix that is guaranteed to be unique across the binary
// or DSO the generated object file will be linked into.
std::string symbol_prefix;
// `qualified_cpp_protobuf_name` is a qualified ("qualified" as in C++
// namespace qualified) protocol buffer name. This is only used in
// CPPShim::expression so relatively qualified names are fine as long as
// they're valid wherever CPPShim::expression is emitted.
std::string qualified_cpp_protobuf_name;
// `message` is the protocol buffer to be embedded. It is allowed to be
// nullptr, in which case the generated C++ shim expression is just `nullptr`,
// and the generated object file does not define any symbols.
const ::tensorflow::protobuf::MessageLite* message;
};
// Embeds a sequence of protocol buffers into an object file.
//
// `target_triple` is the target triple for the target architecture for the
// generated object file.
//
// `protobufs_to_embed` describes the protocol buffers to embed into the
// resulting object file. The C++ shim for protobufs_to_embed[i] is
// cpp_shims[i] in the returned EmbeddedProtocolBuffers instance. The contents
// of all the protocol buffers are embedded into a single .o file whose content
// is stored in the object_file_data field in the returned
// EmbeddedProtocolBuffers instance.
absl::StatusOr<EmbeddedProtocolBuffers> CreateEmbeddedProtocolBuffers(
absl::string_view target_triple,
absl::Span<const ProtobufToEmbed> protobufs_to_embed);
} // namespace tfcompile
} // namespace tensorflow
#endif // TENSORFLOW_COMPILER_AOT_EMBEDDED_PROTOCOL_BUFFERS_H_
tensorflow/compiler/aot/flags.cc
+101
View File
@@ -0,0 +1,101 @@
/* Copyright 2017 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/compiler/aot/flags.h"
namespace tensorflow {
namespace tfcompile {
void AppendMainFlags(std::vector<Flag>* flag_list, MainFlags* flags) {
const std::vector<Flag> tmp = {
{"graph", &flags->graph,
"Input GraphDef file. If the file ends in '.pbtxt' it is expected to "
"be in the human-readable proto text format, otherwise it is expected "
"to be in the proto binary format."},
{"debug_info", &flags->debug_info,
"Graph debug info file. If the file ends in '.pbtxt' it is expected to "
"be in the human-readable proto text format, otherwise it is expected "
"to be in the proto binary format."},
{"debug_info_path_begin_marker", &flags->debug_info_path_begin_marker,
"If not none, only keep the file path in the debug information after the"
" marker. The default value is empty"},
{"config", &flags->config,
"Input file containing Config proto. If the file ends in '.pbtxt' it "
"is expected to be in the human-readable proto text format, otherwise "
"it is expected to be in the proto binary format."},
{"dump_fetch_nodes", &flags->dump_fetch_nodes,
"If set, only flags related to fetches are processed, and the resulting "
"fetch nodes will be dumped to stdout in a comma-separated list. "
"Typically used to format arguments for other tools, e.g. "
"freeze_graph."},
// Flags controlling the XLA ahead-of-time compilation, that correspond to
// the fields of xla::cpu::CpuAotCompilationOptions.
//
// TODO(toddw): The following flags also need to be supported:
// --xla_cpu_llvm_opt_level
// --xla_cpu_llvm_cl_opts
{"target_triple", &flags->target_triple,
"Target platform, similar to the clang -target flag. The general "
"format is <arch><sub>-<vendor>-<sys>-<abi>. "
"http://clang.llvm.org/docs/CrossCompilation.html#target-triple."},
{"target_cpu", &flags->target_cpu,
"Target cpu, similar to the clang -mcpu flag. "
"http://clang.llvm.org/docs/CrossCompilation.html#cpu-fpu-abi"},
{"target_features", &flags->target_features,
"Target features, e.g. +avx2, +neon, etc."},
{"entry_point", &flags->entry_point,
"Name of the generated function. If multiple generated object files "
"will be linked into the same binary, each will need a unique entry "
"point."},
{"cpp_class", &flags->cpp_class,
"Name of the generated C++ class, wrapping the generated function. The "
"syntax of this flag is [[<optional_namespace>::],...]<class_name>. "
"This mirrors the C++ syntax for referring to a class, where multiple "
"namespaces may precede the class name, separated by double-colons. "
"The class will be generated in the given namespace(s), or if no "
"namespaces are given, within the global namespace."},
{"out_function_object", &flags->out_function_object,
"Output object file containing the generated function for the "
"TensorFlow model."},
{"out_header", &flags->out_header, "Output header file name."},
{"out_metadata_object", &flags->out_metadata_object,
"Output object file name containing optional metadata for the generated "
"function."},
{"out_constant_buffers_object", &flags->out_constant_buffers_object,
"Output object file name containing constant buffers for the runtime."},
{"out_session_module", &flags->out_session_module,
"Output session module proto."},
{"mlir_components", &flags->mlir_components,
"The MLIR components to enable. Currently only Bridge is supported."},
{"experimental_quantize", &flags->experimental_quantize,
"If set, quantization passes will run and dump the result before HLO "
"code generation."},
{"sanitize_dataflow", &flags->sanitize_dataflow,
"Enable DataFlow Sanitizer pass."},
{"sanitize_abilists_dataflow", &flags->sanitize_abilists_dataflow,
"Comma separated list of ABIList file paths."},
{"gen_name_to_index", &flags->gen_name_to_index,
"Generate name-to-index data for Lookup{Arg,Result}Index methods."},
{"gen_program_shape", &flags->gen_program_shape,
"Generate program shape data for the ProgramShape method."},
{"use_xla_nanort_runtime", &flags->use_xla_nanort_runtime,
"Use xla cpu nanort runtime, otherwise each thunk execution gets "
"serialized directly into the header."},
};
flag_list->insert(flag_list->end(), tmp.begin(), tmp.end());
}
} // namespace tfcompile
} // namespace tensorflow
tensorflow/compiler/aot/flags.h
+64
View File
@@ -0,0 +1,64 @@
/* Copyright 2017 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_COMPILER_AOT_FLAGS_H_
#define TENSORFLOW_COMPILER_AOT_FLAGS_H_
#include <string>
#include <vector>
#include "tensorflow/core/util/command_line_flags.h"
namespace tensorflow {
namespace tfcompile {
// Flags for the tfcompile binary. See *.cc file for descriptions.
struct MainFlags {
std::string graph;
std::string debug_info;
std::string debug_info_path_begin_marker;
std::string config;
bool dump_fetch_nodes = false;
std::string target_triple;
std::string target_cpu;
std::string target_features;
std::string entry_point;
std::string cpp_class;
std::string out_function_object;
std::string out_metadata_object;
std::string out_header;
std::string out_constant_buffers_object;
std::string out_session_module;
std::string mlir_components;
bool experimental_quantize = false;
// Sanitizer pass options
bool sanitize_dataflow = false;
std::string sanitize_abilists_dataflow;
// C++ codegen options
bool gen_name_to_index = false;
bool gen_program_shape = false;
bool use_xla_nanort_runtime = false;
};
// Appends to flag_list a tensorflow::Flag for each field in MainFlags.
void AppendMainFlags(std::vector<Flag>* flag_list, MainFlags* flags);
} // namespace tfcompile
} // namespace tensorflow
#endif // TENSORFLOW_COMPILER_AOT_FLAGS_H_
tensorflow/compiler/aot/quantize.h
+41
View File
@@ -0,0 +1,41 @@
/* 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.
==============================================================================*/
#ifndef TENSORFLOW_COMPILER_AOT_QUANTIZE_H_
#define TENSORFLOW_COMPILER_AOT_QUANTIZE_H_
#include <functional>
#include <iostream>
#include <ostream>
#include "tensorflow/compiler/tf2xla/tf2xla.pb.h"
#include "xla/hlo/builder/xla_computation.h"
#include "tensorflow/core/platform/errors.h"
#include "tensorflow/core/platform/status.h"
namespace tensorflow {
namespace tfcompile {
using QuantizeXlaFn = std::function<absl::Status(
const tf2xla::Config& config, xla::XlaComputation* computation)>;
// Set the static quantization function to the `fn` if it hasn't been set.
// Return false if the static function has been set.
bool RegisterQuantizeFn(const QuantizeXlaFn& fn);
} // namespace tfcompile
} // namespace tensorflow
#endif // TENSORFLOW_COMPILER_AOT_QUANTIZE_H_
tensorflow/compiler/aot/test.cc
+89
View File
@@ -0,0 +1,89 @@
/* Copyright 2017 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.
==============================================================================*/
// Generated by the tf_library build rule. DO NOT EDIT!
//
// This file contains a test and benchmark for the function generated by
// tfcompile. All tokens of the form `{{TFCOMPILE_*}}` must be rewritten to
// real values before this file can be compiled.
//
// TFCOMPILE_HEADER : Path to the header file generated by tfcompile.
// TFCOMPILE_CPP_CLASS : Name of the C++ class generated by tfcompile.
// TFCOMPILE_NAME : Name for tests and benchmarks.
//
// The tf_library bazel macro in tfcompile.bzl performs the token rewriting, and
// generates a cc_test rule for you.
// These macros must be defined before eigen files are included.
#define EIGEN_USE_THREADS
#define EIGEN_USE_CUSTOM_THREAD_POOL
// clang-format off
#include "{{TFCOMPILE_HEADER}}" // NOLINT(whitespace/braces)
// clang-format on
#include "unsupported/Eigen/CXX11/Tensor" // from @eigen_archive
#include "tensorflow/core/platform/byte_order.h"
#include "tensorflow/core/platform/cpu_info.h"
#include "tensorflow/core/platform/test.h"
#include "tensorflow/core/platform/test_benchmark.h"
// Macros that expand to tokens based on the entry point name.
// clang-format off
#define CPP_CLASS {{TFCOMPILE_CPP_CLASS}} // NOLINT(whitespace/braces)
#define TEST_NAME {{TFCOMPILE_NAME}}Test // NOLINT(whitespace/braces)
#define BM_NAME BM_{{TFCOMPILE_NAME}} // NOLINT(whitespace/braces)
// clang-format on
namespace tensorflow {
namespace tfcompile {
namespace {
void zero_buffers(XlaCompiledCpuFunction* computation) {
for (int i = 0; i < computation->num_args(); ++i) {
memset(computation->arg_data(i), 0, computation->arg_size(i));
}
}
// Trivial test that runs the generated function to ensure it doesn't crash.
TEST(TEST_NAME, NoCrash) {
Eigen::ThreadPool pool(port::MaxParallelism());
Eigen::ThreadPoolDevice device(&pool, pool.NumThreads());
CPP_CLASS computation;
computation.set_thread_pool(&device);
zero_buffers(&computation);
EXPECT_TRUE(computation.Run());
}
// Simple benchmark that repeatedly runs the generated function.
void BM_NAME(benchmark::State& state) {
Eigen::ThreadPool pool(port::MaxParallelism());
Eigen::ThreadPoolDevice device(&pool, pool.NumThreads());
CPP_CLASS computation;
computation.set_thread_pool(&device);
zero_buffers(&computation);
for (auto s : state) {
computation.Run();
}
}
BENCHMARK(BM_NAME);
} // namespace
} // namespace tfcompile
} // namespace tensorflow
tensorflow/compiler/aot/test_google.cc
+89
View File
@@ -0,0 +1,89 @@
/* Copyright 2023 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.
==============================================================================*/
// Generated by the tf_library build rule. DO NOT EDIT!
//
// This file contains a test and benchmark for the function generated by
// tfcompile. All tokens of the form `{{TFCOMPILE_*}}` must be rewritten to
// real values before this file can be compiled.
//
// TFCOMPILE_HEADER : Path to the header file generated by tfcompile.
// TFCOMPILE_CPP_CLASS : Name of the C++ class generated by tfcompile.
// TFCOMPILE_NAME : Name for tests and benchmarks.
//
// The tf_library bazel macro in tfcompile.bzl performs the token rewriting, and
// generates a cc_test rule for you.
// These macros must be defined before eigen files are included.
#define EIGEN_USE_THREADS
#define EIGEN_USE_CUSTOM_THREAD_POOL
// clang-format off
#include "{{TFCOMPILE_HEADER}}" // NOLINT(whitespace/braces)
// clang-format on
#include "testing/base/public/benchmark.h"
#include <gtest/gtest.h>
#include "unsupported/Eigen/CXX11/Tensor" // from @eigen_archive
#include "tensorflow/core/platform/byte_order.h"
#include "tensorflow/core/platform/cpu_info.h"
// Macros that expand to tokens based on the entry point name.
// clang-format off
#define CPP_CLASS {{TFCOMPILE_CPP_CLASS}} // NOLINT(whitespace/braces)
#define TEST_NAME {{TFCOMPILE_NAME}}Test // NOLINT(whitespace/braces)
#define BM_NAME BM_{{TFCOMPILE_NAME}} // NOLINT(whitespace/braces)
// clang-format on
namespace tensorflow {
namespace tfcompile {
namespace {
void zero_buffers(XlaCompiledCpuFunction* computation) {
for (int i = 0; i < computation->num_args(); ++i) {
memset(computation->arg_data(i), 0, computation->arg_size(i));
}
}
// Trivial test that runs the generated function to ensure it doesn't crash.
TEST(TEST_NAME, NoCrash) {
Eigen::ThreadPool pool(port::MaxParallelism());
Eigen::ThreadPoolDevice device(&pool, pool.NumThreads());
CPP_CLASS computation;
computation.set_thread_pool(&device);
zero_buffers(&computation);
EXPECT_TRUE(computation.Run());
}
// Simple benchmark that repeatedly runs the generated function.
void BM_NAME(benchmark::State& state) {
Eigen::ThreadPool pool(port::MaxParallelism());
Eigen::ThreadPoolDevice device(&pool, pool.NumThreads());
CPP_CLASS computation;
computation.set_thread_pool(&device);
zero_buffers(&computation);
for (auto s : state) {
computation.Run();
}
}
BENCHMARK(BM_NAME);
} // namespace
} // namespace tfcompile
} // namespace tensorflow
tensorflow/compiler/aot/test_graph_tfadd.config.pbtxt
+16
View File
@@ -0,0 +1,16 @@
# Text form of tensorflow.tf2xla.Config proto.
feed {
id { node_name: "x_const" }
shape {
dim { size: 1 }
}
}
feed {
id { node_name: "y_reshape" }
shape {
dim { size: 1 }
}
}
fetch {
id { node_name: "x_y_sum" }
}
tensorflow/compiler/aot/test_graph_tfadd.pbtxt
+56
View File
@@ -0,0 +1,56 @@
node {
name : "x_const"
op : "Const"
attr {
key: "value"
value {
tensor { dtype: DT_INT32 tensor_shape { dim { size: 1 } } int_val: 1 }
}
}
attr {
key : "dtype"
value {
type : DT_INT32
}
}
}
node {
name : "y_const"
op : "Const"
attr {
key: "value"
value {
tensor { dtype: DT_INT32 tensor_shape { dim { size: 1 } } int_val: 2 }
}
}
attr {
key: "dtype"
value {
type: DT_INT32
}
}
}
node {
name : "y_reshape"
op : "Reshape"
input : "y_const"
input : "y_shape"
attr { key: "T" value { type: DT_INT32 } }
# Attribute TShape not specified; needs to be set to its default
# by tfcompile.
}
node {
name : "x_y_sum"
op : "Add"
input : "x_const"
input : "y_reshape"
attr {
key : "T"
value {
type: DT_INT32
}
}
}
versions {
producer: 15
}
tensorflow/compiler/aot/test_graph_tfunknownop.config.pbtxt
+16
View File
@@ -0,0 +1,16 @@
# Text form of tensorflow.tf2xla.Config proto.
feed {
id { node_name: "x_const" }
shape {
dim { size: 1 }
}
}
feed {
id { node_name: "y_const" }
shape {
dim { size: 1 }
}
}
fetch {
id { node_name: "x_y_sum" }
}
tensorflow/compiler/aot/test_graph_tfunknownop.pbtxt
+58
View File
@@ -0,0 +1,58 @@
node {
name : "x_const"
op : "Const"
attr {
key: "value"
value {
tensor {
dtype: DT_INT32
tensor_shape { dim { size: 1 } }
int_val: 1
}
}
}
attr { key : "dtype" value { type: DT_INT32 } }
}
node {
name : "y_const"
op : "Const"
attr {
key: "value"
value {
tensor {
dtype: DT_INT32
tensor_shape { dim { size: 1 } }
int_val: 2
}
}
}
attr { key: "dtype" value { type: DT_INT32 } }
}
node {
name : "x_y_sum"
op : "Add"
input : "x_const"
input : "y_const"
attr { key : "T" value { type: DT_INT32 } }
}
node {
name : "z"
op : "SomeUnknownOp"
input : "x_const"
}
node {
name : "z_identity"
op : "Identity"
input : "z:1"
attr { key : "T" value { type: DT_INT32 } }
}
node {
name : "x_z_sum"
op : "Add"
input : "x_const"
input : "z_identity"
attr { key : "T" value { type: DT_INT32 } }
}
versions {
producer: 15
}
tensorflow/compiler/aot/test_graph_tfunknownop2.config.pbtxt
+25
View File
@@ -0,0 +1,25 @@
# Text form of tensorflow.tf2xla.Config proto.
feed {
id { node_name: "x_const" }
shape {
dim { size: 1 }
}
}
feed {
id { node_name: "y_const" }
shape {
dim { size: 1 }
}
}
feed {
id { node_name: "z_identity"}
shape {
dim { size: 1 }
}
}
fetch {
id { node_name: "x_y_sum" }
}
fetch {
id { node_name: "x_z_sum" }
}
tensorflow/compiler/aot/test_graph_tfunknownop3.config.pbtxt
+26
View File
@@ -0,0 +1,26 @@
# Text form of tensorflow.tf2xla.Config proto.
feed {
id { node_name: "x_const" }
shape {
dim { size: 1 }
}
}
feed {
id { node_name: "y_const" }
shape {
dim { size: 1 }
}
}
feed {
id { node_name: "z" output_index: 1}
shape {
dim { size: 1 }
}
type: DT_INT32
}
fetch {
id { node_name: "x_y_sum" }
}
fetch {
id { node_name: "x_z_sum" }
}
tensorflow/compiler/aot/tests/BUILD
+520
View File
@@ -0,0 +1,520 @@
load("@xla//third_party/rules_python/python:py_binary.bzl", "py_binary")
load("//tensorflow:tensorflow.bzl", "tf_cc_test")
load("//tensorflow:tensorflow.default.bzl", "filegroup", "genrule")
load("//tensorflow/compiler/aot:tfcompile.bzl", "tf_library")
load("//tensorflow/compiler/mlir:glob_lit_test.bzl", "glob_lit_tests")
package(
# copybara:uncomment default_applicable_licenses = ["//tensorflow:license"],
default_visibility = ["//visibility:private"],
licenses = ["notice"],
)
glob_lit_tests(
name = "all_lit_tests",
data = [":filecheck_test_utilities"],
driver = "@llvm-project//mlir:run_lit.sh",
tags_override = {
"test_error_message.lit.pbtxt": ["no_oss"], # TODO(b/150957738): to be fixed on oss.
},
test_file_exts = ["lit.pbtxt"],
)
# Bundle together all of the test utilities that are used by tests.
filegroup(
name = "filecheck_test_utilities",
testonly = True,
srcs = [
"test_error_message.lit.pbtxt.config.pbtxt",
"test_error_message.lit.pbtxt.debug.pbtxt",
"test_error_message.lit.pbtxt.fake_py.debug",
],
data = [
"//tensorflow/compiler/aot:tfcompile",
"@llvm-project//llvm:FileCheck",
"@llvm-project//llvm:not",
],
)
# We disable some tfcompile tests in the open source build with the
# "manual" tag to avoid making our OSS users build LLVM twice
# (once for host and once for target).
test_suite(
name = "all_tests",
tags = ["manual"],
tests = [
":test_graph_tfadd_thunks_test",
":test_graph_tfadd_with_ckpt_saver_thunks_test",
":test_graph_tfadd_with_ckpt_thunks_test",
":test_graph_tfassert_eq_thunks_test",
":test_graph_tfcond_thunks_test",
":test_graph_tffunction_thunks_test",
":test_graph_tfgather_thunks_test",
":test_graph_tfmatmul_thunks_test",
":test_graph_tfmatmulandadd_thunks_test",
":test_graph_tfsplits_thunks_test",
":test_graph_tftop_k_thunks_test",
":test_graph_tfvariable_readonly_thunks_test",
":test_graph_tfvariable_sequential_updates_thunks_test",
":test_graph_tfvariable_thunks_test",
":tfcompile_test_thunks",
],
visibility = ["//visibility:public"],
)
py_binary(
name = "make_test_graphs",
testonly = 1,
srcs = ["make_test_graphs.py"],
strict_deps = True,
deps = [
"//tensorflow/core:protos_all_py",
"//tensorflow/python/client:session",
"//tensorflow/python/framework:constant_op",
"//tensorflow/python/framework:dtypes",
"//tensorflow/python/framework:function",
"//tensorflow/python/framework:ops",
"//tensorflow/python/ops:array_ops",
"//tensorflow/python/ops:array_ops_stack",
"//tensorflow/python/ops:cond",
"//tensorflow/python/ops:control_flow_assert",
"//tensorflow/python/ops:control_flow_ops",
"//tensorflow/python/ops:control_flow_util",
"//tensorflow/python/ops:math_ops",
"//tensorflow/python/ops:nn_ops",
"//tensorflow/python/ops:random_ops",
"//tensorflow/python/ops:variable_v1",
"//tensorflow/python/ops:variables",
"//tensorflow/python/training:saver",
"@absl_py//absl:app",
"@six_archive//:six",
],
)
genrule(
name = "gen_test_graphs",
testonly = 1,
outs = [
"test_graph_tfadd.pb",
"test_graph_tfadd_with_ckpt.ckpt",
"test_graph_tfadd_with_ckpt.pb",
"test_graph_tfadd_with_ckpt_saver.ckpt",
"test_graph_tfadd_with_ckpt_saver.pb",
"test_graph_tfadd_with_ckpt_saver.saver",
"test_graph_tfassert_eq.pb",
"test_graph_tfcond.pb",
"test_graph_tffunction.pb",
"test_graph_tfgather.pb",
"test_graph_tfmatmul.pb",
"test_graph_tfmatmulandadd.pb",
"test_graph_tfmatmul_with_constant.pb",
"test_graph_tfsplits.pb",
"test_graph_tftop_k.pb",
"test_graph_tfvariable.pb",
"test_graph_tfvariable_readonly.pb",
"test_graph_tfvariable_sequential_updates.pb",
"test_graph_tfrandom_uniform.pb",
"test_graph_tfscatter.pb",
],
# Set CUDA_VISIBLE_DEVICES='' to prevent the code we launch from using any
# GPUs which might be present. This is important because builds may run
# concurrently with tests, and tests need to be able to assume that they
# have control of the full GPU.
cmd = "CUDA_VISIBLE_DEVICES='' " +
"$(location :make_test_graphs) --out_dir $(@D)",
tags = ["manual"],
tools = [":make_test_graphs"],
)
# This list is used to run tf_compile on the compiled graphs used by
# tfcompile_test. The test is run for several different configurations of
# tf_compile, so this allows the configurations to be specified in one place.
# suffix, mlir_component, xla_flags, tfcompile_flags
tfcompile_test_dep_configs = [
("_thunks", "None", "--xla_cpu_use_thunk_runtime=true", ""),
("_nanort", "None", "--xla_cpu_use_thunk_runtime=true", "--use_xla_nanort_runtime"),
]
[
[
tf_library(
name = "test_graph_tfadd" + suffix,
testonly = 1,
config = "test_graph_tfadd.config.pbtxt",
cpp_class = "AddComp",
graph = "test_graph_tfadd.pb",
# This serves as a test for the list of minimal deps included even when
# include_standard_runtime_deps is False. If this target fails to
# compile but the others in this directory succeed, you may need to
# expand the "required by all tf_library targets" list in tfcompile.bzl.
include_standard_runtime_deps = False,
mlir_components = mlir_component,
tags = [
"manual",
"no_mac", # TODO(b/228273415)
],
tfcompile_flags = tfcompile_flags,
xla_flags = xla_flags,
),
tf_library(
name = "test_graph_tfadd_with_ckpt" + suffix,
testonly = 1,
config = "test_graph_tfadd_with_ckpt.config.pbtxt",
cpp_class = "AddWithCkptComp",
freeze_checkpoint = "test_graph_tfadd_with_ckpt.ckpt",
graph = "test_graph_tfadd_with_ckpt.pb",
mlir_components = mlir_component,
tags = [
"manual",
"no_mac", # TODO(b/228273415)
],
tfcompile_flags = tfcompile_flags,
xla_flags = xla_flags,
),
tf_library(
name = "test_graph_tfadd_with_ckpt_saver" + suffix,
testonly = 1,
config = "test_graph_tfadd_with_ckpt.config.pbtxt",
cpp_class = "AddWithCkptSaverComp",
freeze_checkpoint = "test_graph_tfadd_with_ckpt_saver.ckpt",
freeze_saver = "test_graph_tfadd_with_ckpt_saver.saver",
graph = "test_graph_tfadd_with_ckpt_saver.pb",
mlir_components = mlir_component,
tags = [
"manual",
"no_mac", # TODO(b/228273415)
],
tfcompile_flags = tfcompile_flags,
xla_flags = xla_flags,
),
tf_library(
name = "test_graph_tfassert_eq" + suffix,
testonly = 1,
config = "test_graph_tfassert_eq.config.pbtxt",
cpp_class = "AssertComp",
graph = "test_graph_tfassert_eq.pb",
mlir_components = mlir_component,
tags = [
"manual",
"no_mac", # TODO(b/228273415)
],
tfcompile_flags = tfcompile_flags,
xla_flags = xla_flags,
),
tf_library(
name = "test_graph_tfcond" + suffix,
testonly = 1,
config = "test_graph_tfcond.config.pbtxt",
cpp_class = "CondComp",
graph = "test_graph_tfcond.pb",
mlir_components = mlir_component,
tags = [
"manual",
"no_mac", # TODO(b/228273415)
],
tfcompile_flags = tfcompile_flags,
xla_flags = xla_flags,
),
tf_library(
name = "test_graph_tffunction" + suffix,
testonly = 1,
config = "test_graph_tffunction.config.pbtxt",
cpp_class = "FunctionComp",
graph = "test_graph_tffunction.pb",
mlir_components = mlir_component,
tags = [
"manual",
"no_mac", # TODO(b/228273415)
],
tfcompile_flags = tfcompile_flags,
xla_flags = xla_flags,
),
tf_library(
name = "test_graph_tfgather" + suffix,
testonly = 1,
config = "test_graph_tfgather.config.pbtxt",
cpp_class = "GatherComp",
graph = "test_graph_tfgather.pb",
mlir_components = mlir_component,
tags = [
"manual",
"no_mac", # TODO(b/228273415)
],
tfcompile_flags = tfcompile_flags,
xla_flags = xla_flags,
),
tf_library(
name = "test_graph_tfmatmul" + suffix,
testonly = 1,
config = "test_graph_tfmatmul.config.pbtxt",
cpp_class = "foo::bar::MatMulComp",
graph = "test_graph_tfmatmul.pb",
mlir_components = mlir_component,
tags = [
"manual",
"no_mac", # TODO(b/228273415)
],
tfcompile_flags = tfcompile_flags,
xla_flags = xla_flags,
),
tf_library(
name = "test_graph_tfmatmulandadd" + suffix,
testonly = 1,
config = "test_graph_tfmatmulandadd.config.pbtxt",
cpp_class = "::foo::bar::MatMulAndAddComp",
graph = "test_graph_tfmatmulandadd.pb",
mlir_components = mlir_component,
tags = [
"manual",
"no_mac", # TODO(b/228273415)
],
tfcompile_flags = "--gen_name_to_index --gen_program_shape " + tfcompile_flags,
xla_flags = xla_flags,
),
tf_library(
name = "test_graph_tfmatmul_with_constant" + suffix,
testonly = 1,
config = "test_graph_tfmatmul_with_constant.config.pbtxt",
cpp_class = "foo::bar::MatMulWithConstantComp",
graph = "test_graph_tfmatmul_with_constant.pb",
include_standard_runtime_deps = True,
mlir_components = mlir_component,
tags = [
"manual",
"no_mac", # TODO(b/228273415)
],
tfcompile_flags = tfcompile_flags,
xla_flags = xla_flags,
),
tf_library(
name = "test_graph_tfsplits" + suffix,
testonly = 1,
config = "test_graph_tfsplits.config.pbtxt",
cpp_class = "SplitsComp",
graph = "test_graph_tfsplits.pb",
mlir_components = mlir_component,
tags = [
"manual",
"no_mac", # TODO(b/228273415)
],
tfcompile_flags = tfcompile_flags,
xla_flags = xla_flags,
),
tf_library(
name = "test_graph_tftop_k" + suffix,
testonly = 1,
config = "test_graph_tftop_k.config.pbtxt",
cpp_class = "TopKComp",
graph = "test_graph_tftop_k.pb",
mlir_components = "None",
tags = [
"manual",
"no_mac", # TODO(b/228273415)
],
tfcompile_flags = tfcompile_flags,
xla_flags = xla_flags,
),
tf_library(
name = "test_graph_tfvariable" + suffix,
testonly = 1,
config = "test_graph_tfvariable.config.pbtxt",
cpp_class = "VariableComp",
graph = "test_graph_tfvariable.pb",
mlir_components = mlir_component,
tags = [
"manual",
"no_mac", # TODO(b/228273415)
],
tfcompile_flags = tfcompile_flags,
xla_flags = xla_flags,
),
tf_library(
name = "test_graph_tfvariable_readonly" + suffix,
testonly = 1,
config = "test_graph_tfvariable_readonly.config.pbtxt",
cpp_class = "VariableReadonlyComp",
graph = "test_graph_tfvariable_readonly.pb",
mlir_components = mlir_component,
tags = [
"manual",
"no_mac", # TODO(b/228273415)
],
tfcompile_flags = tfcompile_flags,
xla_flags = xla_flags,
),
tf_library(
name = "test_graph_tfvariable_sequential_updates" + suffix,
testonly = 1,
config = "test_graph_tfvariable_sequential_updates.config.pbtxt",
cpp_class = "VariableSequentialUpdatesComp",
graph = "test_graph_tfvariable_sequential_updates.pb",
mlir_components = mlir_component,
tags = [
"manual",
"no_mac", # TODO(b/228273415)
],
tfcompile_flags = tfcompile_flags,
xla_flags = xla_flags,
),
tf_library(
name = "test_graph_tfrandom_uniform" + suffix,
testonly = 1,
config = "test_graph_tfrandom_uniform.config.pbtxt",
cpp_class = "RandomUniformComp",
graph = "test_graph_tfrandom_uniform.pb",
include_standard_runtime_deps = True,
mlir_components = mlir_component,
tags = [
"manual",
"no_mac", # TODO(b/228273415)
],
tfcompile_flags = tfcompile_flags,
xla_flags = xla_flags,
),
tf_library(
name = "test_graph_tfscatter" + suffix,
testonly = 1,
config = "test_graph_tfscatter.config.pbtxt",
cpp_class = "ScatterComp",
graph = "test_graph_tfscatter.pb",
mlir_components = mlir_component,
tags = [
"manual",
"no_mac", # TODO(b/228273415)
],
tfcompile_flags = tfcompile_flags,
xla_flags = xla_flags,
),
]
for suffix, mlir_component, xla_flags, tfcompile_flags in tfcompile_test_dep_configs
]
tfcompile_bench_tfmatmul_mkn = [
# Intentionally empty to avoid running unnecessary tests.
# Add here your desired (M, K, N) parameters, e.g.
# (1, 1, 256),
# (1, 2, 256),
]
tfcompile_bench_tfmatmul = [
(
"bench_graph_tfmatmul_%sx%sx%s" % (m, k, n),
"bench_graph_tfmatmul_%sx%sx%s.config.pbtxt" % (m, k, n),
"bench_graph_tfmatmul.template.pbtxt",
"-e \"s|<M>|%s|g\" -e \"s|<K>|%s|g\" -e \"s|<N>|%s|g\"" % (m, k, n),
)
for (m, k, n) in tfcompile_bench_tfmatmul_mkn
]
test_suite(
name = "all_tfmatmul_benchmarks",
tags = ["manual"],
tests = [
(":%s_test" % bench_name)
for (bench_name, _, _, _) in tfcompile_bench_tfmatmul
],
visibility = ["//visibility:public"],
)
[[
genrule(
name = "gen_" + config_file,
testonly = 1,
srcs = [template_file],
outs = [config_file],
cmd = ("sed " + sed_replace + " " +
"$(location " + template_file + ") " +
"> $(OUTS)"),
tags = ["manual"],
),
tf_library(
name = bench_name,
testonly = 1,
config = config_file,
cpp_class = "foo::bar::MatMulComp",
graph = "test_graph_tfmatmul.pb",
tags = [
"manual",
"no_mac", # TODO(b/228273415)
],
),
] for (bench_name, config_file, template_file, sed_replace) in tfcompile_bench_tfmatmul]
tf_cc_test(
name = "tfcompile_test_thunks",
srcs = ["tfcompile_test.cc"],
extra_copts = ["-DENABLE_XLA_THUNK_TEST"],
tags = [
"manual",
"no_mac", # TODO(b/228273415)
"not_run:arm",
],
deps = [
":test_graph_tfadd_thunks",
":test_graph_tfadd_with_ckpt_saver_thunks",
":test_graph_tfadd_with_ckpt_thunks",
":test_graph_tfassert_eq_thunks",
":test_graph_tfcond_thunks",
":test_graph_tffunction_thunks",
":test_graph_tfgather_thunks",
":test_graph_tfmatmul_thunks",
":test_graph_tfmatmul_with_constant_thunks",
":test_graph_tfmatmulandadd_thunks",
":test_graph_tfrandom_uniform_thunks",
":test_graph_tfscatter_thunks",
":test_graph_tfsplits_thunks",
":test_graph_tftop_k_thunks",
":test_graph_tfvariable_readonly_thunks",
":test_graph_tfvariable_sequential_updates_thunks",
":test_graph_tfvariable_thunks",
"//tensorflow/compiler/tf2xla:xla_compiled_cpu_function",
"//tensorflow/core:portable_gif_internal",
"//tensorflow/core:test",
"//tensorflow/core:test_main",
"@eigen_archive//:eigen3",
"@xla//xla:shape_util",
"@xla//xla/hlo/testlib:test",
"@xla//xla/tsl/platform:env",
"@xla//xla/tsl/platform:statusor",
],
)
tf_cc_test(
name = "tfcompile_test_nanort",
srcs = ["tfcompile_test.cc"],
extra_copts = ["-DENABLE_XLA_NANORT_TEST"],
tags = [
"manual",
"no_mac", # TODO(b/228273415)
"not_run:arm",
],
deps = [
":test_graph_tfadd_nanort",
":test_graph_tfadd_with_ckpt_nanort",
":test_graph_tfadd_with_ckpt_saver_nanort",
":test_graph_tfassert_eq_nanort",
":test_graph_tfcond_nanort",
":test_graph_tffunction_nanort",
":test_graph_tfgather_nanort",
":test_graph_tfmatmul_nanort",
":test_graph_tfmatmul_with_constant_nanort",
":test_graph_tfmatmulandadd_nanort",
":test_graph_tfrandom_uniform_nanort",
":test_graph_tfscatter_nanort",
":test_graph_tfsplits_nanort",
":test_graph_tftop_k_nanort",
":test_graph_tfvariable_nanort",
":test_graph_tfvariable_readonly_nanort",
":test_graph_tfvariable_sequential_updates_nanort",
"//tensorflow/compiler/tf2xla:xla_compiled_cpu_function",
"//tensorflow/core:portable_gif_internal",
"//tensorflow/core:test",
"//tensorflow/core:test_main",
"@eigen_archive//:eigen3",
"@xla//xla:shape_util",
"@xla//xla/hlo/testlib:test",
"@xla//xla/tsl/platform:env",
"@xla//xla/tsl/platform:statusor",
],
)
tensorflow/compiler/aot/tests/bench_graph_tfmatmul.template.pbtxt
+18
View File
@@ -0,0 +1,18 @@
# Text form of tensorflow.tf2xla.Config proto.
feed {
id { node_name: "x_hold" }
shape {
dim { size: <M> }
dim { size: <K> }
}
}
feed {
id { node_name: "y_hold" }
shape {
dim { size: <K> }
dim { size: <N> }
}
}
fetch {
id { node_name: "x_y_prod" }
}
tensorflow/compiler/aot/tests/make_test_graphs.py
+259
View File
@@ -0,0 +1,259 @@
# Copyright 2017 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 tensorflow graphs for testing tfcompile."""
import argparse
import os
import sys
from absl import app
import six
from six.moves import range
from tensorflow.core.protobuf import saver_pb2
from tensorflow.python.client import session
from tensorflow.python.framework import constant_op
from tensorflow.python.framework import dtypes
from tensorflow.python.framework import function
from tensorflow.python.framework import ops
from tensorflow.python.ops import array_ops
from tensorflow.python.ops import array_ops_stack # pylint: disable=g-direct-tensorflow-import
from tensorflow.python.ops import cond
from tensorflow.python.ops import control_flow_assert
from tensorflow.python.ops import control_flow_ops
from tensorflow.python.ops import control_flow_util
from tensorflow.python.ops import math_ops
from tensorflow.python.ops import nn_ops
from tensorflow.python.ops import random_ops
from tensorflow.python.ops import variable_v1
from tensorflow.python.ops import variables
from tensorflow.python.training import saver as saver_lib
FLAGS = None
def tfadd(_):
x = constant_op.constant([1], name='x_const')
y = constant_op.constant([2], name='y_const')
math_ops.add(x, y, name='x_y_sum')
def tfadd_with_ckpt(out_dir):
x = array_ops.placeholder(dtypes.int32, name='x_hold')
y = variable_v1.VariableV1(constant_op.constant([0]), name='y_saved')
math_ops.add(x, y, name='x_y_sum')
init_op = variables.global_variables_initializer()
saver = saver_lib.Saver(write_version=saver_pb2.SaverDef.V1)
with session.Session() as sess:
sess.run(init_op)
sess.run(y.assign(y + 42))
# Without the checkpoint, the variable won't be set to 42.
ckpt = os.path.join(out_dir, 'test_graph_tfadd_with_ckpt.ckpt')
saver.save(sess, ckpt)
def tfadd_with_ckpt_saver(out_dir):
x = array_ops.placeholder(dtypes.int32, name='x_hold')
y = variable_v1.VariableV1(constant_op.constant([0]), name='y_saved')
math_ops.add(x, y, name='x_y_sum')
init_op = variables.global_variables_initializer()
saver = saver_lib.Saver(name='abcprefix', write_version=saver_pb2.SaverDef.V1)
with session.Session() as sess:
sess.run(init_op)
sess.run(y.assign(y + 42))
# Without the checkpoint, the variable won't be set to 42.
ckpt_file = os.path.join(out_dir, 'test_graph_tfadd_with_ckpt_saver.ckpt')
saver.save(sess, ckpt_file)
# Without the SaverDef, the restore op won't be named correctly.
saver_file = os.path.join(out_dir, 'test_graph_tfadd_with_ckpt_saver.saver')
with open(saver_file, 'wb') as f:
f.write(six.ensure_binary(saver.as_saver_def().SerializeToString()))
def tfassert_eq(_):
x = array_ops.placeholder(dtypes.int32, name='x_hold')
y = array_ops.placeholder(dtypes.int32, name='y_hold')
control_flow_assert.Assert(
math_ops.equal(x, y), ['Expected x == y.'], name='assert_eq'
)
math_ops.add(x, math_ops.negative(y), name='x_y_diff')
def tfcond(_):
p = array_ops.placeholder(dtypes.bool, name='p_hold')
x = array_ops.placeholder(dtypes.int32, name='x_hold')
y = array_ops.placeholder(dtypes.int32, name='y_hold')
z = cond.cond(p, lambda: x, lambda: y)
array_ops.identity(z, name='result')
def tfgather(_):
params = array_ops.placeholder(dtypes.float32, name='params')
indices = array_ops.placeholder(dtypes.int32, name='indices')
array_ops.gather(params, indices, name='gather_output')
def tfmatmul(_):
x = array_ops.placeholder(dtypes.float32, name='x_hold')
y = array_ops.placeholder(dtypes.float32, name='y_hold')
math_ops.matmul(x, y, name='x_y_prod')
def tfmatmul_with_constant(_):
x = array_ops.placeholder(dtypes.float32, name='x_hold')
constant_y = array_ops.ones(
shape=[1024, 256], dtype=dtypes.float32, name='constant_y_hold'
)
math_ops.matmul(x, constant_y, name='x_constant_y_prod')
def tfmatmulandadd(_):
# This tests multiple outputs.
x = array_ops.placeholder(dtypes.float32, name='x_hold')
y = array_ops.placeholder(dtypes.float32, name='y_hold')
math_ops.matmul(x, y, name='x_y_prod')
math_ops.add(x, y, name='x_y_sum')
def tffunction(_):
@function.Defun(dtypes.int32, dtypes.int32)
def test_func(a, b):
return a + b
x = constant_op.constant([1], name='x_const')
y = constant_op.constant([2], name='y_const')
test_func(x, y, name='func_call') # pylint: disable=unexpected-keyword-arg
def tfsplits(_):
"""A more complex graph, including splits."""
x = array_ops.placeholder(dtypes.float32, shape=[2, 2], name='x')
y = array_ops.placeholder(dtypes.float32, shape=[2, 2], name='y')
for _ in range(3):
x0, x1 = array_ops.split(x, 2, 0)
y0, y1 = array_ops.split(y, 2, 0)
x0 += 1
y0 += 1
z = math_ops.matmul(x, y, name='x_y_prod')
a = array_ops.concat([x0, y1], axis=0, name='concat_x0_y1')
b = array_ops.concat([y0, x1], axis=0, name='concat_y0_x1')
x = math_ops.matmul(a, b, name='a_b')
y = math_ops.add(x, z)
array_ops.identity(y, name='result')
def tftop_k(_):
x = array_ops.placeholder(dtypes.int32, shape=[5], name='x')
output = nn_ops.top_k(x, 2, name='values')
array_ops.identity(output[1], name='indices')
def tfvariable_readonly(_):
x = variables.Variable(1000.0, name='x')
unused_y = variables.Variable(1000.0, name='y')
old_x = x.value()
with ops.control_dependencies([old_x]):
new_value = math_ops.add(old_x, 42.0)
array_ops.identity(new_value, name='result')
# TODO(b/147908587): Change x and the two constants back to have a scalar shape
# when the bug is fixed.
def tfvariable(_):
x = variables.Variable([1000.0], name='x', shape=[1])
old_x = x.value()
with ops.control_dependencies([old_x]):
new_x = x.assign_add([42.0])
array_ops_stack.stack([old_x, new_x], name='result')
def tfvariable_sequential_updates(_):
x = variables.Variable(1.0, name='x')
y = variables.Variable(1.0, name='y')
updates = control_flow_ops.no_op()
for _ in range(3):
with ops.control_dependencies([updates]):
x_val = x.read_value() + y
updates = x.assign_sub(0.1 * x_val)
array_ops.identity(updates, name='result')
def tfrandom_uniform(_):
x = random_ops.random_uniform(shape=[1], minval=0.0, maxval=5.0)
array_ops.identity(x, name='result')
def tfscatter(_):
# NOTE(basioli): We run two scatter operations to make sure fusion kernels
# get named correctly in XLA:CPU.
indices0 = array_ops.placeholder(dtypes.int32, name='indices_0')
updates0 = array_ops.placeholder(dtypes.float32, name='updates_0')
indices1 = array_ops.placeholder(dtypes.int32, name='indices_1')
updates1 = array_ops.placeholder(dtypes.float32, name='updates_1')
shape = constant_op.constant([8], name='shape')
array_ops.scatter_nd(indices0, updates0, shape, name='result_0')
array_ops.scatter_nd(indices1, updates1, shape, name='result_1')
def write_graph(build_graph, out_dir):
"""Build a graph using build_graph and write it out."""
g = ops.Graph()
with g.as_default():
build_graph(out_dir)
filename = os.path.join(out_dir, 'test_graph_%s.pb' % build_graph.__name__)
with open(filename, 'wb') as f:
f.write(
six.ensure_binary(
g.as_graph_def().SerializeToString(deterministic=True)
)
)
def main(_):
control_flow_util.enable_control_flow_v2()
write_graph(tfadd, FLAGS.out_dir)
write_graph(tfadd_with_ckpt, FLAGS.out_dir)
write_graph(tfadd_with_ckpt_saver, FLAGS.out_dir)
write_graph(tfassert_eq, FLAGS.out_dir)
write_graph(tfcond, FLAGS.out_dir)
write_graph(tffunction, FLAGS.out_dir)
write_graph(tfgather, FLAGS.out_dir)
write_graph(tfmatmul, FLAGS.out_dir)
write_graph(tfmatmul_with_constant, FLAGS.out_dir)
write_graph(tfmatmulandadd, FLAGS.out_dir)
write_graph(tfsplits, FLAGS.out_dir)
write_graph(tftop_k, FLAGS.out_dir)
write_graph(tfvariable, FLAGS.out_dir)
write_graph(tfvariable_readonly, FLAGS.out_dir)
write_graph(tfvariable_sequential_updates, FLAGS.out_dir)
write_graph(tfrandom_uniform, FLAGS.out_dir)
write_graph(tfscatter, FLAGS.out_dir)
if __name__ == '__main__':
parser = argparse.ArgumentParser()
parser.register('type', 'bool', lambda v: v.lower() == 'true')
parser.add_argument(
'--out_dir',
type=str,
default='',
help='Output directory for graphs, checkpoints and savers.',
)
FLAGS, unparsed = parser.parse_known_args()
app.run(main=main, argv=[sys.argv[0]] + unparsed)
tensorflow/compiler/aot/tests/test_error_message.lit.pbtxt
+69
View File
@@ -0,0 +1,69 @@
# RUN: not tfcompile --graph=%s --config=%s.config.pbtxt --mlir_components=Bridge --debug_info=%s.debug.pbtxt 2>&1 | FileCheck %s
# RUN: not tfcompile --graph=%s --config=%s.config.pbtxt --mlir_components=None 2>&1 | FileCheck -check-prefix=OLD %s
# Checks the error message produced by tfcompile with mlir_component
# Checks that source debug information is used in the output error message and
# the node x_y_sum = Add
# CHECK: INVALID ARGUMENTS: Dimensions must be equal, but are 2 and 3 for 'x_y_sum = Add[T=DT_INT32](aot_feed_0/x, aot_feed_0/y)'
# CHECK: math_ops.add(x, y, name='x_y_sum')
# CHECK: build_graph(out_dir)
# Checks the error message produced by tfcompile without mlir_component
# OLD: INVALID ARGUMENTS: Incompatible shapes: [2] vs. [3]
# OLD: x_y_sum
node: {
name: "x"
op: "Placeholder"
attr: {
key: "shape"
value: {
shape: {
dim: {
size: -1
}
}
}
}
attr: {
key: "dtype"
value: {
type: DT_INT32
}
}
}
node: {
name: "y"
op: "Placeholder"
attr: {
key: "shape"
value: {
shape: {
dim: {
size: -1
}
}
}
}
attr: {
key: "dtype"
value: {
type: DT_INT32
}
}
}
node: {
name: "x_y_sum"
op: "Add"
input: "x"
input: "y"
attr: {
key: "T"
value: {
type: DT_INT32
}
}
}
versions: {
producer: 321
}
tensorflow/compiler/aot/tests/test_error_message.lit.pbtxt.config.pbtxt
+16
View File
@@ -0,0 +1,16 @@
# Text form of tensorflow.tf2xla.Config proto.
feed {
id { node_name: "x" }
shape {
dim { size: 2 }
}
}
feed {
id { node_name: "y" }
shape {
dim { size: 3 }
}
}
fetch {
id { node_name: "x_y_sum" }
}
tensorflow/compiler/aot/tests/test_error_message.lit.pbtxt.debug.pbtxt
+32
View File
@@ -0,0 +1,32 @@
files: "org_tensorflow/tensorflow/compiler/aot/tests/test_error_message.lit.pbtxt.fake_py.debug"
traces: {
key: "x@"
value: {
file_line_cols: {
line: 1
col: 1
}
}
}
traces: {
key: "x_y_sum@"
value: {
file_line_cols: {
line: 3
col: 1
}
file_line_cols: {
line: 4
col: 1
}
}
}
traces: {
key: "y@"
value: {
file_line_cols: {
line: 2
col: 1
}
}
}
tensorflow/compiler/aot/tests/test_error_message.lit.pbtxt.fake_py.debug
+4
View File
@@ -0,0 +1,4 @@
x = value
y = value
math_ops.add(x, y, name='x_y_sum')
build_graph(out_dir)
tensorflow/compiler/aot/tests/test_graph_tfadd.config.pbtxt
+16
View File
@@ -0,0 +1,16 @@
# Text form of tensorflow.tf2xla.Config proto.
feed {
id { node_name: "x_const" }
shape {
dim { size: 1 }
}
}
feed {
id { node_name: "y_const" }
shape {
dim { size: 1 }
}
}
fetch {
id { node_name: "x_y_sum" }
}
tensorflow/compiler/aot/tests/test_graph_tfadd_with_ckpt.config.pbtxt
+10
View File
@@ -0,0 +1,10 @@
# Text form of tensorflow.tf2xla.Config proto.
feed {
id { node_name: "x_hold" }
shape {
dim { size: 1 }
}
}
fetch {
id { node_name: "x_y_sum" }
}
tensorflow/compiler/aot/tests/test_graph_tfassert_eq.config.pbtxt
+16
View File
@@ -0,0 +1,16 @@
# Text form of tensorflow.tf2xla.Config proto.
feed {
id { node_name: "x_hold" }
shape {
dim { size: 1 }
}
}
feed {
id { node_name: "y_hold" }
shape {
dim { size: 1 }
}
}
fetch {
id { node_name: "x_y_diff" }
}
tensorflow/compiler/aot/tests/test_graph_tfcond.config.pbtxt
+20
View File
@@ -0,0 +1,20 @@
# Text form of tensorflow.tf2xla.Config proto.
feed {
id { node_name: "p_hold" }
shape {}
}
feed {
id { node_name: "x_hold" }
shape {
dim { size: 1 }
}
}
feed {
id { node_name: "y_hold" }
shape {
dim { size: 1 }
}
}
fetch {
id { node_name: "result" }
}
tensorflow/compiler/aot/tests/test_graph_tffunction.config.pbtxt
+16
View File
@@ -0,0 +1,16 @@
# Text form of tensorflow.tf2xla.Config proto.
feed {
id { node_name: "x_const" }
shape {
dim { size: 1 }
}
}
feed {
id { node_name: "y_const" }
shape {
dim { size: 1 }
}
}
fetch {
id { node_name: "func_call" }
}
tensorflow/compiler/aot/tests/test_graph_tfgather.config.pbtxt
+16
View File
@@ -0,0 +1,16 @@
# Text form of tensorflow.tf2xla.Config proto.
feed {
id { node_name: "params" }
shape {
dim { size: 4 }
}
}
feed {
id { node_name: "indices" }
shape {
dim { size: 2 }
}
}
fetch {
id { node_name: "gather_output" }
}
tensorflow/compiler/aot/tests/test_graph_tfmatmul.config.pbtxt
+18
View File
@@ -0,0 +1,18 @@
# Text form of tensorflow.tf2xla.Config proto.
feed {
id { node_name: "x_hold" }
shape {
dim { size: 2 }
dim { size: 3 }
}
}
feed {
id { node_name: "y_hold" }
shape {
dim { size: 3 }
dim { size: 2 }
}
}
fetch {
id { node_name: "x_y_prod" }
}
tensorflow/compiler/aot/tests/test_graph_tfmatmul_with_constant.config.pbtxt
+11
View File
@@ -0,0 +1,11 @@
# Text form of tensorflow.tf2xla.Config proto.
feed {
id { node_name: "x_hold" }
shape {
dim { size: 512 }
dim { size: 1024 }
}
}
fetch {
id { node_name: "x_constant_y_prod" }
}
tensorflow/compiler/aot/tests/test_graph_tfmatmulandadd.config.pbtxt
+25
View File
@@ -0,0 +1,25 @@
# Text form of tensorflow.tf2xla.Config proto.
feed {
id { node_name: "x_hold" }
shape {
dim { size: 2 }
dim { size: 2 }
}
name: "x"
}
feed {
id { node_name: "y_hold" }
shape {
dim { size: 2 }
dim { size: 2 }
}
name: "y"
}
fetch {
id { node_name: "x_y_prod" }
name: "x_y_prod"
}
fetch {
id { node_name: "x_y_sum" }
name: "x_y_sum"
}
tensorflow/compiler/aot/tests/test_graph_tfrandom_uniform.config.pbtxt
+4
View File
@@ -0,0 +1,4 @@
# Text form of tensorflow.tf2xla.Config proto.
fetch {
id { node_name: "result" }
}
tensorflow/compiler/aot/tests/test_graph_tfscatter.config.pbtxt
+33
View File
@@ -0,0 +1,33 @@
# Text form of tensorflow.tf2xla.Config proto.
feed {
id { node_name: "indices_0" }
shape {
dim { size: 4 }
dim { size: 1 }
}
}
feed {
id { node_name: "updates_0" }
shape {
dim { size: 4 }
}
}
feed {
id { node_name: "indices_1" }
shape {
dim { size: 4 }
dim { size: 1 }
}
}
feed {
id { node_name: "updates_1" }
shape {
dim { size: 4 }
}
}
fetch {
id { node_name: "result_0" }
}
fetch {
id { node_name: "result_1" }
}
tensorflow/compiler/aot/tests/test_graph_tfsplits.config.pbtxt
+18
View File
@@ -0,0 +1,18 @@
# Text form of tensorflow.tf2xla.Config proto.
feed {
id { node_name: "x" }
shape {
dim { size: 2 }
dim { size: 2 }
}
}
feed {
id { node_name: "y" }
shape {
dim { size: 2 }
dim { size: 2 }
}
}
fetch {
id { node_name: "result" }
}
tensorflow/compiler/aot/tests/test_graph_tftop_k.config.pbtxt
+13
View File
@@ -0,0 +1,13 @@
# Text form of tensorflow.tf2xla.Config proto.
feed {
id { node_name: "x" }
shape {
dim { size: 5 }
}
}
fetch {
id { node_name: "values" }
}
fetch {
id { node_name: "indices" }
}
tensorflow/compiler/aot/tests/test_graph_tfvariable.config.pbtxt
+12
View File
@@ -0,0 +1,12 @@
# Text form of tensorflow.tf2xla.Config proto.
fetch {
id { node_name: "result" }
}
variable {
node_name: "x"
shape {
dim { size: 1 }
}
type: DT_FLOAT
}
tensorflow/compiler/aot/tests/test_graph_tfvariable_readonly.config.pbtxt
+20
View File
@@ -0,0 +1,20 @@
# Text form of tensorflow.tf2xla.Config proto.
fetch {
id { node_name: "result" }
}
variable {
node_name: "x"
shape {
}
type: DT_FLOAT
readonly: true
}
variable {
node_name: "y"
shape {
}
type: DT_FLOAT
readonly: true
}
tensorflow/compiler/aot/tests/test_graph_tfvariable_sequential_updates.config.pbtxt
+15
View File
@@ -0,0 +1,15 @@
# Text form of tensorflow.tf2xla.Config proto.
fetch {
id { node_name: "result" }
}
variable {
node_name: "x"
type: DT_FLOAT
}
variable {
node_name: "y"
type: DT_FLOAT
readonly: true
}
tensorflow/compiler/aot/tests/tfcompile_test.cc
+760
View File
@@ -0,0 +1,760 @@
/* Copyright 2017 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 <cstdint>
#include <cstring>
#include <vector>
#if defined(ENABLE_XLA_THUNK_TEST)
#include "tensorflow/compiler/aot/tests/test_graph_tfadd_thunks.h"
#include "tensorflow/compiler/aot/tests/test_graph_tfadd_with_ckpt_saver_thunks.h"
#include "tensorflow/compiler/aot/tests/test_graph_tfadd_with_ckpt_thunks.h"
#include "tensorflow/compiler/aot/tests/test_graph_tfassert_eq_thunks.h"
#include "tensorflow/compiler/aot/tests/test_graph_tfcond_thunks.h"
#include "tensorflow/compiler/aot/tests/test_graph_tffunction_thunks.h"
#include "tensorflow/compiler/aot/tests/test_graph_tfgather_thunks.h"
#include "tensorflow/compiler/aot/tests/test_graph_tfmatmul_thunks.h"
#include "tensorflow/compiler/aot/tests/test_graph_tfmatmul_with_constant_thunks.h"
#include "tensorflow/compiler/aot/tests/test_graph_tfmatmulandadd_thunks.h"
#include "tensorflow/compiler/aot/tests/test_graph_tfrandom_uniform_thunks.h"
#include "tensorflow/compiler/aot/tests/test_graph_tfscatter_thunks.h"
#include "tensorflow/compiler/aot/tests/test_graph_tfsplits_thunks.h"
#include "tensorflow/compiler/aot/tests/test_graph_tftop_k_thunks.h"
#include "tensorflow/compiler/aot/tests/test_graph_tfvariable_readonly_thunks.h"
#include "tensorflow/compiler/aot/tests/test_graph_tfvariable_sequential_updates_thunks.h"
#include "tensorflow/compiler/aot/tests/test_graph_tfvariable_thunks.h"
#elif defined(ENABLE_XLA_NANORT_TEST)
#include "tensorflow/compiler/aot/tests/test_graph_tfadd_nanort.h"
#include "tensorflow/compiler/aot/tests/test_graph_tfadd_with_ckpt_nanort.h"
#include "tensorflow/compiler/aot/tests/test_graph_tfadd_with_ckpt_saver_nanort.h"
#include "tensorflow/compiler/aot/tests/test_graph_tfassert_eq_nanort.h"
#include "tensorflow/compiler/aot/tests/test_graph_tfcond_nanort.h"
#include "tensorflow/compiler/aot/tests/test_graph_tffunction_nanort.h"
#include "tensorflow/compiler/aot/tests/test_graph_tfgather_nanort.h"
#include "tensorflow/compiler/aot/tests/test_graph_tfmatmul_nanort.h"
#include "tensorflow/compiler/aot/tests/test_graph_tfmatmul_with_constant_nanort.h"
#include "tensorflow/compiler/aot/tests/test_graph_tfmatmulandadd_nanort.h"
#include "tensorflow/compiler/aot/tests/test_graph_tfrandom_uniform_nanort.h"
#include "tensorflow/compiler/aot/tests/test_graph_tfscatter_nanort.h"
#include "tensorflow/compiler/aot/tests/test_graph_tfsplits_nanort.h"
#include "tensorflow/compiler/aot/tests/test_graph_tftop_k_nanort.h"
#include "tensorflow/compiler/aot/tests/test_graph_tfvariable_nanort.h"
#include "tensorflow/compiler/aot/tests/test_graph_tfvariable_readonly_nanort.h"
#include "tensorflow/compiler/aot/tests/test_graph_tfvariable_sequential_updates_nanort.h"
#endif
#include "tensorflow/compiler/tf2xla/xla_compiled_cpu_function.h"
#include "xla/hlo/testlib/test.h"
#include "xla/shape.h"
#include "xla/shape_util.h"
#include "xla/tsl/platform/statusor.h"
#include "xla/tsl/platform/threadpool.h"
#include "tensorflow/core/platform/test.h"
#include "tensorflow/core/platform/types.h"
#define EIGEN_USE_THREADS
#define EIGEN_USE_CUSTOM_THREAD_POOL
#include "unsupported/Eigen/CXX11/Tensor" // from @eigen_archive
namespace tensorflow {
namespace tfcompile {
namespace {
TEST(TFCompileTest, Add) {
AddComp add;
EXPECT_EQ(add.arg0_data(), add.arg_data(0));
EXPECT_EQ(add.arg1_data(), add.arg_data(1));
add.arg0() = 1;
add.arg1() = 2;
EXPECT_TRUE(add.Run());
EXPECT_EQ(add.error_msg(), "");
EXPECT_EQ(add.result0(), 3);
EXPECT_EQ(add.result0_data()[0], 3);
EXPECT_EQ(add.result0_data(), add.result_data(0));
add.arg0_data()[0] = 123;
add.arg1_data()[0] = 456;
EXPECT_TRUE(add.Run());
EXPECT_EQ(add.error_msg(), "");
EXPECT_EQ(add.result0(), 579);
EXPECT_EQ(add.result0_data()[0], 579);
EXPECT_EQ(add.result0_data(), add.result_data(0));
const AddComp& add_const = add;
EXPECT_EQ(add_const.error_msg(), "");
EXPECT_EQ(add_const.arg0(), 123);
EXPECT_EQ(add_const.arg0_data()[0], 123);
EXPECT_EQ(add_const.arg0_data(), add.arg_data(0));
EXPECT_EQ(add_const.arg1(), 456);
EXPECT_EQ(add_const.arg1_data()[0], 456);
EXPECT_EQ(add_const.arg1_data(), add.arg_data(1));
EXPECT_EQ(add_const.result0(), 579);
EXPECT_EQ(add_const.result0_data()[0], 579);
EXPECT_EQ(add_const.result0_data(), add_const.result_data(0));
}
// Run tests that use set_argN_data separately, to avoid accidentally re-using
// non-existent buffers.
TEST(TFCompileTest, Add_SetArg) {
AddComp add(
XlaCompiledCpuFunction::AllocMode::RESULTS_PROFILES_AND_TEMPS_ONLY);
alignas(32) int32_t arg_x = 10;
alignas(32) int32_t arg_y = 32;
add.set_arg0_data(&arg_x);
add.set_arg1_data(&arg_y);
EXPECT_EQ(add.arg0_data(), add.arg_data(0));
EXPECT_EQ(add.arg1_data(), add.arg_data(1));
EXPECT_TRUE(add.Run());
EXPECT_EQ(add.error_msg(), "");
EXPECT_EQ(add.result0(), 42);
EXPECT_EQ(add.result0_data()[0], 42);
EXPECT_EQ(add.result0_data(), add.result_data(0));
}
TEST(TFCompileTest, AddWithCkpt) {
AddWithCkptComp add;
EXPECT_EQ(add.arg0_data(), add.arg_data(0));
add.arg0() = 1;
EXPECT_TRUE(add.Run());
EXPECT_EQ(add.error_msg(), "");
EXPECT_EQ(add.result0(), 43);
EXPECT_EQ(add.result0_data()[0], 43);
EXPECT_EQ(add.result0_data(), add.result_data(0));
add.arg0_data()[0] = 111;
EXPECT_TRUE(add.Run());
EXPECT_EQ(add.error_msg(), "");
EXPECT_EQ(add.result0(), 153);
EXPECT_EQ(add.result0_data()[0], 153);
EXPECT_EQ(add.result0_data(), add.result_data(0));
const AddWithCkptComp& add_const = add;
EXPECT_EQ(add_const.error_msg(), "");
EXPECT_EQ(add_const.arg0(), 111);
EXPECT_EQ(add_const.arg0_data()[0], 111);
EXPECT_EQ(add_const.arg0_data(), add_const.arg_data(0));
EXPECT_EQ(add_const.result0(), 153);
EXPECT_EQ(add_const.result0_data()[0], 153);
EXPECT_EQ(add_const.result0_data(), add_const.result_data(0));
}
TEST(TFCompileTest, AddWithCkptSaver) {
AddWithCkptSaverComp add;
EXPECT_EQ(add.arg0_data(), add.arg_data(0));
add.arg0() = 1;
EXPECT_TRUE(add.Run());
EXPECT_EQ(add.error_msg(), "");
EXPECT_EQ(add.result0(), 43);
EXPECT_EQ(add.result0_data()[0], 43);
EXPECT_EQ(add.result0_data(), add.result_data(0));
add.arg0_data()[0] = 111;
EXPECT_TRUE(add.Run());
EXPECT_EQ(add.error_msg(), "");
EXPECT_EQ(add.result0(), 153);
EXPECT_EQ(add.result0_data()[0], 153);
EXPECT_EQ(add.result0_data(), add.result_data(0));
const AddWithCkptSaverComp& add_const = add;
EXPECT_EQ(add_const.error_msg(), "");
EXPECT_EQ(add_const.arg0(), 111);
EXPECT_EQ(add_const.arg0_data()[0], 111);
EXPECT_EQ(add_const.arg0_data(), add_const.arg_data(0));
EXPECT_EQ(add_const.result0(), 153);
EXPECT_EQ(add_const.result0_data()[0], 153);
EXPECT_EQ(add_const.result0_data(), add_const.result_data(0));
}
TEST(TFCompileTest, Cond) {
CondComp cond;
EXPECT_EQ(cond.arg0_data(), cond.arg_data(0));
EXPECT_EQ(cond.arg1_data(), cond.arg_data(1));
EXPECT_EQ(cond.arg2_data(), cond.arg_data(2));
cond.arg1() = 10;
cond.arg2() = 20;
{
cond.arg0() = true;
const int32_t expected_result = cond.arg1();
EXPECT_TRUE(cond.Run());
EXPECT_EQ(cond.error_msg(), "");
EXPECT_EQ(cond.result0(), expected_result);
EXPECT_EQ(cond.result0_data()[0], expected_result);
EXPECT_EQ(cond.result0_data(), cond.result_data(0));
}
{
cond.arg0() = false;
const int32_t expected_result = cond.arg2();
EXPECT_TRUE(cond.Run());
EXPECT_EQ(cond.error_msg(), "");
EXPECT_EQ(cond.result0(), expected_result);
EXPECT_EQ(cond.result0_data()[0], expected_result);
EXPECT_EQ(cond.result0_data(), cond.result_data(0));
}
}
TEST(TFCompileTest, Gather) {
GatherComp gather;
EXPECT_EQ(gather.arg0_data(), gather.arg_data(0));
EXPECT_EQ(gather.arg1_data(), gather.arg_data(1));
// Successful gather.
{
const float params[4] = {1, 2, 3, 4};
std::copy(params + 0, params + 4, gather.arg0_data());
const int32_t indices[2] = {1, 3};
std::copy(indices + 0, indices + 2, gather.arg1_data());
EXPECT_TRUE(gather.Run());
EXPECT_EQ(gather.error_msg(), "");
const float results[2] = {2, 4};
for (int i = 0; i < 2; ++i) {
EXPECT_EQ(gather.result0(i), results[i]);
EXPECT_EQ(gather.result0_data()[i], results[i]);
}
EXPECT_EQ(gather.result0_data(), gather.result_data(0));
const GatherComp& gather_const = gather;
EXPECT_EQ(gather_const.error_msg(), "");
for (int i = 0; i < 4; ++i) {
EXPECT_EQ(gather_const.arg0(i), params[i]);
EXPECT_EQ(gather_const.arg0_data()[i], params[i]);
}
EXPECT_EQ(gather_const.arg0_data(), gather_const.arg_data(0));
for (int i = 0; i < 2; ++i) {
EXPECT_EQ(gather_const.arg1(i), indices[i]);
EXPECT_EQ(gather_const.arg1_data()[i], indices[i]);
}
EXPECT_EQ(gather_const.arg1_data(), gather_const.arg_data(1));
for (int i = 0; i < 2; ++i) {
EXPECT_EQ(gather_const.result0(i), results[i]);
EXPECT_EQ(gather_const.result0_data()[i], results[i]);
}
EXPECT_EQ(gather_const.result0_data(), gather.result_data(0));
}
}
TEST(TFCompileTest, MatMul2) {
Eigen::ThreadPool tp(2);
Eigen::ThreadPoolDevice device(&tp, tp.NumThreads());
foo::bar::MatMulComp matmul;
matmul.set_thread_pool(&device);
EXPECT_EQ(matmul.arg0_data(), matmul.arg_data(0));
EXPECT_EQ(matmul.arg1_data(), matmul.arg_data(1));
// Test using the argN() methods.
{
matmul.arg0(0, 0) = 1;
matmul.arg0(0, 1) = 2;
matmul.arg0(0, 2) = 3;
matmul.arg0(1, 0) = 4;
matmul.arg0(1, 1) = 5;
matmul.arg0(1, 2) = 6;
matmul.arg1(0, 0) = 7;
matmul.arg1(0, 1) = 8;
matmul.arg1(1, 0) = 9;
matmul.arg1(1, 1) = 10;
matmul.arg1(2, 0) = 11;
matmul.arg1(2, 1) = 12;
EXPECT_TRUE(matmul.Run());
EXPECT_EQ(matmul.error_msg(), "");
const float results[4] = {58, 64, 139, 154};
for (int i = 0; i < 4; ++i) {
EXPECT_EQ(matmul.result0(i / 2, i % 2), results[i]);
EXPECT_EQ(matmul.result0_data()[i], results[i]);
}
EXPECT_EQ(matmul.result0_data(), matmul.result_data(0));
}
// Test using the argN_data() methods.
{
const float args[12] = {10, 20, 30, 40, 50, 60, 70, 80, 90, 100, 110, 120};
std::copy(args + 0, args + 6, matmul.arg0_data());
std::copy(args + 6, args + 12, matmul.arg1_data());
EXPECT_TRUE(matmul.Run());
EXPECT_EQ(matmul.error_msg(), "");
const float results[4] = {5800, 6400, 13900, 15400};
for (int i = 0; i < 4; ++i) {
EXPECT_EQ(matmul.result0(i / 2, i % 2), results[i]);
EXPECT_EQ(matmul.result0_data()[i], results[i]);
}
EXPECT_EQ(matmul.result0_data(), matmul.result_data(0));
const foo::bar::MatMulComp& matmul_const = matmul;
EXPECT_EQ(matmul_const.error_msg(), "");
for (int i = 0; i < 6; ++i) {
EXPECT_EQ(matmul_const.arg0(i / 3, i % 3), args[i]);
EXPECT_EQ(matmul_const.arg0_data()[i], args[i]);
}
EXPECT_EQ(matmul_const.arg0_data(), matmul.arg_data(0));
for (int i = 0; i < 6; ++i) {
EXPECT_EQ(matmul_const.arg1(i / 2, i % 2), args[i + 6]);
EXPECT_EQ(matmul_const.arg1_data()[i], args[i + 6]);
}
EXPECT_EQ(matmul_const.arg1_data(), matmul.arg_data(1));
for (int i = 0; i < 4; ++i) {
EXPECT_EQ(matmul_const.result0(i / 2, i % 2), results[i]);
EXPECT_EQ(matmul_const.result0_data()[i], results[i]);
}
EXPECT_EQ(matmul_const.result0_data(), matmul.result_data(0));
}
}
// Run tests that use set_argN_data separately, to avoid accidentally re-using
// non-existent buffers.
TEST(TFCompileTest, MatMul2_SetArg) {
Eigen::ThreadPool tp(2);
Eigen::ThreadPoolDevice device(&tp, tp.NumThreads());
foo::bar::MatMulComp matmul(
XlaCompiledCpuFunction::AllocMode::RESULTS_PROFILES_AND_TEMPS_ONLY);
matmul.set_thread_pool(&device);
// Test using the set_argN_data() methods.
alignas(32) float arg0[2][3] = {{1, 2, 3}, {4, 5, 6}};
alignas(32) float arg1[3][2] = {{7, 8}, {9, 10}, {11, 12}};
matmul.set_arg0_data(&arg0);
matmul.set_arg1_data(&arg1);
EXPECT_EQ(matmul.arg0_data(), matmul.arg_data(0));
EXPECT_EQ(matmul.arg1_data(), matmul.arg_data(1));
EXPECT_TRUE(matmul.Run());
EXPECT_EQ(matmul.error_msg(), "");
const float results[4] = {58, 64, 139, 154};
for (int i = 0; i < 4; ++i) {
EXPECT_EQ(matmul.result0(i / 2, i % 2), results[i]);
EXPECT_EQ(matmul.result0_data()[i], results[i]);
}
EXPECT_EQ(matmul.result0_data(), matmul.result_data(0));
}
TEST(TFCompileTest, MatMulAndAdd1) {
Eigen::ThreadPool tp(1);
Eigen::ThreadPoolDevice device(&tp, tp.NumThreads());
::foo::bar::MatMulAndAddComp muladd;
muladd.set_thread_pool(&device);
EXPECT_EQ(muladd.arg0_data(), muladd.arg_data(0));
EXPECT_EQ(muladd.arg1_data(), muladd.arg_data(1));
// Test methods with positional args and results.
{
const float args[8] = {1, 2, 3, 4, 5, 6, 7, 8};
std::copy(args + 0, args + 4, muladd.arg0_data());
std::copy(args + 4, args + 8, muladd.arg1_data());
EXPECT_TRUE(muladd.Run());
EXPECT_EQ(muladd.error_msg(), "");
const float results0[4] = {19, 22, 43, 50};
const float results1[4] = {6, 8, 10, 12};
for (int i = 0; i < 4; ++i) {
EXPECT_EQ(muladd.result0(i / 2, i % 2), results0[i]);
EXPECT_EQ(muladd.result0_data()[i], results0[i]);
EXPECT_EQ(muladd.result1(i / 2, i % 2), results1[i]);
EXPECT_EQ(muladd.result1_data()[i], results1[i]);
}
EXPECT_EQ(muladd.result0_data(), muladd.result_data(0));
EXPECT_EQ(muladd.result1_data(), muladd.result_data(1));
const ::foo::bar::MatMulAndAddComp& muladd_const = muladd;
EXPECT_EQ(muladd_const.error_msg(), "");
for (int i = 0; i < 4; ++i) {
EXPECT_EQ(muladd_const.arg0(i / 2, i % 2), args[i]);
EXPECT_EQ(muladd_const.arg0_data()[i], args[i]);
}
EXPECT_EQ(muladd_const.arg0_data(), muladd.arg_data(0));
for (int i = 0; i < 4; ++i) {
EXPECT_EQ(muladd_const.arg1(i / 2, i % 2), args[i + 4]);
EXPECT_EQ(muladd_const.arg1_data()[i], args[i + 4]);
}
EXPECT_EQ(muladd_const.arg1_data(), muladd.arg_data(1));
for (int i = 0; i < 4; ++i) {
EXPECT_EQ(muladd_const.result0(i / 2, i % 2), results0[i]);
EXPECT_EQ(muladd_const.result0_data()[i], results0[i]);
EXPECT_EQ(muladd_const.result1(i / 2, i % 2), results1[i]);
EXPECT_EQ(muladd_const.result1_data()[i], results1[i]);
}
EXPECT_EQ(muladd_const.result0_data(), muladd.result_data(0));
EXPECT_EQ(muladd_const.result1_data(), muladd.result_data(1));
}
// Test methods with named args and results.
{
const float args[8] = {10, 20, 30, 40, 50, 60, 70, 80};
std::copy(args + 0, args + 4, muladd.arg_x_data());
std::copy(args + 4, args + 8, muladd.arg_y_data());
EXPECT_TRUE(muladd.Run());
EXPECT_EQ(muladd.error_msg(), "");
const float results0[4] = {1900, 2200, 4300, 5000};
const float results1[4] = {60, 80, 100, 120};
for (int i = 0; i < 4; ++i) {
EXPECT_EQ(muladd.result_x_y_prod(i / 2, i % 2), results0[i]);
EXPECT_EQ(muladd.result_x_y_prod_data()[i], results0[i]);
EXPECT_EQ(muladd.result_x_y_sum(i / 2, i % 2), results1[i]);
EXPECT_EQ(muladd.result_x_y_sum_data()[i], results1[i]);
}
EXPECT_EQ(muladd.result_x_y_prod_data(), muladd.result_data(0));
EXPECT_EQ(muladd.result_x_y_sum_data(), muladd.result_data(1));
// Test const methods.
const ::foo::bar::MatMulAndAddComp& muladd_const = muladd;
EXPECT_EQ(muladd_const.error_msg(), "");
for (int i = 0; i < 4; ++i) {
EXPECT_EQ(muladd_const.arg_x(i / 2, i % 2), args[i]);
EXPECT_EQ(muladd_const.arg_x_data()[i], args[i]);
}
EXPECT_EQ(muladd_const.arg_x_data(), muladd.arg_data(0));
for (int i = 0; i < 4; ++i) {
EXPECT_EQ(muladd_const.arg_y(i / 2, i % 2), args[i + 4]);
EXPECT_EQ(muladd_const.arg_y_data()[i], args[i + 4]);
}
EXPECT_EQ(muladd_const.arg_y_data(), muladd.arg_data(1));
for (int i = 0; i < 4; ++i) {
EXPECT_EQ(muladd_const.result_x_y_prod(i / 2, i % 2), results0[i]);
EXPECT_EQ(muladd_const.result_x_y_prod_data()[i], results0[i]);
EXPECT_EQ(muladd_const.result_x_y_sum(i / 2, i % 2), results1[i]);
EXPECT_EQ(muladd_const.result_x_y_sum_data()[i], results1[i]);
}
EXPECT_EQ(muladd_const.result_x_y_prod_data(), muladd.result_data(0));
EXPECT_EQ(muladd_const.result_x_y_sum_data(), muladd.result_data(1));
}
}
TEST(TFCompileTest, Function) {
// The function is equivalent to an addition
FunctionComp add_fn;
EXPECT_EQ(add_fn.arg0_data(), add_fn.arg_data(0));
EXPECT_EQ(add_fn.arg1_data(), add_fn.arg_data(1));
add_fn.arg0() = 1;
add_fn.arg1() = 2;
EXPECT_TRUE(add_fn.Run());
EXPECT_EQ(add_fn.error_msg(), "");
EXPECT_EQ(add_fn.result0(), 3);
EXPECT_EQ(add_fn.result0_data()[0], 3);
EXPECT_EQ(add_fn.result0_data(), add_fn.result_data(0));
}
TEST(TFCompileTest, Splits) {
Eigen::ThreadPool tp(1);
Eigen::ThreadPoolDevice device(&tp, tp.NumThreads());
SplitsComp fn;
fn.set_thread_pool(&device);
// x = [[1, 2], [3, 4]]
fn.arg0(0, 0) = 1;
fn.arg0(0, 1) = 2;
fn.arg0(1, 0) = 3;
fn.arg0(1, 1) = 4;
// y = [[10, 20], [30, 40]]
fn.arg1(0, 0) = 10;
fn.arg1(0, 1) = 20;
fn.arg1(1, 0) = 30;
fn.arg1(1, 1) = 40;
EXPECT_TRUE(fn.Run());
EXPECT_EQ(fn.error_msg(), "");
const float expected[] = {7.86375557e+10, 1.34274679e+11, 1.92741717e+12,
3.29964742e+12};
EXPECT_NEAR(expected[0], fn.result0(0, 0), 1e4);
EXPECT_NEAR(expected[1], fn.result0(0, 1), 1e4);
EXPECT_NEAR(expected[2], fn.result0(1, 0), 1e4);
EXPECT_NEAR(expected[3], fn.result0(1, 1), 1e4);
}
TEST(TFCompileTest, TopK) {
Eigen::ThreadPool tp(1);
Eigen::ThreadPoolDevice device(&tp, tp.NumThreads());
TopKComp fn;
fn.set_thread_pool(&device);
// x = [4, 1, 4, 4, 3]
fn.arg0(0) = 4;
fn.arg0(1) = 1;
fn.arg0(2) = 4;
fn.arg0(3) = 4;
fn.arg0(4) = 3;
EXPECT_TRUE(fn.Run());
EXPECT_EQ(fn.error_msg(), "");
const int32_t expected_values[] = {4, 4};
const int32_t expected_indices[] = {0, 2};
EXPECT_EQ(expected_values[0], fn.result0(0));
EXPECT_EQ(expected_values[1], fn.result0(1));
EXPECT_EQ(expected_indices[0], fn.result1(0));
EXPECT_EQ(expected_indices[1], fn.result1(1));
}
TEST(TFCompileTest, VariableReadonly) {
Eigen::ThreadPool tp(1);
Eigen::ThreadPoolDevice device(&tp, tp.NumThreads());
VariableReadonlyComp fn;
alignas(32) float x = 23;
fn.set_var_x_data(&x);
fn.set_thread_pool(&device);
EXPECT_TRUE(fn.Run());
EXPECT_EQ(fn.error_msg(), "");
EXPECT_EQ(fn.result0(), 65);
EXPECT_EQ(fn.var_x(), 23);
}
TEST(TFCompileTest, Variable) {
Eigen::ThreadPool tp(1);
Eigen::ThreadPoolDevice device(&tp, tp.NumThreads());
VariableComp fn;
alignas(32) float x = 23;
fn.set_var_x_data(&x);
fn.set_thread_pool(&device);
EXPECT_TRUE(fn.Run());
EXPECT_EQ(fn.error_msg(), "");
EXPECT_EQ(fn.result0(0, 0), 23);
EXPECT_EQ(fn.result0(1, 0), 65);
EXPECT_EQ(fn.var_x(), 65);
EXPECT_EQ(fn.var_x_data(), &x);
EXPECT_EQ(x, 65);
EXPECT_TRUE(fn.Run());
EXPECT_EQ(fn.error_msg(), "");
EXPECT_EQ(fn.result0(0, 0), 65);
EXPECT_EQ(fn.result0(1, 0), 107);
EXPECT_EQ(fn.var_x(), 107);
}
TEST(TFCompileTest, VariableSequentialUpdates) {
Eigen::ThreadPool tp(1);
Eigen::ThreadPoolDevice device(&tp, tp.NumThreads());
// This implements the recursion:
// x[0] = 2.0
// x[n+1] = x[n] - 0.1*(x[n-1] + y)
VariableSequentialUpdatesComp fn;
fn.var_x() = 2;
*const_cast<float*>(fn.var_y_data()) = 1;
fn.set_thread_pool(&device);
// First calculate x[3]
EXPECT_TRUE(fn.Run());
EXPECT_EQ(fn.error_msg(), "");
EXPECT_NEAR(fn.var_x(), 1.187f, 1e-6);
alignas(32) const float y = 1;
fn.set_var_y_data(&y);
// Now const_cast<float*>(fn.var_y_data()) is not longer legal since we've
// set the buffer to point to a constant location.
// Then calculate x[6]
EXPECT_TRUE(fn.Run());
EXPECT_EQ(fn.error_msg(), "");
EXPECT_NEAR(fn.var_x(), 0.594322f, 1e-6);
}
TEST(TFCompileTest, VariableSequentialUpdatesNoAlloc) {
Eigen::ThreadPool tp(1);
Eigen::ThreadPoolDevice device(&tp, tp.NumThreads());
// This implements the recursion:
// x[0] = 2.0
// x[n+1] = x[n] - 0.1*(x[n-1] + 1.0)
VariableSequentialUpdatesComp fn(
XlaCompiledCpuFunction::AllocMode::RESULTS_PROFILES_AND_TEMPS_ONLY);
alignas(32) float x = 2;
alignas(32) float y = 1;
fn.set_var_x_data(&x);
fn.set_var_y_data(&y);
fn.set_thread_pool(&device);
// First calculate x[3]
EXPECT_TRUE(fn.Run());
EXPECT_EQ(fn.error_msg(), "");
EXPECT_NEAR(x, 1.187f, 1e-6);
// Then calculate x[6]
EXPECT_TRUE(fn.Run());
EXPECT_EQ(fn.error_msg(), "");
EXPECT_NEAR(x, 0.594322f, 1e-6);
}
TEST(TFCompileTest, MatMulWithConstants) {
Eigen::ThreadPool tp(2);
Eigen::ThreadPoolDevice device(&tp, tp.NumThreads());
foo::bar::MatMulWithConstantComp matmul;
matmul.set_thread_pool(&device);
EXPECT_EQ(matmul.arg0_data(), matmul.arg_data(0));
// Test using the argN() methods.
{
for (int i = 0; i < 512; ++i) {
for (int j = 0; j < 1024; ++j) {
matmul.arg0(i, j) = 1;
}
}
EXPECT_TRUE(matmul.Run());
EXPECT_EQ(matmul.error_msg(), "");
std::vector<float> results(512 * 256, 1024);
for (int i = 0; i < results.size(); ++i) {
ASSERT_EQ(matmul.result0(i / 512, i % 256), results[i]);
ASSERT_EQ(matmul.result0_data()[i], results[i]);
}
EXPECT_EQ(matmul.result0_data(), matmul.result_data(0));
}
}
TEST(TFCompileTest, RandomUniform) {
RandomUniformComp random_uniform;
EXPECT_TRUE(random_uniform.Run());
EXPECT_EQ(random_uniform.error_msg(), "");
EXPECT_LE(random_uniform.result0(), 5.0);
EXPECT_GE(random_uniform.result0(), 0.0);
}
TEST(TFCompileTest, Scatter) {
ScatterComp scatter;
const std::vector<int32_t> indices0 = {4, 3, 1, 7};
const std::vector<float> updates0 = {9.0, 10.0, 11.0, 12.0};
const std::vector<int32_t> indices1 = {2, 5, 3, 6};
const std::vector<float> updates1 = {17.0, 2.0, 5.0, -1.0};
std::memcpy(scatter.arg0_data(), indices0.data(),
indices0.size() * sizeof(int32_t));
std::memcpy(scatter.arg1_data(), updates0.data(),
updates0.size() * sizeof(float));
std::memcpy(scatter.arg2_data(), indices1.data(),
indices1.size() * sizeof(int32_t));
std::memcpy(scatter.arg3_data(), updates1.data(),
updates1.size() * sizeof(float));
EXPECT_TRUE(scatter.Run());
EXPECT_EQ(scatter.error_msg(), "");
const std::vector<float> expected0 = {0, 11, 0, 10, 9, 0, 0, 12};
const std::vector<float> expected1 = {0, 0, 17, 5, 0, 2, -1, 0};
// NOTE(basioli): Shape is hardcoded to 8 in tensorflow config.
EXPECT_EQ(scatter.result0_count(), expected0.size());
for (int i = 0; i < scatter.result0_count(); ++i) {
EXPECT_EQ(scatter.result0(i), expected0[i]);
}
// NOTE(basioli): Shape is hardcoded to 8 in tensorflow config.
EXPECT_EQ(scatter.result1_count(), expected1.size());
for (int i = 0; i < scatter.result1_count(); ++i) {
EXPECT_EQ(scatter.result1(i), expected1[i]) << "i: " << i;
}
}
TEST(TFCompileTest, AssertEqAndReturnDiff) {
// Assert is converted into a no-op in XLA, so there is no failure even if
// the two args are different.
AssertComp assert;
EXPECT_EQ(assert.arg0_data(), assert.arg_data(0));
EXPECT_EQ(assert.arg1_data(), assert.arg_data(1));
assert.arg0() = 2;
assert.arg1() = 1;
const int32_t expected_result = assert.arg0() - assert.arg1();
EXPECT_TRUE(assert.Run());
EXPECT_EQ(assert.error_msg(), "");
EXPECT_EQ(assert.result0(), expected_result);
EXPECT_EQ(assert.result0_data()[0], expected_result);
EXPECT_EQ(assert.result0_data(), assert.result_data(0));
}
TEST(TFCompileTest, LookupNameIndex) {
// add doesn't have any names defined in its config.
AddComp add;
EXPECT_FALSE(add.HasNameIndices());
// muladd has names defined for all feeds and fetches.
::foo::bar::MatMulAndAddComp muladd;
EXPECT_TRUE(muladd.HasNameIndices());
EXPECT_EQ(muladd.LookupArgIndex("x"), 0);
EXPECT_EQ(muladd.LookupArgIndex("y"), 1);
EXPECT_EQ(muladd.LookupArgIndex(""), -1);
EXPECT_EQ(muladd.LookupArgIndex("x_hold"), -1);
EXPECT_EQ(muladd.LookupArgIndex("y_hold"), -1);
EXPECT_EQ(muladd.LookupArgIndex("x_y_prod"), -1);
EXPECT_EQ(muladd.LookupArgIndex("x_y_sum"), -1);
EXPECT_EQ(muladd.LookupResultIndex("x_y_prod"), 0);
EXPECT_EQ(muladd.LookupResultIndex("x_y_sum"), 1);
EXPECT_EQ(muladd.LookupResultIndex(""), -1);
EXPECT_EQ(muladd.LookupResultIndex("x"), -1);
EXPECT_EQ(muladd.LookupResultIndex("y"), -1);
EXPECT_EQ(muladd.LookupResultIndex("x_hold"), -1);
EXPECT_EQ(muladd.LookupResultIndex("y_hold"), -1);
}
TEST(TFCompileTest, ProgramShape) {
using xla::ShapeUtil;
const xla::Shape f32_2x2 = ShapeUtil::MakeShape(xla::F32, {2, 2});
// add doesn't have the program shape defined.
AddComp add;
ASSERT_TRUE(add.ProgramShape() == nullptr);
// muladd has the program shape defined.
::foo::bar::MatMulAndAddComp muladd;
const xla::ProgramShapeProto* muladd_shape = muladd.ProgramShape();
ASSERT_TRUE(muladd_shape != nullptr);
ASSERT_EQ(muladd_shape->parameters_size(), 2);
TF_ASSERT_OK_AND_ASSIGN(xla::Shape muladd_arg0,
xla::Shape::FromProto(muladd_shape->parameters(0)));
TF_ASSERT_OK_AND_ASSIGN(xla::Shape muladd_arg1,
xla::Shape::FromProto(muladd_shape->parameters(1)));
EXPECT_TRUE(ShapeUtil::Compatible(muladd_arg0, f32_2x2));
EXPECT_TRUE(ShapeUtil::Compatible(muladd_arg1, f32_2x2));
TF_ASSERT_OK_AND_ASSIGN(xla::Shape muladd_result,
xla::Shape::FromProto(muladd_shape->result()));
ASSERT_EQ(muladd_result.element_type(), xla::TUPLE);
ASSERT_EQ(ShapeUtil::TupleElementCount(muladd_result), 2);
const xla::Shape& muladd_result0 =
ShapeUtil::GetTupleElementShape(muladd_result, 0);
EXPECT_TRUE(ShapeUtil::Compatible(muladd_result0, f32_2x2));
const xla::Shape& muladd_result1 =
ShapeUtil::GetTupleElementShape(muladd_result, 1);
EXPECT_TRUE(ShapeUtil::Compatible(muladd_result1, f32_2x2));
}
} // namespace
} // namespace tfcompile
} // namespace tensorflow
tensorflow/compiler/aot/tfcompile.bzl
+591
View File
@@ -0,0 +1,591 @@
"""Build macro that compiles a TensorFlow graph into a cc_library.
To use from your BUILD file, add the following line to load the macro:
load("//tensorflow/compiler/aot:tfcompile.bzl", "tf_library")
Then call the macro like this:
tf_library(
name = "test_graph_tfmatmul",
config = "test_graph_tfmatmul.config.pbtxt",
cpp_class = "MatMulComp",
graph = ":test_graph_tfmatmul.pb",
)
"""
load("@rules_cc//cc:cc_binary.bzl", "cc_binary")
load("@rules_cc//cc:cc_library.bzl", "cc_library")
load(
"//tensorflow:tensorflow.bzl",
"if_android",
"if_google",
"if_oss",
"tf_cc_test",
"tf_copts",
)
load("//tensorflow:tensorflow.default.bzl", "tfcompile_dfsan_abilists", "tfcompile_dfsan_enabled", "tfcompile_friends", "tfcompile_target_cpu")
visibility(tfcompile_friends())
def _tfcompile_model_library_rule_impl(ctx):
header_file = ctx.outputs.header_out
metadata_object_file = ctx.actions.declare_file("%s_tfcompile_metadata.o" % ctx.attr.model_name)
function_object_file = ctx.actions.declare_file("%s_tfcompile_function.o" % ctx.attr.model_name)
constant_buffers_object_file = ctx.actions.declare_file("%s_tfcompile_constant_buffers.o" % ctx.attr.model_name)
session_module_pb = ctx.actions.declare_file("%s_session_module.pb" % ctx.attr.model_name)
out_files = [header_file, metadata_object_file, function_object_file, constant_buffers_object_file, session_module_pb]
compiler_log_file = None
if ctx.attr.gen_compiler_log:
compiler_log_file = ctx.actions.declare_file("%s_compiler.log" % ctx.attr.model_name)
out_files.append(compiler_log_file)
output_dict = {}
output_dict["header_files"] = [header_file]
output_dict["object_files"] = [metadata_object_file, function_object_file, constant_buffers_object_file]
if compiler_log_file:
output_dict["log_files"] = [compiler_log_file]
output_flags = [
"--out_header=" + header_file.path,
"--out_metadata_object=" + metadata_object_file.path,
"--out_function_object=" + function_object_file.path,
"--out_constant_buffers_object=" + constant_buffers_object_file.path,
"--out_session_module=" + session_module_pb.path,
]
additional_xla_flags = ctx.attr.xla_flags
tfcompile_env = {
"XLA_FLAGS": ("--xla_cpu_enable_fast_math=true " +
"--xla_cpu_fast_math_honor_nans=false " +
"--xla_cpu_fast_math_honor_infs=false " +
"--xla_cpu_fast_math_honor_functions=false " +
"--xla_cpu_fast_math_honor_division=false " +
"--xla_cpu_enable_fast_min_max=true " +
"--xla_cpu_experimental_ynn_fusion_type= " +
additional_xla_flags + " " +
"$${XLA_FLAGS:-}' "),
"CUDA_VISIBLE_DEVICES": "",
}
dfsan_flags = []
dfsan_deps = []
# DFSan is only supported on linux.
if ctx.attr.is_linux and ctx.attr.dfsan:
dfsan_flags = [
"--sanitize_dataflow",
"--sanitize_abilists_dataflow=" + ",".join([f.path for f in ctx.files.dfsan_abilists]),
]
dfsan_deps = ctx.files.dfsan_abilists
cpu_flags = ["--target_cpu=" + ctx.attr.target_cpu] if ctx.attr.target_cpu else []
flags = [
"--graph=" + ctx.file.tfcompile_graph.path,
"--config=" + ctx.file.tfcompile_config.path,
"--entry_point=" + ctx.attr.entry_point,
"--cpp_class=" + ctx.attr.cpp_class,
"--target_triple=" + ctx.attr.target_triple,
] + cpu_flags + output_flags + ctx.attr.extra_flags + dfsan_flags
post_command = ""
if ctx.attr.gen_compiler_log:
post_command += " --vmodule=cpu_compiler=5 2> >(tee -a " + compiler_log_file.path + " >&2) "
full_cmd = (
ctx.executable.tfcompile_tool.path + " " + " ".join(flags) + " " + ctx.attr.flags + post_command
)
ctx.actions.run_shell(
inputs = ctx.files.srcs,
outputs = out_files,
tools = [ctx.executable.tfcompile_tool] + dfsan_deps,
env = tfcompile_env,
command = full_cmd,
progress_message = "tfcompile for model %s (%s)" % (ctx.attr.model_name, ctx.file.tfcompile_graph.path),
mnemonic = "TensorflowCompile",
)
return [
DefaultInfo(
files = depset(out_files),
),
OutputGroupInfo(**output_dict),
]
# Use tf_library macro instead of using this rule directly.
_tfcompile_model_library = rule(
implementation = _tfcompile_model_library_rule_impl,
attrs = {
"model_name": attr.string(),
"srcs": attr.label_list(mandatory = True, allow_files = True),
"header_out": attr.output(),
"cmd": attr.string(),
"tfcompile_tool": attr.label(cfg = "exec", executable = True, allow_files = True),
"tfcompile_graph": attr.label(allow_single_file = True),
"tfcompile_config": attr.label(allow_single_file = True),
"entry_point": attr.string(),
"cpp_class": attr.string(),
"target_triple": attr.string(),
"target_cpu": attr.string(),
# The tfcompile_flags passed into tf_library macro may be a string
# containing multiple flags (and there are cases that do this).
"flags": attr.string(),
# Extra flags are built in the tf_library macro as a list.
"extra_flags": attr.string_list(),
"dfsan": attr.bool(default = False),
"dfsan_abilists": attr.label_list(default = [], allow_files = True),
"is_linux": attr.bool(),
"gen_compiler_log": attr.bool(),
"xla_flags": attr.string(),
},
)
def _tf_library(
name,
graph,
config,
debug_info = None,
freeze_checkpoint = None,
freeze_saver = None,
cpp_class = None,
gen_test = True,
gen_benchmark = True,
gen_compiler_log = False,
visibility = None,
testonly = None,
tfcompile_flags = None,
tfcompile_tool = "//tensorflow/compiler/aot:tfcompile",
include_standard_runtime_deps = True,
enable_xla_hlo_profiling = False,
enable_tracemes = False,
mlir_components = "None",
deps = None,
tags = [],
copts = [],
xla_flags = None):
if not cpp_class:
fail("cpp_class must be specified")
tfcompile_graph = graph
if freeze_checkpoint or freeze_saver:
if not freeze_checkpoint:
fail("freeze_checkpoint must be specified when freeze_saver is " +
"specified")
freeze_name = "freeze_" + name
freeze_file = freeze_name + ".pb"
# First run tfcompile to generate the list of out_nodes.
#
# Here and below, we set CUDA_VISIBLE_DEVICES='' to prevent the code we
# launch from using any GPUs which might be present. This is important
# because builds may run concurrently with tests, and tests need to be
# able to assume that they have control of the full GPU.
out_nodes_file = "out_nodes_" + freeze_name
native.genrule(
name = ("gen_" + out_nodes_file),
srcs = [config],
outs = [out_nodes_file],
cmd = ("CUDA_VISIBLE_DEVICES='' " +
"$(location " + tfcompile_tool + ")" +
" --config=$(location " + config + ")" +
" --dump_fetch_nodes > $@"),
tools = [tfcompile_tool],
# Run tfcompile on the build host, rather than forge, since it's
# typically way faster on the local machine.
local = 1,
tags = tags,
)
# Now run freeze_graph to convert variables into constants.
freeze_args = (
" --input_graph=$(location " + graph + ")" +
" --checkpoint_version=1" +
" --input_binary=" + str(not graph.endswith(".pbtxt")) +
" --input_checkpoint=$(location " + freeze_checkpoint + ")" +
" --output_graph=$(location " + freeze_file + ")" +
" --output_node_names=$$(<$(location " + out_nodes_file +
"))"
)
freeze_saver_srcs = []
if freeze_saver:
freeze_args += " --input_saver=$(location " + freeze_saver + ")"
freeze_saver_srcs.append(freeze_saver)
native.genrule(
name = freeze_name,
srcs = [
graph,
freeze_checkpoint,
out_nodes_file,
] + freeze_saver_srcs,
outs = [freeze_file],
cmd = (
"CUDA_VISIBLE_DEVICES='' " +
"$(location " +
"//tensorflow/python/tools:freeze_graph)" +
freeze_args
),
tools = ["//tensorflow/python/tools:freeze_graph"],
tags = tags,
)
tfcompile_graph = freeze_file
# Rule that runs tfcompile to produce the header and object file.
header_file = name + ".h"
# The XLA backends morph kernel name prefix __ that is not in the form of
# __xla_.
ep = ("__xla_" + native.package_name() + "__" + name).replace("/", "_")
if type(tfcompile_flags) == type(""):
flags = tfcompile_flags
else:
flags = " ".join([
"'" + arg.replace("'", "'\\''") + "'"
for arg in (tfcompile_flags or [])
])
# Do this before we append the `select` into `flags`, because doing so
# transforms `flags` into a variable of type `select`, and we can't call
# `find` on such an object.
need_xla_data_proto = flags and flags.find("--gen_program_shape") != -1
if enable_xla_hlo_profiling:
profiling_flags = ["--xla_hlo_profile"]
else:
profiling_flags = []
if enable_tracemes:
traceme_flags = ["--xla_cpu_enable_xprof_traceme=true"]
else:
traceme_flags = ["--xla_cpu_enable_xprof_traceme=false"]
mlir_flags = ["--mlir_components=" + mlir_components]
srcs = [tfcompile_graph, config]
debug_info_flags = []
if debug_info:
srcs.append(debug_info)
debug_info_flags = ["--debug_info=$(location " + debug_info + ")"]
tfcompile_gen = "gen_" + name
_tfcompile_model_library(
name = tfcompile_gen,
model_name = name,
srcs = srcs,
gen_compiler_log = gen_compiler_log,
header_out = header_file,
tfcompile_tool = tfcompile_tool,
tfcompile_graph = tfcompile_graph,
tfcompile_config = config,
entry_point = ep,
cpp_class = cpp_class,
target_cpu = tfcompile_target_cpu(name),
target_triple = target_llvm_triple(),
flags = flags,
extra_flags = debug_info_flags + profiling_flags + mlir_flags + traceme_flags,
dfsan = tfcompile_dfsan_enabled(),
dfsan_abilists = tfcompile_dfsan_abilists(),
is_linux = select({
"//tensorflow:linux_x86_64": True,
"//conditions:default": False,
}),
visibility = visibility,
testonly = testonly,
tags = tags,
xla_flags = xla_flags,
)
tfcompile_gen_object_files = tfcompile_gen + "_object_files"
native.filegroup(
name = tfcompile_gen_object_files,
srcs = [tfcompile_gen],
output_group = "object_files",
visibility = visibility,
testonly = testonly,
)
use_xla_nanort_runtime = False
if tfcompile_flags and "--use_xla_nanort_runtime" in tfcompile_flags:
use_xla_nanort_runtime = True
# The cc_library rule packaging up the header and object file, and needed
# kernel implementations.
cc_library(
name = name,
srcs = [tfcompile_gen_object_files],
hdrs = [header_file],
visibility = visibility,
testonly = testonly,
deps = [
# These deps are required by all tf_library targets even if
# include_standard_runtime_deps is False. Without them, the
# generated code will fail to compile.
"//third_party/absl/log:check",
"//third_party/absl/synchronization",
"//tensorflow/core:framework_lite",
"//tensorflow/compiler/tf2xla:xla_compiled_cpu_function",
"@xla//xla:types",
"@xla//xla/backends/cpu/runtime:kernel_c_api",
"@xla//xla/backends/cpu/runtime:rng_state_lib",
] + (need_xla_data_proto and [
# If we're generating the program shape, we must depend on the
# proto.
"@xla//xla:xla_data_proto_cc",
] or []) + (enable_xla_hlo_profiling and [
"@xla//xla/service:hlo_profile_printer_data_cc",
] or []) + (include_standard_runtime_deps and [
# TODO(cwhipkey): only depend on kernel code that the model actually
# needed.
"@xla//xla/backends/cpu/runtime:dot_lib",
"@xla//xla/backends/cpu/runtime:sort_lib",
"@xla//xla/backends/cpu/runtime:topk_lib",
"@xla//xla/backends/cpu/runtime:convolution_lib",
"@xla//xla/service/cpu:runtime_matmul",
"@xla//xla/service/cpu:runtime_single_threaded_matmul",
"@eigen_archive//:eigen3",
] or []) + (use_xla_nanort_runtime and [
"//tensorflow/compiler/tf2xla:xla_compiled_cpu_function_thunks",
] or []) + (deps or []),
tags = tags,
copts = copts,
)
# Variables used for gen_test and gen_benchmark.
cpp_class_split = cpp_class.rsplit("::", 2)
if len(cpp_class_split) == 1:
no_ns_name = cpp_class_split[0]
else:
no_ns_name = cpp_class_split[1]
sed_replace = (
"-e \"s|{{TFCOMPILE_HEADER}}|$(location " + header_file + ")|g\" " +
"-e \"s|{{TFCOMPILE_CPP_CLASS}}|" + cpp_class + "|g\" " +
"-e \"s|{{TFCOMPILE_NAME}}|" + no_ns_name + "|g\" "
)
if gen_test:
test_name = name + "_test"
test_file = test_name + ".cc"
template_file = "//tensorflow/compiler/aot:test"
template_file += if_oss("", "_google") + ".cc"
# Rule to rewrite the template_file to produce the test_file.
native.genrule(
name = ("gen_" + test_name),
testonly = 1,
srcs = [
template_file,
header_file,
],
outs = [test_file],
cmd = (
"sed " + sed_replace +
" $(location " + template_file + ") " +
"> $(OUTS)"
),
tags = tags,
)
# The cc_test rule for the generated code. To ensure that this works
# reliably across build configurations, we must use tf_cc_test instead
# of native.cc_test. This is related to how we build
# //tensorflow/core:lib -- see the note in
# tensorflow/core/BUILD for more details.
tf_cc_test(
name = test_name,
srcs = [test_file],
deps = [
":" + name,
"//tensorflow/compiler/aot:tf_library_test_main",
"@xla//xla:executable_run_options",
"@eigen_archive//:eigen3",
] + if_oss([
"//tensorflow/core:lib",
"//tensorflow/core:test",
]) + if_google([
"@com_google_googletest//:gtest",
"//tensorflow/core/platform:byte_order",
"//tensorflow/core/platform:platform_port",
]),
tags = tags,
extra_copts = copts,
visibility = visibility,
)
if gen_benchmark:
benchmark_name = name + "_benchmark"
benchmark_file = benchmark_name + ".cc"
benchmark_main = ("//tensorflow/compiler/aot:" +
"benchmark_main.template")
# Rule to rewrite benchmark.cc to produce the benchmark_file.
native.genrule(
name = ("gen_" + benchmark_name),
srcs = [
benchmark_main,
header_file,
],
testonly = testonly,
outs = [benchmark_file],
cmd = ("sed " + sed_replace +
" $(location " + benchmark_main + ") " +
"> $(OUTS)"),
tags = tags,
)
# The cc_benchmark rule for the generated code. This does not need the
# tf_cc_binary since we (by deliberate design) do not depend on
# //tensorflow/core:lib.
#
# Note: to get smaller size on android for comparison, compile with:
# --copt=-fvisibility=hidden
# --copt=-D_LIBCPP_TYPE_VIS=_LIBCPP_HIDDEN
# --copt=-D_LIBCPP_EXCEPTION_ABI=_LIBCPP_HIDDEN
cc_binary(
name = benchmark_name,
srcs = [benchmark_file],
testonly = testonly,
copts = copts + tf_copts(),
linkopts = if_android(["-pie", "-s"]),
deps = [
":" + name,
"//tensorflow/compiler/aot:benchmark",
"@xla//xla:executable_run_options",
"@eigen_archive//:eigen3",
] + if_android([
"//tensorflow/compiler/aot:benchmark_extra_android",
]),
tags = tags,
visibility = visibility,
)
def tf_library(
name,
graph,
config,
debug_info = None,
freeze_checkpoint = None,
freeze_saver = None,
cpp_class = None,
gen_test = True,
gen_benchmark = True,
gen_compiler_log = False,
visibility = None,
testonly = None,
tfcompile_flags = None,
tfcompile_tool = "//tensorflow/compiler/aot:tfcompile",
include_standard_runtime_deps = True,
enable_xla_hlo_profiling = False,
enable_tracemes = False,
mlir_components = "None",
deps = None,
tags = [],
copts = [],
xla_flags = None):
"""Compiles a TensorFlow graph into an executable with fast math enabled.
Given an invocation of tf_library(name="foo", ...), generates the following
build targets:
foo: A cc_library containing the generated header and
computation.
foo_test: A cc_test with simple tests and benchmarks. Only created if
gen_test=True.
foo_benchmark: A cc_binary that runs a minimal-dependency benchmark,
useful for mobile devices or other platforms that can't
compile the full test libraries. Only created if
gen_benchmark=True.
The output header is called <name>.h.
Args:
name: The name of the build rule.
graph: The TensorFlow GraphDef to compile. If the file ends in '.pbtxt'
it is expected to be in the human-readable proto text format, otherwise
it is expected to be in the proto binary format.
config: File containing tensorflow.tf2xla.Config proto. If the file ends
in '.pbtxt' it is expected to be in the human-readable proto text
format, otherwise it is expected to be in the proto binary format.
debug_info: Debug info to include in the output.
freeze_checkpoint: If provided, run freeze_graph with this checkpoint to
convert variables into constants.
freeze_saver: If provided, run freeze_graph with this saver, in SaverDef
binary form, to convert variables into constants.
cpp_class: The name of the generated C++ class, wrapping the generated
function. The syntax of this flag is
[[<optional_namespace>::],...]<class_name>. This mirrors the C++ syntax
for referring to a class, where multiple namespaces may precede the
class name, separated by double-colons. The class will be generated in
the given namespace(s), or if no namespaces are given, within the global
namespace.
gen_test: If True, also generate a cc_test rule that builds a simple
test and benchmark.
gen_benchmark: If True, also generate a binary with a simple benchmark.
Unlike the output of gen_test, this benchmark can be run on android.
gen_compiler_log: If True, dumps XLA:CPU debug output to a log file.
visibility: Bazel build visibility.
testonly: Bazel testonly attribute.
tfcompile_flags: Extra flags to pass to tfcompile to control compilation.
tfcompile_tool: The tfcompile binary. A non-default can be passed to
use a tfcompile built with extra dependencies.
include_standard_runtime_deps: If True, the standard list of
kernel/runtime deps is added to deps. If False, deps must contain the
full set of deps needed by the generated library.
enable_xla_hlo_profiling: Enable XLA HLO profiling in the generated
program, and emit metadata that lets us pretty-print the gathered
profile counters.
enable_tracemes: Tell tfcompile to generate calls to
TraceMe::Activity{Start|End} around HLO instructions that can be used by
Xprof to construct profiler timelines.
mlir_components: When the value is "None", no components use MLIR. When
the value is "Bridge", use MLIR to translate GraphDef to HLO.
deps: a list of deps to include on the build rules for the generated
library, added to the standard deps if standard_runtime_deps is True.
tags: tags to apply to subsidiary build rules.
copts: list of copts to pass to cc rules.
"""
_tf_library(
name,
graph,
config,
debug_info,
freeze_checkpoint,
freeze_saver,
cpp_class,
gen_test,
gen_benchmark,
gen_compiler_log,
visibility,
testonly,
tfcompile_flags,
tfcompile_tool,
include_standard_runtime_deps,
enable_xla_hlo_profiling,
enable_tracemes,
mlir_components,
deps,
tags,
copts,
xla_flags,
)
def target_llvm_triple():
"""Returns the target LLVM triple to be used for compiling the target."""
# TODO(toddw): Add target_triple for other targets. For details see:
# http://llvm.org/docs/doxygen/html/Triple_8h_source.html
return select({
"//tensorflow:android_armeabi": "armv5-none-android",
"//tensorflow:android_arm": "armv7-none-android",
"//tensorflow:android_arm64": "aarch64-none-android",
"//tensorflow:android_x86": "i686-none-android",
"//tensorflow:ios": "arm64-none-ios",
"//tensorflow:ios_x86_64": "x86_64-apple-ios",
"//tensorflow:linux_ppc64le": "ppc64le-ibm-linux-gnu",
"//tensorflow:linux_aarch64": "aarch64-none-linux-gnu",
"//tensorflow:macos_x86_64": "x86_64-none-darwin",
"//tensorflow:macos_arm64": "aarch64-none-darwin",
"//tensorflow:windows": "x86_64-none-windows",
"//tensorflow:linux_s390x": "systemz-none-linux-gnu",
# internal placeholder,
"//conditions:default": "x86_64-pc-linux",
})
tensorflow/compiler/aot/tfcompile_main.cc
+84
View File
@@ -0,0 +1,84 @@
/* Copyright 2017 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 <vector>
#include "absl/log/check.h"
#include "absl/status/status.h"
#include "absl/strings/string_view.h"
#include "tensorflow/compiler/aot/compile.h"
#include "tensorflow/compiler/aot/flags.h"
#include "tensorflow/compiler/tf2xla/tf2xla.pb.h"
#include "xla/debug_options_flags.h"
#include "xla/tsl/platform/status.h"
#include "tensorflow/core/framework/graph.pb.h"
#include "tensorflow/core/platform/init_main.h"
#include "tensorflow/core/platform/logging.h"
#include "tensorflow/core/util/command_line_flags.h"
namespace tensorflow {
namespace tfcompile {
const char kUsageHeader[] =
"tfcompile performs ahead-of-time compilation of a TensorFlow graph,\n"
"resulting in an object file compiled for your target architecture, and a\n"
"header file that gives access to the functionality in the object file.\n"
"A typical invocation looks like this:\n"
"\n"
" $ tfcompile --graph=mygraph.pb --config=myfile.pbtxt "
"--cpp_class=\"mynamespace::MyComputation\"\n"
"\n";
} // end namespace tfcompile
} // end namespace tensorflow
int main(int argc, char** argv) {
tensorflow::tfcompile::MainFlags flags;
#ifndef __s390x__
flags.target_triple = "x86_64-pc-linux";
#endif
flags.out_function_object = "out_model.o";
flags.out_metadata_object = "out_helper.o";
flags.out_header = "out.h";
flags.entry_point = "entry";
flags.debug_info_path_begin_marker = "";
// Note that tfcompile.bzl's tf_library macro sets fast math flags as that is
// generally the preferred case.
std::vector<tensorflow::Flag> flag_list;
AppendMainFlags(&flag_list, &flags);
xla::AppendDebugOptionsFlags(&flag_list);
std::string usage = tensorflow::tfcompile::kUsageHeader;
usage += tensorflow::Flags::Usage(argv[0], flag_list);
if (argc > 1 && absl::string_view(argv[1]) == "--help") {
std::cerr << usage << "\n";
return 0;
}
bool parsed_flags_ok = tensorflow::Flags::Parse(&argc, argv, flag_list);
QCHECK(parsed_flags_ok) << "\n" << usage;
tensorflow::port::InitMain(usage.c_str(), &argc, &argv);
QCHECK(argc == 1) << "\nERROR: This command does not take any arguments "
"other than flags. See --help.\n\n";
absl::Status status = tensorflow::tfcompile::Main(flags);
if (status.code() == absl::StatusCode::kInvalidArgument) {
std::cerr << "INVALID ARGUMENTS: " << status.message() << "\n\n";
return 1;
} else {
TF_QCHECK_OK(status);
}
return 0;
}
tensorflow/compiler/aot/thunk_proto_execution_deserializer.cc
+918
View File
@@ -0,0 +1,918 @@
/* Copyright 2025 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/compiler/aot/thunk_proto_execution_deserializer.h"
#include <cstddef>
#include <cstdint>
#include <string>
#include <utility>
#include <vector>
#include "absl/numeric/int128.h"
#include "absl/status/status.h"
#include "absl/status/statusor.h"
#include "absl/strings/str_cat.h"
#include "absl/strings/str_join.h"
#include "absl/strings/str_replace.h"
#include "absl/strings/string_view.h"
#include "xla/backends/cpu/runtime/convolution_dims.h"
#include "xla/backends/cpu/runtime/dot_dims.h"
#include "xla/backends/cpu/runtime/kernel_thunk.h"
#include "xla/backends/cpu/runtime/thunk.pb.h"
#include "xla/layout_util.h"
#include "xla/service/cpu/executable.pb.h"
#include "xla/shape.h"
#include "xla/shape_util.h"
#include "xla/status_macros.h"
#include "xla/tsl/platform/statusor.h"
#include "xla/util.h"
#include "xla/xla_data.pb.h"
namespace tensorflow {
namespace tfcompile {
namespace {
std::string GetBufferAllocationString(
const xla::buffer_assignment::BufferAllocationSliceProto& slice) {
return absl::StrCat("reinterpret_cast<std::byte*>(buffer_table[",
slice.buffer_allocation_index(), "]) + ", slice.offset());
}
} // namespace
absl::StatusOr<std::string>
ThunkProtoExecutionDeserializer::GetThunkSpecificRunImpl(
const xla::cpu::CompilationResultProto& proto) && {
return ThunkSpecificRunImplFromThunkSequence(proto.thunk_sequence());
}
absl::StatusOr<std::string>
ThunkProtoExecutionDeserializer::ThunkSpecificRunImplFromThunkSequence(
const xla::cpu::ThunkSequenceProto& thunk_sequence_proto) {
std::vector<std::string> thunk_run_impls;
thunk_run_impls.reserve(thunk_sequence_proto.thunks_size());
for (const auto& thunk : thunk_sequence_proto.thunks()) {
switch (thunk.impl_case()) {
case xla::cpu::ThunkProto::kKernelThunk: {
TF_ASSIGN_OR_RETURN(thunk_run_impls.emplace_back(),
GetKernelThunkRunImpl(thunk));
break;
}
case xla::cpu::ThunkProto::kDotThunk: {
TF_ASSIGN_OR_RETURN(thunk_run_impls.emplace_back(),
GetDotThunkRunImpl(thunk));
break;
}
case xla::cpu::ThunkProto::kConditionalThunk: {
TF_ASSIGN_OR_RETURN(thunk_run_impls.emplace_back(),
GetConditionalThunkRunImpl(thunk));
break;
}
case xla::cpu::ThunkProto::kWhileThunk: {
TF_ASSIGN_OR_RETURN(thunk_run_impls.emplace_back(),
GetWhileThunkRunImpl(thunk));
break;
}
case xla::cpu::ThunkProto::kConvolutionThunk: {
TF_ASSIGN_OR_RETURN(thunk_run_impls.emplace_back(),
GetConvolutionFusionThunkRunImpl(thunk));
break;
}
case xla::cpu::ThunkProto::kRngGetAndUpdateStateThunk: {
TF_ASSIGN_OR_RETURN(thunk_run_impls.emplace_back(),
GetRngGetAndUpdateStateThunkRunImpl(thunk));
break;
}
case xla::cpu::ThunkProto::kCallThunk: {
TF_ASSIGN_OR_RETURN(thunk_run_impls.emplace_back(),
GetCallThunkRunImpl(thunk));
break;
}
case xla::cpu::ThunkProto::kCopyThunk: {
TF_ASSIGN_OR_RETURN(thunk_run_impls.emplace_back(),
GetCopyThunkRunImpl(thunk));
break;
}
case xla::cpu::ThunkProto::kSortThunk: {
TF_ASSIGN_OR_RETURN(thunk_run_impls.emplace_back(),
GetSortThunkRunImpl(thunk));
break;
}
case xla::cpu::ThunkProto::kTopKThunk: {
TF_ASSIGN_OR_RETURN(thunk_run_impls.emplace_back(),
GetTopKThunkRunImpl(thunk));
break;
}
case xla::cpu::ThunkProto::kRngSeedThunk: {
TF_ASSIGN_OR_RETURN(thunk_run_impls.emplace_back(),
GetRngSeedThunkRunImpl(thunk));
break;
}
default: {
return xla::Internal("Unsupported thunk type: %s.", thunk.kind());
}
}
}
return absl::StrJoin(thunk_run_impls, "\n");
}
absl::StatusOr<std::string> ThunkProtoExecutionDeserializer::GetMatmulFunction(
xla::PrimitiveType xla_type) {
switch (xla_type) {
case xla::F16:
return "::xla::cpu::internal::TypedMatMul<Eigen::half, Eigen::half, "
"Eigen::half>";
case xla::F32:
return "::xla::cpu::internal::TypedMatMul<float, float, float>";
case xla::F64:
return "::xla::cpu::internal::TypedMatMul<double, double, double>";
case xla::C64:
return "::xla::cpu::internal::TypedMatMul<std::complex<float>, "
"std::complex<float>, std::complex<float>";
case xla::C128:
return "::xla::cpu::internal::TypedMatMul<std::complex<double>, "
"std::complex<double>, std::complex<double>";
case xla::S32:
return "::xla::cpu::internal::TypedMatMul<int32_t, int32_t, int32_t>";
default:
return xla::Internal("Unsupported xla type: %d", xla_type);
}
}
absl::StatusOr<std::string> ThunkProtoExecutionDeserializer::GetDotThunkRunImpl(
const xla::cpu::ThunkProto& thunk) {
if (!thunk.has_dot_thunk()) {
return xla::Internal(
"Dot thunk was expected when getting thunk run implementation.");
}
const xla::cpu::DotThunkProto& dot_thunk = thunk.dot_thunk();
absl::string_view dot_thunk_invocation_format = R"(
// Dot Thunk
{
absl::BlockingCounter done({{BATCH_SIZE}});
for (int64_t i = 0; i < {{BATCH_SIZE}}; ++i) {
{{MATMUL_FUNCTION}}(
run_options->intra_op_thread_pool(),
{{OUTPUT_PTR}} + {{OUTPUT_STRIDE}} * i,
{{LHS_PTR}} + {{LHS_STRIDE}} * i,
{{RHS_PTR}} + {{RHS_STRIDE}} * i,
{{M}}, {{N}}, {{K}}, {{TRANSPOSE_LHS}}, {{TRANSPOSE_RHS}},
[&done] { done.DecrementCount(); }
);
}
done.Wait();
}
)";
if (!(dot_thunk.lhs_buffer_shape().shape().element_type() ==
dot_thunk.rhs_buffer_shape().shape().element_type() &&
dot_thunk.rhs_buffer_shape().shape().element_type() ==
dot_thunk.out_buffer_shape().shape().element_type())) {
return xla::Internal(
"Dot thunk has mismatched types between lhs, rhs, and out buffers.");
}
TF_ASSIGN_OR_RETURN(
std::string matmul_function,
GetMatmulFunction(dot_thunk.lhs_buffer_shape().shape().element_type()));
TF_ASSIGN_OR_RETURN(std::string data_type,
CppDataTypeFromXlaType(
dot_thunk.lhs_buffer_shape().shape().element_type()));
std::string output_ptr = absl::StrCat(
"reinterpret_cast<", data_type, "*>(",
GetBufferAllocationString(dot_thunk.out_buffer_shape().slice()), ")");
std::string lhs_ptr = absl::StrCat(
"reinterpret_cast<", data_type, "*>(",
GetBufferAllocationString(dot_thunk.lhs_buffer_shape().slice()), ")");
std::string rhs_ptr = absl::StrCat(
"reinterpret_cast<", data_type, "*>(",
GetBufferAllocationString(dot_thunk.rhs_buffer_shape().slice()), ")");
TF_ASSIGN_OR_RETURN(
auto lhs_shape,
xla::Shape::FromProto(dot_thunk.lhs_buffer_shape().shape()));
TF_ASSIGN_OR_RETURN(
auto rhs_shape,
xla::Shape::FromProto(dot_thunk.rhs_buffer_shape().shape()));
TF_ASSIGN_OR_RETURN(
auto out_shape,
xla::Shape::FromProto(dot_thunk.out_buffer_shape().shape()));
TF_ASSIGN_OR_RETURN(xla::cpu::DotShape dot_shape,
xla::cpu::GetDotShape(dot_thunk.dot_dimensions(),
lhs_shape, rhs_shape, out_shape));
TF_ASSIGN_OR_RETURN(
xla::cpu::DotCanonicalDims dot_canonical_dims,
GetDotCanonicalDims(dot_thunk.dot_dimensions(), dot_shape));
size_t m = dot_canonical_dims.m;
size_t k = dot_canonical_dims.k;
size_t n = dot_canonical_dims.n;
// Decide if a transpose is required based on an XOR of the canonical and
// column major flags.
bool transpose_lhs =
(dot_canonical_dims.lhs_canonical != dot_canonical_dims.lhs_column_major);
bool transpose_rhs =
(dot_canonical_dims.rhs_canonical != dot_canonical_dims.rhs_column_major);
if (!dot_canonical_dims.output_column_major) {
std::swap(m, n);
std::swap(lhs_ptr, rhs_ptr);
std::swap(transpose_lhs, transpose_rhs);
transpose_lhs = !transpose_lhs;
transpose_rhs = !transpose_rhs;
}
// NOTE: Don't use byte width because the buffer pointers will already be cast
// to the correct type. Reference xla/backends/cpu/runtime/dot_thunk.cc.
int64_t lhs_stride = m * k;
int64_t rhs_stride = k * n;
int64_t out_stride = m * n;
std::vector<std::pair<std::string, std::string>> rewrites = {
{"{{MATMUL_FUNCTION}}", matmul_function},
{"{{OUTPUT_PTR}}", output_ptr},
{"{{OUTPUT_STRIDE}}", absl::StrCat(out_stride)},
{"{{LHS_PTR}}", lhs_ptr},
{"{{LHS_STRIDE}}", absl::StrCat(lhs_stride)},
{"{{RHS_PTR}}", rhs_ptr},
{"{{RHS_STRIDE}}", absl::StrCat(rhs_stride)},
{"{{M}}", absl::StrCat(m)},
{"{{N}}", absl::StrCat(n)},
{"{{K}}", absl::StrCat(k)},
{"{{TRANSPOSE_LHS}}", transpose_lhs ? "true" : "false"},
{"{{TRANSPOSE_RHS}}", transpose_rhs ? "true" : "false"},
{"{{BATCH_SIZE}}", absl::StrCat(dot_shape.batch_size)}};
return absl::StrReplaceAll(dot_thunk_invocation_format, rewrites);
};
absl::StatusOr<std::string>
ThunkProtoExecutionDeserializer::GetConvolutionFunction(
xla::PrimitiveType xla_type) {
switch (xla_type) {
case xla::F16:
return "xla::cpu::internal::EigenConv2D<Eigen::half>";
case xla::F32:
return "xla::cpu::internal::EigenConv2D<float>";
default:
return xla::Internal("Unsupported xla type: %d", xla_type);
}
}
absl::StatusOr<std::string>
ThunkProtoExecutionDeserializer::GetConvolution2DRunImpl(
const xla::cpu::ConvolutionThunkProto& convolution_thunk,
const xla::cpu::ConvolutionCanonicalDims& canonical_dims) {
TF_ASSIGN_OR_RETURN(
std::string data_type,
CppDataTypeFromXlaType(
convolution_thunk.input_buffer_shape().shape().element_type()));
std::string output_ptr =
absl::StrCat("reinterpret_cast<", data_type, "*>(",
GetBufferAllocationString(
convolution_thunk.output_buffer_shape().slice()),
")");
std::string lhs_ptr = absl::StrCat(
"reinterpret_cast<", data_type, "*>(",
GetBufferAllocationString(convolution_thunk.input_buffer_shape().slice()),
")");
std::string rhs_ptr =
absl::StrCat("reinterpret_cast<", data_type, "*>(",
GetBufferAllocationString(
convolution_thunk.kernel_buffer_shape().slice()),
")");
TF_ASSIGN_OR_RETURN(
std::string convolution_function,
GetConvolutionFunction(
convolution_thunk.input_buffer_shape().shape().element_type()));
absl::string_view convolution_thunk_invocation_format = R"(
// Convolution Thunk
{
absl::Notification done;
{{CONVOLUTION_FUNCTION}}(
run_options->intra_op_thread_pool(),
{{OUTPUT_PTR}}, {{LHS_PTR}}, {{RHS_PTR}}, {{INPUT_BATCH}},
{{INPUT_ROWS}}, {{INPUT_COLS}}, {{INPUT_CHANNELS}}, {{KERNEL_ROWS}},
{{KERNEL_COLS}}, {{KERNEL_CHANNELS}}, {{KERNEL_FILTERS}},
{{OUTPUT_ROWS}}, {{OUTPUT_COLS}}, {{ROW_STRIDE}}, {{COL_STRIDE}},
{{PADDING_TOP}}, {{PADDING_BOTTOM}}, {{PADDING_LEFT}},
{{PADDING_RIGHT}}, {{LHS_ROW_DILATION}}, {{LHS_COL_DILATION}},
{{RHS_ROW_DILATION}}, {{RHS_COL_DILATION}}, {{FEATURE_GROUP_COUNT}},
[&done] { done.Notify(); }
);
done.WaitForNotification();
})";
std::vector<std::pair<std::string, std::string>> rewrites = {
{"{{CONVOLUTION_FUNCTION}}", convolution_function},
{"{{OUTPUT_PTR}}", output_ptr},
{"{{LHS_PTR}}", lhs_ptr},
{"{{RHS_PTR}}", rhs_ptr},
{"{{INPUT_BATCH}}", absl::StrCat(canonical_dims.input_batch)},
{"{{INPUT_ROWS}}", absl::StrCat(canonical_dims.input_dims.x)},
{"{{INPUT_COLS}}", absl::StrCat(canonical_dims.input_dims.y)},
{"{{INPUT_CHANNELS}}", absl::StrCat(canonical_dims.input_channels)},
{"{{KERNEL_ROWS}}", absl::StrCat(canonical_dims.kernel_dims.x)},
{"{{KERNEL_COLS}}", absl::StrCat(canonical_dims.kernel_dims.y)},
{"{{KERNEL_CHANNELS}}", absl::StrCat(canonical_dims.kernel_channels)},
{"{{KERNEL_FILTERS}}", absl::StrCat(canonical_dims.kernel_filters)},
{"{{OUTPUT_ROWS}}", absl::StrCat(canonical_dims.output_dims.x)},
{"{{OUTPUT_COLS}}", absl::StrCat(canonical_dims.output_dims.y)},
{"{{ROW_STRIDE}}", absl::StrCat(canonical_dims.strides.x)},
{"{{COL_STRIDE}}", absl::StrCat(canonical_dims.strides.y)},
{"{{PADDING_TOP}}", absl::StrCat(canonical_dims.padding_before.x)},
{"{{PADDING_BOTTOM}}", absl::StrCat(canonical_dims.padding_after.x)},
{"{{PADDING_LEFT}}", absl::StrCat(canonical_dims.padding_before.y)},
{"{{PADDING_RIGHT}}", absl::StrCat(canonical_dims.padding_after.y)},
{"{{LHS_ROW_DILATION}}", absl::StrCat(canonical_dims.base_dilation.x)},
{"{{LHS_COL_DILATION}}", absl::StrCat(canonical_dims.base_dilation.y)},
{"{{RHS_ROW_DILATION}}", absl::StrCat(canonical_dims.window_dilation.x)},
{"{{RHS_COL_DILATION}}", absl::StrCat(canonical_dims.window_dilation.y)},
{"{{FEATURE_GROUP_COUNT}}",
absl::StrCat(canonical_dims.feature_group_count)}};
return absl::StrReplaceAll(convolution_thunk_invocation_format, rewrites);
}
absl::StatusOr<std::string>
ThunkProtoExecutionDeserializer::GetConvolutionFusionThunkRunImpl(
const xla::cpu::ThunkProto& thunk) {
if (!thunk.has_convolution_thunk()) {
return xla::Internal(
"Convolution thunk was expected when getting thunk run "
"implementation.");
}
const xla::cpu::ConvolutionThunkProto& convolution_thunk =
thunk.convolution_thunk();
TF_ASSIGN_OR_RETURN(
xla::Shape input_buffer_shape,
xla::Shape::FromProto(convolution_thunk.input_buffer_shape().shape()));
TF_ASSIGN_OR_RETURN(
xla::Shape kernel_buffer_shape,
xla::Shape::FromProto(convolution_thunk.kernel_buffer_shape().shape()));
TF_ASSIGN_OR_RETURN(
xla::Shape output_buffer_shape,
xla::Shape::FromProto(convolution_thunk.output_buffer_shape().shape()));
// NOTE(basioli): Slices are not needed here, we only use this class to
// invoke GetConvolutionCanonicalDims.
xla::cpu::ConvolutionSlices slices{
/*input_buffer=*/{},
/*input_shape=*/input_buffer_shape,
/*kernel_buffer=*/{},
/*kernel_shape=*/kernel_buffer_shape,
/*output_buffer=*/{},
/*output_shape=*/output_buffer_shape,
};
TF_ASSIGN_OR_RETURN(
xla::cpu::ConvolutionCanonicalDims canonical_dims,
xla::cpu::GetConvolutionCanonicalDims(
slices, convolution_thunk.dimension_numbers(),
convolution_thunk.window(), convolution_thunk.feature_group_count()));
if (canonical_dims.convolution_rank() == 2) {
return GetConvolution2DRunImpl(convolution_thunk, canonical_dims);
} else {
return xla::Internal("3D convolution is not implemented.");
}
}
absl::StatusOr<std::string>
ThunkProtoExecutionDeserializer::GetRngGetAndUpdateStateThunkRunImpl(
const xla::cpu::ThunkProto& thunk) {
if (!thunk.has_rng_get_and_update_state_thunk()) {
return xla::Internal(
"RngGetAndUpdateState thunk was expected when getting thunk run "
"implementation.");
}
const xla::cpu::RngGetAndUpdateStateThunkProto& rng_thunk =
thunk.rng_get_and_update_state_thunk();
absl::string_view rng_thunk_invocation_format = R"(
// Rng Thunk
{
rng_states[{{RNG_STATE_INDEX}}]->GetAndUpdateState({{RNG_STATE_PTR}});
})";
if (rng_thunk.state_buffer().size() != sizeof(absl::int128)) {
return absl::InvalidArgumentError(
absl::StrCat("Rng state buffer size: ", rng_thunk.state_buffer().size(),
" is not equal to the size of an absl::int128: ",
sizeof(absl::int128)));
}
return absl::StrReplaceAll(
rng_thunk_invocation_format,
{{"{{RNG_STATE_INDEX}}", absl::StrCat(rng_state_index_++)},
{"{{RNG_STATE_PTR}}",
absl::StrCat("reinterpret_cast<uint64_t*>(",
GetBufferAllocationString(rng_thunk.state_buffer()),
")")}});
}
absl::StatusOr<std::string>
ThunkProtoExecutionDeserializer::GetRngSeedThunkRunImpl(
const xla::cpu::ThunkProto& thunk) {
if (!thunk.has_rng_seed_thunk()) {
return xla::Internal(
"RngSeed thunk was expected when getting thunk run implementation.");
}
const xla::cpu::RngSeedThunkProto& rng_seed_thunk = thunk.rng_seed_thunk();
absl::string_view rng_seed_thunk_invocation_format = R"(
// Rng Seed Thunk
{
uint64_t seed = static_cast<uint64_t>(run_options->rng_seed());
if (seed == 0) {
static thread_local std::random_device rd;
static thread_local std::mt19937_64 gen(rd());
std::uniform_int_distribution<uint64_t> distrib(1, std::numeric_limits<uint64_t>::max());
seed = distrib(gen);
}
*reinterpret_cast<uint64_t*>({{RNG_SEED_PTR}}) = seed;
})";
return absl::StrReplaceAll(
rng_seed_thunk_invocation_format,
{{"{{RNG_SEED_PTR}}",
absl::StrCat("reinterpret_cast<uint64_t*>(",
GetBufferAllocationString(rng_seed_thunk.dest_buffer()),
")")}});
}
absl::StatusOr<std::string>
ThunkProtoExecutionDeserializer::GetCallThunkRunImpl(
const xla::cpu::ThunkProto& thunk) {
if (!thunk.has_call_thunk()) {
return xla::Internal(
"Calls thunk was expected when getting thunk run implementation.");
}
const xla::cpu::CallThunkProto& call_thunk = thunk.call_thunk();
absl::string_view call_thunk_invocation_format = R"(
// Call Thunk
{
{{CALL_THUNK_IMPL}}
})";
TF_ASSIGN_OR_RETURN(
std::string call_thunk_impl,
ThunkSpecificRunImplFromThunkSequence(call_thunk.called_sequence()));
return absl::StrReplaceAll(call_thunk_invocation_format,
{{"{{CALL_THUNK_IMPL}}", call_thunk_impl}});
}
absl::StatusOr<std::string>
ThunkProtoExecutionDeserializer::GetCopyThunkRunImpl(
const xla::cpu::ThunkProto& thunk) {
// IMPORTANT(basioli): tfcompiled models should always emit llvm kernels for
// copy thunks. Here we emit just a memcpy. This is done exclusively for copy
// thunks that get created by the sort thunk.
if (!thunk.has_copy_thunk()) {
return xla::Internal(
"Copy thunk was expected when getting thunk run implementation.");
}
const xla::cpu::CopyThunkProto& copy_thunk = thunk.copy_thunk();
TF_ASSIGN_OR_RETURN(
auto input_shape,
xla::Shape::FromProto(copy_thunk.src_buffer_shape().shape()));
TF_ASSIGN_OR_RETURN(
auto output_shape,
xla::Shape::FromProto(copy_thunk.dst_buffer_shape().shape()));
if (input_shape != output_shape) {
return xla::Internal(
"Copy thunk has input shape %s and output shape %s that are not the "
"same.",
input_shape.ToString(true), output_shape.ToString(true));
}
absl::string_view copy_thunk_invocation_format = R"(
// Copy Thunk
{
std::memcpy({{OUTPUT_PTR}}, {{INPUT_PTR}}, {{SIZE}});
})";
return absl::StrReplaceAll(
copy_thunk_invocation_format,
{{"{{OUTPUT_PTR}}",
absl::StrCat(
"reinterpret_cast<char*>(",
GetBufferAllocationString(copy_thunk.dst_buffer_shape().slice()),
")")},
{"{{INPUT_PTR}}",
absl::StrCat(
"reinterpret_cast<char*>(",
GetBufferAllocationString(copy_thunk.src_buffer_shape().slice()),
")")},
{"{{SIZE}}",
absl::StrCat(copy_thunk.src_buffer_shape().slice().size())}});
}
absl::StatusOr<std::string>
ThunkProtoExecutionDeserializer::GetSortThunkRunImpl(
const xla::cpu::ThunkProto& thunk) {
if (!thunk.has_sort_thunk()) {
return xla::Internal(
"Sort thunk was expected when getting thunk run implementation.");
}
const xla::cpu::SortThunkProto& sort_thunk = thunk.sort_thunk();
std::vector<std::string> buffers_to_sort;
buffers_to_sort.reserve(sort_thunk.inputs_shapes_size());
std::vector<int32_t> primitive_sizes;
primitive_sizes.reserve(sort_thunk.inputs_shapes_size());
for (const auto& buffer_proto : sort_thunk.inputs_shapes()) {
buffers_to_sort.push_back(
absl::StrCat("reinterpret_cast<std::byte*>(",
GetBufferAllocationString(buffer_proto.slice()), ")"));
primitive_sizes.push_back(xla::ShapeUtil::ByteSizeOfPrimitiveType(
buffer_proto.shape().element_type()));
}
absl::string_view sort_thunk_invocation_format = R"(
// Sort Thunk
{
std::vector<std::byte*> values = {
{{BUFFERS_TO_SORT}}
};
std::vector<size_t> primitive_sizes = {
{{VALUES_PRIMITIVE_TYPE_SIZE_IN_BYTES}}
};
// Type alias compatible with `FunctionLibrary::Comparator`.
using Comparator = void(bool* result, const void* run_options,
const void** params, const void* buffer_table,
const void* status, const void* prof_counters);
Comparator* comparator = reinterpret_cast<Comparator*>(
{{SORT_FUNCTION_NAME}});
absl::AnyInvocable<bool(const void** data)> less_than =
[comparator](const void** data) {
bool result;
(*comparator)(&result, nullptr, data, nullptr, nullptr, nullptr);
ABSL_ANNOTATE_MEMORY_IS_INITIALIZED(&result, sizeof(result));
return result;
};
xla::cpu::internal::SortInplace(
{
{{HIGHER_DIMENSIONS}},
{{SORT_DIMENSION_ELEMENTS}},
{{LOWER_DIMENSIONS}}
},
values, primitive_sizes, {{IS_STABLE}}, &less_than);
})";
TF_ASSIGN_OR_RETURN(
auto keys_shape,
xla::Shape::FromProto(sort_thunk.inputs_shapes(0).shape()));
// Normalize the shape and the dimension to sort.
xla::Shape normalized_keys_shape =
xla::ShapeUtil::MakeShapeWithDescendingLayoutAndSamePhysicalLayout(
keys_shape);
auto logical_to_physical =
xla::LayoutUtil::MakeLogicalToPhysical(keys_shape.layout());
TF_RET_CHECK(sort_thunk.dimension() < logical_to_physical.size());
int64_t physical_dimension_to_sort =
logical_to_physical[sort_thunk.dimension()];
int64_t sort_dimension_elements =
normalized_keys_shape.dimensions(physical_dimension_to_sort);
int64_t higher_dimensions = 1;
for (int64_t i = 0; i < physical_dimension_to_sort; ++i) {
higher_dimensions *= normalized_keys_shape.dimensions(i);
}
int64_t lower_dimensions = 1;
for (int64_t i = normalized_keys_shape.dimensions().size() - 1;
i > physical_dimension_to_sort; --i) {
lower_dimensions *= normalized_keys_shape.dimensions(i);
}
return absl::StrReplaceAll(
sort_thunk_invocation_format,
{
{"{{HIGHER_DIMENSIONS}}", absl::StrCat(higher_dimensions)},
{"{{SORT_DIMENSION_ELEMENTS}}",
absl::StrCat(sort_dimension_elements)},
{"{{LOWER_DIMENSIONS}}", absl::StrCat(lower_dimensions)},
{"{{SORT_FUNCTION_NAME}}", sort_thunk.comparator_name()},
{"{{BUFFERS_TO_SORT}}", absl::StrJoin(buffers_to_sort, ", ")},
{"{{VALUES_PRIMITIVE_TYPE_SIZE_IN_BYTES}}",
absl::StrJoin(primitive_sizes, ", ")},
{"{{IS_STABLE}}", sort_thunk.is_stable() ? "true" : "false"},
});
}
absl::StatusOr<std::string>
ThunkProtoExecutionDeserializer::GetTopKThunkRunImpl(
const xla::cpu::ThunkProto& thunk) {
if (!thunk.has_top_k_thunk()) {
return xla::Internal(
"TopK thunk was expected when getting thunk run implementation.");
}
const xla::cpu::TopKThunkProto& topk_thunk_proto = thunk.top_k_thunk();
absl::string_view topk_thunk_invocation_format = R"(
// TopK Thunk
{
::xla::cpu::internal::TopK({{BATCH_SIZE}}, {{INPUT_SIZE}}, {{K}},
reinterpret_cast<const float*>({{VALUES_PTR}}),
reinterpret_cast<float*>({{OUTPUT_PTR}}),
reinterpret_cast<int32_t*>({{INDICES_PTR}}));
})";
return absl::StrReplaceAll(
topk_thunk_invocation_format,
{
{"{{BATCH_SIZE}}", absl::StrCat(topk_thunk_proto.batch_size())},
{"{{INPUT_SIZE}}", absl::StrCat(topk_thunk_proto.input_size())},
{"{{K}}", absl::StrCat(topk_thunk_proto.k())},
{"{{VALUES_PTR}}",
GetBufferAllocationString(topk_thunk_proto.values_buffer())},
{"{{OUTPUT_PTR}}",
GetBufferAllocationString(topk_thunk_proto.output_buffer())},
{"{{INDICES_PTR}}",
GetBufferAllocationString(topk_thunk_proto.indices_buffer())},
});
}
absl::StatusOr<std::string>
ThunkProtoExecutionDeserializer::GetKernelThunkRunImpl(
const xla::cpu::ThunkProto& thunk) {
if (!thunk.has_kernel_thunk()) {
return xla::Internal(
"Kernel thunk was expected when getting thunk run implementation.");
}
const xla::cpu::KernelThunkProto& kernel_thunk = thunk.kernel_thunk();
auto get_args_initializer_as_string =
[](const xla::cpu::KernelThunkProto& kernel_thunk) -> std::string {
std::vector<std::string> args_initializer;
for (const auto& buffer_proto : kernel_thunk.arguments_buffers()) {
args_initializer.push_back(absl::StrCat(
"XLA_CPU_KernelArg{", GetBufferAllocationString(buffer_proto.slice()),
", ", buffer_proto.slice().size(), "}"));
}
for (const auto& buffer_proto : kernel_thunk.results_buffers()) {
args_initializer.push_back(absl::StrCat(
"XLA_CPU_KernelArg{", GetBufferAllocationString(buffer_proto.slice()),
", ", buffer_proto.slice().size(), "}"));
}
return absl::StrCat("{", absl::StrJoin(args_initializer, ", "), "}");
};
// Execute in block so we don't have to worry about naming for now
absl::string_view kernel_invocation_format = R"(
// Kernel Thunk
{
std::array<XLA_CPU_KernelArg, {{NUM_ARGS}}> args = {{ARGS_INITIALIZER}};
XLA_CPU_NumWorkGroups kernel_thread_dims = {
{{THREAD_DIM_X}},
{{THREAD_DIM_Y}},
{{THREAD_DIM_Z}},
};
for (uint64_t z = 0; z < {{THREAD_DIM_Z}}; ++z) {
for (uint64_t y = 0; y < {{THREAD_DIM_Y}}; ++y) {
for (uint64_t x = 0; x < {{THREAD_DIM_X}}; ++x) {
XLA_CPU_WorkGroupId kernel_thread = {x, y, z};
XLA_CPU_KernelCallFrame call_frame = {
&kernel_thread_dims, &kernel_thread, args.size(), args.data()};
XLA_CPU_KernelError* error = (*{{KERNEL_NAME}})(&call_frame);
if (ABSL_PREDICT_FALSE(error != nullptr)) {
return false;
}
}
}
}
})";
return absl::StrReplaceAll(
kernel_invocation_format,
{
{"{{NUM_ARGS}}",
absl::StrCat(kernel_thunk.arguments_buffers().size() +
kernel_thunk.results_buffers().size())},
{"{{ARGS_INITIALIZER}}",
get_args_initializer_as_string(kernel_thunk)},
{"{{THREAD_DIM_X}}", absl::StrCat(kernel_thunk.num_workgroups().x())},
{"{{THREAD_DIM_Y}}", absl::StrCat(kernel_thunk.num_workgroups().y())},
{"{{THREAD_DIM_Z}}", absl::StrCat(kernel_thunk.num_workgroups().z())},
{"{{KERNEL_NAME}}", kernel_thunk.kernel_name()},
});
}
absl::StatusOr<std::string>
ThunkProtoExecutionDeserializer::GetConditionalThunkRunImpl(
const xla::cpu::ThunkProto& thunk) {
if (!thunk.has_conditional_thunk()) {
return xla::Internal(
"Conditional thunk was expected when getting thunk run "
"implementation.");
}
const xla::cpu::ConditionalThunkProto& conditional_thunk =
thunk.conditional_thunk();
std::vector<std::string> conditional_thunk_branches;
conditional_thunk_branches.reserve(conditional_thunk.branch_sequences_size());
for (const auto& branch_sequence : conditional_thunk.branch_sequences()) {
TF_ASSIGN_OR_RETURN(conditional_thunk_branches.emplace_back(),
ThunkSpecificRunImplFromThunkSequence(branch_sequence));
}
absl::string_view branch_execution_format = R"(
case {{CASE_INDEX}}: {
{{BRANCH_EXECUTION}}
break;
}
)";
std::vector<std::string> branch_execution_impls;
branch_execution_impls.reserve(conditional_thunk_branches.size());
for (size_t i = 0; i < conditional_thunk_branches.size(); ++i) {
branch_execution_impls.push_back(absl::StrReplaceAll(
branch_execution_format,
{
{"{{CASE_INDEX}}", absl::StrCat(i)},
{"{{BRANCH_EXECUTION}}", conditional_thunk_branches[i]},
}));
}
absl::string_view conditional_thunk_invocation_format = R"(
// Conditional Thunk
{
size_t branch_index = {{BRANCH_INDEX}};
CHECK(branch_index < {{NUM_BRANCHES}}) << "branch_index is out of bounds";
switch (branch_index) {
{{BRANCH_EXECUTIONS}}
}
})";
auto get_branch_index =
[](const xla::buffer_assignment::BufferAllocationSliceProto&
branch_index_buffer) -> absl::StatusOr<std::string> {
if (branch_index_buffer.size() == sizeof(bool)) {
return absl::StrCat("*reinterpret_cast<bool*>(",
GetBufferAllocationString(branch_index_buffer),
") ? 0 : 1");
}
if (branch_index_buffer.size() == sizeof(int32_t)) {
return absl::StrCat("*reinterpret_cast<int32_t*>(",
GetBufferAllocationString(branch_index_buffer), ")");
}
return xla::Internal("Unsupported branch index buffer size %d",
branch_index_buffer.size());
};
TF_ASSIGN_OR_RETURN(
std::string branch_index,
get_branch_index(conditional_thunk.branch_index_buffer()));
return absl::StrReplaceAll(
conditional_thunk_invocation_format,
{
{"{{BRANCH_INDEX}}", branch_index},
{"{{NUM_BRANCHES}}", absl::StrCat(branch_execution_impls.size())},
{"{{BRANCH_EXECUTIONS}}",
absl::StrJoin(branch_execution_impls, "\n")},
});
}
absl::StatusOr<std::string>
ThunkProtoExecutionDeserializer::GetForLoopThunkRunImpl(
const xla::cpu::WhileThunkProto& while_thunk) {
if (!while_thunk.trip_count().contains_value()) {
return xla::Internal("While thunk is missing trip count.");
}
int64_t trip_count = while_thunk.trip_count().value();
absl::string_view for_loop_thunk_invocation_format = R"(
// For Loop Thunk
{
for (int64_t loop_counter = 0; loop_counter < {{TRIP_COUNT}}; ++loop_counter) {
{{BODY_EXECUTION}};
}
}
)";
TF_ASSIGN_OR_RETURN(
std::string body_execution,
ThunkSpecificRunImplFromThunkSequence(while_thunk.body_sequence()));
return absl::StrReplaceAll(for_loop_thunk_invocation_format,
{
{"{{TRIP_COUNT}}", absl::StrCat(trip_count)},
{"{{BODY_EXECUTION}}", body_execution},
});
}
absl::StatusOr<std::string>
ThunkProtoExecutionDeserializer::GetWhileLoopThunkRunImpl(
const xla::cpu::WhileThunkProto& while_thunk) {
if (while_thunk.trip_count().contains_value()) {
return xla::Internal("While loop path should not have trip count.");
}
absl::string_view while_loop_thunk_invocation_format = R"(
// While Loop Thunk
{
{{CONDITION_EXECUTION}};
while ({{CONDITION}}) {
{{BODY_EXECUTION}};
{{CONDITION_EXECUTION}};
}
}
)";
TF_ASSIGN_OR_RETURN(
std::string body_execution,
ThunkSpecificRunImplFromThunkSequence(while_thunk.body_sequence()));
TF_ASSIGN_OR_RETURN(
std::string condition_execution,
ThunkSpecificRunImplFromThunkSequence(while_thunk.cond_sequence()));
return absl::StrReplaceAll(
while_loop_thunk_invocation_format,
{
{"{{CONDITION}}",
absl::StrCat("*reinterpret_cast<bool*>(",
GetBufferAllocationString(while_thunk.cond_buffer()),
")")},
{"{{BODY_EXECUTION}}", body_execution},
{"{{CONDITION_EXECUTION}}", condition_execution},
});
}
absl::StatusOr<std::string>
ThunkProtoExecutionDeserializer::GetWhileThunkRunImpl(
const xla::cpu::ThunkProto& thunk) {
if (!thunk.has_while_thunk()) {
return xla::Internal(
"While thunk was expected when getting thunk run implementation.");
}
const xla::cpu::WhileThunkProto& while_thunk = thunk.while_thunk();
if (!while_thunk.trip_count().contains_value()) {
return GetWhileLoopThunkRunImpl(while_thunk);
}
return GetForLoopThunkRunImpl(while_thunk);
}
absl::StatusOr<std::string>
ThunkProtoExecutionDeserializer::CppDataTypeFromXlaType(
xla::PrimitiveType xla_type) {
switch (xla_type) {
case xla::F16:
return "Eigen::half";
case xla::F32:
return "float";
case xla::F64:
return "double";
case xla::C64:
return "std::complex<float>";
case xla::C128:
return "std::complex<double>";
case xla::S32:
return "int32_t";
default:
return xla::Internal("Unsupported xla type: %d", xla_type);
}
}
} // namespace tfcompile
} // namespace tensorflow
tensorflow/compiler/aot/thunk_proto_execution_deserializer.h
+109
View File
@@ -0,0 +1,109 @@
/* Copyright 2025 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_COMPILER_AOT_THUNK_PROTO_EXECUTION_DESERIALIZER_H_
#define TENSORFLOW_COMPILER_AOT_THUNK_PROTO_EXECUTION_DESERIALIZER_H_
#include <cstdint>
#include <string>
#include "absl/status/statusor.h"
#include "xla/backends/cpu/runtime/convolution_dims.h"
#include "xla/backends/cpu/runtime/thunk.pb.h"
#include "xla/debug_options_flags.h"
#include "xla/service/cpu/executable.pb.h"
#include "xla/xla_data.pb.h"
namespace tensorflow {
namespace tfcompile {
// Helper class for deserializing the contents of specific thunks into C++ code
// that is used to codegen the `Run` method of the tfcompiled models.
class ThunkProtoExecutionDeserializer {
public:
explicit ThunkProtoExecutionDeserializer()
: xla_cpu_multi_thread_eigen_(
xla::GetDebugOptionsFromFlags().xla_cpu_multi_thread_eigen()) {}
absl::StatusOr<std::string> GetThunkSpecificRunImpl(
const xla::cpu::CompilationResultProto& proto) &&;
absl::StatusOr<std::string> ThunkSpecificRunImplFromThunkSequence(
const xla::cpu::ThunkSequenceProto& thunk_sequence_proto);
protected:
absl::StatusOr<std::string> GetMatmulFunction(xla::PrimitiveType xla_type);
absl::StatusOr<std::string> GetDotThunkRunImpl(
const xla::cpu::ThunkProto& thunk);
absl::StatusOr<std::string> GetConvolutionFunction(
xla::PrimitiveType xla_type);
absl::StatusOr<std::string> GetConvolution2DRunImpl(
const xla::cpu::ConvolutionThunkProto& convolution_thunk,
const xla::cpu::ConvolutionCanonicalDims& canonical_dims);
absl::StatusOr<std::string> GetConvolutionFusionThunkRunImpl(
const xla::cpu::ThunkProto& thunk);
absl::StatusOr<std::string> GetRngGetAndUpdateStateThunkRunImpl(
const xla::cpu::ThunkProto& thunk);
absl::StatusOr<std::string> GetRngSeedThunkRunImpl(
const xla::cpu::ThunkProto& thunk);
absl::StatusOr<std::string> GetCallThunkRunImpl(
const xla::cpu::ThunkProto& thunk);
absl::StatusOr<std::string> GetKernelThunkRunImpl(
const xla::cpu::ThunkProto& thunk);
absl::StatusOr<std::string> GetCopyThunkRunImpl(
const xla::cpu::ThunkProto& thunk);
absl::StatusOr<std::string> GetConditionalThunkRunImpl(
const xla::cpu::ThunkProto& thunk);
absl::StatusOr<std::string> GetForLoopThunkRunImpl(
const xla::cpu::WhileThunkProto& while_thunk);
absl::StatusOr<std::string> GetWhileLoopThunkRunImpl(
const xla::cpu::WhileThunkProto& while_thunk);
absl::StatusOr<std::string> GetWhileThunkRunImpl(
const xla::cpu::ThunkProto& thunk);
absl::StatusOr<std::string> GetSortThunkRunImpl(
const xla::cpu::ThunkProto& thunk);
absl::StatusOr<std::string> GetTopKThunkRunImpl(
const xla::cpu::ThunkProto& thunk);
absl::StatusOr<std::string> CppDataTypeFromXlaType(
xla::PrimitiveType xla_type);
private:
// The index of the next rng state to use when deserializing the rng state
// from the ThunkProto.
int64_t rng_state_index_ = 0;
bool xla_cpu_multi_thread_eigen_;
};
} // namespace tfcompile
} // namespace tensorflow
#endif // TENSORFLOW_COMPILER_AOT_THUNK_PROTO_EXECUTION_DESERIALIZER_H_
tensorflow/compiler/jit/BUILD
+2115
View File
File diff suppressed because it is too large Load Diff
tensorflow/compiler/jit/build_xla_ops_pass.cc
+624
View File
@@ -0,0 +1,624 @@
/* Copyright 2017 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/compiler/jit/build_xla_ops_pass.h"
#include "absl/algorithm/container.h"
#include "absl/strings/str_cat.h"
#include "absl/strings/str_join.h"
#include "tensorflow/cc/framework/ops.h"
#include "tensorflow/cc/framework/scope_internal.h"
#include "tensorflow/cc/ops/array_ops.h"
#include "tensorflow/cc/ops/const_op.h"
#include "tensorflow/cc/ops/control_flow_ops.h"
#include "tensorflow/cc/ops/functional_ops.h"
#include "tensorflow/cc/ops/logging_ops.h"
#include "tensorflow/compiler/jit/defs.h"
#include "tensorflow/compiler/jit/device_util.h"
#include "tensorflow/compiler/jit/encapsulate_subgraphs_pass.h"
#include "tensorflow/compiler/jit/flags.h"
#include "tensorflow/compiler/jit/xla_cluster_util.h"
#include "tensorflow/compiler/tf2xla/cc/ops/xla_jit_ops.h"
#include "tensorflow/compiler/tf2xla/xla_op_registry.h"
#include "xla/status_macros.h"
#include "tensorflow/core/common_runtime/function.h"
#include "tensorflow/core/common_runtime/graph_constructor.h"
#include "tensorflow/core/common_runtime/optimization_registry.h"
#include "tensorflow/core/framework/graph_def_util.h"
#include "tensorflow/core/framework/memory_types.h"
#include "tensorflow/core/framework/node_def_builder.h"
#include "tensorflow/core/framework/node_def_util.h"
#include "tensorflow/core/graph/algorithm.h"
#include "tensorflow/core/graph/graph.h"
#include "tensorflow/core/lib/core/status.h"
#include "tensorflow/core/lib/hash/hash.h"
#include "tensorflow/core/public/version.h"
#include "tensorflow/core/util/dump_graph.h"
namespace tensorflow {
namespace {
struct DebuggingOpts {
// If true, insert Print nodes to print every output from an XLA cluster.
bool print_outputs;
// If true, insert CheckNumerics nodes for every floating point typed input to
// an XLA cluster.
bool check_input_numerics;
// If true, insert CheckNumerics nodes for every floating point typed output
// from an XLA cluster.
bool check_output_numerics;
};
void MoveOutgoingEdges(Graph* g, Node* old_node, Node* new_node) {
std::vector<const Edge*> out_edges(old_node->out_edges().begin(),
old_node->out_edges().end());
for (const Edge* edge : out_edges) {
// TODO(sanjoy): This does not update NodeDef inputs. To be able to update
// NodeDef inputs we first need to fix encapsulate_subgraphs_pass to fix up
// the NodeDef inputs to the function call nodes.
g->AddEdge(new_node, edge->src_output(), edge->dst(), edge->dst_input());
g->RemoveEdge(edge);
}
}
// Returns a data value that is dead iff `control` is dead.
Output ControlToData(const Scope& scope, Node* control) {
// The choice of data type here is important.
//
// We implement a "control merge", which is a control edge that is alive if
// either of two nodes (denoted as A and B below) are alive, in the following
// manner:
//
// A --ctrl--> Const0 --data--> Merge --data--> Identity
// ^ |
// | ctrl
// B --ctrl--> Const1 --data-----+ |
// v
// ***
//
// where *** denotes the merged control output.
//
// We want everything starting from Const{0/1} to Identity to either wholly
// live on the host or wholly live on device so we need to pick a data type
// that is either consistently assigned to the device (e.g. float) or
// consistently assigned to the host (e.g. int32). We should *not* pick a
// data type that partly placed on the host and partly on the device
// (e.g. bool constants are placed on the device but bool Identity is placed
// on the host).
Output data = ops::Const(scope.WithOpName("ctrl_as_data"),
Tensor(DT_INT32, TensorShape({0})));
scope.graph()->AddControlEdge(control, data.node());
return Output(data.node());
}
// Returns an operation that can be control-depended on that is dead iff `data`
// is dead.
Operation DataToControl(const Scope& scope, Output data) {
return Operation(
ops::Identity(scope.WithOpName("data_as_ctrl"), data).node());
}
// Replaces each outgoing edge from `old_node` with a merge node that merges in
// the corresponding output from `new_node`.
void MergeOutgoingDataEdges(const Scope& s, Node* old_node, Node* new_node,
absl::string_view cluster_name,
const DebuggingOpts& debugging_opts) {
if (!s.status().ok()) {
return;
}
std::vector<Output> merged_outputs(old_node->num_outputs(), Output(nullptr));
std::vector<const Edge*> data_edges;
absl::c_copy_if(old_node->out_edges(), std::back_inserter(data_edges),
[](const Edge* e) { return !e->IsControlEdge(); });
for (const Edge* e : data_edges) {
int oidx = e->src_output();
Output merged_output = merged_outputs[oidx];
if (merged_output.node() == nullptr) {
Output new_output(new_node, oidx);
if (debugging_opts.print_outputs) {
std::string cpu_device = "/job:localhost/replica:0/task:0/device:CPU:0";
ops::Print print_op(s.WithOpName("print_", oidx)
.WithDevice(cpu_device)
.WithAssignedDevice(cpu_device),
new_output, {new_output},
ops::Print::Attrs{}
.Message(absl::StrCat("output ", oidx, " from ",
old_node->name(), " is "))
.FirstN(1000)
.Summarize(-1));
new_output = print_op;
}
if (debugging_opts.check_output_numerics &&
DataTypeIsFloating(new_output.type())) {
ops::CheckNumerics check_numerics_op(
s.WithOpName("check_output_", oidx)
.WithDevice(new_node->requested_device())
.WithAssignedDevice(new_node->assigned_device_name()),
new_output,
absl::StrCat("CheckNumerics failed for output ", oidx, "(",
new_output.name(), ") from cluster ", cluster_name));
new_output = check_numerics_op;
}
ops::_XlaMerge xla_merge_op(s.WithOpName("merge_oidx_", oidx),
Output(old_node, oidx), new_output);
merged_output = merged_outputs[oidx] = xla_merge_op.output;
}
Node* dst = e->dst();
int dst_idx = e->dst_input();
s.graph()->RemoveEdge(e);
s.graph()->AddEdge(merged_output.node(), merged_output.index(), dst,
dst_idx);
}
}
// Replaces each control successor of `old_node` to execute whenever either
// `old_node` or `new_node` is executed.
void MergeOutgoingControlEdges(const Scope& s, Node* old_node, Node* new_node) {
if (!s.status().ok()) {
return;
}
std::vector<const Edge*> ctrl_edges;
absl::c_copy_if(old_node->out_edges(), std::back_inserter(ctrl_edges),
[](const Edge* e) { return e->IsControlEdge(); });
if (ctrl_edges.empty()) {
return;
}
if (ctrl_edges.size() == 1 && ctrl_edges.front()->dst()->IsSink()) {
// Avoid creating a Merge node if we can just add an edge to _SINK
// instead.
s.graph()->AddControlEdge(new_node, s.graph()->sink_node());
return;
}
// We can't merge control edges directly so we instead first "convert" them to
// normal values that can be merged, merge the values and then "convert" the
// merged value back into control.
//
// NB! We need to copy out the outgoing control edges before constructing
// old_ctrl_as_data otherwise the control edge from old_node to the constant
// in ControlToData will be present in ctrl_edges.
Output old_ctrl_as_data = ControlToData(s, old_node);
Output new_ctrl_as_data = ControlToData(s, new_node);
ops::Merge ctrl_merge_as_data(s.WithOpName("ctrl_merge"),
{old_ctrl_as_data, new_ctrl_as_data});
Operation ctrl_merge = DataToControl(s, ctrl_merge_as_data.output);
for (const Edge* e : ctrl_edges) {
s.graph()->AddControlEdge(ctrl_merge.node(), e->dst());
s.graph()->RemoveControlEdge(e);
}
}
struct XlaClusterInfo {
std::vector<Output> constant_inputs;
std::vector<Output> non_constant_inputs;
std::vector<Output> resource_inputs;
NameAttrList function;
};
Output IncomingEdgeAsOutput(const Edge* e) {
return Output(e->src(), e->src_output());
}
absl::Status GetXlaClusterInfo(Node* n, XlaClusterInfo* result) {
int num_constant_inputs, num_resource_inputs;
TF_RETURN_IF_ERROR(
GetNodeAttr(n->attrs(), kXlaNumConstantArgsAttr, &num_constant_inputs));
TF_RETURN_IF_ERROR(
GetNodeAttr(n->attrs(), kXlaNumResourceArgsAttr, &num_resource_inputs));
if (num_constant_inputs < 0 || num_resource_inputs < 0 ||
num_constant_inputs + num_resource_inputs > n->num_inputs()) {
return absl::InvalidArgumentError(
"Invalid number of constant/resource arguments to XLA kernel.");
}
int num_non_constant_inputs =
n->num_inputs() - num_constant_inputs - num_resource_inputs;
std::vector<const Edge*> input_edges_vector;
TF_RETURN_IF_ERROR(n->input_edges(&input_edges_vector));
absl::Span<const Edge*> input_edges(input_edges_vector);
absl::c_transform(input_edges.subspan(0, num_constant_inputs),
std::back_inserter(result->constant_inputs),
IncomingEdgeAsOutput);
absl::c_transform(
input_edges.subspan(num_constant_inputs, num_non_constant_inputs),
std::back_inserter(result->non_constant_inputs), IncomingEdgeAsOutput);
absl::c_transform(
input_edges.subspan(num_constant_inputs + num_non_constant_inputs,
num_resource_inputs),
std::back_inserter(result->resource_inputs), IncomingEdgeAsOutput);
result->function.set_name(n->type_string());
*result->function.mutable_attr() = n->def().attr();
return absl::OkStatus();
}
absl::Status CopyIncomingControlEdges(Graph* g, Node* from, Node* to) {
for (const Edge* e : from->in_edges()) {
if (e->IsControlEdge()) {
g->AddControlEdge(e->src(), to);
}
}
return absl::OkStatus();
}
void RemoveAllIncomingControlEdges(Graph* g, Node* n) {
std::vector<const Edge*> incoming_ctrl_edges;
absl::c_copy_if(n->in_edges(), std::back_inserter(incoming_ctrl_edges),
[](const Edge* e) { return e->IsControlEdge(); });
for (const Edge* e : incoming_ctrl_edges) {
g->RemoveControlEdge(e);
}
}
// Returns true (into `result`) if a node placed on `device` must be compiled.
absl::Status DeviceRequiresCompilation(
const jit::DeviceInfoCache& device_info_cache, jit::DeviceId device,
bool* result) {
const XlaOpRegistry::DeviceRegistration* registration =
device_info_cache.GetCompilationDevice(device);
*result = registration->autoclustering_policy ==
XlaOpRegistry::AutoclusteringPolicy::kAlways;
return absl::OkStatus();
}
// Replaces `n` with a `PartitionedCall` op that calls the same function.
absl::StatusOr<Node*> ReplaceFunctionCallWithPartitionedCall(
const GraphOptimizationPassOptions& options,
const FunctionLibraryDefinition& flib_def, Node* n, Graph* g,
const NameAttrList& func, const Scope& root) {
std::string config_string =
options.session_options->config.SerializeAsString();
int input_count = absl::c_count_if(
n->in_edges(), [](const Edge* e) { return !e->IsControlEdge(); });
std::vector<Output> args(input_count);
for (const Edge* e : n->in_edges()) {
if (!e->IsControlEdge()) {
args[e->dst_input()] = Output(e->src(), e->src_output());
}
}
// In theory we can use PartitionedCall if the XLA cluster does not have any
// stateful operations. However, for now we choose to be conservative since
// we don't have any evidence that choosing a stateless partitioned call helps
// for performance.
ops::StatefulPartitionedCall call(
root.WithOpName("stateful_partitioned_call"), args, n->output_types(),
func, ops::StatefulPartitionedCall::Attrs{}.ConfigProto(config_string));
for (const Edge* e : n->in_edges()) {
if (e->IsControlEdge()) {
g->AddControlEdge(e->src(), call.operation.node());
}
}
std::vector<const Edge*> edges_to_delete;
for (const Edge* e : n->out_edges()) {
edges_to_delete.push_back(e);
if (e->IsControlEdge()) {
g->AddControlEdge(call.operation.node(), e->dst());
} else {
g->AddEdge(call.operation.node(), e->src_output(), e->dst(),
e->dst_input());
}
}
for (const Edge* e : edges_to_delete) {
g->RemoveEdge(e);
}
g->RemoveNode(n);
return call.operation.node();
}
absl::StatusOr<jit::DeviceId> InferDeviceForCluster(
jit::DeviceInfoCache* device_info_cache, Node* n,
const std::string& function_name,
const FunctionLibraryDefinition& flib_def) {
const FunctionDef* func_def = flib_def.Find(function_name);
TF_RET_CHECK(func_def) << "Could not find " << function_name;
jit::DeviceSet device_set;
for (const NodeDef& ndef : func_def->node_def()) {
VLOG(3) << ndef.DebugString();
if (!ndef.device().empty()) {
TF_ASSIGN_OR_RETURN(jit::DeviceId device_id,
device_info_cache->GetIdFor(ndef.device()));
device_set.Insert(device_id);
}
}
if (!n->assigned_device_name().empty()) {
// TODO(sanjoy): We need this because EncapsulateSubgraphsPass drops device
// assignment when constant folding. We should fix EncapsulateSubgraphsPass
// instead.
TF_ASSIGN_OR_RETURN(jit::DeviceId device_id,
device_info_cache->GetIdFor(n->assigned_device_name()));
device_set.Insert(device_id);
}
TF_ASSIGN_OR_RETURN(jit::DeviceId result,
PickDeviceForXla(*device_info_cache, device_set,
/*allow_mixing_unknown_and_cpu=*/true));
VLOG(2) << "For " << function_name << " PickDeviceForXla("
<< device_info_cache->DebugString(device_set) << ") -> "
<< device_info_cache->GetNameFor(result);
return result;
}
std::vector<Output> GetXlaRunArgs(const Scope& s,
const XlaClusterInfo& cluster_info,
const DebuggingOpts& debugging_opts) {
std::vector<Output> xla_run_args;
xla_run_args.reserve(cluster_info.non_constant_inputs.size() +
cluster_info.resource_inputs.size());
int input_idx = 0;
for (const Output& o : cluster_info.non_constant_inputs) {
if (debugging_opts.check_input_numerics && DataTypeIsFloating(o.type())) {
ops::CheckNumerics check_numerics_op(
s.WithOpName("check_input_", input_idx), o,
absl::StrCat("CheckNumerics failed for input ", input_idx, "(",
o.name(), ") into ", cluster_info.function.name()));
xla_run_args.push_back(check_numerics_op);
} else {
xla_run_args.push_back(o);
}
input_idx++;
}
absl::c_copy(cluster_info.resource_inputs, std::back_inserter(xla_run_args));
return xla_run_args;
}
absl::StatusOr<MemoryTypeVector> GetOutputMemoryTypes(const Scope& root,
Node* n) {
MemoryTypeVector input_mtypes, output_mtypes;
DeviceType device_type("");
TF_RETURN_IF_ERROR(
DeviceNameToDeviceType(n->assigned_device_name(), &device_type));
TF_RETURN_IF_ERROR(MemoryTypesForNode(root.graph()->op_registry(),
device_type, n->def(), &input_mtypes,
&output_mtypes));
return output_mtypes;
}
// Predicate INT32 typed inputs to `n` on the deadness of
// `predicate_as_control`.
//
// This is a performance optimization. Since INT32 arguments to a
// PartitionedCall are placed on the host, a producer that produces them on the
// device will incur a D2H copy, even if the PartitionedCall is not executed
// (i.e. even if we choose to execute the XLA compiled computation via _XlaRun).
// To prevent this, we add control dependencies to make the int32 input edges
// into the PartitionedCall dead. With this change the D2H copy only happens if
// the PartitionedCall is actually executed.
absl::Status PredicateInt32Inputs(const Scope& root, Node* n,
Operation predicate_as_control) {
std::vector<Output> int32_inputs;
std::vector<int> int32_inputs_input_idxs;
for (const Edge* e : n->in_edges()) {
if (e->IsControlEdge()) {
continue;
}
if (e->src()->output_type(e->src_output()) == DT_INT32) {
TF_ASSIGN_OR_RETURN(MemoryTypeVector source_output_mem_types,
GetOutputMemoryTypes(root, e->src()));
if (source_output_mem_types[e->src_output()] == DEVICE_MEMORY) {
int32_inputs.push_back(Output(e->src(), e->src_output()));
int32_inputs_input_idxs.push_back(e->dst_input());
}
}
}
if (int32_inputs.empty()) {
return absl::OkStatus();
}
// Create a single IdentityN that is dead if and only if
// `predicate_as_control` is dead.
//
// IdentityN is also special in that, unlike `Identity`, it does not place
// int32 inputs in host memory. Placing int32 inputs in host memory would
// defeat the purpose of adding this indirection.
ops::IdentityN identity_n(root.WithOpName("int32_id_n"), int32_inputs);
root.graph()->AddControlEdge(predicate_as_control.node(),
identity_n.operation.node());
for (int i = 0, end = int32_inputs.size(); i < end; i++) {
TF_RETURN_IF_ERROR(root.graph()->UpdateEdge(identity_n[i].node(), i, n,
int32_inputs_input_idxs[i]));
}
return absl::OkStatus();
}
absl::Status ReplaceNodeWithXlaCompileAndXlaRun(
jit::DeviceInfoCache* device_info_cache,
const GraphOptimizationPassOptions& options,
const FunctionLibraryDefinition& flib_def, bool lazy_compilation_enabled,
const DebuggingOpts& debugging_opts, Graph* g, Node* n) {
XlaClusterInfo cluster_info;
TF_RETURN_IF_ERROR(GetXlaClusterInfo(n, &cluster_info));
TF_ASSIGN_OR_RETURN(
jit::DeviceId device,
InferDeviceForCluster(device_info_cache, n, cluster_info.function.name(),
flib_def));
bool requires_compilation;
TF_RETURN_IF_ERROR(DeviceRequiresCompilation(*device_info_cache, device,
&requires_compilation));
if (!lazy_compilation_enabled) {
requires_compilation = true;
}
std::string device_name_str =
std::string(device_info_cache->GetNameFor(device));
absl::Status status;
Scope root = NewInternalScope(g, &status, /*refiner=*/nullptr)
.NewSubScope(n->name())
.WithDevice(n->requested_device())
.WithAssignedDevice(device_name_str);
ops::_XlaCompile xla_compile(root.WithOpName("xla_compile"),
/*constants=*/cluster_info.constant_inputs,
/*args=*/cluster_info.non_constant_inputs,
/*resources=*/cluster_info.resource_inputs,
/*must_compile=*/requires_compilation,
cluster_info.function);
bool has_ref_attr;
TF_RETURN_IF_ERROR(
GetNodeAttr(n->attrs(), kXlaHasReferenceVarsAttr, &has_ref_attr));
xla_compile.operation.node()->AddAttr(kXlaHasReferenceVarsAttr, has_ref_attr);
TF_RETURN_IF_ERROR(
CopyIncomingControlEdges(g, /*from=*/n, /*to=*/xla_compile.key.node()));
std::vector<Output> xla_run_args =
GetXlaRunArgs(root, cluster_info, debugging_opts);
if (requires_compilation) {
// "Strict" compilation: every _XlaCompile invocation must compile the
// cluster.
ops::_XlaRun xla_run(root.WithOpName("xla_run"), xla_run_args,
xla_compile.key, n->output_types());
MoveOutgoingEdges(g, /*old_node=*/n,
/*new_node=*/xla_run.operation.node());
g->RemoveNode(n);
} else {
// "Lazy" compilation: an _XlaCompile invocation may decide not to compile
// the cluster based on profitability heuristics.
// We generate the following graph:
//
// (use_tf_call, use_xla_run) =
// Switch(pred=xla_compile.compilation_successful,
// value=xla_compile.key)
//
// tf_call_outputs = cluster_N(..., ^use_tf_call)
// xla_run_outputs = _XlaRun(..., key=use_xla_run)
// outputs = Merge(tf_call_outputs, xla_run_outputs).
ops::Switch s(root.WithOpName("predicated_compilation_key"),
xla_compile.key, xla_compile.compilation_successful);
Output predicated_compilation_key = s.output_true;
Output inverse_predicated_compilation_key = s.output_false;
ops::_XlaRun xla_run(root.WithOpName("xla_run"), xla_run_args,
predicated_compilation_key, n->output_types());
MergeOutgoingControlEdges(root, /*old_node=*/n,
/*new_node=*/xla_run.operation.node());
MergeOutgoingDataEdges(root, /*old_node=*/n,
/*new_node=*/xla_run.operation.node(),
cluster_info.function.name(), debugging_opts);
TF_RETURN_IF_ERROR(root.status());
// We already have a TensorFlow function call into the cluster -- the
// original node we set out to rewrite. We just wire in the correct control
// deps and we're done.
RemoveAllIncomingControlEdges(g, n);
Operation inverse_predicate_as_control =
DataToControl(root, inverse_predicated_compilation_key);
g->AddControlEdge(inverse_predicate_as_control.node(), n);
n->ClearAttr(kXlaCompiledKernelAttr);
TF_ASSIGN_OR_RETURN(Node* const pco, ReplaceFunctionCallWithPartitionedCall(
options, flib_def, n, g,
cluster_info.function, root));
TF_RETURN_IF_ERROR(
PredicateInt32Inputs(root, pco, inverse_predicate_as_control));
}
return absl::OkStatus();
}
} // namespace
absl::Status BuildXlaOpsPass::Run(const GraphOptimizationPassOptions& options) {
Graph* graph = options.graph->get();
// Copy out the nodes we want to rewrite to avoid modifying the graph while we
// iterate on graph->op_nodes().
std::vector<Node*> xla_compiled_kernels;
absl::c_copy_if(graph->op_nodes(), std::back_inserter(xla_compiled_kernels),
[](const Node* n) {
if (n->IsSend() || n->IsRecv() || n->IsControlFlow()) {
return false;
}
// Only compile nodes that are marked for compilation by the
// compilation-marking pass (via 'attr_name').
return IsXlaCompiledKernel(*n);
});
bool lazy_compilation_enabled =
enable_lazy_compilation_
? *enable_lazy_compilation_
: GetBuildXlaOpsPassFlags()->tf_xla_enable_lazy_compilation;
jit::DeviceInfoCache device_info_cache;
const BuildXlaOpsPassFlags& flags = *GetBuildXlaOpsPassFlags();
DebuggingOpts debugging_opts;
debugging_opts.print_outputs = flags.tf_xla_print_cluster_outputs;
debugging_opts.check_input_numerics =
flags.tf_xla_check_cluster_input_numerics;
debugging_opts.check_output_numerics =
flags.tf_xla_check_cluster_output_numerics;
VLOG(1) << "print_outputs = " << debugging_opts.print_outputs;
VLOG(1) << "check_input_numerics = " << debugging_opts.check_input_numerics;
VLOG(1) << "check_output_numerics = " << debugging_opts.check_output_numerics;
for (Node* n : xla_compiled_kernels) {
TF_RETURN_IF_ERROR(ReplaceNodeWithXlaCompileAndXlaRun(
&device_info_cache, options, *options.flib_def,
lazy_compilation_enabled, debugging_opts, graph, n));
}
if (VLOG_IS_ON(1)) {
DumpGraphToFile("build_xla_ops", *graph, options.flib_def);
}
return absl::OkStatus();
}
} // namespace tensorflow
tensorflow/compiler/jit/build_xla_ops_pass.h
+45
View File
@@ -0,0 +1,45 @@
/* Copyright 2017 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_COMPILER_JIT_BUILD_XLA_OPS_PASS_H_
#define TENSORFLOW_COMPILER_JIT_BUILD_XLA_OPS_PASS_H_
#include "absl/types/optional.h"
#include "tensorflow/core/common_runtime/optimization_registry.h"
#include "tensorflow/core/lib/core/status.h"
namespace tensorflow {
// Replaces TF function calls marked with `_XlaCompiledKernel` with _XlaCompile
// and _XlaRun nodes (which compile and launch, respectively, the corresponding
// HLO module).
class BuildXlaOpsPass : public GraphOptimizationPass {
public:
// If enable_lazy_compilation is not nullopt then *enable_lazy_compilation
// overrides --tf_xla_enable_lazy_compilation flag in deciding whether lazy
// compilation is enabled.
explicit BuildXlaOpsPass(
std::optional<bool> enable_lazy_compilation = std::nullopt)
: enable_lazy_compilation_(enable_lazy_compilation) {}
absl::Status Run(const GraphOptimizationPassOptions& options) override;
private:
std::optional<bool> enable_lazy_compilation_;
};
} // namespace tensorflow
#endif // TENSORFLOW_COMPILER_JIT_BUILD_XLA_OPS_PASS_H_
tensorflow/compiler/jit/build_xla_ops_pass_test.cc
+328
View File
@@ -0,0 +1,328 @@
/* 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/compiler/jit/build_xla_ops_pass.h"
#include "tensorflow/cc/framework/ops.h"
#include "tensorflow/cc/ops/array_ops.h"
#include "tensorflow/cc/ops/resource_variable_ops.h"
#include "tensorflow/cc/ops/standard_ops.h"
#include "tensorflow/compiler/jit/defs.h"
#include "tensorflow/compiler/jit/encapsulate_subgraphs_pass.h"
#include "tensorflow/compiler/jit/node_matchers.h"
#include "tensorflow/compiler/jit/test_util.h"
#include "tensorflow/core/common_runtime/device_factory.h"
#include "tensorflow/core/graph/algorithm.h"
#include "tensorflow/core/lib/core/status_test_util.h"
#include "tensorflow/core/platform/logging.h"
#include "tensorflow/core/platform/test.h"
#include "tensorflow/core/public/session_options.h"
namespace tensorflow {
namespace {
class BuildXlaOpsTest : public ::testing::Test {
protected:
void SetUp() override {
// This is needed to register the XLA_* devices.
CHECK(DeviceFactory::AddDevices(
SessionOptions(), "/job:localhost/replica:0/task:0", &devices_)
.ok());
}
private:
std::vector<std::unique_ptr<Device>> devices_;
};
using ::tensorflow::testing::FindNodeByName;
using ::tensorflow::testing::matchers::Attr;
using ::tensorflow::testing::matchers::CtrlDeps;
using ::tensorflow::testing::matchers::Inputs;
using ::tensorflow::testing::matchers::NodeWith;
using ::tensorflow::testing::matchers::Op;
using ::tensorflow::testing::matchers::Out;
using ::testing::_;
absl::Status BuildXlaOps(const Scope& s, const FunctionDefLibrary& fdef_lib,
std::unique_ptr<Graph>* result) {
auto graph = std::make_unique<Graph>(OpRegistry::Global());
TF_RETURN_IF_ERROR(s.ToGraph(graph.get()));
FunctionLibraryDefinition flib_def(graph->op_registry(), fdef_lib);
// Assign all nodes to the CPU device.
static const char* kCpuDevice = "/job:localhost/replica:0/task:0/cpu:0";
for (Node* n : graph->nodes()) {
if (n->requested_device().empty()) {
n->set_assigned_device_name(kCpuDevice);
} else {
n->set_assigned_device_name(n->requested_device());
}
}
FixupSourceAndSinkEdges(graph.get());
GraphOptimizationPassWrapper wrapper;
GraphOptimizationPassOptions opt_options =
wrapper.CreateGraphOptimizationPassOptions(&graph);
opt_options.flib_def = &flib_def;
BuildXlaOpsPass pass(/*enable_lazy_compilation=*/true);
TF_RETURN_IF_ERROR(pass.Run(opt_options));
VLOG(3) << graph->ToGraphDefDebug().DebugString();
*result = std::move(graph);
return absl::OkStatus();
}
absl::Status MakeXlaCompiledKernel(Graph* graph, const std::string& callee_name,
const std::string& node_name,
int num_constant_args, int num_resource_args,
Node** result) {
NodeDef call_node;
call_node.set_name(node_name);
call_node.set_op(callee_name);
AddNodeAttr(kXlaCompiledKernelAttr, true, &call_node);
AddNodeAttr(kXlaNumConstantArgsAttr, num_constant_args, &call_node);
AddNodeAttr(kXlaNumResourceArgsAttr, num_resource_args, &call_node);
TF_ASSIGN_OR_RETURN(*result, graph->AddNode(call_node));
return absl::OkStatus();
}
absl::Status MakeXlaCompiledKernel(Graph* graph, const std::string& callee_name,
const std::string& node_name,
Node** result) {
return MakeXlaCompiledKernel(graph, callee_name, node_name,
/*num_constant_args=*/0, /*num_resource_args=*/0,
result);
}
Node* MakeWrite(const Scope& scope, Output value_to_write,
const std::string& id) {
Output var_handle = ops::VarHandleOp(scope.WithOpName("Var_" + id), DT_FLOAT,
TensorShape({}));
ops::AssignVariableOp assign_op(scope.WithOpName("Assignee_" + id),
var_handle, value_to_write);
return assign_op.operation.node();
}
Node* MakeWrite(const Scope& scope, const std::string& id) {
return MakeWrite(
scope, ops::Const(scope.WithOpName("ValueToAssign" + id), 1.0f), id);
}
FunctionDefLibrary CreateFunctionDefLibWithConstFunction(
const std::string& name) {
FunctionDefLibrary fdef_lib;
FunctionDef func = FunctionDefHelper::Create(
/*function_name=*/name, /*in_def=*/{}, /*out_def=*/{"out: float"},
/*attr_def*/
{}, /*node_def=*/{FunctionDefHelper::Const("one", 1.0f)},
/*ret_def=*/{{"out", "out:output:0"}});
*fdef_lib.add_function() = std::move(func);
return fdef_lib;
}
TEST_F(BuildXlaOpsTest, ControlDepsPreserved) {
const char* kXlaDeviceName = "/job:worker/replica:0/task:0/device:XLA_CPU:0";
Scope root = Scope::NewRootScope().WithDevice(kXlaDeviceName).ExitOnError();
FunctionDefLibrary fdef_lib =
CreateFunctionDefLibWithConstFunction("cluster_0");
TF_ASSERT_OK(root.graph()->AddFunctionLibrary(fdef_lib));
Node* call;
TF_ASSERT_OK(MakeXlaCompiledKernel(root.graph(), "cluster_0", "C", &call));
call->AddAttr(kXlaHasReferenceVarsAttr, false);
call->set_requested_device(kXlaDeviceName);
Node* write_op = MakeWrite(root, "write");
write_op->AddAttr(kXlaHasReferenceVarsAttr, false);
root.graph()->AddControlEdge(call, write_op);
std::unique_ptr<Graph> graph;
TF_ASSERT_OK(BuildXlaOps(root, fdef_lib, &graph));
Node* write_op_new = FindNodeByName(graph.get(), write_op->name());
ASSERT_NE(write_op_new, nullptr);
EXPECT_THAT(write_op_new, NodeWith(CtrlDeps(NodeWith(Op("_XlaRun")))));
}
TEST_F(BuildXlaOpsTest, CleanFailureOnBogusAttr) {
Scope root = Scope::NewRootScope().ExitOnError();
FunctionDefLibrary fdef_lib =
CreateFunctionDefLibWithConstFunction("cluster_0");
TF_ASSERT_OK(root.graph()->AddFunctionLibrary(fdef_lib));
Node* call;
TF_ASSERT_OK(
MakeXlaCompiledKernel(root.graph(), "cluster_0", "C", 100, 100, &call));
Node* write_op = MakeWrite(root, "write");
root.graph()->AddControlEdge(call, write_op);
std::unique_ptr<Graph> graph;
absl::Status failure_status = BuildXlaOps(root, fdef_lib, &graph);
ASSERT_FALSE(failure_status.ok());
EXPECT_EQ(failure_status.code(), error::INVALID_ARGUMENT);
}
TEST_F(BuildXlaOpsTest, OnNonXlaDevice) {
Scope root = Scope::NewRootScope().ExitOnError();
FunctionDefLibrary fdef_lib =
CreateFunctionDefLibWithConstFunction("cluster_0");
TF_ASSERT_OK(root.graph()->AddFunctionLibrary(fdef_lib));
Node* call;
TF_ASSERT_OK(MakeXlaCompiledKernel(root.graph(), "cluster_0", "C", &call));
TF_ASSERT_OK(root.DoShapeInference(call));
call->AddAttr(kXlaHasReferenceVarsAttr, false);
Node* write_op = MakeWrite(root, Output(call), "write_result");
write_op->AddAttr(kXlaHasReferenceVarsAttr, false);
auto xla_compile = NodeWith(Op("_XlaCompile"), Attr("must_compile", false));
auto predicated_compilation_key =
NodeWith(Op("Switch"), Inputs(Out(0, xla_compile), Out(1, xla_compile)));
auto xla_run =
NodeWith(Op("_XlaRun"), Inputs(Out(1, predicated_compilation_key)));
auto tf_call =
NodeWith(Op("StatefulPartitionedCall"),
CtrlDeps(NodeWith(Op("Identity"),
Inputs(Out(0, predicated_compilation_key)))));
auto merge = NodeWith(Op("_XlaMerge"), Inputs(Out(tf_call), Out(xla_run)));
auto assign_var = NodeWith(Op("AssignVariableOp"), Inputs(_, Out(merge)));
std::unique_ptr<Graph> graph;
TF_ASSERT_OK(BuildXlaOps(root, fdef_lib, &graph));
Node* write_op_new = FindNodeByName(graph.get(), write_op->name());
ASSERT_NE(write_op_new, nullptr);
EXPECT_THAT(write_op_new, assign_var);
}
TEST_F(BuildXlaOpsTest, OnXlaDevice) {
const char* kXlaDeviceName = "/job:worker/replica:0/task:0/device:XLA_CPU:0";
Scope root = Scope::NewRootScope().WithDevice(kXlaDeviceName).ExitOnError();
FunctionDefLibrary fdef_lib =
CreateFunctionDefLibWithConstFunction("cluster_0");
TF_ASSERT_OK(root.graph()->AddFunctionLibrary(fdef_lib));
Node* call;
TF_ASSERT_OK(MakeXlaCompiledKernel(root.graph(), "cluster_0", "C", &call));
call->set_requested_device(kXlaDeviceName);
TF_ASSERT_OK(root.DoShapeInference(call));
call->AddAttr(kXlaHasReferenceVarsAttr, false);
Node* write_op = MakeWrite(root, Output(call), "write_result");
write_op->AddAttr(kXlaHasReferenceVarsAttr, false);
std::unique_ptr<Graph> graph;
TF_ASSERT_OK(BuildXlaOps(root, fdef_lib, &graph));
auto xla_op =
NodeWith(Op("_XlaRun"), Inputs(Out(NodeWith(Op("_XlaCompile")))));
auto assign_var =
NodeWith(Op("AssignVariableOp"), Inputs(Out(NodeWith()), Out(xla_op)));
Node* write_op_new = FindNodeByName(graph.get(), write_op->name());
ASSERT_NE(write_op_new, nullptr);
EXPECT_THAT(write_op_new, assign_var);
}
TEST_F(BuildXlaOpsTest, NoExtraMergeForEdgeToSink) {
Scope root = Scope::NewRootScope().ExitOnError();
FunctionDefLibrary fdef_lib =
CreateFunctionDefLibWithConstFunction("cluster_0");
TF_ASSERT_OK(root.graph()->AddFunctionLibrary(fdef_lib));
Node* call;
TF_ASSERT_OK(MakeXlaCompiledKernel(root.graph(), "cluster_0", "C", &call));
call->AddAttr(kXlaHasReferenceVarsAttr, false);
std::unique_ptr<Graph> graph;
TF_ASSERT_OK(BuildXlaOps(root, fdef_lib, &graph));
Node* sink_node = graph->sink_node();
EXPECT_THAT(sink_node,
NodeWith(CtrlDeps(NodeWith(Op("_XlaRun")),
NodeWith(Op("StatefulPartitionedCall")),
NodeWith(Op("NoOp")))));
}
#ifdef GOOGLE_CUDA
FunctionDefLibrary CreateFunctionDefLibWithInt32Input(const std::string& name) {
FunctionDefLibrary fdef_lib;
FunctionDef func = FunctionDefHelper::Create(
/*function_name=*/name, /*in_def=*/{"in: int32"},
/*out_def=*/{"out: int32"},
/*attr_def=*/{}, /*node_def=*/{{{"out"}, "Identity", {"in"}}},
/*ret_def=*/{{"out", "out:output:0"}});
*fdef_lib.add_function() = std::move(func);
return fdef_lib;
}
// This tests a rewrite that only makes sense and is active in a CUDA-enabled
// build. Specifically we check that we insert an IdentityN op to avoid extra
// device-to-host copies.
TEST_F(BuildXlaOpsTest, NoDeviceToHostCopiesForClustersWithInt32Inputs) {
const char* kXlaDeviceName = "/job:worker/replica:0/task:0/device:GPU:0";
Scope root = Scope::NewRootScope()
.WithDevice(kXlaDeviceName)
.WithAssignedDevice(kXlaDeviceName)
.ExitOnError();
FunctionDefLibrary fdef_lib =
CreateFunctionDefLibWithInt32Input("cluster_int32");
TF_ASSERT_OK(root.graph()->AddFunctionLibrary(fdef_lib));
Node* call;
TF_ASSERT_OK(
MakeXlaCompiledKernel(root.graph(), "cluster_int32", "C", &call));
call->set_requested_device(kXlaDeviceName);
call->AddAttr(kXlaHasReferenceVarsAttr, false);
auto var =
ops::VarHandleOp(root.WithOpName("var"), DT_INT32, TensorShape({}));
auto int32_on_device =
ops::ReadVariableOp(root.WithOpName("int32_on_device"), var, DT_INT32);
root.graph()->AddEdge(int32_on_device.node(), 0, call, 0);
std::unique_ptr<Graph> graph;
TF_ASSERT_OK(BuildXlaOps(root, fdef_lib, &graph));
Node* stateful_partitioned_call_op = nullptr;
for (Node* n : graph->op_nodes()) {
if (n->type_string() == "StatefulPartitionedCall") {
ASSERT_EQ(stateful_partitioned_call_op, nullptr);
stateful_partitioned_call_op = n;
}
}
ASSERT_NE(stateful_partitioned_call_op, nullptr);
auto xla_compile = NodeWith(Op("_XlaCompile"));
auto switch_on_compilation_pred =
NodeWith(Op("Switch"), Inputs(Out(0, xla_compile), Out(1, xla_compile)));
auto ctrl_dep =
NodeWith(Op("Identity"), Inputs(Out(0, switch_on_compilation_pred)));
// Check that we pipe int32 inputs through an IdentityN to avoid extra D2H
// copies.
EXPECT_THAT(
stateful_partitioned_call_op,
NodeWith(Inputs(Out(NodeWith(Op("IdentityN"), CtrlDeps(ctrl_dep))))));
}
#endif
} // namespace
} // namespace tensorflow
tensorflow/compiler/jit/clone_constants_for_better_clustering.cc
+225
View File
@@ -0,0 +1,225 @@
/* Copyright 2019 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/compiler/jit/clone_constants_for_better_clustering.h"
#include <string>
#include "absl/algorithm/container.h"
#include "tensorflow/compiler/jit/xla_cluster_util.h"
#include "xla/status_macros.h"
#include "tensorflow/core/framework/node_def.pb.h"
#include "tensorflow/core/framework/tensor.pb.h"
#include "tensorflow/core/util/dump_graph.h"
namespace tensorflow {
using tsl::StatusOr;
class CloneConstantsForBetterClusteringPassImpl {
public:
explicit CloneConstantsForBetterClusteringPassImpl(Graph* graph)
: graph_(graph), unique_name_counter_(0) {}
absl::Status Run();
private:
absl::Status CloneSmallConstantInputs(
const absl::flat_hash_set<std::string>& name_set, Node* n);
std::string GenerateUniqueName(
const absl::flat_hash_set<std::string>& name_set,
absl::string_view prefix);
absl::StatusOr<Node*> CloneNode(
const absl::flat_hash_set<std::string>& name_set, Node* n);
Graph* graph_;
int unique_name_counter_;
};
std::string CloneConstantsForBetterClusteringPassImpl::GenerateUniqueName(
const absl::flat_hash_set<std::string>& name_set,
absl::string_view prefix) {
std::string candidate;
do {
candidate = absl::StrCat(prefix, "/clone_", unique_name_counter_++);
} while (name_set.contains(candidate));
return candidate;
}
absl::StatusOr<Node*> CloneConstantsForBetterClusteringPassImpl::CloneNode(
const absl::flat_hash_set<std::string>& name_set, Node* n) {
NodeDef new_in_def = n->def();
new_in_def.clear_input();
new_in_def.set_name(GenerateUniqueName(name_set, new_in_def.name()));
TF_ASSIGN_OR_RETURN(Node * new_in, graph_->AddNode(new_in_def));
for (const Edge* e : n->in_edges()) {
if (e->IsControlEdge()) {
graph_->AddControlEdge(e->src(), new_in);
} else {
graph_->AddEdge(e->src(), e->src_output(), new_in, e->dst_input());
}
}
new_in->set_assigned_device_name(n->assigned_device_name());
return new_in;
}
namespace {
absl::StatusOr<bool> IsConstantSmall(Node* n) {
const TensorProto* proto = nullptr;
TF_RETURN_IF_ERROR(GetNodeAttr(n->def(), "value", &proto));
int64_t total_elements = 1;
for (const auto& dim : proto->tensor_shape().dim()) {
if (dim.size() < 0) {
return absl::InternalError(absl::StrCat(
"Unknown dimension size in constant tensor ", n->name()));
}
total_elements *= dim.size();
}
// TODO(sanjoy): It may make sense to combine this threshold with XLA's "large
// constant" threshold, if there is one.
const int kSmallTensorThreshold = 16;
return total_elements < kSmallTensorThreshold;
}
// We only clone small constants since we want to avoid increasing memory
// pressure on GPUs.
absl::StatusOr<bool> IsSmallConstant(Node* n) {
if (!n->IsConstant()) {
return false;
}
return IsConstantSmall(n);
}
bool IsInPlaceOp(absl::string_view op_name) {
return op_name == "InplaceUpdate" || op_name == "InplaceAdd" ||
op_name == "InplaceSub";
}
} // namespace
absl::Status
CloneConstantsForBetterClusteringPassImpl::CloneSmallConstantInputs(
const absl::flat_hash_set<std::string>& name_set, Node* n) {
std::vector<const Edge*> in_edges;
// Get the edges and sort them so we clone in a deterministic order.
absl::c_copy(n->in_edges(), std::back_inserter(in_edges));
absl::c_stable_sort(in_edges, [](const Edge* e1, const Edge* e2) {
return e1->id() < e2->id();
});
for (const Edge* e : in_edges) {
Node* input = e->src();
TF_ASSIGN_OR_RETURN(bool is_small_constant, IsSmallConstant(input));
if (is_small_constant && input->out_edges().size() != 1) {
VLOG(2) << "Cloning small constant " << input->name();
TF_ASSIGN_OR_RETURN(Node* const input_cloned, CloneNode(name_set, input));
if (e->IsControlEdge()) {
graph_->AddControlEdge(input_cloned, e->dst());
} else {
int dst_input = e->dst_input();
TF_RET_CHECK(e->src_output() == 0)
<< "expected constant to have exactly one non-control output, but "
"found output index = "
<< e->src_output();
graph_->RemoveEdge(e);
graph_->AddEdge(input_cloned, 0, n, dst_input);
}
}
}
return absl::OkStatus();
}
absl::Status CloneConstantsForBetterClusteringPassImpl::Run() {
absl::flat_hash_set<std::string> name_set;
absl::c_transform(graph_->nodes(), std::inserter(name_set, name_set.begin()),
[](Node* n) { return n->name(); });
std::vector<Node*> nodes;
for (Node* n : graph_->nodes()) {
// We rely on the immutability of Tensors to safely clone Const operations.
// However, "in place" ops do not respect the immutability of Tensors so we
// avoid this transformation when such ops are present in the graph.
//
// In-place operations are problematic because they break the semantic
// illusion that tensorflow::Tensor instances are immutable. For instance
// if we have the following graph:
//
// digraph {
// SRC -> Const
// SRC -> I
// SRC -> V
// Const -> Identity
// Const -> InplaceAdd [label="x"]
// I -> InplaceAdd [label="i"]
// V -> InplaceAdd [label="v"]
// InplaceAdd -> Identity [style=dotted]
// }
//
// then the value produced by `Identity` is Const+I*V since InplaceAdd
// modifies the tensor in place. However, if we clone `Const` and turn the
// graph into:
//
// digraph {
// SRC -> "Const/clone_1"
// SRC -> "Const/clone_2"
// SRC -> I
// SRC -> V
// "Const/clone_1" -> Identity
// "Const/clone_2" -> InplaceAdd [label="x"]
// I -> InplaceAdd [label="i"]
// V -> InplaceAdd [label="v"]
// InplaceAdd -> Identity [style=dotted]
// }
//
// then `Identity` no longer produces Const+I*V because the InplaceAdd
// operation only modifies Const/clone_2 in place.
if (IsInPlaceOp(n->type_string())) {
return absl::OkStatus();
}
nodes.push_back(n);
}
// Iterate over a copy of the nodes to avoid iterating over g->nodes() while
// creating more nodes.
for (Node* n : nodes) {
TF_RETURN_IF_ERROR(CloneSmallConstantInputs(name_set, n));
}
return absl::OkStatus();
}
absl::Status CloneConstantsForBetterClusteringPass::Run(
const GraphOptimizationPassOptions& options) {
if (GetGlobalJitLevelForGraph(options) == OptimizerOptions::OFF) {
return absl::OkStatus();
}
Graph* g = options.graph->get();
if (VLOG_IS_ON(1)) {
DumpGraphToFile("before_clone_constants_for_better_clustering", *g);
}
TF_RETURN_IF_ERROR(CloneConstantsForBetterClusteringPassImpl{g}.Run());
if (VLOG_IS_ON(1)) {
DumpGraphToFile("after_clone_constants_for_better_clustering", *g);
}
return absl::OkStatus();
}
} // namespace tensorflow
tensorflow/compiler/jit/clone_constants_for_better_clustering.h
+62
View File
@@ -0,0 +1,62 @@
/* Copyright 2019 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_COMPILER_JIT_CLONE_CONSTANTS_FOR_BETTER_CLUSTERING_H_
#define TENSORFLOW_COMPILER_JIT_CLONE_CONSTANTS_FOR_BETTER_CLUSTERING_H_
#include "absl/container/flat_hash_set.h"
#include "tensorflow/core/common_runtime/optimization_registry.h"
namespace tensorflow {
// Clones small host constants in the graph to make it easier to form larger
// clusters.
//
// This helps us in two ways:
//
// - It reduces dependencies between clusters. Let's say a constant C is used
// by nodes X and Y. If X and Y are put in different clusters (for whatever
// reason) Y's cluster now has to wait for all the operations in X's cluster
// to finish before it starts running.
//
// - It lets us create bigger clusters in multi-GPU benchmarks. Consider the
// following graph:
//
// digraph {
// Const -> GPU_1
// Const -> GPU_0_Y
// GPU_0_X -> GPU_0_Y
// }
//
// We'd cluster Const and GPU_1 together (and place it on GPU_1), and this
// will block us from clustering GPU_0_X and GPU_0_Y together since that
// would increase the amount of work on GPU 0 waiting on work on GPU 1.
// However, cloning Const into two copies, one for GPU_0_Y and one for GPU_1
// will let us create one cluster containing {Const/copy_0, GPU_1} and
// another containing {Const/copy_1, GPU_0_X, GPU_0_Y}.
//
// We only clone small host constants now to avoid increasing memory consumption
// too much. Moreover, in practice the constants we have to duplicate are
// things like the `perm` input to `Transpose` and the `size` input to `Slice`
// which tend to be small anyway.
class CloneConstantsForBetterClusteringPass : public GraphOptimizationPass {
public:
CloneConstantsForBetterClusteringPass() = default;
absl::Status Run(const GraphOptimizationPassOptions& options) override;
};
} // namespace tensorflow
#endif // TENSORFLOW_COMPILER_JIT_CLONE_CONSTANTS_FOR_BETTER_CLUSTERING_H_
tensorflow/compiler/jit/clone_constants_for_better_clustering_test.cc
+232
View File
@@ -0,0 +1,232 @@
/* Copyright 2019 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/compiler/jit/clone_constants_for_better_clustering.h"
#include "tensorflow/cc/framework/ops.h"
#include "tensorflow/cc/ops/array_ops.h"
#include "tensorflow/cc/ops/const_op.h"
#include "tensorflow/cc/ops/math_ops.h"
#include "tensorflow/compiler/jit/node_matchers.h"
#include "tensorflow/core/lib/core/status_test_util.h"
#include "tensorflow/core/platform/test.h"
#include "tensorflow/core/public/session_options.h"
namespace tensorflow {
namespace {
using ::tensorflow::testing::FindNodeByName;
absl::Status CloneConstantsForBetterClustering(const Scope& s,
std::unique_ptr<Graph>* result) {
auto graph = std::make_unique<Graph>(OpRegistry::Global());
SessionOptions session_options;
session_options.config.mutable_graph_options()
->mutable_optimizer_options()
->set_global_jit_level(OptimizerOptions::ON_2);
GraphOptimizationPassOptions options;
options.graph = &graph;
options.session_options = &session_options;
// Scope::ToGraph seems to drop assigned devices, probably because it goes
// through a GraphDef. So explicitly maintain the device assignment.
// std::unordered_map<string, string> assigned_device_names;
// for (Node* n : s.graph()->nodes()) {
// assigned_device_names[n->name()] = n->assigned_device_name();
// }
GraphConstructorOptions opts;
opts.expect_device_spec = true;
TF_RETURN_IF_ERROR(s.ToGraph(graph.get(), opts));
// for (Node* n : graph->nodes()) {
// n->set_assigned_device_name(assigned_device_names[n->name()]);
// }
CloneConstantsForBetterClusteringPass rewriter;
TF_RETURN_IF_ERROR(rewriter.Run(options));
*result = std::move(graph);
return absl::OkStatus();
}
const char* kCPU = "/job:localhost/replica:0/task:0/device:CPU:0";
const char* kGPU = "/job:localhost/replica:0/task:0/device:GPU:0";
TEST(CloneConstantsForBetterClusteringTest, ScalarConstantPlacedOnGpu) {
Scope root = Scope::NewRootScope().ExitOnError();
Scope on_gpu = root.WithAssignedDevice(kGPU).WithDevice(kGPU);
Output in = ops::Placeholder(on_gpu.WithOpName("in"), DT_FLOAT);
Output c = ops::Const(on_gpu.WithOpName("const"), 1.0f, {});
Output add1 = ops::AddV2(on_gpu.WithOpName("add1"), in, c);
Output add2 = ops::AddV2(on_gpu.WithOpName("add2"), add1, c);
std::unique_ptr<Graph> result;
TF_ASSERT_OK(CloneConstantsForBetterClustering(root, &result));
OutputTensor add1_operand;
TF_ASSERT_OK(
FindNodeByName(result.get(), "add1")->input_tensor(1, &add1_operand));
OutputTensor add2_operand;
TF_ASSERT_OK(
FindNodeByName(result.get(), "add2")->input_tensor(1, &add2_operand));
EXPECT_NE(add1_operand.node, add2_operand.node);
}
TEST(CloneConstantsForBetterClusteringTest, HostConstantPlacedOnCpu) {
Scope root = Scope::NewRootScope().ExitOnError();
Scope on_gpu = root.WithAssignedDevice(kGPU).WithDevice(kGPU);
Scope on_cpu = root.WithAssignedDevice(kCPU).WithDevice(kCPU);
Output in0 = ops::Placeholder(on_gpu.WithOpName("in0"), DT_FLOAT);
Output in1 = ops::Placeholder(on_gpu.WithOpName("in1"), DT_FLOAT);
Output perm = ops::Const(on_cpu.WithOpName("perm"), {3, 1, 2, 0});
{
Output tr0 = ops::Transpose(on_gpu.WithOpName("tr0"), in0, perm);
Output tr1 = ops::Transpose(on_gpu.WithOpName("tr1"), in1, perm);
}
std::unique_ptr<Graph> result;
TF_ASSERT_OK(CloneConstantsForBetterClustering(root, &result));
OutputTensor tr0_perm;
TF_ASSERT_OK(FindNodeByName(result.get(), "tr0")->input_tensor(1, &tr0_perm));
OutputTensor tr1_perm;
TF_ASSERT_OK(FindNodeByName(result.get(), "tr1")->input_tensor(1, &tr1_perm));
EXPECT_NE(tr0_perm.node, tr1_perm.node);
}
TEST(CloneConstantsForBetterClusteringTest, HostConstantPlacedOnGpu) {
Scope root = Scope::NewRootScope().ExitOnError();
Scope on_gpu = root.WithAssignedDevice(kGPU).WithDevice(kGPU);
Output in0 = ops::Placeholder(on_gpu.WithOpName("in0"), DT_FLOAT);
Output in1 = ops::Placeholder(on_gpu.WithOpName("in1"), DT_FLOAT);
Output perm = ops::Const(on_gpu.WithOpName("perm"), {3, 1, 2, 0});
{
Output tr0 = ops::Transpose(on_gpu.WithOpName("tr0"), in0, perm);
Output tr1 = ops::Transpose(on_gpu.WithOpName("tr1"), in1, perm);
}
std::unique_ptr<Graph> result;
TF_ASSERT_OK(CloneConstantsForBetterClustering(root, &result));
OutputTensor tr0_perm;
TF_ASSERT_OK(FindNodeByName(result.get(), "tr0")->input_tensor(1, &tr0_perm));
OutputTensor tr1_perm;
TF_ASSERT_OK(FindNodeByName(result.get(), "tr1")->input_tensor(1, &tr1_perm));
EXPECT_NE(tr0_perm.node, tr1_perm.node);
}
TEST(CloneConstantsForBetterClusteringTest, CloneSmallDeviceConstants) {
Scope root = Scope::NewRootScope().ExitOnError();
Scope on_gpu = root.WithAssignedDevice(kGPU).WithDevice(kGPU);
Output in0 = ops::Placeholder(on_gpu.WithOpName("in0"), DT_FLOAT);
Output in1 = ops::Placeholder(on_gpu.WithOpName("in1"), DT_FLOAT);
Output perm_f32 = ops::Const(on_gpu.WithOpName("perm"), {3.0, 1.0, 2.0, 0.0});
Output perm_int0 =
ops::Cast(on_gpu.WithOpName("perm_cast_0"), perm_f32, DT_INT32);
Output perm_int1 =
ops::Cast(on_gpu.WithOpName("perm_cast_1"), perm_f32, DT_INT32);
{
Output tr0 = ops::Transpose(on_gpu.WithOpName("tr0"), in0, perm_int0);
Output tr1 = ops::Transpose(on_gpu.WithOpName("tr1"), in1, perm_int1);
}
std::unique_ptr<Graph> result;
TF_ASSERT_OK(CloneConstantsForBetterClustering(root, &result));
OutputTensor tr0_perm;
TF_ASSERT_OK(
FindNodeByName(result.get(), "perm_cast_0")->input_tensor(0, &tr0_perm));
OutputTensor tr1_perm;
TF_ASSERT_OK(
FindNodeByName(result.get(), "perm_cast_1")->input_tensor(0, &tr1_perm));
EXPECT_NE(tr0_perm.node, tr1_perm.node);
}
TEST(CloneConstantsForBetterClusteringTest, DontCloneLargeConstants) {
Scope root = Scope::NewRootScope().ExitOnError();
Scope on_gpu = root.WithAssignedDevice(kGPU).WithDevice(kGPU);
Scope on_cpu = root.WithAssignedDevice(kCPU).WithDevice(kCPU);
Output in0 = ops::Placeholder(on_gpu.WithOpName("in0"), DT_FLOAT);
Output in1 = ops::Placeholder(on_gpu.WithOpName("in1"), DT_FLOAT);
Output perm = ops::Const(
on_cpu.WithOpName("perm"),
{17, 16, 15, 14, 13, 12, 11, 10, 9, 8, 7, 6, 5, 4, 3, 2, 1, 0});
{
Output tr0 = ops::Transpose(on_gpu.WithOpName("tr0"), in0, perm);
Output tr1 = ops::Transpose(on_gpu.WithOpName("tr1"), in1, perm);
}
std::unique_ptr<Graph> result;
TF_ASSERT_OK(CloneConstantsForBetterClustering(root, &result));
OutputTensor tr0_perm;
TF_ASSERT_OK(FindNodeByName(result.get(), "tr0")->input_tensor(1, &tr0_perm));
OutputTensor tr1_perm;
TF_ASSERT_OK(FindNodeByName(result.get(), "tr1")->input_tensor(1, &tr1_perm));
EXPECT_EQ(tr0_perm.node, tr1_perm.node);
}
TEST(CloneConstantsForBetterClusteringTest, InplaceOps) {
Scope root = Scope::NewRootScope().ExitOnError();
Scope on_gpu = root.WithAssignedDevice(kGPU).WithDevice(kGPU);
Scope on_cpu = root.WithAssignedDevice(kCPU).WithDevice(kCPU);
Output in0 = ops::Placeholder(on_gpu.WithOpName("in0"), DT_FLOAT);
Output in1 = ops::Placeholder(on_gpu.WithOpName("in1"), DT_FLOAT);
Output perm = ops::Const(on_cpu.WithOpName("perm"), {3, 1, 2, 0});
{
Output tr0 = ops::Transpose(on_gpu.WithOpName("tr0"), in0, perm);
Output tr1 = ops::Transpose(on_gpu.WithOpName("tr1"), in1, perm);
}
Output in_place_add =
ops::InplaceAdd(on_cpu.WithOpName("tr0"), perm,
ops::Placeholder(on_cpu.WithOpName("i"), DT_INT32), perm);
std::unique_ptr<Graph> result;
TF_ASSERT_OK(CloneConstantsForBetterClustering(root, &result));
OutputTensor tr0_perm;
TF_ASSERT_OK(FindNodeByName(result.get(), "tr0")->input_tensor(1, &tr0_perm));
OutputTensor tr1_perm;
TF_ASSERT_OK(FindNodeByName(result.get(), "tr1")->input_tensor(1, &tr1_perm));
EXPECT_EQ(tr0_perm.node, tr1_perm.node);
}
} // namespace
} // namespace tensorflow
tensorflow/compiler/jit/cluster_scoping_pass.cc
+164
View File
@@ -0,0 +1,164 @@
/* Copyright 2019 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/compiler/jit/cluster_scoping_pass.h"
#include "absl/algorithm/container.h"
#include "absl/container/flat_hash_set.h"
#include "absl/strings/str_cat.h"
#include "tensorflow/compiler/jit/defs.h"
#include "tensorflow/compiler/jit/xla_cluster_util.h"
#include "tensorflow/core/framework/node_def_util.h"
#include "tensorflow/core/graph/algorithm.h"
namespace tensorflow {
namespace {
class ClusterScopingPassImpl {
public:
ClusterScopingPassImpl(Graph* graph,
OptimizerOptions::GlobalJitLevel global_jit_level)
: graph_(graph),
global_jit_level_(global_jit_level),
unique_scope_id_(0) {}
absl::Status Run();
private:
absl::Status ScopingForPipelineStages();
size_t GetUniqueScopeId() { return unique_scope_id_++; }
void AddScopeToAllTransitivePredecessors(Node* start);
void AddScopeToAllTransitiveSuccessors(Node* start);
private:
Graph* graph_;
OptimizerOptions::GlobalJitLevel global_jit_level_;
size_t unique_scope_id_;
};
std::optional<std::string> GetXlaInternalScope(Node* node) {
std::string scope;
if (GetNodeAttr(node->attrs(), kXlaInternalScopeAttr, &scope).ok()) {
return scope;
}
return std::nullopt;
}
void SetXlaInternalScope(Node* node, absl::string_view scope) {
node->AddAttr(kXlaInternalScopeAttr, scope);
}
// NB! We append a new scope as suffix to the _XlaInternalScope attribute
// instead of overriding the old value. In other words, appending scope B to
// scope A creates the conjunction of the scopes A and B (i.e, A & B) and,
// in effect, the node gets both the old and new scopes. As a unique scope
// disallows a node being merged with nodes in other scopes, the scope
// conjunction preserves the semantic of the old scope (i.e., the node still
// cannot be merged with the previously incompatible nodes.)
//
// For example, the below case should be rare in practice but can serve for the
// purpose of discussion. After adding scopes for both Stage and Unstage,
// Node_Y will receive both scopes "unstage" and "stage", while Node_X receives
// only scope "stage". The semantic of scope "unstage" is preserved although
// scope "stage" is later appended. As a result, Node_X and Node_Y will be put
// into different clusters.
//
// Unstage -> Node_Y (scope "unstage & stage")
// |
// V
// Node_X (scope "stage") -> Stage
//
void AddOrAppendXlaInternalScope(Node* node, absl::string_view suffix) {
std::string updated_scope;
std::optional<std::string> cur_scope = GetXlaInternalScope(node);
if (cur_scope == std::nullopt) {
updated_scope = std::string(suffix);
} else {
updated_scope = absl::StrCat(cur_scope.value(), "&", suffix);
}
SetXlaInternalScope(node, updated_scope);
}
void ClusterScopingPassImpl::AddScopeToAllTransitivePredecessors(Node* start) {
const std::string unique_suffix = absl::StrCat("_", GetUniqueScopeId());
std::vector<Node*> starts;
starts.push_back(start);
auto enter = [&](Node* n) { AddOrAppendXlaInternalScope(n, unique_suffix); };
ReverseDFSFrom(*graph_, starts, enter, /*leave=*/nullptr,
/*stable_comparator=*/NodeComparatorName());
}
void ClusterScopingPassImpl::AddScopeToAllTransitiveSuccessors(Node* start) {
const std::string unique_suffix = absl::StrCat("_", GetUniqueScopeId());
std::vector<Node*> starts;
starts.push_back(start);
auto enter = [&](Node* n) { AddOrAppendXlaInternalScope(n, unique_suffix); };
DFSFrom(*graph_, starts, enter, /*leave=*/nullptr,
/*stable_comparator=*/NodeComparatorName(),
// Do not filter any edges to better capture the semantics of
// transitive closure of successors. We may revisit this when
// we see more cases needing cluster scoping in the future.
/*edge_filter=*/nullptr);
}
// This preserves the parallelism between pipeline stages. For example, below
// is a typical pattern of input pipelining in Tensorflow and this heuristic
// ensures Node_X and Node_Y are put into different clusters. Without the
// heuristic, they may be put into the same cluster and it can introduce
// artificial dependencies and incur great performance loss. In this example,
// Node_Y becomes dependent on IteratorGetNext and the latencies add up if
// Node_X and Node_Y are in the same cluster.
//
// IteratorGetNext -> Node_X -> Stage
//
// Unstage -> Node_Y
//
absl::Status ClusterScopingPassImpl::ScopingForPipelineStages() {
for (Node* n : graph_->nodes()) {
DCHECK(n);
if (n->type_string() == "Unstage") {
AddScopeToAllTransitiveSuccessors(n);
}
if (n->type_string() == "Stage") {
AddScopeToAllTransitivePredecessors(n);
}
}
return absl::OkStatus();
}
absl::Status ClusterScopingPassImpl::Run() {
if (global_jit_level_ == OptimizerOptions::OFF) {
return absl::OkStatus();
}
return ScopingForPipelineStages();
}
} // namespace
absl::Status ClusterScopingPass::Run(
const GraphOptimizationPassOptions& options) {
Graph* graph = options.graph->get();
return ClusterScopingPassImpl{graph, GetGlobalJitLevelForGraph(options)}
.Run();
}
} // namespace tensorflow
tensorflow/compiler/jit/cluster_scoping_pass.h
+38
View File
@@ -0,0 +1,38 @@
/* Copyright 2019 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_COMPILER_JIT_CLUSTER_SCOPING_PASS_H_
#define TENSORFLOW_COMPILER_JIT_CLUSTER_SCOPING_PASS_H_
#include "tensorflow/core/common_runtime/optimization_registry.h"
namespace tensorflow {
// This pass adds scopes to nodes in the _XlaInternalScope attribute to guide
// the later clustering passes. A major reason to do this is to prevent the
// clustering from losing critical parallelism in the Tensorflow graph, which
// can incur great performance degradation.
//
// This pass must be run before MarkForCompilationPass, as it stores the
// scoping information that MarkForCompilationPass will need to respect for
// clustering decision.
class ClusterScopingPass : public GraphOptimizationPass {
public:
absl::Status Run(const GraphOptimizationPassOptions& options) override;
};
} // namespace tensorflow
#endif // TENSORFLOW_COMPILER_JIT_CLUSTER_SCOPING_PASS_H_
tensorflow/compiler/jit/cluster_scoping_pass_test.cc
+180
View File
@@ -0,0 +1,180 @@
/* Copyright 2019 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/compiler/jit/cluster_scoping_pass.h"
#include "absl/container/flat_hash_map.h"
#include "tensorflow/compiler/jit/defs.h"
#include "tensorflow/compiler/jit/test_util.h"
#include "tensorflow/core/common_runtime/graph_constructor.h"
#include "tensorflow/core/common_runtime/graph_def_builder_util.h"
#include "tensorflow/core/framework/node_def_util.h"
#include "tensorflow/core/framework/op.h"
#include "tensorflow/core/graph/algorithm.h"
#include "tensorflow/core/graph/graph_def_builder.h"
#include "tensorflow/core/lib/core/status_test_util.h"
#include "tensorflow/core/platform/test.h"
#include "tensorflow/core/public/session_options.h"
namespace tensorflow {
namespace {
absl::Status ClusterScoping(std::unique_ptr<Graph>* graph) {
FixupSourceAndSinkEdges(graph->get());
GraphOptimizationPassWrapper wrapper;
wrapper.session_options.config.mutable_graph_options()
->mutable_optimizer_options()
->set_global_jit_level(OptimizerOptions::ON_2);
GraphOptimizationPassOptions opt_options =
wrapper.CreateGraphOptimizationPassOptions(graph);
ClusterScopingPass pass;
return pass.Run(opt_options);
}
absl::flat_hash_map<std::string, std::string> GetXlaInternalScopes(
const Graph& graph) {
absl::flat_hash_map<std::string, std::string> scopes;
for (Node* node : graph.nodes()) {
std::string scope;
if (GetNodeAttr(node->attrs(), kXlaInternalScopeAttr, &scope).ok()) {
scopes[node->name()] = scope;
}
}
if (VLOG_IS_ON(2)) {
VLOG(2) << "_XlaInternalScopes:";
for (const auto& p : scopes) {
VLOG(2) << " " << p.first << " -> " << p.second;
}
}
return scopes;
}
Node* BuildStageNode(GraphDefBuilder& builder, std::string name,
std::initializer_list<DataType> dtypes,
absl::Span<const ops::NodeOut> values) {
auto opts = builder.opts()
.WithName(std::move(name))
.WithAttr("dtypes", std::move(dtypes));
if (opts.HaveError()) {
return nullptr;
}
NodeBuilder node_builder(name, "Stage", opts.op_registry());
node_builder.Input(values);
return opts.FinalizeBuilder(&node_builder);
}
TEST(XlaCompilationTest, StagePipelinePreserved) {
std::unique_ptr<Graph> graph(new Graph(OpRegistry::Global()));
{
// Graph:
// b
// |
// v
// a -> add0 (ClusterX) -> relu0 (ClusterX) -> stage
//
// b
// |
// v
// unstage -> add1 (ClusterY) -> relu1 (ClusterY)
GraphDefBuilder builder(GraphDefBuilder::kFailImmediately);
Node* a = ops::SourceOp("Const", builder.opts()
.WithName("a")
.WithAttr("dtype", DT_FLOAT)
.WithAttr("value", Tensor()));
Node* b = ops::SourceOp("Const", builder.opts()
.WithName("b")
.WithAttr("dtype", DT_FLOAT)
.WithAttr("value", Tensor()));
Node* unstage = ops::SourceOp(
"Unstage",
builder.opts().WithName("unstage").WithAttr("dtypes", {DT_FLOAT}));
Node* add0 = ops::BinaryOp("Add", a, b, builder.opts().WithName("add0"));
Node* add1 =
ops::BinaryOp("Add", unstage, b, builder.opts().WithName("add1"));
Node* relu0 = ops::UnaryOp("Relu", add0, builder.opts().WithName("relu0"));
ops::UnaryOp("Relu", add1, builder.opts().WithName("relu1"));
BuildStageNode(builder, "stage", {DT_FLOAT}, {relu0});
TF_EXPECT_OK(GraphDefBuilderToGraph(builder, graph.get()));
}
TF_ASSERT_OK(ClusterScoping(&graph));
auto scopes = GetXlaInternalScopes(*graph);
EXPECT_NE(scopes["add0"], scopes["add1"]);
EXPECT_EQ(scopes["add0"], scopes["relu0"]);
EXPECT_EQ(scopes["add1"], scopes["relu1"]);
}
TEST(XlaCompilationTest, StagePipelinePreservedAndInitialScopesRespected) {
std::unique_ptr<Graph> graph(new Graph(OpRegistry::Global()));
{
// Graph:
// b
// |
// v
// a -> add0 (ClusterA) -> relu0 (ClusterB) -> stage
//
// b
// |
// v
// unstage -> add1 (ClusterC) -> relu1 (ClusterD)
GraphDefBuilder builder(GraphDefBuilder::kFailImmediately);
Node* a = ops::SourceOp("Const", builder.opts()
.WithName("a")
.WithAttr("dtype", DT_FLOAT)
.WithAttr("value", Tensor()));
Node* b = ops::SourceOp("Const", builder.opts()
.WithName("b")
.WithAttr("dtype", DT_FLOAT)
.WithAttr("value", Tensor()));
Node* unstage = ops::SourceOp(
"Unstage",
builder.opts().WithName("unstage").WithAttr("dtypes", {DT_FLOAT}));
// Intentionally give add0 and add1 the same initial scope but they should
// be separated by the ClusterScopingPass.
Node* add0 = ops::BinaryOp("Add", a, b,
builder.opts().WithName("add0").WithAttr(
kXlaInternalScopeAttr, "ClusterA"));
Node* add1 = ops::BinaryOp("Add", unstage, b,
builder.opts().WithName("add1").WithAttr(
kXlaInternalScopeAttr, "ClusterA"));
Node* relu0 = ops::UnaryOp("Relu", add0,
builder.opts().WithName("relu0").WithAttr(
kXlaInternalScopeAttr, "ClusterB"));
ops::UnaryOp("Relu", add1,
builder.opts().WithName("relu1").WithAttr(
kXlaInternalScopeAttr, "ClusterD"));
BuildStageNode(builder, "stage", {DT_FLOAT}, {relu0});
TF_EXPECT_OK(GraphDefBuilderToGraph(builder, graph.get()));
}
TF_ASSERT_OK(ClusterScoping(&graph));
auto scopes = GetXlaInternalScopes(*graph);
EXPECT_NE(scopes["add0"], scopes["add1"]);
EXPECT_NE(scopes["add0"], scopes["relu0"]);
EXPECT_NE(scopes["add1"], scopes["relu1"]);
}
} // namespace
} // namespace tensorflow
tensorflow/compiler/jit/compilability_check_util.cc
+758
View File
@@ -0,0 +1,758 @@
/* Copyright 2019 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/compiler/jit/compilability_check_util.h"
#include <cstddef>
#include <iterator>
#include <string>
#include <utility>
#include <vector>
#include "absl/algorithm/container.h"
#include "absl/container/flat_hash_set.h"
#include "absl/log/check.h"
#include "absl/log/log.h"
#include "absl/status/status.h"
#include "absl/status/statusor.h"
#include "absl/strings/str_cat.h"
#include "absl/strings/string_view.h"
#include "absl/types/span.h"
#include "tensorflow/compiler/jit/defs.h"
#include "tensorflow/compiler/jit/xla_activity.pb.h"
#include "tensorflow/compiler/jit/xla_activity_listener.h"
#include "tensorflow/compiler/jit/xla_cluster_util.h"
#include "tensorflow/compiler/tf2xla/const_analysis.h"
#include "tensorflow/compiler/tf2xla/tf2xla_defs.h"
#include "tensorflow/compiler/tf2xla/xla_op_registry.h"
#include "xla/tsl/platform/errors.h"
#include "tensorflow/core/common_runtime/function_body.h"
#include "tensorflow/core/common_runtime/function_utils.h"
#include "tensorflow/core/framework/attr_value.pb.h"
#include "tensorflow/core/framework/graph_def_util.h"
#include "tensorflow/core/framework/memory_types.h"
#include "tensorflow/core/framework/node_def.pb.h"
#include "tensorflow/core/framework/node_def_util.h"
#include "tensorflow/core/framework/op_kernel.h"
#include "tensorflow/core/framework/types.h"
#include "tensorflow/core/graph/graph.h"
#include "tensorflow/core/lib/gtl/cleanup.h"
namespace tensorflow {
namespace {
bool HasResourceInput(const Node& node) {
return absl::c_count(node.input_types(), DT_RESOURCE) != 0;
}
void LogNotCompilable(const Node& node, absl::string_view reason = "") {
VLOG(3) << "Found uncompilable node " << node.name() << " (op "
<< node.type_string() << ")" << (reason.empty() ? "" : ": ")
<< reason;
}
bool IsInOutsideCompilationCluster(const Node& n) {
return n.attrs().Find(kXlaOutsideCompilationAttr) != nullptr;
}
absl::Status MakeCallNodeFromAttribute(const Node& node,
const std::string& attr_name,
NodeDef* node_def) {
const NameAttrList* name_attr;
TF_RETURN_IF_ERROR(GetNodeAttr(node.attrs(), attr_name, &name_attr));
node_def->set_op(name_attr->name());
*(node_def->mutable_attr()) = name_attr->attr();
return absl::OkStatus();
}
absl::StatusOr<std::vector<NodeDef>> MakeCallNodesFromAttribute(
const Node& node, absl::string_view attr_name,
absl::string_view call_name) {
std::vector<NameAttrList> attr_lists;
TF_RETURN_IF_ERROR(GetNodeAttr(node.attrs(), attr_name, &attr_lists));
std::vector<NodeDef> out;
out.reserve(attr_lists.size());
for (int i = 0; i < attr_lists.size(); i++) {
out.emplace_back();
NodeDef& inserted = out.back();
inserted.set_name(absl::StrCat(call_name, "_", i));
inserted.set_op(attr_lists[i].name());
*inserted.mutable_attr() = attr_lists[i].attr();
}
return out;
}
// Utility which searches for values in a sorted list by scanning over it once.
// No matter how many times ScanForValue is called, the list is scanned at most
// once. However, if a call to ScanForValue skips over a value, that value is
// not revisited in future calls to ScanForValue, so callers must take
// care to order their calls.
//
// Useful for merging multiple sorted lists in O(n) time.
class SinglePassSearch {
public:
// Creates a SinglePassSearch object that can be used to search in `values`.
// Does not take ownership of `values`. `values` must outlive this.
// `values` must be sorted.
explicit SinglePassSearch(absl::Span<int const> values)
: current_index_(0), values_(values) {}
// Scans forward in the vector looking for "value", updating the internal
// position in to the vector.
// Returns true iff the vector contains the given value at or after current
// position.
// Not thread-safe.
bool ScanForValue(int value) {
while (current_index_ < values_.size() &&
values_[current_index_] <= value) {
if (values_[current_index_] == value) {
current_index_++;
return true;
}
current_index_++;
}
return false;
}
private:
int current_index_;
const absl::Span<int const> values_;
};
} // anonymous namespace
RecursiveCompilabilityChecker::UncompilableNodesMap
RecursiveCompilabilityChecker::FindUncompilableNodes(
const Node& node, FunctionLibraryRuntime* lib_runtime,
const std::vector<RecursiveCompilabilityChecker::StackFrame>*
node_stack_trace) const {
std::vector<StackFrameView> stack_trace;
// If `node_stack_trace` is provided, that means `node` is inside
// a function body, and therefore, arg nodes and retval nodes are
// not considered uncompilable.
if (node_stack_trace != nullptr) {
for (const auto& frame : *node_stack_trace) {
stack_trace.emplace_back(
StackFrameView{frame.name, frame.function_name, frame.stack_trace});
}
}
stack_trace.emplace_back(
StackFrameView{node.name(), "", node.GetStackTrace()});
RecursiveCompilabilityChecker::UncompilableNodesMap uncompilable_nodes;
IsCompilableNode(node, lib_runtime, &stack_trace,
/*encapsulating_function=*/nullptr, &uncompilable_nodes);
return uncompilable_nodes;
}
bool RecursiveCompilabilityChecker::HasXLAKernel(
const Node& node, std::string* uncompilable_reason) const {
// There is a SymbolicGradient kernel on the XLA_JIT device, but the gradient
// is really a kind of function call and will be handled by
// IsCompilableCall().
if (node.type_string() == "SymbolicGradient") {
*uncompilable_reason =
"SymbolicGradient should be handled by IsCompilableCall().";
return false;
}
if (node.type_string() == "Const") {
const AttrValue* attr = node.attrs().Find("dtype");
if (!op_filter_.allow_string_consts && attr != nullptr &&
attr->type() == DT_STRING) {
*uncompilable_reason =
"Const op with type DT_STRING is not supported by XLA.";
return false;
}
}
// XLA does not offer guaranteed aliasing between the input and output of the
// XLA cluster so it can't implement the forward-tensor-ref semantic. Leave
// such nodes out of XLA clusters.
if (HasForwardedRefInput(node)) {
VLOG(2) << "Rejecting " << node.name() << ": Identity with unsafe cast.";
*uncompilable_reason = "Identity with unsafe cast.";
return false;
}
absl::Status s =
FindKernelDef(jit_device_type_, node.def(), nullptr, nullptr);
if (!s.ok()) {
*uncompilable_reason = s.message();
return false;
}
return true;
}
// Tests whether 'if_node' is compilable. Every operator in the then_branch and
// else_branch functions must be compilable for 'if_node' to be compilable.
bool RecursiveCompilabilityChecker::IsCompilableIf(
const Node& if_node, FunctionLibraryRuntime* lib_runtime,
std::vector<StackFrameView>* stack_trace,
NameAttrList* encapsulating_function,
RecursiveCompilabilityChecker::UncompilableNodesMap* uncompilable_nodes)
const {
bool is_compilable = true;
is_compilable &= ExtractNodeDefAndCheckCompilability(
if_node, "then_branch", "if_then", encapsulating_function, lib_runtime,
stack_trace, uncompilable_nodes);
if (!uncompilable_nodes && !is_compilable) return is_compilable;
is_compilable &= ExtractNodeDefAndCheckCompilability(
if_node, "else_branch", "if_else", encapsulating_function, lib_runtime,
stack_trace, uncompilable_nodes);
return is_compilable;
}
bool RecursiveCompilabilityChecker::IsCompilableCase(
const Node& case_node, FunctionLibraryRuntime* lib_runtime,
std::vector<StackFrameView>* stack_trace,
NameAttrList* encapsulating_function,
RecursiveCompilabilityChecker::UncompilableNodesMap* uncompilable_nodes)
const {
absl::StatusOr<std::vector<NodeDef>> calls =
MakeCallNodesFromAttribute(case_node, "branches", "branch");
if (!calls.ok()) {
VLOG(2) << "Rejecting node " << case_node.name() << ": "
<< "missing attribute 'branches'";
return false;
}
bool is_compilable = true;
for (const NodeDef& call : *calls) {
is_compilable &=
IsCompilableCall(call, lib_runtime, stack_trace, encapsulating_function,
uncompilable_nodes);
}
return is_compilable;
}
// Tests whether 'while_node' is a completely compilable loop.
// Every operator in the condition and body functions must be compilable for a
// while loop to be compilable.
bool RecursiveCompilabilityChecker::IsCompilableWhile(
const Node& while_node, FunctionLibraryRuntime* lib_runtime,
std::vector<StackFrameView>* stack_trace,
NameAttrList* encapsulating_function,
RecursiveCompilabilityChecker::UncompilableNodesMap* uncompilable_nodes)
const {
bool is_compilable = true;
is_compilable &= ExtractNodeDefAndCheckCompilability(
while_node, "cond", "while_cond", encapsulating_function, lib_runtime,
stack_trace, uncompilable_nodes);
if (!uncompilable_nodes && !is_compilable) return is_compilable;
is_compilable &= ExtractNodeDefAndCheckCompilability(
while_node, "body", "while_body", encapsulating_function, lib_runtime,
stack_trace, uncompilable_nodes);
return is_compilable;
}
bool RecursiveCompilabilityChecker::ExtractNodeDefAndCheckCompilability(
const Node& node, const std::string& attr_name,
const std::string& call_name, NameAttrList* encapsulating_function,
FunctionLibraryRuntime* lib_runtime,
std::vector<StackFrameView>* stack_trace,
RecursiveCompilabilityChecker::UncompilableNodesMap* uncompilable_nodes)
const {
NodeDef call;
call.set_name(call_name);
if (!MakeCallNodeFromAttribute(node, attr_name, &call).ok()) {
const auto uncompilable_reason = absl::StrCat(
"missing '", attr_name, "' attribute from node", node.name());
MaybeMarkUncompilableNode(uncompilable_reason, *stack_trace,
encapsulating_function, uncompilable_nodes);
VLOG(2) << "Rejecting node " << node.name() << ": " << uncompilable_reason
<< ".";
return false;
}
if (!IsCompilableCall(call, lib_runtime, stack_trace, encapsulating_function,
uncompilable_nodes)) {
VLOG(2) << "Rejecting node " << node.name()
<< ": can't compile : " << call.op();
return false;
}
return true;
}
// Tests whether 'call_def' is a call to a completely compilable function.
// Every operator in the function must be compilable for a function to be
// compilable.
bool RecursiveCompilabilityChecker::IsCompilableCall(
const NodeDef& call_def, FunctionLibraryRuntime* lib_runtime,
std::vector<StackFrameView>* stack_trace,
NameAttrList* encapsulating_function,
RecursiveCompilabilityChecker::UncompilableNodesMap* uncompilable_nodes)
const {
if (stack_trace->size() > kMaxRecursionDepth) {
std::string uncompilable_reason = "function depth limit exceeded";
MaybeMarkUncompilableNode(uncompilable_reason, *stack_trace,
encapsulating_function, uncompilable_nodes);
VLOG(2) << "Rejecting " << call_def.op() << ": " << uncompilable_reason
<< ".";
return false;
}
FunctionLibraryRuntime::Handle handle;
absl::Status s;
NameAttrList function;
s = NameAndAttrsFromFunctionCall(call_def, &function);
if (s.ok()) {
s = lib_runtime->Instantiate(function.name(), AttrSlice(&function.attr()),
&handle);
}
if (!s.ok()) {
std::string uncompilable_reason =
absl::StrCat("could not instantiate call: '", function.name(), "'");
MaybeMarkUncompilableNode(uncompilable_reason, *stack_trace,
encapsulating_function, uncompilable_nodes);
VLOG(2) << "Rejecting " << call_def.DebugString() << ": "
<< uncompilable_reason << " : " << s;
return false;
}
auto release_handle_on_return =
gtl::MakeCleanup([&] { CHECK_OK(lib_runtime->ReleaseHandle(handle)); });
const FunctionBody* fbody = lib_runtime->GetFunctionBody(handle);
bool is_compilable = true;
for (const Node* node : fbody->graph->op_nodes()) {
stack_trace->emplace_back(
StackFrameView{node->name(), function.name(), node->GetStackTrace()});
is_compilable &= IsCompilableNode(*node, lib_runtime, stack_trace,
&function, uncompilable_nodes);
stack_trace->pop_back();
if (!uncompilable_nodes && !is_compilable) return is_compilable;
}
return is_compilable;
}
bool RecursiveCompilabilityChecker::OpIsInaccurate(const Node& node) const {
// b/127344411: SelfAdjointEigV2 and Svd precision issues.
return node.type_string() == "SelfAdjointEigV2" ||
node.type_string() == "Svd";
}
bool RecursiveCompilabilityChecker::OpIsSlow(const Node& node) const {
// b/128001705: SelfAdjointEigV2 and Svd performance issues.
// b/135640736: MatrixInverse performance issues.
// b/111271662: MatrixSolve performance issues.
// https://github.com/tensorflow/tensorflow/pull/31012:
// ResizeNearestNeighbor, ResizeBilinear, and ResizeBilinearGrad sometimes
// create convolutions too large for CuDNN to handle.
// NonMaxSuppressionV3/V4 in XLA runs significantly slower than TF kernel in
// object detection models, specially when there are a lot of proposed
// bounding boxes.
return node.type_string() == "SelfAdjointEigV2" ||
node.type_string() == "Svd" || node.type_string() == "Qr" ||
node.type_string() == "MatrixInverse" ||
node.type_string() == "MatrixSolve" ||
node.type_string() == "ResizeBilinearGrad" ||
node.type_string() == "NonMaxSuppressionV3" ||
node.type_string() == "NonMaxSuppressionV4";
}
bool RecursiveCompilabilityChecker::IsCompilableNode(
const Node& node, FunctionLibraryRuntime* lib_runtime,
std::vector<StackFrameView>* stack_trace,
NameAttrList* encapsulating_function,
RecursiveCompilabilityChecker::UncompilableNodesMap* uncompilable_nodes)
const {
auto stack_depth = stack_trace->size();
if (op_filter_.allow_outside_compiled && IsInOutsideCompilationCluster(node))
return true;
if (node.IsSource() || node.IsSink()) {
absl::string_view uncompilable_reason = "source or sink node";
MaybeMarkUncompilableNode(uncompilable_reason, *stack_trace,
encapsulating_function, uncompilable_nodes);
LogNotCompilable(node, uncompilable_reason);
return false;
}
// _Arg nodes in a top-level function represent feeds and _Retval nodes in a
// top-level function represent fetches.
if (stack_depth == 1 &&
(node.type_string() == "_Arg" || node.type_string() == "_Retval")) {
absl::string_view uncompilable_reason = "top level _Arg or _Retval";
MaybeMarkUncompilableNode(uncompilable_reason, *stack_trace,
encapsulating_function, uncompilable_nodes);
LogNotCompilable(node, uncompilable_reason);
return false;
}
if (node.attrs().Find("_scoped_allocator") ||
node.attrs().Find("_forward_from")) {
// TODO(b/128858118): XLA does not support _scoped_allocator and
// _forward_from.
absl::string_view uncompilable_reason =
"_scoped_allocator or _forward_from attribute";
MaybeMarkUncompilableNode(uncompilable_reason, *stack_trace,
encapsulating_function, uncompilable_nodes);
LogNotCompilable(node, uncompilable_reason);
return false;
}
std::string uncompilable_reason;
if (IsFunctionCall(*lib_runtime->GetFunctionLibraryDefinition(), node)) {
if (!IsCompilableCall(node.def(), lib_runtime, stack_trace,
encapsulating_function, uncompilable_nodes)) {
LogNotCompilable(node, "unsupported function");
return false;
}
} else if (!HasXLAKernel(node, &uncompilable_reason)) {
MaybeMarkUncompilableNode(
absl::StrCat("unsupported op: ", uncompilable_reason), *stack_trace,
encapsulating_function, uncompilable_nodes);
LogNotCompilable(node, uncompilable_reason);
return false;
}
if (node.IsWhileNode() &&
!IsCompilableWhile(node, lib_runtime, stack_trace, encapsulating_function,
uncompilable_nodes)) {
LogNotCompilable(node, "unsupported while");
return false;
}
if (node.IsIfNode() &&
!IsCompilableIf(node, lib_runtime, stack_trace, encapsulating_function,
uncompilable_nodes)) {
LogNotCompilable(node, "unsupported if");
return false;
}
if (op_filter_.require_always_compilable && node.IsCaseNode() &&
!IsCompilableCase(node, lib_runtime, stack_trace, encapsulating_function,
uncompilable_nodes)) {
LogNotCompilable(node, "unsupported case");
return false;
}
if (!op_filter_.allow_stateful_rng_ops &&
IsStatefulRandomOp(node.type_string())) {
absl::string_view uncompilable_reason = "stateful random op";
MaybeMarkUncompilableNode(uncompilable_reason, *stack_trace,
encapsulating_function, uncompilable_nodes);
LogNotCompilable(node, uncompilable_reason);
return false;
}
if (!op_filter_.allow_control_trigger && node.IsControlTrigger()) {
absl::string_view uncompilable_reason = "not allowed control trigger";
MaybeMarkUncompilableNode(uncompilable_reason, *stack_trace,
encapsulating_function, uncompilable_nodes);
LogNotCompilable(node, uncompilable_reason);
return false;
}
if (!op_filter_.allow_eliding_assert_and_checknumerics_ops &&
IsAssertOrCheckNumerics(node.type_string())) {
absl::string_view uncompilable_reason = "Assert or CheckNumerics";
MaybeMarkUncompilableNode(uncompilable_reason, *stack_trace,
encapsulating_function, uncompilable_nodes);
LogNotCompilable(node, uncompilable_reason);
return false;
}
if (!op_filter_.allow_collective_reduce_v2 &&
node.type_string() == "CollectiveReduceV2") {
absl::string_view uncompilable_reason = "Collective op";
MaybeMarkUncompilableNode(uncompilable_reason, *stack_trace,
encapsulating_function, uncompilable_nodes);
LogNotCompilable(node, uncompilable_reason);
return false;
}
if (!op_filter_.allow_where_op && node.type_string() == "Where") {
absl::string_view uncompilable_reason = "Where op";
MaybeMarkUncompilableNode(uncompilable_reason, *stack_trace,
encapsulating_function, uncompilable_nodes);
LogNotCompilable(node, uncompilable_reason);
return false;
}
if (!op_filter_.allow_unique_op && node.type_string() == "Unique") {
absl::string_view uncompilable_reason = "Unique op";
MaybeMarkUncompilableNode(uncompilable_reason, *stack_trace,
encapsulating_function, uncompilable_nodes);
LogNotCompilable(node, uncompilable_reason);
return false;
}
if (!op_filter_.allow_ops_producing_or_consuming_variant &&
OpProducesOrConsumesVariant(node)) {
absl::string_view uncompilable_reason = "DT_VARIANT producer/consumer";
MaybeMarkUncompilableNode(uncompilable_reason, *stack_trace,
encapsulating_function, uncompilable_nodes);
LogNotCompilable(node, uncompilable_reason);
return false;
}
if (!op_filter_.allow_stack_ops && IsStackOp(node)) {
absl::string_view uncompilable_reason = "Stack op";
MaybeMarkUncompilableNode(uncompilable_reason, *stack_trace,
encapsulating_function, uncompilable_nodes);
LogNotCompilable(node, uncompilable_reason);
return false;
}
if (!op_filter_.allow_tensor_array_ops && IsTensorArrayOp(node)) {
absl::string_view uncompilable_reason = "TensorArray op";
MaybeMarkUncompilableNode(uncompilable_reason, *stack_trace,
encapsulating_function, uncompilable_nodes);
LogNotCompilable(node, uncompilable_reason);
return false;
}
if (!op_filter_.allow_resource_ops_in_called_functions && stack_depth > 1 &&
HasResourceInput(node)) {
absl::string_view uncompilable_reason =
"resource variable op in called function";
MaybeMarkUncompilableNode(uncompilable_reason, *stack_trace,
encapsulating_function, uncompilable_nodes);
LogNotCompilable(node, uncompilable_reason);
return false;
}
if (!op_filter_.allow_inaccurate_ops && OpIsInaccurate(node)) {
absl::string_view uncompilable_reason =
"operation with numerical accuracy issues";
BroadcastOptimizationRemark(XlaOptimizationRemark::INACCURATE_OPERATION,
node.DebugString())
.IgnoreError();
MaybeMarkUncompilableNode(uncompilable_reason, *stack_trace,
encapsulating_function, uncompilable_nodes);
LogNotCompilable(node, uncompilable_reason);
return false;
}
if (!op_filter_.allow_slow_ops && OpIsSlow(node)) {
absl::string_view uncompilable_reason = "slow operation";
BroadcastOptimizationRemark(XlaOptimizationRemark::SLOW_OPERATION,
node.DebugString())
.IgnoreError();
MaybeMarkUncompilableNode(uncompilable_reason, *stack_trace,
encapsulating_function, uncompilable_nodes);
LogNotCompilable(node, uncompilable_reason);
return false;
}
return true;
}
RecursiveCompilabilityChecker::OperationFilter CreateOperationFilter(
const XlaOpRegistry::DeviceRegistration& registration) {
RecursiveCompilabilityChecker::OperationFilter op_filter;
op_filter.allow_resource_ops_in_called_functions =
registration.cluster_resource_variable_ops_unsafely;
op_filter.allow_stack_ops = registration.cluster_stack_ops;
op_filter.allow_tensor_array_ops = registration.cluster_tensor_array_ops;
op_filter.allow_stateful_rng_ops = registration.cluster_stateful_rng_ops;
op_filter.allow_control_trigger = registration.cluster_control_trigger;
op_filter.allow_eliding_assert_and_checknumerics_ops =
registration.elide_assert_and_checknumerics;
op_filter.allow_ops_producing_or_consuming_variant =
registration.cluster_variant_ops;
op_filter.allow_slow_ops = registration.cluster_slow_ops;
op_filter.allow_inaccurate_ops = registration.cluster_inaccurate_ops;
return op_filter;
}
/*static*/ void RecursiveCompilabilityChecker::MaybeMarkUncompilableNode(
const absl::string_view reason,
const std::vector<StackFrameView>& stack_trace,
NameAttrList* encapsulating_function,
RecursiveCompilabilityChecker::UncompilableNodesMap* uncompilable_nodes) {
if (!uncompilable_nodes) return;
UncompilableNodeInfo node_info;
node_info.uncompilable_reason = std::string(reason);
absl::c_transform(stack_trace, std::back_inserter(node_info.stack_trace),
[](const StackFrameView& stack_element) {
return StackFrame{
std::string(stack_element.name),
std::string(stack_element.function_name),
stack_element.stack_trace};
});
node_info.name = std::string(stack_trace.back().name);
auto function =
encapsulating_function ? *encapsulating_function : NameAttrList();
auto function_identifier = function.ShortDebugString();
auto it = uncompilable_nodes->find(function_identifier);
if (it == uncompilable_nodes->end()) {
std::vector<RecursiveCompilabilityChecker::UncompilableNodeInfo>
uncompilable_node_info{std::move(node_info)};
uncompilable_nodes->emplace(
std::move(function_identifier),
std::make_pair(function, std::move(uncompilable_node_info)));
} else {
it->second.second.emplace_back(std::move(node_info));
}
}
// Returns `true` iff node has a given `attr` set to `true`. Returns `false`
// both for the missing attr, and the attr set to `false`.
static bool HasBoolAttr(const NodeDef& node, const char* attr) {
const auto& it = node.attr().find(attr);
return it != node.attr().end() && it->second.b();
}
bool CanCreateXlaKernel(const NodeDef& node_def) {
return HasBoolAttr(node_def, kXlaMustCompileAttr);
}
absl::Status GetBodyAndConstantsAndResources(
FunctionLibraryRuntime* flr, const NameAttrList& function,
const FunctionBody** fbody, std::vector<int>* constant_arg_indices,
std::vector<int>* resource_arg_indices) {
FunctionLibraryRuntime::Handle handle;
TF_RETURN_IF_ERROR(
flr->Instantiate(function.name(), AttrSlice(&function.attr()), &handle));
*fbody = flr->GetFunctionBody(handle);
CHECK(*fbody); // Can't be nullptr since we just instantiated it.
const DataTypeVector& arg_types = (*fbody)->arg_types;
std::vector<bool> const_args(arg_types.size());
// If we can't analyze the const args. Bail out.
TF_RETURN_IF_ERROR(
BackwardsConstAnalysis(*((*fbody)->graph), &const_args,
/*compile_time_const_nodes=*/nullptr, flr));
for (size_t i = 0; i < const_args.size(); ++i) {
if (const_args[i]) {
constant_arg_indices->push_back(i);
}
}
// There can be hundreds of resource variables. Reserve the space for them.
// We don't reserve for constants above as they are usually few.
resource_arg_indices->reserve(arg_types.size());
for (size_t i = 0; i < arg_types.size(); ++i) {
if (arg_types[i] == DT_RESOURCE) {
resource_arg_indices->push_back(i);
}
}
return absl::OkStatus();
}
tensorflow::MemoryTypeVector GetInputMemoryTypes(
const tensorflow::FunctionBody* fbody,
absl::Span<int const> constant_arg_indices,
absl::Span<int const> resource_arg_indices) {
// Set input and output memory types.
tensorflow::MemoryTypeVector input_memory_types(fbody->arg_types.size(),
tensorflow::DEVICE_MEMORY);
// These indices are used only for optimization purposes. They allow us
// to loop over constant_arg_indices and resource_arg_indices only once
// while iterating over all the function arguments checking if it is a
// resource or a constant.
// The reason we optimized this code is because functions can have a lot of
// captured arguments. For example, the backward pass of ResNet50 takes in all
// 214 variables and a similar number of activations.
SinglePassSearch constants_search(constant_arg_indices);
SinglePassSearch resources_search(resource_arg_indices);
for (size_t i = 0; i < fbody->arg_types.size(); ++i) {
if (resources_search.ScanForValue(i) || constants_search.ScanForValue(i)) {
// Compile-time constants and resource handles are expected to be in
// host memory.
input_memory_types[i] = tensorflow::HOST_MEMORY;
}
}
return input_memory_types;
}
tensorflow::MemoryTypeVector GetOutputMemoryTypes(
const tensorflow::FunctionBody* fbody) {
tensorflow::MemoryTypeVector output_memory_types(fbody->ret_types.size(),
tensorflow::DEVICE_MEMORY);
for (size_t i = 0; i < fbody->ret_types.size(); ++i) {
if (fbody->ret_types[i] == tensorflow::DT_RESOURCE) {
output_memory_types[i] = tensorflow::HOST_MEMORY;
}
}
return output_memory_types;
}
static auto const ops_triggering_xla_compilation =
new absl::flat_hash_set<std::string>{"XlaBroadcastHelper",
"XlaCallModule",
"XlaConv",
"XlaConvV2",
"XlaDequantize",
"XlaDot",
"XlaDotV2",
"XlaDynamicSlice",
"XlaDynamicUpdateSlice",
"XlaEinsum",
"XlaGather",
"XlaIf",
"XlaKeyValueSort",
"XlaPad",
"XlaRecv",
"XlaReduce",
"XlaReduceWindow",
"XlaReplicaId",
"XlaRngBitGenerator",
"XlaScatter",
"XlaSelectAndScatter",
"XlaSelfAdjointEig",
"XlaSend",
"XlaSharding",
"XlaSort",
"XlaSpmdFullToShardShape",
"XlaSpmdShardToFullShape",
"XlaSvd",
"XlaVariadicReduceV2",
"XlaVariadicSort",
"XlaWhile"};
bool NodeCanTriggerXlaCompilation(const NodeDef& node) {
return node.attr().find(kXlaClusterIdAttr) != node.attr().end() ||
HasBoolAttr(node, kXlaMustCompileAttr) ||
HasBoolAttr(node, kXlaCompileAttr) ||
HasBoolAttr(node, kXlaScopeAttr) ||
HasBoolAttr(node, kXlaInternalScopeAttr) ||
ops_triggering_xla_compilation->count(node.op());
}
bool CanTriggerXlaCompilation(const GraphDef& graph) {
for (const FunctionDef& function : graph.library().function()) {
for (const NodeDef& node : function.node_def()) {
if (NodeCanTriggerXlaCompilation(node)) {
return true;
}
}
}
for (const NodeDef& node : graph.node()) {
if (NodeCanTriggerXlaCompilation(node)) {
return true;
}
}
return false;
}
} // namespace tensorflow
tensorflow/compiler/jit/compilability_check_util.h
+340
View File
@@ -0,0 +1,340 @@
/* Copyright 2019 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_COMPILER_JIT_COMPILABILITY_CHECK_UTIL_H_
#define TENSORFLOW_COMPILER_JIT_COMPILABILITY_CHECK_UTIL_H_
#include <string>
#include "absl/algorithm/container.h"
#include "absl/status/statusor.h"
#include "absl/strings/string_view.h"
#include "absl/types/optional.h"
#include "tensorflow/compiler/jit/defs.h"
#include "tensorflow/compiler/jit/device_util.h"
#include "tensorflow/compiler/jit/flags.h"
#include "tensorflow/compiler/jit/resource_operation_safety_analysis.h"
#include "tensorflow/compiler/tf2xla/const_analysis.h"
#include "tensorflow/compiler/tf2xla/resource_operation_table.h"
#include "tensorflow/compiler/tf2xla/xla_op_registry.h"
#include "xla/service/graphcycles/graphcycles.h"
#include "xla/union_find.h"
#include "xla/util.h"
#include "tensorflow/core/common_runtime/function.h"
#include "tensorflow/core/common_runtime/graph_constructor.h"
#include "tensorflow/core/framework/attr_value.pb.h"
#include "tensorflow/core/framework/bounds_check.h"
#include "tensorflow/core/framework/function.h"
#include "tensorflow/core/framework/graph_def_util.h"
#include "tensorflow/core/framework/memory_types.h"
#include "tensorflow/core/framework/node_def.pb.h"
#include "tensorflow/core/framework/op_kernel.h"
#include "tensorflow/core/framework/types.h"
#include "tensorflow/core/framework/types.pb.h"
#include "tensorflow/core/graph/algorithm.h"
#include "tensorflow/core/graph/control_flow.h"
#include "tensorflow/core/graph/graph.h"
#include "tensorflow/core/lib/gtl/cleanup.h"
#include "tensorflow/core/lib/strings/stringprintf.h"
#include "tensorflow/core/public/version.h"
#include "tensorflow/core/util/dump_graph.h"
namespace tensorflow {
// Checks whether a TF node can be compiled or not. "Recursive" as in for call
// and functional while nodes it recursively checks whether the callee functions
// can be compiled.
class RecursiveCompilabilityChecker {
public:
// Contains node name and function name. If the node is not inside a function
// body, function name is an empty string.
struct StackFrame {
std::string name;
std::string function_name;
std::shared_ptr<AbstractStackTrace> stack_trace;
};
// Contains information about uncompilable node inside a function body.
struct UncompilableNodeInfo {
std::string name;
// A list representing a stacktrace from the highest level node in
// increasing call depth to immediate node that fails the
// compilability checker.
std::vector<StackFrame> stack_trace;
std::string uncompilable_reason;
};
// Aggregates information about what kinds of ops are allowed.
struct OperationFilter { // TODO(lzr): Add AllowEverything() helper.
// Whether resource variable ops are allowed are allowed in callees. We do
// not allow resource variable ops in called functions (either as direct TF
// calls or as higher order control flow ops) because we do not yet model
// their memory effects in jit/resource_operation_safety_analysis.
bool allow_resource_ops_in_called_functions = false;
// Whether Stack operations are allowed. We avoid auto-clustering Stack
// operations in general because we do not support snapshotting them.
//
// TODO(b/112837194): This restriction can be lifted with some work.
bool allow_stack_ops = false;
// Whether TensorArray operations are allowed. We avoid auto-clustering
// TensorArray operations in general because we do not support snapshotting
// them.
//
// TODO(b/112837194): This restriction can be lifted with some work.
bool allow_tensor_array_ops = false;
// Whether stateful RNG ops are allowed. XLA's RNG does not have the same
// seeding behavior as TensorFlow's RNG (b/34749654). So we avoid
// auto-clustering stateful RNG ops.
bool allow_stateful_rng_ops = false;
// TODO(b/118970344): Whether ControlTrigger ops are allowed. It is unsound
// to cluster ControlTrigger because of how we use deadness analysis.
bool allow_control_trigger = false;
// Whether it is okay to "cluster" Assert and CheckNumerics by simply
// removing them (they're not removed during clustering, but their
// XlaOpKernel is a no-op kernel). We avoid auto-clustering these ops so
// that the user is not surprised when XLA is implicitly enabled. If the
// user explicitly specifies to use XLA, it is fine to resort to a dummy
// implementation. Currently Assert and CheckNumerics ops have dummy XLA
// implementations.
bool allow_eliding_assert_and_checknumerics_ops = false;
// Whether ops that produce or consume DT_VARIANT values are allowed. We
// don't auto-cluster these ops because we don't yet support live-in or
// live-out DT_VARIANT values.
bool allow_ops_producing_or_consuming_variant = false;
// Whether ops known to be slow on XLA-GPU should be considered compilable.
bool allow_slow_ops = false;
// Whether ops known to have numerical accuracy issues should be considered
// compilable..
bool allow_inaccurate_ops = false;
// Require the function to be always compilable, regardless whether some
// control flow branches might be dead for a given input.
bool require_always_compilable = false;
// Whether string constants are compilable.
bool allow_string_consts = true;
// Whether to allow the compilation of CollectiveReduceV2Op.
bool allow_collective_reduce_v2 = true;
// Whether to allow the compilation of WhereOp.
bool allow_where_op = true;
// Whether to allow the compilation of UniqueOp. Compilation of the UniqueOp
// generates output with bounded dynamic shape that may cause failures with
// auto clustering.
// TODO(b/209813421): Enable tf.unique during
// autoclustering once all failures are rfixed.
bool allow_unique_op = true;
// Whether ops that are marked as outside compiled are always considered
// compilable.
// TODO(b/191502757): Make this behavior true by default and remove this
// option once inference converter supports outside compilation.
bool allow_outside_compiled = false;
};
RecursiveCompilabilityChecker(OperationFilter op_filter,
DeviceType jit_device_type)
: op_filter_(std::move(op_filter)),
jit_device_type_(std::move(jit_device_type)) {}
using UncompilableNodesMap =
std::map<std::string,
std::pair<NameAttrList, std::vector<UncompilableNodeInfo>>>;
// Returns a map where the key is the function identifier(short debug
// string) of the function encapsulating the uncompilable nodes, and the
// value is a pair of NameAttrList of the function and a vector of
// uncompilable node info. When uncompilable node is not inside any
// function call nodes, then key is a ShortDebugString() of an empty
// NameAttrList.
//
// Also, when `node` is inside a function body, users can set
// `node_stack_trace` to provide an additional context for `node`'s
// placement within the outer most graph.
UncompilableNodesMap FindUncompilableNodes(
const Node& node, FunctionLibraryRuntime* lib_runtime,
const std::vector<StackFrame>* node_stack_trace = nullptr) const;
// Returns true if `node` can be compiled by XLA.
bool IsCompilableNode(const Node& node,
FunctionLibraryRuntime* lib_runtime) const {
std::vector<StackFrameView> stack_trace;
stack_trace.emplace_back(StackFrameView{node.name(), ""});
return IsCompilableNode(node, lib_runtime, &stack_trace);
}
// Returns true if XLA supports this Op, but we don't want to cluster it (ie:
// due to performance or correctness concerns).
bool OpIsInaccurate(const Node& node) const;
bool OpIsSlow(const Node& node) const;
private:
struct StackFrameView {
absl::string_view name;
absl::string_view function_name;
std::shared_ptr<AbstractStackTrace> stack_trace;
};
bool IsCompilableNode(
const Node& node, FunctionLibraryRuntime* lib_runtime,
std::vector<StackFrameView>* stack_trace,
NameAttrList* encapsulating_function = nullptr,
UncompilableNodesMap* uncompilable_nodes = nullptr) const;
bool IsCompilableCall(
const NodeDef& call_def, FunctionLibraryRuntime* lib_runtime,
std::vector<StackFrameView>* stack_trace,
NameAttrList* encapsulating_function = nullptr,
UncompilableNodesMap* uncompilable_nodes = nullptr) const;
bool IsCompilableIf(const Node& if_node, FunctionLibraryRuntime* lib_runtime,
std::vector<StackFrameView>* stack_trace,
NameAttrList* encapsulating_function,
UncompilableNodesMap* uncompilable_nodes) const;
bool IsCompilableWhile(const Node& while_node,
FunctionLibraryRuntime* lib_runtime,
std::vector<StackFrameView>* stack_trace,
NameAttrList* encapsulating_function,
UncompilableNodesMap* uncompilable_nodes) const;
// Tests whether 'case_node' is compilable. Every operator in all branches
// must be compilable.
bool IsCompilableCase(const Node& case_node,
FunctionLibraryRuntime* lib_runtime,
std::vector<StackFrameView>* stack_trace,
NameAttrList* encapsulating_function,
UncompilableNodesMap* uncompilable_nodes) const;
// Returns compilability of node def retrieved from `node`'s attribute with
// name `attr_name`.
bool ExtractNodeDefAndCheckCompilability(
const Node& node, const std::string& attr_name,
const std::string& call_name, NameAttrList* encapsulating_function,
FunctionLibraryRuntime* lib_runtime,
std::vector<StackFrameView>* stack_trace,
UncompilableNodesMap* uncompilable_nodes) const;
bool IsStackOp(const Node& node) const {
const XlaResourceOpInfo* op_info =
GetResourceOpInfoForOp(node.type_string());
return op_info && op_info->resource_kind() == XlaResourceKind::kStack;
}
bool IsTensorArrayOp(const Node& node) const {
const XlaResourceOpInfo* op_info =
GetResourceOpInfoForOp(node.type_string());
return op_info && op_info->resource_kind() == XlaResourceKind::kTensorArray;
}
bool IsAssertOrCheckNumerics(absl::string_view op_name) const {
return op_name == "Assert" || op_name == "CheckNumerics";
}
bool IsStatefulRandomOp(absl::string_view op_name) const {
return op_name == "RandomUniform" || op_name == "RandomShuffle" ||
op_name == "RandomUniformInt" || op_name == "RandomStandardNormal" ||
op_name == "TruncatedNormal" || op_name == "Multinomial";
}
bool OpProducesOrConsumesVariant(const Node& node) const {
auto is_variant = [](DataType dtype) { return dtype == DT_VARIANT; };
return absl::c_any_of(node.input_types(), is_variant) ||
absl::c_any_of(node.output_types(), is_variant);
}
bool HasXLAKernel(const Node& node,
std::string* uncompilable_reason = nullptr) const;
static void MaybeMarkUncompilableNode(
const absl::string_view reason,
const std::vector<StackFrameView>& stack_trace,
NameAttrList* encapsulating_function,
UncompilableNodesMap* uncompilable_nodes_map);
// Make sure we don't recurse infinitely on recursive functions.
const size_t kMaxRecursionDepth = 50;
const OperationFilter op_filter_;
const DeviceType jit_device_type_;
};
RecursiveCompilabilityChecker::OperationFilter CreateOperationFilter(
const XlaOpRegistry::DeviceRegistration& registration);
// Given a FunctionLibraryRuntime and a `function`, returns this function's body
// in `fbody` as well as the indices of its constant and resource arguments.
// `fbody` is owned by `flr`.
// `constant_arg_indices` and `resource_arg_indices` should be empty vector.
// They are sorted in ascending order on this function's return.
absl::Status GetBodyAndConstantsAndResources(
FunctionLibraryRuntime* flr, const NameAttrList& function,
const FunctionBody** fbody, std::vector<int>* constant_arg_indices,
std::vector<int>* resource_arg_indices);
// Given a NodeDef `node_def` returns true iff `node_def` has kXlaCompileAttr
// set.
bool CanCreateXlaKernel(const NodeDef& node_def);
// Returns memory types for the input.
// `constant_arg_indices` and `resource_arg_indices` are sorted arrays of
// indices corresponding to constant and resource arguments respectively.
//
// One might wonder, about the case where a compile-time constant argument
// (which must be in host memory) is also used as an input into an op,
// e.g. `Add`, that expects its inputs in device memory. Here is how it
// works now.
// First, what do we mean by "op expects an input in XYZ memory"?
// There are two types of "ops" here: the tf2xla kernel and the HLO
// computation it builds. The tf2xla kernel needs to retrieve the actual
// numeric value of the compile-time constant tensors, so it really expects
// them to be on in host memory. However, for other inputs, it refers to them
// using xla::ComputationDataHandle, which is just a symbolic handle that
// xla::ComputationBuilder assigns. How does this handle gets assigned for
// constant arguments? Even constant arguments get an _Arg node in the graph
// instantiated for Function compilation. The tf2xla kernel for constant _Arg
// nodes takes the constant value, converts it to XlaLiteral, and feeds it
// to xla::ComputationBuilder.ConstantLiteral, which returns the handle. This
// constant XlaLiteral is included in the HLO graph, and subsequently, in
// the actual executable, which is copied to the device before being
// executed. Thus, when this executable runs, the constant is available in
// device memory.
tensorflow::MemoryTypeVector GetInputMemoryTypes(
const tensorflow::FunctionBody* fbody,
absl::Span<int const> constant_arg_indices,
absl::Span<int const> resource_arg_indices);
// Returns output memory types.
//
// XlaLaunch kernel keeps all outputs (including constants, which it copies),
// in device memory except for resources.
tensorflow::MemoryTypeVector GetOutputMemoryTypes(
const tensorflow::FunctionBody* fbody);
// Check whether graph can trigger XLA compilation.
bool CanTriggerXlaCompilation(const GraphDef& graph);
// Returns true iff the node can trigger XLA compilation.
bool NodeCanTriggerXlaCompilation(const NodeDef& node);
} // namespace tensorflow
#endif // TENSORFLOW_COMPILER_JIT_COMPILABILITY_CHECK_UTIL_H_
tensorflow/compiler/jit/compilability_check_util_test.cc
+656
View File
@@ -0,0 +1,656 @@
/* Copyright 2019 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/compiler/jit/compilability_check_util.h"
#include <algorithm>
#include <memory>
#include <vector>
#include "absl/log/check.h"
#include "absl/strings/match.h"
#include "absl/types/span.h"
#include "tensorflow/cc/framework/ops.h"
#include "tensorflow/cc/framework/scope.h"
#include "tensorflow/cc/ops/array_ops.h"
#include "tensorflow/cc/ops/functional_ops.h"
#include "tensorflow/compiler/jit/defs.h"
#include "tensorflow/compiler/tf2xla/xla_op_kernel.h"
#include "tensorflow/compiler/tf2xla/xla_op_registry.h"
#include "xla/tsl/lib/core/status_test_util.h"
#include "tensorflow/core/common_runtime/graph_def_builder_util.h"
#include "tensorflow/core/framework/attr_value.pb.h"
#include "tensorflow/core/framework/function.h"
#include "tensorflow/core/framework/node_def_util.h"
#include "tensorflow/core/framework/op.h"
#include "tensorflow/core/framework/op_kernel.h"
#include "tensorflow/core/framework/types.h"
#include "tensorflow/core/graph/graph_def_builder.h"
#include "tensorflow/core/platform/env.h"
#include "tensorflow/core/platform/test.h"
#include "tensorflow/core/public/version.h"
namespace tensorflow {
namespace {
AttrValue FuncListAttr(const absl::Span<const char* const> names) {
AttrValue attr;
for (const char* name : names) {
attr.mutable_list()->add_func()->set_name(name);
}
return attr;
}
constexpr char kFunctionalIfNodeName[] = "If";
constexpr char kFunctionalCaseNodeName[] = "Case";
constexpr char kFunctionalWhileNodeName[] = "While";
constexpr char kCompilableFunctionName[] = "CompilableFn";
constexpr char kCompilableFunctionNodeName[] = "n_c";
constexpr char kUncompilableFunctionName[] = "UncompilableFn";
constexpr char kUncompilableFunctionNodeName[] = "n_c_uncompilable";
constexpr char kUncompilableFunctionTwoName[] = "UncompilableFnTwo";
constexpr char kUncompilableFunctionNodeTwoName[] = "n_d_uncompilable";
constexpr char kNonMaxSuppressionNodeName[] = "NonMaxSuppression";
// A dummy OpKernel for testing.
class DummyCompilableOp : public XlaOpKernel {
public:
explicit DummyCompilableOp(OpKernelConstruction* ctx) : XlaOpKernel(ctx) {}
void Compile(XlaOpKernelContext* ctx) override {
ctx->SetOutput(0, ctx->Input(0));
}
};
// Register the DummyCompilableOp kernel for CPU.
REGISTER_OP("InputFloatOp").Output("o: float");
REGISTER_OP("InputInt32Op").Output("o: int32");
REGISTER_OP("CompilableOp").Input("i: float").Output("o: float");
REGISTER_XLA_OP(Name("CompilableOp").Device(DEVICE_CPU_XLA_JIT),
DummyCompilableOp);
// Dummy op that is uncompilable in CPU.
REGISTER_OP("MissingKernel").Input("i: float").Output("o: float");
class CompilabilityCheckUtilTest : public ::testing::Test {
protected:
void SetUp() override {
XlaOpRegistry::RegisterCompilationKernels();
op_filter_.allow_resource_ops_in_called_functions = false;
op_filter_.allow_stack_ops = false;
op_filter_.allow_tensor_array_ops = false;
op_filter_.allow_stateful_rng_ops = false;
op_filter_.allow_control_trigger = false;
op_filter_.allow_eliding_assert_and_checknumerics_ops = false;
op_filter_.allow_ops_producing_or_consuming_variant = false;
op_filter_.allow_inaccurate_ops = false;
op_filter_.allow_slow_ops = false;
op_filter_.allow_outside_compiled = false;
checker_ = CreateCompilabilityChecker();
}
std::unique_ptr<RecursiveCompilabilityChecker> CreateCompilabilityChecker() {
return std::make_unique<RecursiveCompilabilityChecker>(op_filter_,
device_type_);
}
FunctionLibraryRuntime* GetFunctionLibraryRuntime() {
OptimizerOptions opts;
pflr_ = std::make_unique<ProcessFunctionLibraryRuntime>(
nullptr, Env::Default(), /*config=*/nullptr, TF_GRAPH_DEF_VERSION,
flib_def_.get(), opts);
return pflr_->GetFLR(ProcessFunctionLibraryRuntime::kDefaultFLRDevice);
}
RecursiveCompilabilityChecker::OperationFilter op_filter_;
DeviceType device_type_ = DeviceType(DEVICE_CPU_XLA_JIT);
std::unique_ptr<FunctionDefLibrary> func_library_ =
std::make_unique<FunctionDefLibrary>();
std::unique_ptr<FunctionLibraryDefinition> flib_def_ =
std::make_unique<FunctionLibraryDefinition>(OpRegistry::Global(),
*func_library_);
std::unique_ptr<RecursiveCompilabilityChecker> checker_;
std::unique_ptr<ProcessFunctionLibraryRuntime> pflr_;
};
TEST_F(CompilabilityCheckUtilTest, CheckNonFunctionalNodes) {
GraphDefBuilder builder(GraphDefBuilder::kFailImmediately);
auto opts = builder.opts();
Node* const0 = ops::SourceOp("InputFloatOp", opts);
Node* compilable_op = ops::UnaryOp("CompilableOp", const0, opts);
Node* uncompilable_op = ops::UnaryOp("MissingKernel", compilable_op, opts);
GraphDef graph_def;
TF_EXPECT_OK(builder.ToGraphDef(&graph_def));
auto* flib_runtime = GetFunctionLibraryRuntime();
// Source node is not compilable.
EXPECT_FALSE(checker_->IsCompilableNode(*const0, flib_runtime));
EXPECT_TRUE(checker_->IsCompilableNode(*compilable_op, flib_runtime));
// Uncompilable as we are only checking compilability in CPU device type.
EXPECT_FALSE(checker_->IsCompilableNode(*uncompilable_op, flib_runtime));
const auto uncompilable_nodes =
checker_->FindUncompilableNodes(*uncompilable_op, flib_runtime);
ASSERT_EQ(1, uncompilable_nodes.size());
auto node_info_it =
uncompilable_nodes.find(NameAttrList().ShortDebugString());
ASSERT_NE(uncompilable_nodes.end(), node_info_it);
const auto& uncompilable_nodes_inside_function = node_info_it->second.second;
ASSERT_EQ(1, uncompilable_nodes_inside_function.size());
const auto& uncompilable_node_info = uncompilable_nodes_inside_function.at(0);
EXPECT_TRUE(absl::StrContains(uncompilable_node_info.uncompilable_reason,
"unsupported op"));
ASSERT_EQ(1, uncompilable_node_info.stack_trace.size());
ASSERT_EQ("", uncompilable_node_info.stack_trace.at(0).function_name);
}
TEST_F(CompilabilityCheckUtilTest, CheckOutsideCompiledNode) {
GraphDefBuilder builder(GraphDefBuilder::kFailImmediately);
auto opts = builder.opts();
Node* const0 = ops::SourceOp("InputFloatOp", opts);
Node* uncompilable_op = ops::UnaryOp("MissingKernel", const0, opts);
uncompilable_op->AddAttr("_xla_outside_compilation", "0");
GraphDef graph_def;
TF_EXPECT_OK(builder.ToGraphDef(&graph_def));
auto* flib_runtime = GetFunctionLibraryRuntime();
// Outside compiled ops are considered by default..
EXPECT_FALSE(checker_->IsCompilableNode(*uncompilable_op, flib_runtime));
const auto uncompilable_nodes =
checker_->FindUncompilableNodes(*uncompilable_op, flib_runtime);
ASSERT_EQ(1, uncompilable_nodes.size());
op_filter_.allow_outside_compiled = true;
checker_ = CreateCompilabilityChecker();
// With filter option outside compiled ops are ignored and considered
// compilable.
EXPECT_TRUE(checker_->IsCompilableNode(*uncompilable_op, flib_runtime));
const auto uncompilable_nodes2 =
checker_->FindUncompilableNodes(*uncompilable_op, flib_runtime);
ASSERT_EQ(0, uncompilable_nodes2.size());
}
TEST_F(CompilabilityCheckUtilTest, CheckSimpleFunctionNode) {
FunctionDefLibrary flib;
*flib.add_function() = FunctionDefHelper::Define(
/*Function*/ kUncompilableFunctionName,
/*Inputs*/ {"n_a:float"},
/*Outputs*/ {"n_c_uncompilable:float"},
/*Attributes*/ {},
// Node info
{{{kUncompilableFunctionNodeName}, "MissingKernel", {"n_a"}}});
flib_def_.reset(new FunctionLibraryDefinition(OpRegistry::Global(), flib));
GraphDefBuilder builder(GraphDefBuilder::kFailImmediately, flib_def_.get());
std::unique_ptr<Graph> graph(new Graph(flib_def_.get()));
Node* const0 = ops::SourceOp("InputFloatOp", builder.opts());
Node* functional_node = ops::UnaryOp(kUncompilableFunctionName, const0,
builder.opts().WithName("D"));
TF_EXPECT_OK(GraphDefBuilderToGraph(builder, graph.get()));
auto* flib_runtime = GetFunctionLibraryRuntime();
EXPECT_FALSE(checker_->IsCompilableNode(*functional_node, flib_runtime));
const auto uncompilable_nodes =
checker_->FindUncompilableNodes(*functional_node, flib_runtime);
EXPECT_EQ(1, uncompilable_nodes.size());
NameAttrList function;
function.set_name(kUncompilableFunctionName);
const auto node_info_it =
uncompilable_nodes.find(function.ShortDebugString());
ASSERT_NE(uncompilable_nodes.end(), node_info_it);
const auto& uncompilable_node_list = node_info_it->second.second;
ASSERT_EQ(1, uncompilable_node_list.size());
const auto& node_info = uncompilable_node_list.at(0);
const auto& node_stack = node_info.stack_trace;
ASSERT_EQ(2, node_stack.size());
EXPECT_EQ("D", node_stack.at(0).name);
EXPECT_EQ(kUncompilableFunctionNodeName, node_stack.at(1).name);
EXPECT_EQ(kUncompilableFunctionNodeName, node_info.name);
EXPECT_TRUE(
absl::StrContains(node_info.uncompilable_reason, "unsupported op"));
}
TEST_F(CompilabilityCheckUtilTest, CheckFunctionalWhileNode) {
FunctionDefLibrary flib;
*flib.add_function() = FunctionDefHelper::Define(
/*Function*/ kCompilableFunctionName,
/*Inputs*/ {"n_a:float", "n_b:float"},
/*Outputs*/ {"n_c:float"},
/*Attribute*/ {},
// Node info
{{{kCompilableFunctionNodeName},
"Add",
{"n_a", "n_b"},
{{"T", DT_FLOAT}}}});
*flib.add_function() = FunctionDefHelper::Define(
/*Function*/ kUncompilableFunctionName,
/*Inputs*/ {"n_a:float"},
/*Outputs*/ {"n_c_uncompilable:float"},
/*Attributes*/ {},
// Node info
{{{kUncompilableFunctionNodeName}, "MissingKernel", {"n_a"}}});
flib_def_.reset(new FunctionLibraryDefinition(OpRegistry::Global(), flib));
GraphDefBuilder builder(GraphDefBuilder::kFailImmediately, flib_def_.get());
Node* const0 = ops::SourceOp("InputFloatOp", builder.opts());
Node* input_node = ops::UnaryOp("CompilableOp", const0, builder.opts());
NameAttrList compilable;
compilable.set_name(kCompilableFunctionName);
NameAttrList uncompilable;
uncompilable.set_name(kUncompilableFunctionName);
NodeBuilder while_builder(kFunctionalWhileNodeName, "While",
builder.opts().op_registry());
while_builder.Input({input_node, input_node})
.Attr("cond", compilable)
.Attr("body", uncompilable);
builder.opts().FinalizeBuilder(&while_builder);
GraphDef graph_def;
TF_EXPECT_OK(builder.ToGraphDef(&graph_def));
std::unique_ptr<Graph> graph(new Graph(flib_def_.get()));
CHECK_OK(GraphDefBuilderToGraph(builder, graph.get()));
auto while_node_it = std::find_if(
graph->nodes().begin(), graph->nodes().end(),
[&](const Node* n) { return n->name() == kFunctionalWhileNodeName; });
EXPECT_NE(while_node_it, graph->nodes().end());
auto* flib_runtime = GetFunctionLibraryRuntime();
EXPECT_FALSE(checker_->IsCompilableNode(**while_node_it, flib_runtime));
const auto uncompilable_nodes =
checker_->FindUncompilableNodes(**while_node_it, flib_runtime);
ASSERT_EQ(1, uncompilable_nodes.size());
NameAttrList function;
function.set_name(kUncompilableFunctionName);
const auto node_info_it =
uncompilable_nodes.find(function.ShortDebugString());
ASSERT_NE(uncompilable_nodes.end(), node_info_it);
const auto& uncompilable_node_list = node_info_it->second.second;
ASSERT_EQ(1, uncompilable_node_list.size());
const auto& node_info = uncompilable_node_list.at(0);
const auto& node_stack = node_info.stack_trace;
ASSERT_EQ(2, node_stack.size());
const auto& stacktrace_first_node_info = node_stack.at(0);
EXPECT_EQ(kFunctionalWhileNodeName, stacktrace_first_node_info.name);
EXPECT_EQ("", stacktrace_first_node_info.function_name);
const auto& stacktrace_second_node_info = node_stack.at(1);
EXPECT_EQ(kUncompilableFunctionNodeName, stacktrace_second_node_info.name);
EXPECT_EQ(kUncompilableFunctionName,
stacktrace_second_node_info.function_name);
EXPECT_EQ(kUncompilableFunctionNodeName, node_info.name);
EXPECT_TRUE(
absl::StrContains(node_info.uncompilable_reason, "unsupported op"));
}
TEST_F(CompilabilityCheckUtilTest, CheckFunctionalIfNode) {
FunctionDefLibrary flib;
*flib.add_function() = FunctionDefHelper::Define(
/*Function*/ kUncompilableFunctionName,
/*Inputs*/ {"n_a:float"},
/*Outputs*/ {"n_c_uncompilable:float"},
/*Attributes*/ {},
// Node info
{{{kUncompilableFunctionNodeName}, "MissingKernel", {"n_a"}}});
*flib.add_function() = FunctionDefHelper::Define(
/*Function*/ kUncompilableFunctionTwoName,
/*Inputs*/ {"n_a:float"},
/*Outputs*/ {"n_d_uncompilable:float"},
/*Attribute*/ {},
// Node info
{{{kUncompilableFunctionNodeTwoName}, "MissingKernel", {"n_a"}}});
NameAttrList uncompilable_fn1_attr;
uncompilable_fn1_attr.set_name(kUncompilableFunctionName);
NameAttrList uncompilable_fn2_attr;
uncompilable_fn2_attr.set_name(kUncompilableFunctionTwoName);
Scope root = Scope::NewRootScope().ExitOnError();
TF_ASSERT_OK(root.graph()->AddFunctionLibrary(flib));
auto predicate = ops::Placeholder(root.WithOpName("pred"), DT_BOOL);
auto placeholder = ops::Placeholder(root.WithOpName("A"), DT_INT32);
std::vector<NodeBuilder::NodeOut> if_inputs(
{NodeBuilder::NodeOut(placeholder.node())});
Node* if_node;
TF_ASSERT_OK(
NodeBuilder(kFunctionalIfNodeName, "If", &root.graph()->flib_def())
.Input(predicate.node())
.Input(if_inputs)
.Attr("then_branch", uncompilable_fn1_attr)
.Attr("else_branch", uncompilable_fn2_attr)
.Attr("Tout", {DT_INT32})
.Finalize(root.graph(), &if_node));
std::unique_ptr<Graph> graph(new Graph(OpRegistry::Global()));
TF_ASSERT_OK(root.ToGraph(graph.get()));
flib_def_.reset(new FunctionLibraryDefinition(OpRegistry::Global(), flib));
auto if_node_it = std::find_if(
graph->nodes().begin(), graph->nodes().end(),
[&](const Node* n) { return n->name() == kFunctionalIfNodeName; });
EXPECT_NE(if_node_it, graph->nodes().end());
auto* flib_runtime = GetFunctionLibraryRuntime();
EXPECT_FALSE(checker_->IsCompilableNode(**if_node_it, flib_runtime));
const auto uncompilable_nodes =
checker_->FindUncompilableNodes(**if_node_it, flib_runtime);
ASSERT_EQ(2, uncompilable_nodes.size());
NameAttrList function_one;
function_one.set_name(kUncompilableFunctionName);
auto it = uncompilable_nodes.find(function_one.ShortDebugString());
ASSERT_NE(uncompilable_nodes.end(), it);
const auto& uncompilable_node_list = it->second.second;
ASSERT_EQ(1, uncompilable_node_list.size());
const auto& uncompilable_node_one = uncompilable_node_list.at(0);
const auto& node_one_stack = uncompilable_node_one.stack_trace;
ASSERT_EQ(2, node_one_stack.size());
const auto& node_one_stacktrace_first_node = node_one_stack.at(0);
EXPECT_EQ(kFunctionalIfNodeName, node_one_stacktrace_first_node.name);
EXPECT_EQ("", node_one_stacktrace_first_node.function_name);
const auto& stacktrace_second_node_info = node_one_stack.at(1);
EXPECT_EQ(kUncompilableFunctionNodeName, stacktrace_second_node_info.name);
EXPECT_EQ(kUncompilableFunctionName,
stacktrace_second_node_info.function_name);
EXPECT_EQ(kUncompilableFunctionNodeName, uncompilable_node_one.name);
EXPECT_TRUE(absl::StrContains(uncompilable_node_one.uncompilable_reason,
"unsupported op"));
NameAttrList function_two;
function_two.set_name(kUncompilableFunctionTwoName);
it = uncompilable_nodes.find(function_two.ShortDebugString());
ASSERT_NE(uncompilable_nodes.end(), it);
const auto& uncompilable_node_two_list = it->second.second;
ASSERT_EQ(1, uncompilable_node_two_list.size());
const auto& uncompilable_node_two = uncompilable_node_two_list.at(0);
const auto& node_two_stack = uncompilable_node_two.stack_trace;
ASSERT_EQ(2, node_two_stack.size());
const auto& node_two_stacktrace_first_node = node_two_stack.at(0);
EXPECT_EQ(kFunctionalIfNodeName, node_two_stacktrace_first_node.name);
EXPECT_EQ("", node_two_stacktrace_first_node.function_name);
const auto& node_two_stacktrace_second_node = node_two_stack.at(1);
EXPECT_EQ(kUncompilableFunctionNodeTwoName,
node_two_stacktrace_second_node.name);
EXPECT_EQ(kUncompilableFunctionTwoName,
node_two_stacktrace_second_node.function_name);
EXPECT_EQ(kUncompilableFunctionNodeTwoName, uncompilable_node_two.name);
EXPECT_TRUE(absl::StrContains(uncompilable_node_one.uncompilable_reason,
"unsupported op"));
}
TEST_F(CompilabilityCheckUtilTest, CheckFunctionalCaseNode) {
FunctionDefLibrary flib;
*flib.add_function() = FunctionDefHelper::Define(
/*Function*/ kUncompilableFunctionName,
/*Inputs*/ {"n_a:float"},
/*Outputs*/ {"n_c_uncompilable:float"},
/*Attributes*/ {},
// Node info
{{{kUncompilableFunctionNodeName}, "MissingKernel", {"n_a"}}});
*flib.add_function() = FunctionDefHelper::Define(
/*Function*/ kUncompilableFunctionTwoName,
/*Inputs*/ {"n_a:float"},
/*Outputs*/ {"n_d_uncompilable:float"},
/*Attribute*/ {},
// Node info
{{{kUncompilableFunctionNodeTwoName}, "MissingKernel", {"n_a"}}});
Scope root = Scope::NewRootScope().ExitOnError();
TF_ASSERT_OK(root.graph()->AddFunctionLibrary(flib));
auto branch_index = ops::Placeholder(root.WithOpName("pred"), DT_INT32);
auto placeholder = ops::Placeholder(root.WithOpName("A"), DT_INT32);
std::vector<NodeBuilder::NodeOut> inputes(
{NodeBuilder::NodeOut(placeholder.node())});
Node* case_node;
TF_ASSERT_OK(
NodeBuilder(kFunctionalCaseNodeName, "Case", &root.graph()->flib_def())
.Input(branch_index.node())
.Input(inputes)
.Attr("branches", FuncListAttr({kUncompilableFunctionName,
kUncompilableFunctionTwoName}))
.Attr("Tout", {DT_INT32})
.Finalize(root.graph(), &case_node));
std::unique_ptr<Graph> graph(new Graph(OpRegistry::Global()));
TF_ASSERT_OK(root.ToGraph(graph.get()));
flib_def_.reset(new FunctionLibraryDefinition(OpRegistry::Global(), flib));
auto case_node_it = std::find_if(
graph->nodes().begin(), graph->nodes().end(),
[&](const Node* n) { return n->name() == kFunctionalCaseNodeName; });
EXPECT_NE(case_node_it, graph->nodes().end());
auto* flib_runtime = GetFunctionLibraryRuntime();
op_filter_.require_always_compilable = false;
checker_ = CreateCompilabilityChecker();
EXPECT_TRUE(checker_->IsCompilableNode(**case_node_it, flib_runtime));
op_filter_.require_always_compilable = true;
checker_ = CreateCompilabilityChecker();
EXPECT_FALSE(checker_->IsCompilableNode(**case_node_it, flib_runtime));
}
TEST_F(CompilabilityCheckUtilTest, TestCanNotTriggerXlaCompilation) {
GraphDefBuilder b(GraphDefBuilder::kFailImmediately);
Scope root = Scope::NewRootScope().ExitOnError();
FunctionDefLibrary library;
FunctionDef identity_func = FunctionDefHelper::Create(
"IdentityFunc",
/*in_def=*/{"x:float"},
/*out_def=*/{"res:float"},
/*attr_def=*/{},
/*node_def=*/{{{"t0"}, "Identity", {"x"}, {{"T", DT_FLOAT}}}},
/*ret_def*/ {{"res", "t0:output"}});
*library.add_function() = identity_func;
Output in = ops::Placeholder(root, DT_FLOAT);
NameAttrList b_name_attr;
b_name_attr.set_name("IdentityFunc");
ops::PartitionedCall call(root.WithOpName("call"), {in}, {DT_FLOAT},
b_name_attr);
GraphDef graph_def;
TF_ASSERT_OK(root.graph()->AddFunctionLibrary(library));
TF_ASSERT_OK(root.ToGraphDef(&graph_def));
EXPECT_FALSE(CanTriggerXlaCompilation(graph_def));
}
TEST_F(CompilabilityCheckUtilTest, TestXlaOpsCanTriggerXlaCompilation) {
GraphDefBuilder b(GraphDefBuilder::kFailImmediately);
Scope root = Scope::NewRootScope().ExitOnError();
FunctionDefLibrary library;
FunctionDef sort_func = FunctionDefHelper::Create(
"SortFunc",
/*in_def=*/{"x:float"},
/*out_def=*/{"res:float"},
/*attr_def=*/{},
/*node_def=*/{{{"t0"}, "XlaSort", {"x"}, {{"T", DT_FLOAT}}}},
/*ret_def*/ {{"res", "t0:output"}});
*library.add_function() = sort_func;
Output in = ops::Placeholder(root, DT_FLOAT);
NameAttrList b_name_attr;
b_name_attr.set_name("SortFunc");
ops::PartitionedCall call(root.WithOpName("call"), {in}, {DT_FLOAT},
b_name_attr);
GraphDef graph_def;
TF_ASSERT_OK(root.graph()->AddFunctionLibrary(library));
TF_ASSERT_OK(root.ToGraphDef(&graph_def));
EXPECT_TRUE(CanTriggerXlaCompilation(graph_def));
}
TEST_F(CompilabilityCheckUtilTest, TestCanTriggerXlaCompilation) {
GraphDefBuilder b(GraphDefBuilder::kFailImmediately);
Scope root = Scope::NewRootScope().ExitOnError();
FunctionDefLibrary library;
AttrValue true_attribute;
true_attribute.set_b(true);
FunctionDef identity_func = FunctionDefHelper::Create(
"IdentityFunc",
/*in_def=*/{"x:float"},
/*out_def=*/{"res:float"},
/*attr_def=*/{},
/*node_def=*/{{{"t0"}, "Identity", {"x"}, {{"T", DT_FLOAT}}}},
/*ret_def*/ {{"res", "t0:output"}});
(*identity_func.mutable_attr())[kXlaMustCompileAttr] = true_attribute;
FunctionDef call_identity = FunctionDefHelper::Create(
"CallIdentity",
/*in_def=*/{"x:float"},
/*out_def=*/{"z:float"}, /*attr_def=*/{},
/*node_def=*/
{{{"func_call"},
"PartitionedCall",
{"x"},
{{"Tin", DataTypeSlice({DT_FLOAT})},
{"Tout", DataTypeSlice({DT_FLOAT})},
{"f",
FunctionDefHelper::FunctionRef("IdentityRef", {{"T", DT_FLOAT}})},
{kXlaMustCompileAttr, true}}}},
/*ret_def=*/{{"z", "func_call:output:0"}});
*library.add_function() = identity_func;
*library.add_function() = call_identity;
Output in = ops::Placeholder(root, DT_FLOAT);
NameAttrList b_name_attr;
b_name_attr.set_name("CallIdentity");
ops::PartitionedCall call(root.WithOpName("call"), {in}, {DT_FLOAT},
b_name_attr);
GraphDef graph_def;
TF_ASSERT_OK(root.graph()->AddFunctionLibrary(library));
TF_ASSERT_OK(root.ToGraphDef(&graph_def));
EXPECT_TRUE(CanTriggerXlaCompilation(graph_def));
}
TEST_F(CompilabilityCheckUtilTest, CheckNonMaxSuppressionV3UncompilableSlowOp) {
GraphDefBuilder builder(GraphDefBuilder::kFailImmediately);
auto opts = builder.opts();
Node* boxes = ops::SourceOp("InputFloatOp", opts);
Node* scores = ops::SourceOp("InputFloatOp", opts);
Node* max_output_size = ops::SourceOp("InputInt32Op", opts);
Node* iou_threshold = ops::SourceOp("InputFloatOp", opts);
Node* score_threshold = ops::SourceOp("InputFloatOp", opts);
NodeBuilder non_max_suppression_builder(
kNonMaxSuppressionNodeName, "NonMaxSuppressionV3", opts.op_registry());
non_max_suppression_builder.Input(boxes)
.Input(scores)
.Input(max_output_size)
.Input(iou_threshold)
.Input(score_threshold)
.Attr("T", DT_FLOAT);
Node* non_max_suppression;
non_max_suppression =
builder.opts().FinalizeBuilder(&non_max_suppression_builder);
GraphDef graph_def;
TF_EXPECT_OK(builder.ToGraphDef(&graph_def));
auto* flib_runtime = GetFunctionLibraryRuntime();
EXPECT_FALSE(checker_->IsCompilableNode(*non_max_suppression, flib_runtime));
const auto uncompilable_nodes =
checker_->FindUncompilableNodes(*non_max_suppression, flib_runtime);
ASSERT_EQ(1, uncompilable_nodes.size());
auto node_info_it =
uncompilable_nodes.find(NameAttrList().ShortDebugString());
ASSERT_NE(uncompilable_nodes.end(), node_info_it);
const auto& uncompilable_nodes_inside_function = node_info_it->second.second;
ASSERT_EQ(1, uncompilable_nodes_inside_function.size());
const auto& uncompilable_node_info = uncompilable_nodes_inside_function.at(0);
EXPECT_TRUE(absl::StrContains(uncompilable_node_info.uncompilable_reason,
"slow operation"));
}
TEST_F(CompilabilityCheckUtilTest, CheckNonMaxSuppressionV4UncompilableSlowOp) {
GraphDefBuilder builder(GraphDefBuilder::kFailImmediately);
auto opts = builder.opts();
Node* boxes = ops::SourceOp("InputFloatOp", opts);
Node* scores = ops::SourceOp("InputFloatOp", opts);
Node* max_output_size = ops::SourceOp("InputInt32Op", opts);
Node* iou_threshold = ops::SourceOp("InputFloatOp", opts);
Node* score_threshold = ops::SourceOp("InputFloatOp", opts);
NodeBuilder non_max_suppression_v4_builder(
kNonMaxSuppressionNodeName, "NonMaxSuppressionV4", opts.op_registry());
non_max_suppression_v4_builder.Input(boxes)
.Input(scores)
.Input(max_output_size)
.Input(iou_threshold)
.Input(score_threshold)
.Attr("T", DT_FLOAT);
Node* non_max_suppression_v4;
non_max_suppression_v4 =
builder.opts().FinalizeBuilder(&non_max_suppression_v4_builder);
GraphDef graph_def;
TF_EXPECT_OK(builder.ToGraphDef(&graph_def));
auto* flib_runtime = GetFunctionLibraryRuntime();
EXPECT_FALSE(
checker_->IsCompilableNode(*non_max_suppression_v4, flib_runtime));
const auto uncompilable_nodes =
checker_->FindUncompilableNodes(*non_max_suppression_v4, flib_runtime);
ASSERT_EQ(1, uncompilable_nodes.size());
auto node_info_it =
uncompilable_nodes.find(NameAttrList().ShortDebugString());
ASSERT_NE(uncompilable_nodes.end(), node_info_it);
const auto& uncompilable_nodes_inside_function = node_info_it->second.second;
ASSERT_EQ(1, uncompilable_nodes_inside_function.size());
const auto& uncompilable_node_info = uncompilable_nodes_inside_function.at(0);
EXPECT_TRUE(absl::StrContains(uncompilable_node_info.uncompilable_reason,
"slow operation"));
}
} // namespace
} // namespace tensorflow
tensorflow/compiler/jit/compilation_passes_test_main.cc
+58
View File
@@ -0,0 +1,58 @@
/* Copyright 2019 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 <memory>
#include "absl/strings/match.h"
#include "absl/strings/string_view.h"
#include "tensorflow/compiler/jit/flags.h"
#include "tensorflow/core/platform/logging.h"
#include "tensorflow/core/platform/test.h"
#include "tensorflow/core/platform/test_benchmark.h"
GTEST_API_ int main(int real_argc, char** real_argv) {
std::vector<tensorflow::Flag> flag_list;
tensorflow::AppendMarkForCompilationPassFlags(&flag_list);
auto usage = tensorflow::Flags::Usage(real_argv[0], flag_list);
std::vector<char*> args;
args.reserve(real_argc + 1);
for (int i = 0; i < real_argc; i++) {
args.push_back(real_argv[i]);
}
struct FreeDeleter {
void operator()(char* ptr) { free(ptr); }
};
std::unique_ptr<char, FreeDeleter> enable_global_jit_arg(
strdup("--tf_xla_cpu_global_jit=true"));
args.push_back(enable_global_jit_arg.get());
std::unique_ptr<char, FreeDeleter> reduce_min_cluster_size_arg(
strdup("--tf_xla_min_cluster_size=2"));
args.push_back(reduce_min_cluster_size_arg.get());
int argc = args.size();
if (!tensorflow::Flags::Parse(&argc, &args.front(), flag_list)) {
LOG(ERROR) << "\n" << usage;
return 2;
}
testing::InitGoogleTest(&argc, &args.front());
return RUN_ALL_TESTS();
}
tensorflow/compiler/jit/deadness_analysis.cc
+1610
View File
File diff suppressed because it is too large Load Diff
tensorflow/compiler/jit/deadness_analysis.h
+99
View File
@@ -0,0 +1,99 @@
/* 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_COMPILER_JIT_DEADNESS_ANALYSIS_H_
#define TENSORFLOW_COMPILER_JIT_DEADNESS_ANALYSIS_H_
#include "tensorflow/core/graph/graph.h"
namespace tensorflow {
// This analyzes a TensorFlow graph to identify nodes which may have partially
// dead inputs (i.e. these nodes may have some dead inputs and some alive
// inputs).
//
// For example, the ADD node in the following graph
//
// V0 PRED0 V1 PRED1
// | | | |
// v v v v
// SWITCH SWITCH
// | |
// +---+ + ---+
// | |
// v v
// ADD
//
// can have its inputs independently dead or alive based on the runtime values
// of PRED0 and PRED1.
//
// It is tempting to call this a liveness analysis but I avoided that because
// "liveness" already has other connotations.
class DeadnessAnalysis {
public:
// An opaque representation of a predicate. DeadnessPredicate
// instances that compare equal via operator== represent predicates
// that always evaluate to the same value.
struct DeadnessPredicate {
public:
DeadnessPredicate(const DeadnessPredicate&) = default;
DeadnessPredicate(DeadnessPredicate&&) = default;
DeadnessPredicate& operator=(const DeadnessPredicate&) = default;
DeadnessPredicate& operator=(DeadnessPredicate&&) = default;
bool operator==(const DeadnessPredicate& other) const {
return other.pred_ == pred_;
}
bool operator!=(const DeadnessPredicate& other) const {
return other.pred_ != pred_;
}
private:
explicit DeadnessPredicate(void* pred) : pred_(pred) {}
// This is really a Predicate*, but we don't want to expose that
// implementation detail to our clients. `pred_` has pointer equality so we
// can just compare the pointer in operator== and operator!=.
void* pred_;
friend class DeadnessAnalysis;
};
virtual absl::StatusOr<DeadnessPredicate> GetPredicateFor(Node* n,
int oidx) const = 0;
// Prints out the internal state of this instance. For debugging purposes
// only.
virtual void Print() const = 0;
virtual ~DeadnessAnalysis();
std::string DebugString(DeadnessPredicate predicate) const;
// Run the deadness analysis over `graph` and returns an error or a populated
// instance of DeadnessAnalysis in `result`.
static absl::Status Run(const Graph& graph,
std::unique_ptr<DeadnessAnalysis>* result);
protected:
static DeadnessPredicate MakeDeadnessPredicate(void* pred) {
return DeadnessPredicate(pred);
}
};
} // namespace tensorflow
#endif // TENSORFLOW_COMPILER_JIT_DEADNESS_ANALYSIS_H_
tensorflow/compiler/jit/deadness_analysis_internal.h
+36
View File
@@ -0,0 +1,36 @@
/* 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_COMPILER_JIT_DEADNESS_ANALYSIS_INTERNAL_H_
#define TENSORFLOW_COMPILER_JIT_DEADNESS_ANALYSIS_INTERNAL_H_
#include "absl/container/flat_hash_map.h"
#include "tensorflow/core/graph/tensor_id.h"
namespace tensorflow {
namespace deadness_analysis_internal {
// Returns a map describing the predicate each Tensor was mapped to. For
// testing purposes only.
using PredicateMapTy =
absl::flat_hash_map<TensorId, std::string, TensorId::Hasher>;
absl::Status ComputePredicates(const Graph& graph,
PredicateMapTy* out_predicate_map,
bool enable_optimistic = true);
} // namespace deadness_analysis_internal
} // namespace tensorflow
#endif // TENSORFLOW_COMPILER_JIT_DEADNESS_ANALYSIS_INTERNAL_H_
tensorflow/compiler/jit/deadness_analysis_test.cc
+1313
View File
File diff suppressed because it is too large Load Diff
tensorflow/compiler/jit/defs.cc
+45
View File
@@ -0,0 +1,45 @@
/* Copyright 2017 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/compiler/jit/defs.h"
#include <atomic>
namespace tensorflow {
const char* const kXlaMustCompileAttr = "_XlaMustCompile";
const char* const kXlaCompileAttr = "_XlaCompile";
// User-provided through jit_scope APIs. Effective only when auto_jit is OFF.
const char* const kXlaScopeAttr = "_XlaScope";
// Automatically inserted by auto_jit to guide clustering results. Effective
// only when auto_jit is ON.
const char* const kXlaInternalScopeAttr = "_XlaInternalScope";
const char* const kXlaClusterIdAttr = "_xla_compile_id";
static std::atomic<bool> xla_devices_creation_required(false);
// Request XLA:GPU and XLA:CPU device creation. Deprecated, only used by XRT
// backend.
void RequestXlaDevicesCreation() { xla_devices_creation_required = true; }
// Check whether XLA:GPU and XLA:CPU device creation was requested. Deprecated,
// only used by XRT backend.
bool XlaDevicesCreationRequired() { return xla_devices_creation_required; }
} // namespace tensorflow
tensorflow/compiler/jit/defs.h
+49
View File
@@ -0,0 +1,49 @@
/* Copyright 2017 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.
==============================================================================*/
// Provides definitions needed for use of the TensorFlow XLA
// device.
#ifndef TENSORFLOW_COMPILER_JIT_DEFS_H_
#define TENSORFLOW_COMPILER_JIT_DEFS_H_
namespace tensorflow {
// Name of attribute used to tag operators for compilation with XLA
// Implies must-compile semantics: either it will be compiled
// with XLA, or an error will be thrown.
extern const char* const kXlaMustCompileAttr; // "_XlaMustCompile"
// Implies auto-clustering: tagged nodes will be clustered and compiled with XLA
// on a best-effort basis.
extern const char* const kXlaCompileAttr; // "_XlaCompile"
// Implies auto-clustering within the given scope.
extern const char* const kXlaScopeAttr; // "_XlaScope"
extern const char* const kXlaInternalScopeAttr; // "_XlaInternalScope"
// The id of the compiled cluster.
extern const char* const kXlaClusterIdAttr; // "_xla_compile_id"
[[deprecated("XLA:CPU/GPU devices are deprecated")]] void
RequestXlaDevicesCreation();
[[deprecated("XLA:CPU/GPU devices are deprecated")]] bool
XlaDevicesCreationRequired();
} // namespace tensorflow
#endif // TENSORFLOW_COMPILER_JIT_DEFS_H_
tensorflow/compiler/jit/device_compilation_cache.h
+280
View File
@@ -0,0 +1,280 @@
/* Copyright 2022 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_COMPILER_JIT_DEVICE_COMPILATION_CACHE_H_
#define TENSORFLOW_COMPILER_JIT_DEVICE_COMPILATION_CACHE_H_
#include <cstdint>
#include <memory>
#include <optional>
#include <string>
#include <utility>
#include "absl/base/nullability.h"
#include "absl/container/flat_hash_map.h"
#include "absl/strings/str_cat.h"
#include "tensorflow/compiler/jit/device_compilation_cluster_signature.h"
#include "tensorflow/compiler/jit/xla_compile_util.h"
#include "tensorflow/compiler/tf2xla/xla_compiler.h"
#include "xla/client/local_client.h"
#include "xla/pjrt/pjrt_client.h"
#include "tensorflow/core/platform/mutex.h"
namespace tensorflow {
namespace device_compilation_cache_internal {
template <typename ExecutableType>
int64_t ExecutableSize(const ExecutableType* executable) {
return 0;
}
template <>
inline int64_t ExecutableSize<xla::LocalExecutable>(
const xla::LocalExecutable* executable) {
if (executable != nullptr && executable->executable() != nullptr) {
return executable->executable()->SizeOfGeneratedCodeInBytes();
}
return 0;
}
template <>
inline int64_t ExecutableSize<xla::PjRtLoadedExecutable>(
const xla::PjRtLoadedExecutable* executable) {
if (executable != nullptr) {
return executable->SizeOfGeneratedCodeInBytes();
}
return 0;
}
} // namespace device_compilation_cache_internal
// Cache to store compiled HLO, executables and related metadata keyed by
// `DeviceCompilationClusterSignature`. The cache owns the stored
// CompilationResults and Executables.
// Currently no cache eviction policy is implemented and the cache grows without
// bound.
template <typename ExecutableType>
class DeviceCompilationCache {
public:
DeviceCompilationCache() = default;
~DeviceCompilationCache() = default;
using Key = DeviceCompilationClusterSignature;
struct Value {
DeviceCompileState compile_state = DeviceCompileState::kUncompiled;
absl::Status compilation_status;
int64_t request_count = 0;
const XlaCompiler::CompilationResult* compilation_result = nullptr;
ExecutableType* executable = nullptr;
};
// Returns std::nullopt if value for the supplied key is not found. If a value
// is found, `request_count` is incremented before returning the value.
std::optional<Value> Lookup(const Key& key) const;
// Inserts an empty value if value is not found and returns it. If a value is
// found, `request_count` is incremented before returning the value.
Value LookupOrCreate(const Key& key);
// Caches `compile_state`, `compilation_status`, `compilation_result` and
// `executable` and associates them with the provided `key`. Takes ownership
// of `compilation_result` and `executable`. Does not increment the
// corresponding `request_count`. Only arguments that are not std::nullopt are
// updated in the cache.
void Store(const Key& key, std::optional<DeviceCompileState> compile_state,
std::optional<absl::Status> compilation_status,
std::optional<std::unique_ptr<XlaCompiler::CompilationResult>>
compilation_result,
std::optional<std::unique_ptr<ExecutableType>> executable);
std::string DebugString() const;
// Erase any cache entries that have a null entry and releases all references
// to `xla::XlaComputation` in the non-null cache entries.
void Finalize() {
const mutex_lock lock(compile_cache_mu_);
absl::erase_if(
cache_,
[&](std::pair<const Key, absl_nullable std::unique_ptr<Entry>>& kv) {
Entry* absl_nullable const entry = kv.second.get();
if (entry == nullptr) {
return true;
}
const mutex_lock entry_lock(entry->mu);
if (entry->compilation_result != nullptr) {
entry->compilation_result->computation.reset();
}
return false;
});
};
private:
// The value associated with a cache entry.
struct Entry {
mutable mutex mu;
// The current compilation state for this entry.
DeviceCompileState compile_state TF_GUARDED_BY(mu) =
DeviceCompileState::kUncompiled;
// The number of times a compilation with this signature has been requested.
int64_t request_count TF_GUARDED_BY(mu) = 0;
// Did compilation succeed?
absl::Status compilation_status TF_GUARDED_BY(mu);
// Output of the XlaCompiler.
std::unique_ptr<XlaCompiler::CompilationResult> compilation_result
TF_GUARDED_BY(mu);
// The XLA executable compiled from <computation>. May be null if no
// executable has been built.
std::unique_ptr<ExecutableType> executable TF_GUARDED_BY(mu);
std::string DebugString() const {
mutex_lock lock(mu);
int64_t executable_size =
device_compilation_cache_internal::ExecutableSize<ExecutableType>(
executable.get());
int64_t hlo_module_size = 0;
if (compilation_result != nullptr &&
compilation_result->computation != nullptr) {
hlo_module_size =
compilation_result->computation->proto().ByteSizeLong();
}
return absl::StrCat(
"{compile_state: ", compile_state, ", request_count: ", request_count,
", compilation_status: ", compilation_status.ToString(),
", compilation_result?: ", compilation_result != nullptr,
", hlo_module_size: ", hlo_module_size, " bytes",
", executable?: ", executable != nullptr,
", executable_size: ", executable_size, " bytes}");
}
};
mutable mutex compile_cache_mu_;
absl::flat_hash_map<Key, std::unique_ptr<Entry>, Key::Hash> cache_
TF_GUARDED_BY(compile_cache_mu_);
DeviceCompilationCache(const DeviceCompilationCache&) = delete;
void operator=(const DeviceCompilationCache&) = delete;
};
template <typename ExecutableType>
std::optional<typename DeviceCompilationCache<ExecutableType>::Value>
DeviceCompilationCache<ExecutableType>::Lookup(const Key& key) const {
// The outer lock protects the existence of the cache entry. It does not
// protect the contents of the cache entry.
Entry* entry;
{
mutex_lock lock(compile_cache_mu_);
// Find cache entry.
auto it = cache_.find(key);
if (it == cache_.cend()) {
return std::nullopt;
}
entry = it->second.get();
}
mutex_lock lock(entry->mu);
Value value = {/*compile_state=*/entry->compile_state,
/*compilation_status=*/entry->compilation_status,
/*request_count=*/++entry->request_count,
/*compilation_result=*/entry->compilation_result.get(),
/*executable=*/entry->executable.get()};
return value;
}
template <typename ExecutableType>
typename DeviceCompilationCache<ExecutableType>::Value
DeviceCompilationCache<ExecutableType>::LookupOrCreate(const Key& key) {
// The outer lock protects the existence of the cache entry. It does not
// protect the contents of the cache entry.
Entry* entry;
{
mutex_lock lock(compile_cache_mu_);
// Emplace empty cache entry if not found.
auto it = cache_.emplace(key, std::make_unique<Entry>()).first;
entry = it->second.get();
}
mutex_lock lock(entry->mu);
Value value = {/*compile_state=*/entry->compile_state,
/*compilation_status=*/entry->compilation_status,
/*request_count=*/++entry->request_count,
/*compilation_result=*/entry->compilation_result.get(),
/*executable=*/entry->executable.get()};
return value;
}
template <typename ExecutableType>
void DeviceCompilationCache<ExecutableType>::Store(
const Key& key, std::optional<DeviceCompileState> compile_state,
std::optional<absl::Status> compilation_status,
std::optional<std::unique_ptr<XlaCompiler::CompilationResult>>
compilation_result,
std::optional<std::unique_ptr<ExecutableType>> executable) {
Entry* entry;
{
mutex_lock lock(compile_cache_mu_);
// Emplace empty cache entry if not found.
auto it = cache_.emplace(key, std::make_unique<Entry>()).first;
entry = it->second.get();
}
{
mutex_lock lock(entry->mu);
if (compile_state.has_value()) {
entry->compile_state = *compile_state;
}
if (compilation_status.has_value()) {
entry->compilation_status = *compilation_status;
}
if (compilation_result.has_value()) {
entry->compilation_result = std::move(*compilation_result);
}
if (executable.has_value()) {
entry->executable = std::move(*executable);
}
}
VLOG(4) << "Added/updated cache entry: key=" << key.HumanString()
<< ", entry=" << entry->DebugString();
}
template <typename ExecutableType>
std::string DeviceCompilationCache<ExecutableType>::DebugString() const {
std::string s = "DeviceCompilationCache<ExecutableType> {\n";
{
mutex_lock lock(compile_cache_mu_);
for (const auto& [key, entry] : cache_) {
absl::StrAppend(&s, key.HumanString(), " : ", entry->DebugString(),
",\n");
}
}
absl::StrAppend(&s, "}");
return s;
}
} // namespace tensorflow
#endif // TENSORFLOW_COMPILER_JIT_DEVICE_COMPILATION_CACHE_H_
tensorflow/compiler/jit/device_compilation_cache_test.cc
+277
View File
@@ -0,0 +1,277 @@
/* Copyright 2022 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/compiler/jit/device_compilation_cache.h"
#include <memory>
#include <optional>
#include <string>
#include <utility>
#include <vector>
#include <gtest/gtest.h>
#include "absl/status/status.h"
#include "absl/status/statusor.h"
#include "tensorflow/compiler/jit/xla_compile_util.h"
#include "tensorflow/compiler/tf2xla/xla_compiler.h"
#include "xla/hlo/builder/xla_computation.h"
#include "xla/tsl/protobuf/error_codes.pb.h"
#include "tensorflow/core/platform/errors.h"
#include "tensorflow/core/platform/test.h"
namespace tensorflow {
namespace {
struct FakeExecutable {
std::string data;
explicit FakeExecutable(const std::string& s) : data(s) {}
};
using Cache = DeviceCompilationCache<FakeExecutable>;
using Signature = DeviceCompilationClusterSignature;
absl::StatusOr<Signature> BuildSampleSignature(const std::string& fn_name) {
NameAttrList fn;
fn.set_name(fn_name);
std::vector<XlaCompiler::Argument> args(1);
args[0].kind = XlaCompiler::Argument::kConstant;
args[0].type = DT_INT32;
args[0].shape = TensorShape({4, 0});
args[0].constant_value = Tensor(DT_INT32, {4, 0});
return Signature::Build(fn, args);
}
TEST(DeviceCompilationCacheTest, LookupEntryDoesntExist) {
auto cache = std::make_unique<Cache>();
TF_ASSERT_OK_AND_ASSIGN(auto key, BuildSampleSignature("foo"));
auto cache_value = cache->Lookup(key);
EXPECT_FALSE(cache_value.has_value());
}
TEST(DeviceCompilationCacheTest, LookupOrCreateEntryDoesntExist) {
auto cache = std::make_unique<Cache>();
TF_ASSERT_OK_AND_ASSIGN(auto key, BuildSampleSignature("foo"));
Cache::Value cache_value = cache->LookupOrCreate(key);
EXPECT_EQ(cache_value.compile_state, DeviceCompileState::kUncompiled);
EXPECT_EQ(cache_value.request_count, 1);
EXPECT_EQ(cache_value.compilation_result, nullptr);
EXPECT_EQ(cache_value.executable, nullptr);
}
TEST(DeviceCompilationCacheTest, IncrementRequestCountOnLookup) {
auto cache = std::make_unique<Cache>();
TF_ASSERT_OK_AND_ASSIGN(auto key, BuildSampleSignature("foo"));
Cache::Value cache_value = cache->LookupOrCreate(key);
EXPECT_EQ(cache_value.request_count, 1);
cache_value = cache->LookupOrCreate(key);
EXPECT_EQ(cache_value.request_count, 2);
cache_value = cache->LookupOrCreate(key);
EXPECT_EQ(cache_value.request_count, 3);
}
TEST(DeviceCompilationCacheTest, RequestCountUnchangedOnStore) {
auto cache = std::make_unique<Cache>();
TF_ASSERT_OK_AND_ASSIGN(auto key, BuildSampleSignature("foo"));
Cache::Value cache_value = cache->LookupOrCreate(key);
EXPECT_EQ(cache_value.request_count, 1);
cache_value = cache->LookupOrCreate(key);
EXPECT_EQ(cache_value.request_count, 2);
cache_value = cache->LookupOrCreate(key);
EXPECT_EQ(cache_value.request_count, 3);
auto compilation_result = std::make_unique<XlaCompiler::CompilationResult>();
cache->Store(key, DeviceCompileState::kCompiled, absl::OkStatus(),
std::move(compilation_result), std::nullopt);
cache_value = cache->LookupOrCreate(key);
EXPECT_EQ(cache_value.request_count, 4);
}
TEST(DeviceCompilationCacheTest, StoreLookup) {
auto cache = std::make_unique<Cache>();
TF_ASSERT_OK_AND_ASSIGN(auto key, BuildSampleSignature("foo"));
auto compilation_result = std::make_unique<XlaCompiler::CompilationResult>();
auto executable = std::make_unique<FakeExecutable>("foo_exe");
cache->Store(key, DeviceCompileState::kCompiled, absl::OkStatus(),
std::move(compilation_result), std::move(executable));
auto cache_value = cache->Lookup(key);
EXPECT_EQ(cache_value->compile_state, DeviceCompileState::kCompiled);
EXPECT_EQ(cache_value->request_count, 1);
EXPECT_TRUE(cache_value->compilation_status.ok());
EXPECT_TRUE(cache_value->compilation_result != nullptr);
EXPECT_TRUE(cache_value->executable != nullptr);
EXPECT_EQ(cache_value->executable->data, "foo_exe");
}
TEST(DeviceCompilationCacheTest, StoreLookupOrCreate) {
auto cache = std::make_unique<Cache>();
TF_ASSERT_OK_AND_ASSIGN(auto key, BuildSampleSignature("foo"));
auto compilation_result = std::make_unique<XlaCompiler::CompilationResult>();
auto executable = std::make_unique<FakeExecutable>("foo_exe");
cache->Store(key, DeviceCompileState::kCompiled, absl::OkStatus(),
std::move(compilation_result), std::move(executable));
auto cache_value = cache->LookupOrCreate(key);
EXPECT_EQ(cache_value.compile_state, DeviceCompileState::kCompiled);
EXPECT_EQ(cache_value.request_count, 1);
EXPECT_TRUE(cache_value.compilation_status.ok());
EXPECT_TRUE(cache_value.compilation_result != nullptr);
EXPECT_TRUE(cache_value.executable != nullptr);
EXPECT_EQ(cache_value.executable->data, "foo_exe");
}
TEST(DeviceCompilationCacheTest, StoreOptionalArgs) {
auto cache = std::make_unique<Cache>();
TF_ASSERT_OK_AND_ASSIGN(auto key, BuildSampleSignature("foo"));
auto compilation_result = std::make_unique<XlaCompiler::CompilationResult>();
auto executable = std::make_unique<FakeExecutable>("foo_exe");
cache->Store(key, DeviceCompileState::kCompiled, std::nullopt, std::nullopt,
std::nullopt);
auto cache_value = cache->Lookup(key);
EXPECT_EQ(cache_value->compile_state, DeviceCompileState::kCompiled);
EXPECT_TRUE(cache_value->compilation_status.ok());
EXPECT_TRUE(cache_value->compilation_result == nullptr);
EXPECT_TRUE(cache_value->executable == nullptr);
cache->Store(key, std::nullopt,
absl::InvalidArgumentError("Couldn't compile."), std::nullopt,
std::nullopt);
cache_value = cache->Lookup(key);
EXPECT_EQ(cache_value->compile_state, DeviceCompileState::kCompiled);
EXPECT_EQ(cache_value->compilation_status.code(), error::INVALID_ARGUMENT);
EXPECT_TRUE(cache_value->compilation_result == nullptr);
EXPECT_TRUE(cache_value->executable == nullptr);
cache->Store(key, std::nullopt, std::nullopt, std::move(compilation_result),
std::nullopt);
cache_value = cache->Lookup(key);
EXPECT_EQ(cache_value->compile_state, DeviceCompileState::kCompiled);
EXPECT_EQ(cache_value->compilation_status.code(), error::INVALID_ARGUMENT);
EXPECT_TRUE(cache_value->compilation_result != nullptr);
EXPECT_TRUE(cache_value->executable == nullptr);
cache->Store(key, std::nullopt, std::nullopt, std::nullopt,
std::move(executable));
cache_value = cache->Lookup(key);
EXPECT_EQ(cache_value->compile_state, DeviceCompileState::kCompiled);
EXPECT_EQ(cache_value->compilation_status.code(), error::INVALID_ARGUMENT);
EXPECT_TRUE(cache_value->compilation_result != nullptr);
EXPECT_TRUE(cache_value->executable != nullptr);
EXPECT_EQ(cache_value->executable->data, "foo_exe");
}
TEST(DeviceCompilationCacheTest, StoreMultipleEntries) {
auto cache = std::make_unique<Cache>();
TF_ASSERT_OK_AND_ASSIGN(auto key1, BuildSampleSignature("foo"));
TF_ASSERT_OK_AND_ASSIGN(auto key2, BuildSampleSignature("bar"));
auto compilation_result1 = std::make_unique<XlaCompiler::CompilationResult>();
auto compilation_result2 = std::make_unique<XlaCompiler::CompilationResult>();
auto executable1 = std::make_unique<FakeExecutable>("foo_exe");
auto executable2 = std::make_unique<FakeExecutable>("bar_exe");
cache->Store(key1, DeviceCompileState::kCompiled,
absl::InvalidArgumentError("Invalid argument."),
std::move(compilation_result1), std::move(executable1));
cache->Store(key2, DeviceCompileState::kCompiling, absl::OkStatus(),
std::move(compilation_result2), std::move(executable2));
auto cache_value_1 = cache->Lookup(key1);
auto cache_value_2 = cache->Lookup(key2);
EXPECT_EQ(cache_value_1->compile_state, DeviceCompileState::kCompiled);
EXPECT_EQ(cache_value_1->compilation_status.code(), error::INVALID_ARGUMENT);
EXPECT_TRUE(cache_value_1->compilation_result != nullptr);
EXPECT_TRUE(cache_value_1->executable != nullptr);
EXPECT_EQ(cache_value_1->executable->data, "foo_exe");
EXPECT_EQ(cache_value_2->compile_state, DeviceCompileState::kCompiling);
EXPECT_TRUE(cache_value_2->compilation_status.ok());
EXPECT_TRUE(cache_value_2->compilation_result != nullptr);
EXPECT_TRUE(cache_value_2->executable != nullptr);
EXPECT_EQ(cache_value_2->executable->data, "bar_exe");
}
TEST(DeviceCompilationCacheTest, Finalize) {
auto cache = std::make_unique<Cache>();
TF_ASSERT_OK_AND_ASSIGN(auto key, BuildSampleSignature("foo"));
auto computation = std::make_shared<xla::XlaComputation>();
auto compilation_result = std::make_unique<XlaCompiler::CompilationResult>();
compilation_result->computation = computation;
cache->Store(key, /*compile_state=*/DeviceCompileState::kCompiled,
/*compilation_status=*/absl::OkStatus(),
std::move(compilation_result),
/*executable=*/std::make_unique<FakeExecutable>("foo_exe"));
std::optional<Cache::Value> cache_value = cache->Lookup(key);
ASSERT_TRUE(cache_value.has_value());
ASSERT_TRUE(cache_value->compilation_result != nullptr);
EXPECT_EQ(cache_value->compilation_result->computation, computation);
cache->Finalize();
cache_value = cache->Lookup(key);
ASSERT_TRUE(cache_value.has_value());
ASSERT_TRUE(cache_value->compilation_result != nullptr);
EXPECT_TRUE(cache_value->compilation_result->computation == nullptr);
}
TEST(DeviceCompilationCache, FinalizeWithNullComputation) {
auto cache = std::make_unique<Cache>();
TF_ASSERT_OK_AND_ASSIGN(auto key, BuildSampleSignature("foo"));
cache->Store(
key, /*compile_state=*/DeviceCompileState::kCompiled,
/*compilation_status=*/absl::OkStatus(),
/*compilation_result=*/std::make_unique<XlaCompiler::CompilationResult>(),
/*executable=*/std::make_unique<FakeExecutable>("foo_exe"));
std::optional<Cache::Value> cache_value = cache->Lookup(key);
ASSERT_TRUE(cache_value.has_value());
ASSERT_TRUE(cache_value->compilation_result != nullptr);
EXPECT_TRUE(cache_value->compilation_result->computation == nullptr);
cache->Finalize();
cache_value = cache->Lookup(key);
ASSERT_TRUE(cache_value.has_value());
ASSERT_TRUE(cache_value->compilation_result != nullptr);
EXPECT_TRUE(cache_value->compilation_result->computation == nullptr);
}
} // namespace
} // namespace tensorflow
tensorflow/compiler/jit/device_compilation_cluster_signature.cc
+158
View File
@@ -0,0 +1,158 @@
/* Copyright 2022 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/compiler/jit/device_compilation_cluster_signature.h"
#include <string>
#include <utility>
#include <variant>
#include "absl/status/statusor.h"
#include "absl/types/span.h"
#include "tensorflow/compiler/tf2xla/xla_compiler.h"
#include "tensorflow/core/framework/function.h"
#include "tensorflow/core/framework/node_def_util.h"
namespace tensorflow {
namespace {
using Signature = DeviceCompilationClusterSignature;
using TensorTypeAndShape = Signature::TensorTypeAndShape;
// Functor that converts a Signature's arg to a human readable string.
struct SignatureHumanStringAppender {
explicit SignatureHumanStringAppender(std::string* dest) : dest(dest) {}
std::string* dest;
void operator()(const Tensor& arg) {
absl::StrAppend(dest, "; ", arg.DebugString());
}
void operator()(const TensorTypeAndShape& arg) {
absl::StrAppend(dest, ",", DataTypeString(arg.first));
absl::StrAppend(dest, " [", absl::StrJoin(arg.second, ","), "]");
}
};
// Functor that compares the arg values of two different signatures. Returns
// true when the args are not equal.
struct SignatureNotEqual {
bool operator()(const Tensor& arg, const Tensor& other) {
return arg.dtype() != other.dtype() || arg.shape() != other.shape() ||
arg.tensor_data() != other.tensor_data();
}
bool operator()(const TensorTypeAndShape& arg,
const TensorTypeAndShape& other) {
return arg.first != other.first || arg.second != other.second;
}
bool operator()(const Tensor& arg, const TensorTypeAndShape& other) {
return true;
}
bool operator()(const TensorTypeAndShape& arg, const Tensor& other) {
return true;
}
};
// Functor that incrementally computes a Signature's hash given its current hash
// and one of its args.
struct SignatureHashCombiner {
explicit SignatureHashCombiner(const uint64_t h) : h(h) {}
uint64_t h;
uint64_t operator()(const Tensor& arg) {
h = Hash64Combine(h, std::hash<int>()(static_cast<int>(arg.dtype())));
h = Hash64Combine(
h, Hash64(arg.tensor_data().data(), arg.tensor_data().size()));
for (int dim = 0; dim < arg.dims(); ++dim) {
h = Hash64Combine(h, std::hash<int>()(arg.dim_size(dim)));
}
return h;
}
uint64_t operator()(const TensorTypeAndShape& arg) {
h = Hash64Combine(h, std::hash<int>()(static_cast<int>(arg.first)));
h = Hash64Combine(h, std::hash<int>()(arg.second.size()));
for (int dim : arg.second) {
h = Hash64Combine(h, std::hash<int>()(dim));
}
return h;
}
};
} // namespace
// Compute a string signature which encodes the shapes of the
// arguments in the supplied list.
std::string Signature::HumanString() const {
std::string result = name;
for (const auto& arg : args) {
std::visit(SignatureHumanStringAppender(&result), arg);
}
return result;
}
bool Signature::operator==(const Signature& other) const {
if (name != other.name) return false;
if (args.size() != other.args.size()) return false;
for (int i = 0, end = args.size(); i < end; ++i) {
if (std::visit(SignatureNotEqual(), args[i], other.args[i])) {
return false;
}
}
return true;
}
uint64_t Signature::Hash::operator()(const Signature& signature) const {
uint64_t h = std::hash<std::string>()(signature.name);
for (const auto& arg : signature.args) {
h = std::visit(SignatureHashCombiner(h), arg);
}
return h;
}
static absl::StatusOr<Signature> AppendArguments(
Signature signature, absl::Span<const XlaCompiler::Argument> args) {
signature.args.reserve(args.size());
for (const XlaCompiler::Argument& arg : args) {
switch (arg.kind) {
case XlaCompiler::Argument::kConstant:
case XlaCompiler::Argument::kConstantResource:
signature.args.push_back(arg.constant_value);
break;
case XlaCompiler::Argument::kParameter:
case XlaCompiler::Argument::kResource:
signature.args.push_back(
TensorTypeAndShape(arg.type, arg.DimensionSizesAsInlinedVector()));
break;
default:
return absl::InvalidArgumentError(
absl::StrCat("Unhandled argument kind in XlaCompilationCache: ",
arg.HumanString()));
}
}
return signature;
}
absl::StatusOr<Signature> Signature::Build(
const NameAttrList& function,
absl::Span<const XlaCompiler::Argument> args) {
Signature signature;
signature.name = Canonicalize(function.name(), AttrSlice(&function.attr()));
return AppendArguments(std::move(signature), args);
}
absl::StatusOr<DeviceCompilationClusterSignature> Signature::Build(
const DeviceCompilationCanonicalFunction& canonical_function,
absl::Span<const XlaCompiler::Argument> args) {
Signature signature;
signature.name = canonical_function.canonical;
return AppendArguments(std::move(signature), args);
}
} // namespace tensorflow
tensorflow/compiler/jit/device_compilation_cluster_signature.h
+80
View File
@@ -0,0 +1,80 @@
/* Copyright 2022 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_COMPILER_JIT_DEVICE_COMPILATION_CLUSTER_SIGNATURE_H_
#define TENSORFLOW_COMPILER_JIT_DEVICE_COMPILATION_CLUSTER_SIGNATURE_H_
#include <cstdint>
#include <string>
#include <utility>
#include <variant>
#include "absl/container/inlined_vector.h"
#include "absl/status/statusor.h"
#include "absl/types/span.h"
#include "tensorflow/compiler/tf2xla/xla_compiler.h"
#include "tensorflow/core/framework/function.h"
#include "tensorflow/core/framework/node_def_util.h"
#include "tensorflow/core/framework/tensor.h"
#include "tensorflow/core/platform/types.h"
namespace tensorflow {
// Pre-computed device compilation canonical function.
struct DeviceCompilationCanonicalFunction {
std::string canonical;
};
inline DeviceCompilationCanonicalFunction Canonicalize(
const NameAttrList& function) {
return {Canonicalize(function.name(), AttrSlice(&function.attr()))};
}
// Describes the types, shapes and any compile-time constant arguments
// to a kernel. Key that uniquely identifies a compilation output.
struct DeviceCompilationClusterSignature {
// Name of the cluster, built from the function name and it's attributes.
std::string name;
// List of args (either as a TensorTypeAndShape or as a Tensor value)
// for compile-time constant arguments to the compilation, ordered by
// argument number. Tensors must be in host memory.
using TensorTypeAndShape =
std::pair<DataType, absl::InlinedVector<int64_t, 4>>;
absl::InlinedVector<std::variant<Tensor, TensorTypeAndShape>, 8> args;
bool operator==(const DeviceCompilationClusterSignature& other) const;
struct Hash {
uint64_t operator()(
const DeviceCompilationClusterSignature& signature) const;
};
// Returns a human-readable description of the signature.
std::string HumanString() const;
// Builds the signature for a compilation.
static absl::StatusOr<DeviceCompilationClusterSignature> Build(
const NameAttrList& function,
absl::Span<const XlaCompiler::Argument> args);
static absl::StatusOr<DeviceCompilationClusterSignature> Build(
const DeviceCompilationCanonicalFunction& canonical_function,
absl::Span<const XlaCompiler::Argument> args);
};
} // namespace tensorflow
#endif // TENSORFLOW_COMPILER_JIT_DEVICE_COMPILATION_CLUSTER_SIGNATURE_H_
tensorflow/compiler/jit/device_compilation_cluster_signature_test.cc
+123
View File
@@ -0,0 +1,123 @@
/* Copyright 2022 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/compiler/jit/device_compilation_cluster_signature.h"
#include <utility>
#include <vector>
#include "tensorflow/compiler/jit/flags.h"
#include "tensorflow/compiler/tf2xla/shape_util.h"
#include "xla/client/client_library.h"
#include "tensorflow/core/platform/test.h"
#include "tensorflow/core/platform/test_benchmark.h"
namespace tensorflow {
namespace {
using SignatureHash = DeviceCompilationClusterSignature::Hash;
TEST(DeviceCompilationClusterSignatureTest, SignatureEquality) {
NameAttrList fn;
fn.set_name("afunction");
std::vector<XlaCompiler::Argument> args(1);
args[0].kind = XlaCompiler::Argument::kConstant;
args[0].type = DT_INT32;
args[0].shape = TensorShape({4, 0});
args[0].constant_value = Tensor(DT_INT32, {4, 0});
TF_ASSERT_OK_AND_ASSIGN(DeviceCompilationClusterSignature s1,
DeviceCompilationClusterSignature::Build(fn, args));
args[0].type = DT_FLOAT;
args[0].constant_value = Tensor(DT_FLOAT, {4, 0});
TF_ASSERT_OK_AND_ASSIGN(DeviceCompilationClusterSignature s2,
DeviceCompilationClusterSignature::Build(fn, args));
args[0].shape = TensorShape({0, 4});
args[0].constant_value = Tensor(DT_FLOAT, {0, 4});
TF_ASSERT_OK_AND_ASSIGN(DeviceCompilationClusterSignature s3,
DeviceCompilationClusterSignature::Build(fn, args));
std::vector<DeviceCompilationClusterSignature> signatures = {s1, s2, s3};
for (int i = 0; i < signatures.size(); ++i) {
for (int j = 0; j < signatures.size(); ++j) {
EXPECT_EQ(i == j, signatures[i] == signatures[j])
<< "s1: " << signatures[i].HumanString() << "\n"
<< "s2: " << signatures[j].HumanString();
EXPECT_EQ(i == j,
signatures[i].HumanString() == signatures[j].HumanString())
<< "s1: " << signatures[i].HumanString() << "\n"
<< "s2: " << signatures[j].HumanString();
EXPECT_EQ(i == j, SignatureHash()(signatures[i]) ==
SignatureHash()(signatures[j]))
<< "s1: " << signatures[i].HumanString() << "\n"
<< "s1_hash: " << SignatureHash()(signatures[i]) << "\n"
<< "s2: " << signatures[j].HumanString() << "\n"
<< "s2_hash: " << SignatureHash()(signatures[j]);
}
}
}
TEST(DeviceCompilationClusterSignatureTest, SignatureUniqueness) {
NameAttrList fn;
fn.set_name("afunction");
std::vector<XlaCompiler::Argument> args(2);
args[0].kind = XlaCompiler::Argument::kConstant;
args[0].type = DT_INT32;
args[0].constant_value = Tensor(DT_INT32, {4, 0});
args[1].kind = XlaCompiler::Argument::kParameter;
args[1].type = DT_INT32;
args[1].shape = TensorShape({4, 0});
TF_ASSERT_OK_AND_ASSIGN(DeviceCompilationClusterSignature s1,
DeviceCompilationClusterSignature::Build(fn, args));
using std::swap; // go/using-std-swap
swap(args[0], args[1]);
TF_ASSERT_OK_AND_ASSIGN(DeviceCompilationClusterSignature s2,
DeviceCompilationClusterSignature::Build(fn, args));
EXPECT_NE(s1.HumanString(), s2.HumanString());
EXPECT_NE(SignatureHash()(s1), SignatureHash()(s2));
EXPECT_FALSE(s1 == s2);
}
void BM_BuildSignature(::testing::benchmark::State& state) {
const int n_args = state.range(0);
NameAttrList fn;
fn.set_name("afunction");
for (int i = 0; i < n_args; i++) {
(*fn.mutable_attr())[absl::StrCat("T", i)].set_type(DT_FLOAT);
}
std::vector<XlaCompiler::Argument> args(n_args);
for (int i = 0; i < n_args; i++) {
args[i].kind = (((i % 3) == 0) ? XlaCompiler::Argument::kConstant
: XlaCompiler::Argument::kParameter);
args[i].type = DT_INT32;
args[i].shape = TensorShape({4, 0});
args[i].constant_value = Tensor(DT_INT32, {4, 0});
}
for (auto i : state) {
auto s = DeviceCompilationClusterSignature::Build(fn, args);
CHECK(s.ok());
DeviceCompilationClusterSignature sig = std::move(s.value());
}
}
BENCHMARK(BM_BuildSignature)->Arg(0)->Arg(1)->Arg(2)->Arg(5)->Arg(10);
} // namespace
} // namespace tensorflow
tensorflow/compiler/jit/device_compilation_profiler.cc
+233
View File
@@ -0,0 +1,233 @@
/* Copyright 2022 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/compiler/jit/device_compilation_profiler.h"
#include <cstdint>
#include <optional>
#include <string>
#include <utility>
#include "absl/strings/str_cat.h"
#include "tensorflow/compiler/jit/xla_activity.pb.h"
#include "tensorflow/compiler/jit/xla_activity_listener.h"
#include "tensorflow/core/framework/attr_value.pb.h"
#include "tensorflow/core/framework/metrics.h"
#include "tensorflow/core/platform/errors.h"
#include "tensorflow/core/platform/status.h"
#include "tsl/platform/mutex.h"
namespace tensorflow {
namespace {
bool ShouldBeMegamorphic(int64_t compile_count, int64_t execution_count) {
const int64_t kCompileThreshold = 10;
const int64_t kMinExecutionsPerCompile = 50;
// This heuristic is trying to capture the following property: have we sunk a
// certain minimum amount of compile time into the cluster that didn't quite
// "pay off"?
return compile_count > kCompileThreshold &&
execution_count < kMinExecutionsPerCompile * compile_count;
}
void RegisterExecutionForCluster(
const NameAttrList& function,
DeviceCompilationProfiler::ClusterCompileStats* stats) {
++stats->execution_count;
// The is_megamorphic bit is "sticky". We assume clusters that have been
// observed to be megamorphic once stay megamorphic forever.
if (!stats->is_megamorphic &&
ShouldBeMegamorphic(stats->compile_count, stats->execution_count)) {
VLOG(1) << "Marking " << function.name()
<< " as megamorphic, compile_count=" << stats->compile_count
<< " execution_count=" << stats->execution_count;
stats->is_megamorphic = true;
}
}
// The number of times a lazy compilation must be requested for a specific
// signature before we attempt to compile it.
constexpr int64_t kDefaultCompilationThreshold = 2;
// Maximum number of ongoing compilations.
constexpr int64_t kMaxNumOngoingCompilations = kNumAsyncDeviceCompilerThreads;
} // namespace
DeviceCompilationProfiler::~DeviceCompilationProfiler() {
mutex_lock lock(mu_);
cluster_compile_stats_.clear();
}
absl::StatusOr<DeviceCompilationProfiler::ClusterCompileStats>
DeviceCompilationProfiler::GetCompileStats(const NameAttrList& function) const {
mutex_lock lock(mu_);
if (auto it = cluster_compile_stats_.find(function.name());
it != cluster_compile_stats_.end()) {
return it->second;
}
return absl::NotFoundError(absl::StrCat(
"Couldn't find compilation stats for cluster: ", function.name()));
}
void DeviceCompilationProfiler::RegisterExecution(
const NameAttrList& function) {
mutex_lock lock(mu_);
auto it =
cluster_compile_stats_.emplace(function.name(), ClusterCompileStats{})
.first;
RegisterExecutionForCluster(function, &it->second);
}
absl::Status DeviceCompilationProfiler::RegisterCompilation(
const NameAttrList& function, int64_t compile_time_us,
bool used_persistent_cache) {
metrics::UpdateXlaCompilationTime(compile_time_us);
const std::string& function_name = function.name();
mutex_lock lock(mu_);
// Create a stats entry if it doesn't already exist.
auto it =
cluster_compile_stats_.emplace(function.name(), ClusterCompileStats{})
.first;
const uint64_t compile_time_s = compile_time_us / 1.0e6;
it->second.compile_count++;
it->second.cumulative_compile_time_us += compile_time_us;
VLOG(1) << "Compiled " << function_name << " " << it->second.compile_count
<< " times, compile time: " << compile_time_us
<< " us, cumulative: " << it->second.cumulative_compile_time_us
<< " us ("
<< tensorflow::strings::HumanReadableElapsedTime(compile_time_s)
<< " / "
<< tensorflow::strings::HumanReadableElapsedTime(
it->second.cumulative_compile_time_us / 1.0e6)
<< ")";
XlaJitCompilationActivity jit_compilation_activity;
jit_compilation_activity.set_cluster_name(function_name);
jit_compilation_activity.set_compile_count(it->second.compile_count);
jit_compilation_activity.set_compile_time_us(compile_time_us);
jit_compilation_activity.set_cumulative_compile_time_us(
it->second.cumulative_compile_time_us);
jit_compilation_activity.set_used_persistent_cache(used_persistent_cache);
return BroadcastXlaActivity(std::move(jit_compilation_activity));
}
bool DeviceCompilationProfiler::ShouldCompileCluster(
const NameAttrList& function, DeviceCompileMode compile_mode,
int64_t current_request_count) {
std::optional<int64_t> compile_threshold;
if (compile_mode == DeviceCompileMode::kLazy) {
compile_threshold = kDefaultCompilationThreshold;
} else if (compile_mode == DeviceCompileMode::kAsync) {
compile_threshold = 0; // for now, always compile right away.
}
if (compile_mode == DeviceCompileMode::kStrict) {
// Lazy compilation is disabled.
return true;
}
mutex_lock lock(mu_);
// Create a stats entry if one isn't found and register an execution.
// Determine eligibility assuming this is the first execution of the cluster
// and this cluster has never been compiled before.
auto [it, cluster_not_found] =
cluster_compile_stats_.emplace(function.name(), ClusterCompileStats{});
if (cluster_not_found) {
RegisterExecutionForCluster(function, &it->second);
}
// We avoid compiling clusters that have "gone megamorphic" i.e. have an
// excessive amount of shape dynamism.
if (it->second.is_megamorphic) {
BroadcastOptimizationRemark(XlaOptimizationRemark::MEGAMORPHIC_FUNCTION,
function.name())
.IgnoreError();
VLOG(2) << "Not compiling cluster " << function.name()
<< " because it is megamorphic.";
return false;
}
// TODO(b/255826209): Figure out if Lazy compilation is still needed given
// that we always compile a cluster the first time it is executed (explained
// below) regardless of compilation mode. If it is not, clean up the related
// logic.
// We always compile a cluster the very first time it is executed. This is an
// optimistic guess that pays off for statically shaped TensorFlow graphs
// (since they get the benefit of XLA right away without waiting for warmup)
// and doesn't hurt much for dynamically shaped TensorFlow graphs (we "pay" at
// most one cluster-compilation's worth of compile time).
if (it->second.execution_count == 1) {
return true;
}
if (compile_mode == DeviceCompileMode::kAsync) {
// Asynchronous compilation is enabled.
if (num_ongoing_compilations_ >= kMaxNumOngoingCompilations) {
VLOG(2) << "Not asynchronously compiling cluster " << function.name()
<< " because of too many ongoing compilations.";
return false;
}
}
bool reached_compile_threshold = current_request_count >= *compile_threshold;
if (!reached_compile_threshold) {
VLOG(2) << "Not compiling cluster " << function.name()
<< " because it has not reached compile threshold; threshold is "
<< *compile_threshold << " execution count "
<< current_request_count << ".";
}
return reached_compile_threshold;
}
void DeviceCompilationProfiler::IncrementOngoingAsyncCompilations() {
mutex_lock lock(mu_);
num_ongoing_compilations_++;
}
void DeviceCompilationProfiler::DecrementOngoingAsyncCompilations() {
mutex_lock lock(mu_);
num_ongoing_compilations_--;
}
int64_t DeviceCompilationProfiler::GetNumOngoingAsyncCompilations() const {
mutex_lock lock(mu_);
return num_ongoing_compilations_;
}
std::string DeviceCompilationProfiler::DebugString() const {
std::string debug_string =
"DeviceCompilationProfiler {\ncluster_compile_stats: {\n";
{
mutex_lock lock(mu_);
for (const auto& [key, stats] : cluster_compile_stats_) {
absl::StrAppend(&debug_string, key, ": ", stats.DebugString(), "\n");
}
}
absl::StrAppend(&debug_string, "}\nnum_ongoing_compilations=",
GetNumOngoingAsyncCompilations(), "\n}\n");
return debug_string;
}
} // namespace tensorflow
tensorflow/compiler/jit/device_compilation_profiler.h
+101
View File
@@ -0,0 +1,101 @@
/* Copyright 2022 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_COMPILER_JIT_DEVICE_COMPILATION_PROFILER_H_
#define TENSORFLOW_COMPILER_JIT_DEVICE_COMPILATION_PROFILER_H_
#include <cstdint>
#include <string>
#include "tensorflow/compiler/jit/xla_compile_util.h"
#include "tensorflow/core/framework/attr_value.pb.h"
namespace tensorflow {
// Tracks statistics for device compilation and uses these to determine whether
// the given cluster should be compiled or not.
class DeviceCompilationProfiler : public ResourceBase {
public:
DeviceCompilationProfiler() = default;
~DeviceCompilationProfiler() override;
struct ClusterCompileStats {
// Number of times the cluster has been (re-)compiled.
int64_t compile_count = 0;
// The number of times this cluster has been executed.
int64_t execution_count = 0;
// Cumulative time spent compiling the cluster.
int64_t cumulative_compile_time_us = 0;
// True if we have decided that this cluster is too dynamic (i.e. its shapes
// change too frequently) to profitably JIT compile. Once a cluster is
// tagged megamorphic, it stays megamorphic forever.
bool is_megamorphic = false;
std::string DebugString() const {
return absl::StrCat(
"DeviceCompilationProfiler::ClusterCompileStats {compile_count=",
compile_count, ", execution_count=", execution_count,
", cumulative_compile_time_us=", cumulative_compile_time_us,
", is_megamorphic=", is_megamorphic, "}");
}
};
// Returns the compilation statistics for the given cluster.
absl::StatusOr<ClusterCompileStats> GetCompileStats(
const NameAttrList& function) const;
// Determines whether the cluster should be compiled. Creates and inserts an
// entry into stats (also calls `RegisterExecution`) for `function` if it
// doesn't already exist.
virtual bool ShouldCompileCluster(const NameAttrList& function,
DeviceCompileMode compile_mode,
int64_t current_request_count);
// Registers a cluster execution. Increments the execution count for the given
// cluster and also determines whether the cluster has gone megamorphic (and
// sets the megamorphic bit accordingly).
void RegisterExecution(const NameAttrList& function);
// Registers a cluster compilation. Increments the compilation count and
// accumulates the compile time for the given cluster. Also broadcasts an
// XlaJitCompilationActivity.
virtual absl::Status RegisterCompilation(const NameAttrList& function,
int64_t compile_time_us,
bool used_persistent_cache);
void IncrementOngoingAsyncCompilations();
void DecrementOngoingAsyncCompilations();
int64_t GetNumOngoingAsyncCompilations() const;
std::string DebugString() const override;
private:
mutable mutex mu_;
// Maps cluster names to compilation statistics for said cluster.
absl::flat_hash_map<std::string, ClusterCompileStats> cluster_compile_stats_
TF_GUARDED_BY(mu_);
int64_t num_ongoing_compilations_ TF_GUARDED_BY(mu_) = 0;
DeviceCompilationProfiler(const DeviceCompilationProfiler&) = delete;
void operator=(const DeviceCompilationProfiler&) = delete;
};
} // namespace tensorflow
#endif // TENSORFLOW_COMPILER_JIT_DEVICE_COMPILATION_PROFILER_H_
tensorflow/compiler/jit/device_compilation_profiler_test.cc
+243
View File
@@ -0,0 +1,243 @@
/* Copyright 2022 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/compiler/jit/device_compilation_profiler.h"
#include <memory>
#include <utility>
#include <vector>
#include <gmock/gmock.h>
#include <gtest/gtest.h>
#include "tensorflow/compiler/jit/tests/device_compiler_test_helper.h"
#include "tensorflow/compiler/jit/xla_activity.pb.h"
#include "tensorflow/core/framework/attr_value.pb.h"
namespace tensorflow {
namespace {
TEST(DeviceCompilationProfilerTest, RegisterExecution) {
DeviceCompilationProfiler* profiler = new DeviceCompilationProfiler();
core::ScopedUnref profiler_ref(profiler);
NameAttrList function;
function.set_name("TestFunc");
for (int i = 0; i < 5; ++i) {
profiler->RegisterExecution(function);
}
TF_ASSERT_OK_AND_ASSIGN(auto stats, profiler->GetCompileStats(function));
EXPECT_EQ(stats.execution_count, 5);
}
TEST(DeviceCompilationProfilerTest, RegisterCompilation) {
DeviceCompilationProfiler* profiler = new DeviceCompilationProfiler();
core::ScopedUnref profiler_ref(profiler);
auto listener = std::make_unique<JitCompilationListener>();
auto listener_ptr = listener.get();
RegisterXlaActivityListener(std::move(listener));
NameAttrList function;
function.set_name("TestFunc");
std::vector<XlaJitCompilationActivity> expected_activities;
for (int i = 0; i < 5; ++i) {
EXPECT_TRUE(profiler->RegisterCompilation(function, 4, false).ok());
TF_ASSERT_OK_AND_ASSIGN(auto stats, profiler->GetCompileStats(function));
XlaJitCompilationActivity expected_activity;
expected_activity.set_cluster_name(function.name());
expected_activity.set_compile_count(stats.compile_count);
expected_activity.set_compile_time_us(4);
expected_activity.set_cumulative_compile_time_us(
stats.cumulative_compile_time_us);
expected_activity.set_used_persistent_cache(false);
expected_activities.push_back(expected_activity);
}
TF_ASSERT_OK_AND_ASSIGN(auto stats, profiler->GetCompileStats(function));
EXPECT_EQ(stats.compile_count, 5);
EXPECT_EQ(stats.cumulative_compile_time_us, 5 * 4);
// TODO(b/255826209): Use ::testing::EqualsProto once b/135192747 is fixed.
const auto& actual_activities = listener_ptr->GetListenerHistory();
EXPECT_EQ(actual_activities.size(), expected_activities.size());
for (size_t i = 0; i < actual_activities.size(); ++i) {
EXPECT_EQ(actual_activities[i].SerializeAsString(),
expected_activities[i].SerializeAsString());
}
}
TEST(DeviceCompilationProfilerTest, OngoingAsyncCompilations) {
DeviceCompilationProfiler* profiler = new DeviceCompilationProfiler();
core::ScopedUnref profiler_ref(profiler);
for (int i = 0; i < 5; ++i) {
profiler->IncrementOngoingAsyncCompilations();
}
EXPECT_EQ(profiler->GetNumOngoingAsyncCompilations(), 5);
for (int i = 0; i < 5; ++i) {
profiler->DecrementOngoingAsyncCompilations();
}
EXPECT_EQ(profiler->GetNumOngoingAsyncCompilations(), 0);
for (int i = 0; i < 5; ++i) {
profiler->IncrementOngoingAsyncCompilations();
profiler->DecrementOngoingAsyncCompilations();
}
EXPECT_EQ(profiler->GetNumOngoingAsyncCompilations(), 0);
}
TEST(DeviceCompilationProfilerTest, ShouldCompileClusterNotFound) {
DeviceCompilationProfiler* profiler = new DeviceCompilationProfiler();
core::ScopedUnref profiler_ref(profiler);
NameAttrList function;
function.set_name("TestFunc");
EXPECT_TRUE(
profiler->ShouldCompileCluster(function, DeviceCompileMode::kAsync, 0));
EXPECT_TRUE(
profiler->ShouldCompileCluster(function, DeviceCompileMode::kLazy, 0));
EXPECT_TRUE(
profiler->ShouldCompileCluster(function, DeviceCompileMode::kStrict, 0));
}
TEST(DeviceCompilationProfilerTest, ShouldCompileClusterFirstExecution) {
DeviceCompilationProfiler* profiler = new DeviceCompilationProfiler();
core::ScopedUnref profiler_ref(profiler);
NameAttrList function;
function.set_name("TestFunc");
profiler->RegisterExecution(function);
EXPECT_TRUE(
profiler->ShouldCompileCluster(function, DeviceCompileMode::kAsync, 0));
EXPECT_TRUE(
profiler->ShouldCompileCluster(function, DeviceCompileMode::kLazy, 0));
}
TEST(DeviceCompilationProfilerTest, ShouldCompileClusterMegamorphic) {
DeviceCompilationProfiler* profiler = new DeviceCompilationProfiler();
core::ScopedUnref profiler_ref(profiler);
NameAttrList function;
function.set_name("TestFunc");
const int64_t kCompileThreshold = 10;
const int64_t kMinExecutionsPerCompile = 50;
// Register compilation enough times (without registering executions enough
// times) so that the function is marked megamorphic.
for (int i = 0; i < kCompileThreshold + 1; ++i) {
EXPECT_TRUE(profiler->RegisterCompilation(function, 1, false).ok());
}
profiler->RegisterExecution(function);
// Shouldn't compile cluster since it has gone megamorphic.
EXPECT_FALSE(
profiler->ShouldCompileCluster(function, DeviceCompileMode::kAsync, 0));
EXPECT_FALSE(
profiler->ShouldCompileCluster(function, DeviceCompileMode::kLazy, 0));
TF_ASSERT_OK_AND_ASSIGN(auto stats, profiler->GetCompileStats(function));
EXPECT_TRUE(stats.is_megamorphic);
// Always compile for strict compile mode.
EXPECT_TRUE(
profiler->ShouldCompileCluster(function, DeviceCompileMode::kStrict, 0));
// Once a cluster has gone megamorphic, it remains megamorphic (even though
// it's being executed more frequently now) and shouldn't be compiled again.
for (int i = 0; i < kCompileThreshold * kMinExecutionsPerCompile + 1; ++i) {
profiler->RegisterExecution(function);
}
EXPECT_FALSE(
profiler->ShouldCompileCluster(function, DeviceCompileMode::kAsync, 0));
EXPECT_FALSE(
profiler->ShouldCompileCluster(function, DeviceCompileMode::kLazy, 0));
TF_ASSERT_OK_AND_ASSIGN(stats, profiler->GetCompileStats(function));
EXPECT_TRUE(stats.is_megamorphic);
// Always compile for strict compile mode.
EXPECT_TRUE(
profiler->ShouldCompileCluster(function, DeviceCompileMode::kStrict, 0));
}
TEST(DeviceCompilationProfilerTest, ShouldCompileClusterAsync) {
DeviceCompilationProfiler* profiler = new DeviceCompilationProfiler();
core::ScopedUnref profiler_ref(profiler);
NameAttrList function;
function.set_name("TestFunc");
const int64_t kMaxNumOngoingCompilations = 10;
for (int i = 0; i < kMaxNumOngoingCompilations; ++i) {
profiler->IncrementOngoingAsyncCompilations();
}
// Should allow compilation since this is the first execution.
profiler->RegisterExecution(function);
EXPECT_TRUE(
profiler->ShouldCompileCluster(function, DeviceCompileMode::kAsync, 0));
// Should not allow compilation since this is not the first execution and
// we've already reached the maximum number of ongoing compilations allowed.
profiler->RegisterExecution(function);
EXPECT_FALSE(
profiler->ShouldCompileCluster(function, DeviceCompileMode::kAsync, 0));
profiler->DecrementOngoingAsyncCompilations();
// Should allow compilation since we've decremented the number of ongoing
// compilations.
EXPECT_TRUE(
profiler->ShouldCompileCluster(function, DeviceCompileMode::kAsync, 0));
}
TEST(DeviceCompilationProfilerTest, ShouldCompileClusterLazy) {
DeviceCompilationProfiler* profiler = new DeviceCompilationProfiler();
core::ScopedUnref profiler_ref(profiler);
NameAttrList function;
function.set_name("TestFunc");
constexpr int64_t kDefaultCompilationThreshold = 2;
// Should allow compilation since this is the first execution.
profiler->RegisterExecution(function);
EXPECT_TRUE(
profiler->ShouldCompileCluster(function, DeviceCompileMode::kLazy, 0));
// Shouldn't allow compilation until compilation has been requested at least
// kDefaultCompilationThreshold times.
profiler->RegisterExecution(function);
for (int current_request_count = 0;
current_request_count < kDefaultCompilationThreshold;
++current_request_count) {
EXPECT_FALSE(profiler->ShouldCompileCluster(
function, DeviceCompileMode::kLazy, current_request_count));
}
EXPECT_TRUE(profiler->ShouldCompileCluster(function, DeviceCompileMode::kLazy,
kDefaultCompilationThreshold));
}
} // namespace
} // namespace tensorflow
tensorflow/compiler/jit/device_compiler.h
+573
View File
@@ -0,0 +1,573 @@
/* Copyright 2017 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_COMPILER_JIT_DEVICE_COMPILER_H_
#define TENSORFLOW_COMPILER_JIT_DEVICE_COMPILER_H_
#include <memory>
#include <optional>
#include <string>
#include <utility>
#include <vector>
#include "absl/algorithm/container.h"
#include "absl/base/call_once.h"
#include "absl/base/nullability.h"
#include "absl/container/flat_hash_map.h"
#include "absl/status/status.h"
#include "absl/types/span.h"
#include "tensorflow/compiler/jit/device_compilation_cache.h"
#include "tensorflow/compiler/jit/device_compilation_cluster_signature.h"
#include "tensorflow/compiler/jit/device_compilation_profiler.h"
#include "tensorflow/compiler/jit/device_compiler_client.h"
#include "tensorflow/compiler/jit/device_executable_persistor.h"
#include "tensorflow/compiler/jit/flags.h"
#include "tensorflow/compiler/jit/tf_graph_to_hlo_compiler.h"
#include "tensorflow/compiler/jit/xla_compile_util.h"
#include "tensorflow/compiler/tf2xla/xla_compiler.h"
#include "xla/tsl/platform/statusor.h"
#include "tensorflow/core/framework/metrics.h"
#include "tensorflow/core/framework/op_kernel.h"
#include "tensorflow/core/framework/resource_base.h"
#include "tensorflow/core/lib/core/threadpool.h"
#include "tensorflow/core/platform/mutex.h"
#include "tensorflow/core/platform/thread_annotations.h"
namespace tensorflow {
// Compiles/lowers a given Tensorflow graph/function/cluster into a compiled XLA
// compilation (HLO) using the XlaCompiler and compiles the resulting
// XlaCompilationResult into an `ExecutableType` (eg. xla::LocalExecutable) by
// calling `ClientType` (eg. xla::LocalClient).
//
// Caches the compiled XlaCompilationResult and Executable using a
// DeviceCompilationCache. Compilation is done only when there's a cache miss.
//
// Uses the DeviceExecutablePersistor class for persistence and tries to load a
// serialized executable from disk upon a request for compilation. If the
// appropriate executable isn't found on disk, compiles the given Tensorflow
// function/graph/cluster into an XlaCompilationResult (HLO) and
// `ExecutableType` and tries saving/persisting the compiled HLO and executable
// to disk.
//
// Since XLA computations must have static shapes, DeviceCompiler generates a
// new XLA computation for each new set of input shapes.
// TODO(b/255826209): De-templatize once we've moved to Device API completely.
template <typename ExecutableType, typename ClientType>
class DeviceCompiler : public ResourceBase {
public:
DeviceCompiler(
std::unique_ptr<DeviceExecutablePersistor<ExecutableType, ClientType>>
persistor,
std::unique_ptr<DeviceCompilerClient<ExecutableType, ClientType>>
compiler_client);
~DeviceCompiler() override;
enum class CompileScope {
kOp,
kFunction,
};
// Compiles a function into a XlaCompiler::CompilationResult that can be used
// to execute an XLA Computation. Compilation results are cached. Compilation
// is skipped if there is a cache hit. `function` is the name of a Tensorflow
// function to compile. `args` is a description of the arguments to the
// computation.
//
// `compile_mode` controls the behavior of the compilation cache on a cache
// miss. If `compile_mode` is `kLazy` then, based on some profitability
// heuristics, the compilation cache may decide not to compile the cluster at
// this time. In this case it returns null into both `out_compilation_result`
// and `out_executable`. If `compile_mode` is `kStrict` then the compilation
// cache always attempts the compilation on a cache miss. If compilation mode
// is 'kAsync' compilation of the cluster happens in the background while the
// fallback path executes.
//
// The result of compilation is written to `*out_compilation_result`, which
// must be non-null. If `out_executable` is non-null, also builds an
// `ExecutableType` and sets `out_executable` to point to it. The
// resulting executable pointer may be null if the computation has no
// non-constant outputs.
absl::Status CompileIfNeeded(
const XlaCompiler::Options& options, const NameAttrList& function,
const std::vector<XlaCompiler::Argument>& args,
const XlaCompiler::CompileOptions& compile_options,
DeviceCompileMode compile_mode, DeviceCompilationProfiler* profiler,
const XlaCompiler::CompilationResult** out_compilation_result,
ExecutableType** out_executable);
// As above, but for a single op.
absl::Status CompileSingleOpIfNeeded(
const XlaCompiler::Options& options,
const std::vector<XlaCompiler::Argument>& args,
const XlaCompiler::CompileOptions& compile_options, OpKernelContext* ctx,
DeviceCompilationProfiler* profiler,
const XlaCompiler::CompilationResult** out_compilation_result,
ExecutableType** out_executable);
// An override that allows the caller to specify the function explicitly.
absl::Status CompileSingleOpIfNeeded(
const XlaCompiler::Options& options, const NameAttrList& function,
const DeviceCompilationCanonicalFunction& canonical_function,
const std::vector<XlaCompiler::Argument>& args,
const XlaCompiler::CompileOptions& compile_options, OpKernelContext* ctx,
DeviceCompilationProfiler* profiler,
const XlaCompiler::CompilationResult** out_compilation_result,
ExecutableType** out_executable);
ClientType* client() const { return compiler_client_->client(); }
const DeviceType& device_type() const { return persistor_->device_type(); }
DeviceCompilationCache<ExecutableType>* cache() { return cache_.get(); }
DeviceExecutablePersistor<ExecutableType, ClientType>* persistor() {
return persistor_.get();
}
DeviceCompilerClient<ExecutableType, ClientType>* compiler_client() {
return compiler_client_.get();
}
std::string DebugString() const override;
private:
// Common implementation of Compile and CompileSingleOp. The `OpKernelContext`
// parameter is always null for the former.
absl::Status CompileImpl(
const XlaCompiler::CompileOptions& compile_options,
const XlaCompiler::Options& options, const NameAttrList& function,
const DeviceCompilationCanonicalFunction& canonical_function,
const std::vector<XlaCompiler::Argument>& args, CompileScope scope,
DeviceCompileMode compile_mode, OpKernelContext* ctx,
DeviceCompilationProfiler* profiler,
const XlaCompiler::CompilationResult** out_compilation_result,
ExecutableType** out_executable);
StatusOr<typename DeviceCompilationCache<ExecutableType>::Value>
CompileStrict(
const DeviceCompilationClusterSignature& sig,
const XlaCompiler::CompileOptions& compile_options,
const XlaCompiler::Options& options,
const std::vector<XlaCompiler::Argument>& args,
const NameAttrList& function,
typename DeviceCompilationCache<ExecutableType>::Value cache_value,
CompileScope scope, OpKernelContext* ctx,
DeviceCompilationProfiler* profiler, mutex* mu)
TF_EXCLUSIVE_LOCKS_REQUIRED(*mu);
absl::Status CompileAsynchronous(
const DeviceCompilationClusterSignature& sig,
const XlaCompiler::CompileOptions& compile_options,
const XlaCompiler::Options& options,
const std::vector<XlaCompiler::Argument>& args,
const NameAttrList& function, CompileScope scope, OpKernelContext* ctx,
DeviceCompilationProfiler* profiler);
// Releases all references held to `std::shared_ptr<xla::XlaComputation>`
// held by the cache.
//
// This is to be called during session finalization, after all compilation
// has completed and computations no longer need to be accessed through the
// cache.
void Finalize() override;
std::unique_ptr<DeviceExecutablePersistor<ExecutableType, ClientType>>
persistor_;
std::unique_ptr<DeviceCompilerClient<ExecutableType, ClientType>>
compiler_client_;
std::unique_ptr<DeviceCompilationCache<ExecutableType>> cache_;
// Pool of threads for asynchronous compilations.
std::unique_ptr<thread::ThreadPool> async_compiler_threads_;
mutex cluster_mutexes_mu_;
absl::flat_hash_map<DeviceCompilationClusterSignature, std::unique_ptr<mutex>,
DeviceCompilationClusterSignature::Hash>
cluster_mutexes_ TF_GUARDED_BY(cluster_mutexes_mu_);
DeviceCompiler(const DeviceCompiler&) = delete;
void operator=(const DeviceCompiler&) = delete;
};
namespace device_compiler_internal {
// Print something that users can search for to definitively ascertain that XLA
// was used for their TF model.
// Prints only once to avoid spamming LOG(INFO).
inline void LogOnceXlaCompiledFirstCluster() {
static absl::once_flag log_once;
absl::call_once(log_once, [] {
LOG(INFO) << "Compiled cluster using XLA! This line is logged at most "
"once for the lifetime of the process.";
});
}
template <typename ExecutableType>
inline absl::Status EligibleToPersist(DeviceCompileState compile_state,
const ExecutableType* executable) {
if (compile_state != DeviceCompileState::kCompiled) {
return absl::FailedPreconditionError(
"Cache entry to serialize is not compiled.");
}
if (executable == nullptr) {
return absl::FailedPreconditionError(
"LocalExecutable not found for cache entry to serialize.");
}
return absl::OkStatus();
}
} // namespace device_compiler_internal
template <typename ExecutableType, typename ClientType>
DeviceCompiler<ExecutableType, ClientType>::DeviceCompiler(
std::unique_ptr<DeviceExecutablePersistor<ExecutableType, ClientType>>
persistor,
std::unique_ptr<DeviceCompilerClient<ExecutableType, ClientType>>
compiler_client)
: persistor_(std::move(persistor)),
compiler_client_(std::move(compiler_client)) {
cache_ = std::make_unique<DeviceCompilationCache<ExecutableType>>();
async_compiler_threads_ = std::make_unique<tensorflow::thread::ThreadPool>(
tensorflow::Env::Default(), "async_compiler_threads",
kNumAsyncDeviceCompilerThreads);
}
template <typename ExecutableType, typename ClientType>
DeviceCompiler<ExecutableType, ClientType>::~DeviceCompiler() {
// Since programs are owned by the cache, ensure any use of our programs have
// completed by waiting for all stream executors to complete.
compiler_client_->WaitForProgramsToFinish();
// Wait for all outstanding compilations to finish.
// Resetting the pointer explicitly in the top level destructor.
// Without this, the pointer would be reset when the AsyncCompilationState
// is destructed, which is dependent on the order of the members in the
// DeviceCompiler class, which is error prone if the order changes.
async_compiler_threads_.reset();
// TODO(b/110813685): Think about the program ownership model. Programs are
// currently owned by the compilation cache which means we must wait for
// program completion in the destructor. There are multiple compilation caches
// around, which complicates things a little. Perhaps having programs be
// shared_ptrs (an invasive change) would make the model easier to reason
// about?
}
template <typename ExecutableType, typename ClientType>
std::string DeviceCompiler<ExecutableType, ClientType>::DebugString() const {
return "DeviceCompiler";
}
template <typename ExecutableType, typename ClientType>
absl::Status DeviceCompiler<ExecutableType, ClientType>::CompileIfNeeded(
const XlaCompiler::Options& options, const NameAttrList& function,
const std::vector<XlaCompiler::Argument>& args,
const XlaCompiler::CompileOptions& compile_options,
DeviceCompileMode compile_mode, DeviceCompilationProfiler* profiler,
const XlaCompiler::CompilationResult** out_compilation_result,
ExecutableType** out_executable) {
return CompileImpl(compile_options, options, function, Canonicalize(function),
args, CompileScope::kFunction, compile_mode,
/*ctx=*/nullptr, profiler, out_compilation_result,
out_executable);
}
inline NameAttrList GetDeviceCompilerFunction(const NodeDef& def) {
NameAttrList function;
function.set_name(def.op());
*function.mutable_attr() = def.attr();
// Remove the "_class" attribute from the attribute set used to create the
// compilation cache key. This attribute is information for the colocator
// and causes false uniqueness between nodes.
function.mutable_attr()->erase("_class");
return function;
}
template <typename ExecutableType, typename ClientType>
absl::Status
DeviceCompiler<ExecutableType, ClientType>::CompileSingleOpIfNeeded(
const XlaCompiler::Options& options,
const std::vector<XlaCompiler::Argument>& args,
const XlaCompiler::CompileOptions& compile_options, OpKernelContext* ctx,
DeviceCompilationProfiler* profiler,
const XlaCompiler::CompilationResult** out_compilation_result,
ExecutableType** out_executable) {
const NodeDef& def = ctx->op_kernel().def();
const NameAttrList function = GetDeviceCompilerFunction(def);
return CompileSingleOpIfNeeded(options, function, Canonicalize(function),
args, compile_options, ctx, profiler,
out_compilation_result, out_executable);
}
template <typename ExecutableType, typename ClientType>
absl::Status
DeviceCompiler<ExecutableType, ClientType>::CompileSingleOpIfNeeded(
const XlaCompiler::Options& options, const NameAttrList& function,
const DeviceCompilationCanonicalFunction& canonical_function,
const std::vector<XlaCompiler::Argument>& args,
const XlaCompiler::CompileOptions& compile_options, OpKernelContext* ctx,
DeviceCompilationProfiler* profiler,
const XlaCompiler::CompilationResult** out_compilation_result,
ExecutableType** out_executable) {
return CompileImpl(compile_options, options, function, canonical_function,
args, CompileScope::kOp, DeviceCompileMode::kStrict, ctx,
profiler, out_compilation_result, out_executable);
}
template <typename ExecutableType, typename ClientType>
StatusOr<typename DeviceCompilationCache<ExecutableType>::Value>
DeviceCompiler<ExecutableType, ClientType>::CompileStrict(
const DeviceCompilationClusterSignature& sig,
const XlaCompiler::CompileOptions& compile_options,
const XlaCompiler::Options& options,
const std::vector<XlaCompiler::Argument>& args,
const NameAttrList& function,
typename DeviceCompilationCache<ExecutableType>::Value cache_value,
CompileScope scope, OpKernelContext* ctx,
DeviceCompilationProfiler* profiler, mutex* mu) {
tensorflow::Env* env = tensorflow::Env::Default();
const uint64_t compile_start_us = env->NowMicros();
TfGraphToHloCompiler compiler(options);
cache_value.compile_state = DeviceCompileState::kCompiled;
std::unique_ptr<ExecutableType> out_executable;
auto out_compilation_result =
std::make_unique<XlaCompiler::CompilationResult>();
if (scope == CompileScope::kOp) {
cache_value.compilation_status = compiler.CompileSingleOp(
compile_options, ctx, args, out_compilation_result.get());
} else {
CHECK(scope == CompileScope::kFunction); // Crash OK
cache_value.compilation_status = compiler.Compile(
compile_options, function, args, out_compilation_result.get());
}
TF_RETURN_IF_ERROR(cache_value.compilation_status);
TF_RET_CHECK(cache_value.executable == nullptr);
TF_RET_CHECK(out_compilation_result->computation != nullptr);
auto loaded_executable = persistor_->TryToLoadExecutable(
DeviceCompilationClusterSignature::Hash()(sig), sig.HumanString(),
options, *out_compilation_result, compiler_client_.get());
if (loaded_executable.has_value()) {
cache_value.compilation_status = loaded_executable->status();
if (loaded_executable->ok()) {
out_executable = *std::move(*loaded_executable);
metrics::UpdatePersistentCacheLoadCount();
}
} else {
auto built_executable =
compiler_client_->BuildExecutable(options, *out_compilation_result);
TF_RETURN_IF_ERROR(built_executable.status());
out_executable = *std::move(built_executable);
TF_RETURN_IF_ERROR(
device_compiler_internal::EligibleToPersist<ExecutableType>(
cache_value.compile_state, out_executable.get()));
TF_RETURN_IF_ERROR(persistor_->TryToPersistExecutable(
DeviceCompilationClusterSignature::Hash()(sig), sig.HumanString(),
options, *out_compilation_result, *out_executable,
compiler_client_.get()));
}
cache_value.compilation_result = out_compilation_result.get();
cache_value.executable = out_executable.get();
cache_->Store(sig, cache_value.compile_state, cache_value.compilation_status,
std::move(out_compilation_result), std::move(out_executable));
// Finalize the cache to release the XlaComputation after it was compiled.
cache_->Finalize();
const uint64_t compile_end_us = env->NowMicros();
const uint64_t compile_time_us = compile_end_us - compile_start_us;
device_compiler_internal::LogOnceXlaCompiledFirstCluster();
TF_RETURN_IF_ERROR(profiler->RegisterCompilation(
function, compile_time_us, loaded_executable.has_value()));
return cache_value;
}
template <typename ExecutableType, typename ClientType>
absl::Status DeviceCompiler<ExecutableType, ClientType>::CompileAsynchronous(
const DeviceCompilationClusterSignature& signature,
const XlaCompiler::CompileOptions& compile_options,
const XlaCompiler::Options& options,
const std::vector<XlaCompiler::Argument>& args,
const NameAttrList& function, CompileScope scope, OpKernelContext* ctx,
DeviceCompilationProfiler* profiler) {
// Explicitly capture all required data by value for async compilation.
// Update compilation state in cache.
cache_->Store(signature, DeviceCompileState::kCompiling, std::nullopt,
std::nullopt, std::nullopt);
profiler->IncrementOngoingAsyncCompilations();
// Don't move the above code into the thread function as it synchronously
// updates the async compilation state!
// When the ThreadPool for the compilation cache is destroyed, it waits for
// compilations to have finished. This means that both 'entry' and 'this' will
// be alive for the duration of the compilation.
// !!Pay attention when additional variables must be captured by this lambda!!
// All values are captured by value. Make sure that all pointer values (like
// entry) do not get freed until the lambda has finished.
const std::string& function_name = function.name();
async_compiler_threads_->Schedule([=] {
VLOG(2) << "Starting asynchronous compilation of cluster " << function_name
<< '.';
// We don't need to lock mu, but do it anyway to satisfy thread safety
// analysis.
mutex mu;
mutex_lock lock(mu);
auto cache_value = typename DeviceCompilationCache<ExecutableType>::Value();
auto s = CompileStrict(signature, compile_options, options, args, function,
cache_value, scope, ctx, profiler, &mu);
VLOG(2) << "Finished asynchronous compililation of cluster "
<< function_name << '.';
profiler->DecrementOngoingAsyncCompilations();
// Update compilation status in cache.
if (!s.ok()) {
cache_->Store(signature, std::nullopt, s.status(), std::nullopt,
std::nullopt);
}
});
return absl::OkStatus();
}
template <typename ExecutableType, typename ClientType>
void DeviceCompiler<ExecutableType, ClientType>::Finalize() {
const mutex_lock lock(cluster_mutexes_mu_);
std::vector<mutex* absl_nonnull> cluster_mutexes;
cluster_mutexes.reserve(cluster_mutexes_.size());
for (auto& [_, mutex] : cluster_mutexes_) {
if (mutex != nullptr) {
cluster_mutexes.push_back(mutex.get());
}
}
// Sort the mutexes before locking to ensure that this happens in a
// deterministic order, consistent between resizes of the `cluster_mutexes_`
// map.
absl::c_sort(cluster_mutexes);
std::vector<mutex_lock> cluster_mutex_locks;
cluster_mutex_locks.reserve(cluster_mutexes.size());
for (mutex* absl_nonnull const mutex : cluster_mutexes) {
cluster_mutex_locks.emplace_back(*mutex);
}
cache_->Finalize();
}
template <typename ExecutableType, typename ClientType>
absl::Status DeviceCompiler<ExecutableType, ClientType>::CompileImpl(
const XlaCompiler::CompileOptions& compile_options,
const XlaCompiler::Options& options, const NameAttrList& function,
const DeviceCompilationCanonicalFunction& canonical_function,
const std::vector<XlaCompiler::Argument>& args, CompileScope scope,
DeviceCompileMode compile_mode, OpKernelContext* ctx,
DeviceCompilationProfiler* profiler,
const XlaCompiler::CompilationResult** out_compilation_result,
ExecutableType** out_executable) {
DCHECK_NE(out_executable, nullptr);
VLOG(2) << "DeviceCompiler::Compile " << DebugString();
if (VLOG_IS_ON(2)) {
VLOG(2) << "num_inputs=" << args.size();
for (int i = 0, end = args.size(); i < end; i++) {
VLOG(3) << i << ": " << args[i].HumanString();
}
}
TF_ASSIGN_OR_RETURN(auto signature, DeviceCompilationClusterSignature::Build(
canonical_function, args));
// The outer lock protects the existence of the mutex in the map.
mutex* cluster_mutex;
{
mutex_lock lock(cluster_mutexes_mu_);
auto it =
cluster_mutexes_.emplace(signature, std::make_unique<mutex>()).first;
cluster_mutex = it->second.get();
}
profiler->RegisterExecution(function);
std::string human_signature;
if (VLOG_IS_ON(2)) {
human_signature = VLOG_IS_ON(3) ? signature.HumanString() : function.name();
VLOG(2) << "DeviceCompilationClusterSignature: " << human_signature;
}
// Acquire the cache entry lock and compile, if necessary.
// TODO(phawkins): this locking will need to be restructured when we implement
// cache eviction.
mutex_lock cluster_compile_lock(*cluster_mutex);
auto cache_value = cache_->LookupOrCreate(signature);
int64_t current_request_count = cache_value.request_count;
VLOG(2) << "Compilation cache entry hit: "
<< static_cast<int>(cache_value.compile_state)
<< " signature: " << human_signature << " with request count "
<< current_request_count;
DeviceCompileState state = cache_value.compile_state;
*out_compilation_result = nullptr;
*out_executable = nullptr;
// Check if the requested entry is uncompiled and return an error if
// compilation is disabled. This will raise an error for kLazy even if we have
// not yet hit the compilation threshold and no compilation happens this
// round. This is to avoid non-determanism of when compilation is disallowed,
// for example by changing the threshold.
if (state == DeviceCompileState::kUncompiled && FailOnXlaCompilation()) {
VLOG(1) << "XLA compilation disabled: " << function.name() << "\n"
<< absl::StrJoin(
args, "\n",
[](std::string* out, const XlaCompiler::Argument& arg) {
absl::StrAppend(out, " arg: ", arg.HumanString());
});
return absl::InternalError("XLA compilation disabled");
}
if (state == DeviceCompileState::kUncompiled) {
XLA_SCOPED_LOGGING_TIMER("Compilation of XLA executable");
if (!profiler->ShouldCompileCluster(function, compile_mode,
current_request_count)) {
VLOG(2) << "Not compiling for signature: " << human_signature;
return absl::OkStatus();
} else if (compile_mode == DeviceCompileMode::kAsync) {
VLOG(2) << "Queueing asynchronous compilation for signature: "
<< human_signature;
TF_RETURN_IF_ERROR(CompileAsynchronous(signature, compile_options,
options, args, function, scope,
ctx, profiler));
return absl::OkStatus();
} else {
VLOG(2) << "Instantly compiling for signature: " << human_signature;
TF_ASSIGN_OR_RETURN(
cache_value,
CompileStrict(signature, compile_options, options, args, function,
cache_value, scope, ctx, profiler, cluster_mutex));
}
} else if (state == DeviceCompileState::kCompiling) {
VLOG(2) << "Ongoing asynchronous compilation for signature: "
<< human_signature;
return absl::OkStatus();
} else if (state == DeviceCompileState::kCompiled) {
VLOG(2) << "Already Compiled for signature: " << human_signature;
}
TF_RETURN_IF_ERROR(cache_value.compilation_status);
*out_compilation_result = cache_value.compilation_result;
*out_executable = cache_value.executable;
return absl::OkStatus();
}
} // namespace tensorflow
#endif // TENSORFLOW_COMPILER_JIT_DEVICE_COMPILER_H_
tensorflow/compiler/jit/device_compiler_client.cc
+46
View File
@@ -0,0 +1,46 @@
/* Copyright 2022 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/compiler/jit/device_compiler_client.h"
#include "tensorflow/compiler/tf2xla/xla_compiler.h"
#include "tensorflow/core/util/determinism.h"
namespace tensorflow {
xla::ExecutableBuildOptions GetExecutableBuildOptions(
const XlaCompiler::Options& options,
const XlaCompiler::CompilationResult& result, int default_device_ordinal) {
xla::ExecutableBuildOptions build_options;
if (result.collective_info) {
build_options.set_num_replicas(result.collective_info->group_size);
}
if (options.device_ordinal != -1) {
build_options.set_device_ordinal(options.device_ordinal);
} else if (default_device_ordinal != -1) {
build_options.set_device_ordinal(default_device_ordinal);
}
build_options.set_result_layout(result.xla_output_shape);
build_options.set_device_allocator(options.device_allocator.get());
build_options.set_alias_passthrough_params(options.alias_passthrough_params);
build_options.mutable_debug_options()->set_xla_detailed_logging(
options.detailed_logging);
if (tensorflow::OpDeterminismRequired()) {
build_options.mutable_debug_options()->set_xla_gpu_deterministic_ops(true);
}
return build_options;
}
} // namespace tensorflow
tensorflow/compiler/jit/device_compiler_client.h
+76
View File
@@ -0,0 +1,76 @@
/* Copyright 2022 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_COMPILER_JIT_DEVICE_COMPILER_CLIENT_H_
#define TENSORFLOW_COMPILER_JIT_DEVICE_COMPILER_CLIENT_H_
#include <optional>
#include <string>
#include <variant>
#include "tensorflow/compiler/tf2xla/xla_compiler.h"
#include "xla/client/executable_build_options.h"
namespace tensorflow {
template <typename ExecutableType, typename ClientType>
class DeviceCompilerClient {
public:
DeviceCompilerClient() = default;
virtual ~DeviceCompilerClient() = default;
// Compiles `result` (HLO) to an `ExecutableType` using `ClientType` and
// returns it.
virtual StatusOr<std::unique_ptr<ExecutableType>> BuildExecutable(
const XlaCompiler::Options& options,
const XlaCompiler::CompilationResult& result) = 0;
// Serializes an available `executable` to string using `ClientType` and
// returns it.
virtual absl::StatusOr<std::string> SerializeExecutable(
const ExecutableType& executable) = 0;
// Compiles `result` (HLO) to a serializable executable (eg.
// xla::AotCompilationResult) using `ClientType`, serializes it to string and
// returns it.
virtual absl::StatusOr<std::string> BuildSerializedExecutable(
const XlaCompiler::Options& options,
const XlaCompiler::CompilationResult& result) = 0;
// Loads `serialized_executable` into an `ExecutableType` using `ClientType`.
virtual StatusOr<std::unique_ptr<ExecutableType>> LoadExecutable(
const XlaCompiler::Options& options,
const XlaCompiler::CompilationResult& result,
const std::string& serialized_executable) = 0;
// Waits for the underlying `ClientType` backend's programs to finish
// executing before returning.
virtual void WaitForProgramsToFinish() = 0;
virtual ClientType* client() const = 0;
private:
DeviceCompilerClient(const DeviceCompilerClient&) = delete;
void operator=(const DeviceCompilerClient&) = delete;
};
// Generates the ExecutableBuildOptions for compilation from HLO to
// executable.
xla::ExecutableBuildOptions GetExecutableBuildOptions(
const XlaCompiler::Options& options,
const XlaCompiler::CompilationResult& result, int default_device_ordinal);
} // namespace tensorflow
#endif // TENSORFLOW_COMPILER_JIT_DEVICE_COMPILER_CLIENT_H_
tensorflow/compiler/jit/device_compiler_client_test.cc
+76
View File
@@ -0,0 +1,76 @@
/* Copyright 2022 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/compiler/jit/device_compiler_client.h"
#include <gtest/gtest.h>
namespace tensorflow {
namespace {
TEST(GetExecutableOptionTest, Basic) {
XlaCompiler::Options options;
options.device_ordinal = 0;
options.alias_passthrough_params = true;
options.detailed_logging = true;
XlaCompiler::CompilationResult result;
xla::Shape xla_output_shape;
result.xla_output_shape = xla_output_shape;
auto build_option =
GetExecutableBuildOptions(options, result, /*default_device_ordinal=*/-1);
EXPECT_EQ(build_option.device_ordinal(), 0);
EXPECT_EQ(build_option.result_layout()->ToString(),
xla_output_shape.ToString());
EXPECT_EQ(build_option.alias_passthrough_params(), true);
EXPECT_EQ(build_option.debug_options().xla_detailed_logging(), true);
EXPECT_EQ(build_option.debug_options().xla_enable_dumping(), true);
}
TEST(GetExecutableOptionTest, DefaultDeviceOrdinal) {
XlaCompiler::Options options;
XlaCompiler::CompilationResult result;
auto build_option =
GetExecutableBuildOptions(options, result, /*default_device_ordinal=*/0);
EXPECT_EQ(build_option.device_ordinal(), 0);
}
TEST(GetExecutableOptionTest, DeviceOrdinalNotSet) {
XlaCompiler::Options options;
XlaCompiler::CompilationResult result;
auto build_option =
GetExecutableBuildOptions(options, result, /*default_device_ordinal=*/-1);
EXPECT_EQ(build_option.device_ordinal(), -1);
}
TEST(GetExecutableOptionTest, DumpingWithoutDetailedLogging) {
XlaCompiler::Options options;
options.detailed_logging = false;
XlaCompiler::CompilationResult result;
auto build_option =
GetExecutableBuildOptions(options, result, /*default_device_ordinal=*/-1);
EXPECT_FALSE(build_option.debug_options().xla_detailed_logging());
EXPECT_TRUE(build_option.debug_options().xla_enable_dumping());
}
} // namespace
} // namespace tensorflow
tensorflow/compiler/jit/device_compiler_disable_test.cc
+89
View File
@@ -0,0 +1,89 @@
/* 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 <memory>
#include <utility>
#include <vector>
#include "tensorflow/compiler/jit/device_compilation_profiler.h"
#include "tensorflow/compiler/jit/flags.h"
#include "tensorflow/compiler/jit/device_compiler.h"
#include "tensorflow/compiler/jit/xla_device_compiler_client.h"
#include "tensorflow/compiler/tf2xla/xla_compiler.h"
#include "xla/client/client_library.h"
#include "tensorflow/core/platform/test.h"
namespace tensorflow {
namespace {
// This test is kept separate because it disables XLA compilation globally.
TEST(DeviceCompilerTest, TestDisabledXlaCompilation) {
NameAttrList fn;
fn.set_name("afunction");
// Create mock arguments so we see them in the VLOG when compilation fails.
std::vector<XlaCompiler::Argument> args(2);
for (int i = 0; i < 2; ++i) {
args[i].kind = XlaCompiler::Argument::kParameter;
args[i].type = DT_INT32;
args[i].shape = TensorShape({2, i + 1});
args[i].name = absl::StrCat("arg", i);
}
DisableXlaCompilation();
xla::LocalClient* client = xla::ClientLibrary::LocalClientOrDie();
DeviceType device_type = DeviceType(DEVICE_CPU_XLA_JIT);
const XlaCompiler::CompilationResult* compilation_result;
xla::LocalExecutable* executable;
using XlaDeviceExecutablePersistor =
DeviceExecutablePersistor<xla::LocalExecutable, xla::LocalClient>;
auto persistor = std::make_unique<XlaDeviceExecutablePersistor>(
XlaDeviceExecutablePersistor::Config(), device_type);
auto compiler_client = std::make_unique<XlaDeviceCompilerClient>(client);
auto xla_device_compiler =
new DeviceCompiler<xla::LocalExecutable, xla::LocalClient>(
std::move(persistor), std::move(compiler_client));
core::ScopedUnref xla_device_compiler_ref(xla_device_compiler);
auto profiler = new DeviceCompilationProfiler();
core::ScopedUnref profiler_ref(profiler);
// Check that strict compilation is disallowed.
absl::Status status = xla_device_compiler->CompileIfNeeded(
XlaCompiler::Options{}, fn, args, XlaCompiler::CompileOptions{},
DeviceCompileMode::kStrict, profiler, &compilation_result, &executable);
EXPECT_FALSE(status.ok());
EXPECT_TRUE(absl::StrContains(status.message(), "XLA compilation disabled"));
// Check that async compilation is disallowed.
status = xla_device_compiler->CompileIfNeeded(
XlaCompiler::Options{}, fn, args, XlaCompiler::CompileOptions{},
DeviceCompileMode::kAsync, profiler, &compilation_result, &executable);
EXPECT_FALSE(status.ok());
EXPECT_TRUE(absl::StrContains(status.message(), "XLA compilation disabled"));
// Check that lazy compilation is disallowed.
status = xla_device_compiler->CompileIfNeeded(
XlaCompiler::Options{}, fn, args, XlaCompiler::CompileOptions{},
DeviceCompileMode::kLazy, profiler, &compilation_result, &executable);
EXPECT_FALSE(status.ok());
EXPECT_TRUE(absl::StrContains(status.message(), "XLA compilation disabled"));
}
} // namespace
} // namespace tensorflow
tensorflow/compiler/jit/device_compiler_test.cc
+673
View File
@@ -0,0 +1,673 @@
/* Copyright 2022 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/compiler/jit/device_compiler.h"
#include <iostream>
#include <iterator>
#include <memory>
#include <optional>
#include <string>
#include <utility>
#include <vector>
#include <gmock/gmock.h>
#include <gtest/gtest.h>
#include "absl/algorithm/container.h"
#include "absl/log/check.h"
#include "absl/strings/str_cat.h"
#include "absl/strings/string_view.h"
#include "absl/synchronization/notification.h"
#include "absl/time/clock.h"
#include "absl/time/time.h"
#include "tensorflow/cc/framework/scope.h"
#include "tensorflow/cc/ops/function_ops.h"
#include "tensorflow/cc/ops/math_ops.h"
#include "tensorflow/compiler/jit/device_compilation_cluster_signature.h"
#include "tensorflow/compiler/jit/device_compiler_client.h"
#include "tensorflow/compiler/jit/tests/device_compiler_test_helper.h"
#include "tensorflow/compiler/jit/tf_graph_to_hlo_compiler.pb.h"
#include "tensorflow/compiler/jit/xla_compile_util.h"
#include "tensorflow/compiler/jit/xla_device_compiler_client.h"
#include "tensorflow/compiler/tf2xla/xla_argument.h"
#include "tensorflow/compiler/tf2xla/xla_compiler.h"
#include "xla/client/client_library.h"
#include "xla/client/local_client.h"
#include "xla/hlo/builder/xla_computation.h"
#include "xla/stream_executor/platform_manager.h"
#include "xla/tsl/lib/core/status_test_util.h"
#include "xla/tsl/lib/strings/proto_serialization.h"
#include "xla/tsl/platform/statusor.h"
#include "tensorflow/core/framework/fake_input.h"
#include "tensorflow/core/framework/function.h"
#include "tensorflow/core/framework/graph_to_functiondef.h"
#include "tensorflow/core/framework/node_def_builder.h"
#include "tensorflow/core/framework/resource_base.h"
#include "tensorflow/core/kernels/ops_testutil.h"
#include "tensorflow/core/lib/core/status_test_util.h"
#include "tensorflow/core/platform/env.h"
#include "tensorflow/core/platform/errors.h"
#include "tensorflow/core/platform/status.h"
#include "tensorflow/core/platform/status_matchers.h"
#include "tensorflow/core/platform/statusor.h"
#include "tensorflow/core/platform/test.h"
#include "tsl/platform/path.h"
namespace tensorflow {
namespace {
using ::testing::_;
using ::testing::ElementsAreArray;
using ::testing::EndsWith;
using ::testing::Eq;
using ::testing::EqualsProto;
using ::testing::IsTrue;
using ::testing::NotNull;
using ::testing::Return;
using ::testing::SizeIs;
using ::testing::StrEq;
using ::testing::UnorderedElementsAre;
using ::testing::status::IsOk;
using XlaDeviceCompiler =
DeviceCompiler<xla::LocalExecutable, xla::LocalClient>;
using XlaDeviceExecutablePersistor =
DeviceExecutablePersistor<xla::LocalExecutable, xla::LocalClient>;
using Signature = DeviceCompilationClusterSignature;
xla::LocalClient* GetLocalClient() {
// TODO(b/255826209): Figure out how to run this test with the CPU client as
// well.
auto platform = se::PlatformManager::PlatformWithName("cuda").value();
return xla::ClientLibrary::GetOrCreateLocalClient(platform).value();
}
XlaDeviceCompiler* CreateXlaDeviceCompiler(bool enable_persistence = false) {
auto xla_compiler_client =
std::make_unique<XlaDeviceCompilerClient>(GetLocalClient());
auto xla_persistor = std::make_unique<XlaDeviceExecutablePersistor>(
XlaDeviceExecutablePersistor::Config{
enable_persistence ? testing::TmpDir() : "", false, "xla"},
DeviceType(DEVICE_GPU_XLA_JIT));
return new XlaDeviceCompiler(std::move(xla_persistor),
std::move(xla_compiler_client));
}
absl::StatusOr<std::unique_ptr<Graph>> SampleGraphAddXY() {
std::unique_ptr<Graph> graph(new Graph(OpRegistry::Global()));
Scope scope = Scope::NewRootScope().ExitOnError();
auto a = ops::_Arg(scope.WithOpName("A"), DT_INT32, 0);
auto b = ops::_Arg(scope.WithOpName("B"), DT_INT32, 1);
auto c = ops::Add(scope.WithOpName("C"), a, b);
auto d = ops::_Retval(scope.WithOpName("D"), c, 0);
TF_RETURN_IF_ERROR(scope.ToGraph(graph.get()));
return graph;
}
absl::StatusOr<FunctionDef> SampleFuntionAddXY(const std::string& name) {
TF_ASSIGN_OR_RETURN(auto graph, SampleGraphAddXY());
FunctionDef fdef;
TF_RETURN_IF_ERROR(GraphToFunctionDef(*graph, name, &fdef));
return fdef;
}
std::vector<XlaCompiler::Argument> SampleArgsForAddXY() {
std::vector<XlaCompiler::Argument> args(2);
args[0].kind = XlaCompiler::Argument::kParameter;
args[0].type = DT_INT32;
args[0].shape = TensorShape({2});
args[1].kind = XlaCompiler::Argument::kParameter;
args[1].type = DT_INT32;
args[1].shape = TensorShape({2});
return args;
}
class MockXlaDeviceExecutablePersistor
: public DeviceExecutablePersistor<xla::LocalExecutable, xla::LocalClient> {
public:
MockXlaDeviceExecutablePersistor()
: DeviceExecutablePersistor<xla::LocalExecutable, xla::LocalClient>(
Config{testing::TmpDir(), false, "xla"},
DeviceType(DEVICE_CPU_XLA_JIT)) {}
MOCK_METHOD(absl::Status, TryToPersistExecutable,
(uint64_t, const std::string&, const XlaCompiler::Options&,
const XlaCompiler::CompilationResult&,
const xla::LocalExecutable&,
(DeviceCompilerClient<xla::LocalExecutable, xla::LocalClient>*)),
(const, override));
};
class MockDeviceCompilationProfiler : public DeviceCompilationProfiler {
public:
MOCK_METHOD(bool, ShouldCompileCluster,
(const NameAttrList& function, DeviceCompileMode compile_mode,
int64_t current_request_count),
(override));
MOCK_METHOD(absl::Status, RegisterCompilation,
(const NameAttrList& function, int64_t compile_time_us,
bool used_persistent_cache),
(override));
};
class DeviceCompilerTest : public ::testing::Test {
protected:
void SetUp() override {
flib_def_ = std::make_unique<FunctionLibraryDefinition>(
OpRegistry::Global(), FunctionDefLibrary());
TF_ASSERT_OK_AND_ASSIGN(auto fdef, SampleFuntionAddXY("foo"));
TF_ASSERT_OK(flib_def_->AddFunctionDef(fdef));
profiler_ = new DeviceCompilationProfiler();
profiler_ref_ = std::make_unique<core::ScopedUnref>(profiler_);
mock_profiler_ = new MockDeviceCompilationProfiler();
mock_profiler_ref_ = std::make_unique<core::ScopedUnref>(mock_profiler_);
xla_device_compiler_ = CreateXlaDeviceCompiler();
xla_device_compiler_ref_ =
std::make_unique<core::ScopedUnref>(xla_device_compiler_);
auto listener = std::make_unique<JitCompilationListener>();
listener_ = listener.get();
RegisterXlaActivityListener(std::move(listener));
}
XlaCompiler::Options GetDefaultXlaOptions() {
XlaCompiler::Options options;
options.device_type = DeviceType(DEVICE_GPU_XLA_JIT);
options.client = xla_device_compiler_->client();
options.flib_def = flib_def_.get();
return options;
}
absl::StatusOr<std::unique_ptr<xla::LocalExecutable>>
BuildSampleXlaExecutable() {
TF_ASSIGN_OR_RETURN(auto graph, SampleGraphAddXY());
auto args = SampleArgsForAddXY();
// Compiles the graph.
XlaCompiler compiler(GetDefaultXlaOptions());
XlaCompiler::CompilationResult compilation_result;
TF_RETURN_IF_ERROR(compiler.CompileGraph(XlaCompiler::CompileOptions(),
"graph", std::move(graph), args,
&compilation_result));
return xla_device_compiler_->compiler_client()->BuildExecutable(
GetDefaultXlaOptions(), compilation_result);
}
std::unique_ptr<FunctionLibraryDefinition> flib_def_;
JitCompilationListener* listener_;
DeviceCompilationProfiler* profiler_;
std::unique_ptr<core::ScopedUnref> profiler_ref_;
MockDeviceCompilationProfiler* mock_profiler_;
std::unique_ptr<core::ScopedUnref> mock_profiler_ref_;
XlaDeviceCompiler* xla_device_compiler_;
std::unique_ptr<core::ScopedUnref> xla_device_compiler_ref_;
};
TEST_F(DeviceCompilerTest, CompileStrictSuccess) {
const XlaCompiler::CompilationResult* compilation_result = nullptr;
xla::LocalExecutable* xla_executable = nullptr;
XlaCompiler::Options options = GetDefaultXlaOptions();
NameAttrList fn;
fn.set_name("foo");
TF_EXPECT_OK(xla_device_compiler_->CompileIfNeeded(
options, fn, SampleArgsForAddXY(), XlaCompiler::CompileOptions{},
DeviceCompileMode::kStrict, profiler_, &compilation_result,
&xla_executable));
EXPECT_TRUE(compilation_result != nullptr);
EXPECT_TRUE(xla_executable != nullptr);
}
TEST_F(DeviceCompilerTest, CompileShouldCompileClusterFalse) {
const XlaCompiler::CompilationResult* compilation_result = nullptr;
xla::LocalExecutable* xla_executable = nullptr;
XlaCompiler::Options options = GetDefaultXlaOptions();
NameAttrList fn;
fn.set_name("foo");
// Using a mock here since it's difficult to have ShouldCompileCluster()
// return false.
EXPECT_CALL(*mock_profiler_,
ShouldCompileCluster(_, DeviceCompileMode::kLazy, 1))
.WillOnce(Return(false));
TF_EXPECT_OK(xla_device_compiler_->CompileIfNeeded(
options, fn, SampleArgsForAddXY(), XlaCompiler::CompileOptions{},
DeviceCompileMode::kLazy, mock_profiler_, &compilation_result,
&xla_executable));
EXPECT_TRUE(compilation_result == nullptr);
EXPECT_TRUE(xla_executable == nullptr);
}
TEST_F(DeviceCompilerTest, CompileCacheHit) {
const XlaCompiler::CompilationResult* compilation_result = nullptr;
xla::LocalExecutable* xla_executable = nullptr;
XlaCompiler::Options options = GetDefaultXlaOptions();
NameAttrList fn;
fn.set_name("foo");
TF_EXPECT_OK(xla_device_compiler_->CompileIfNeeded(
options, fn, SampleArgsForAddXY(), XlaCompiler::CompileOptions{},
DeviceCompileMode::kStrict, profiler_, &compilation_result,
&xla_executable));
EXPECT_TRUE(compilation_result != nullptr);
EXPECT_TRUE(xla_executable != nullptr);
const XlaCompiler::CompilationResult* new_compilation_result = nullptr;
xla::LocalExecutable* new_xla_executable = nullptr;
// Request compiling the same function again.
TF_EXPECT_OK(xla_device_compiler_->CompileIfNeeded(
options, fn, SampleArgsForAddXY(), XlaCompiler::CompileOptions{},
DeviceCompileMode::kStrict, profiler_, &new_compilation_result,
&new_xla_executable));
// new_compilation_result and new_xla_executable should point to the
// compilation_result and executable returned after the first compilation
// request.
EXPECT_EQ(compilation_result, new_compilation_result);
EXPECT_EQ(xla_executable, new_xla_executable);
}
TEST_F(DeviceCompilerTest, CompileAsyncSuccess) {
const XlaCompiler::CompilationResult* compilation_result = nullptr;
xla::LocalExecutable* xla_executable = nullptr;
XlaCompiler::Options options = GetDefaultXlaOptions();
NameAttrList fn;
fn.set_name("foo");
// Using a mock here to determine when the async compilation finishes. This is
// to avoid using absl::SleepFor().
// `RegisterCompilation` is the last call that happens just before the async
// compilation completes. We use the completion of this call to determine when
// the compilation finshes to verify expected behavior.
absl::Notification done;
EXPECT_CALL(*mock_profiler_,
ShouldCompileCluster(_, DeviceCompileMode::kAsync, 1))
.WillOnce(Return(true));
EXPECT_CALL(*mock_profiler_, RegisterCompilation(_, _, false))
.WillOnce([&done] {
done.Notify();
return absl::OkStatus();
});
auto args = SampleArgsForAddXY();
TF_EXPECT_OK(xla_device_compiler_->CompileIfNeeded(
options, fn, args, XlaCompiler::CompileOptions{},
DeviceCompileMode::kAsync, mock_profiler_, &compilation_result,
&xla_executable));
// compilation_result and xla_executable aren't available immediately after
// requesting compilation in asynchronous mode.
EXPECT_TRUE(compilation_result == nullptr);
EXPECT_TRUE(xla_executable == nullptr);
// Check if an appropriate entry is made in xla_cache.
auto xla_cache = xla_device_compiler_->cache();
TF_ASSERT_OK_AND_ASSIGN(auto signature, Signature::Build(fn, args));
auto cache_value = xla_cache->Lookup(signature);
EXPECT_TRUE(cache_value);
EXPECT_TRUE(cache_value->compile_state != DeviceCompileState::kUncompiled);
// Wait for async compilation to complete.
done.WaitForNotification();
cache_value = xla_cache->Lookup(signature);
EXPECT_TRUE(cache_value);
EXPECT_TRUE(cache_value->compile_state == DeviceCompileState::kCompiled);
EXPECT_TRUE(cache_value->compilation_result != nullptr);
EXPECT_TRUE(cache_value->executable != nullptr);
EXPECT_TRUE(cache_value->compilation_status.ok());
}
TEST_F(DeviceCompilerTest, CompilePersistentCacheEnabled) {
auto xla_device_compiler =
CreateXlaDeviceCompiler(/*enable_persistence=*/true);
core::ScopedUnref xla_device_compiler_ref(xla_device_compiler);
NameAttrList fn;
fn.set_name("foo");
auto args = SampleArgsForAddXY();
XlaCompiler::Options options = GetDefaultXlaOptions();
const XlaCompiler::CompilationResult* compilation_result = nullptr;
xla::LocalExecutable* xla_executable = nullptr;
TF_EXPECT_OK(xla_device_compiler->CompileIfNeeded(
options, fn, args, XlaCompiler::CompileOptions{},
DeviceCompileMode::kStrict, profiler_, &compilation_result,
&xla_executable));
EXPECT_TRUE(compilation_result != nullptr);
EXPECT_TRUE(xla_executable != nullptr);
// Check if device_compiler was able to load the executable from the
// persistent cache.
std::vector<XlaJitCompilationActivity> activity_history =
listener_->GetListenerHistory();
EXPECT_EQ(activity_history.size(), 1);
EXPECT_EQ(activity_history[0].cluster_name(), fn.name());
EXPECT_EQ(activity_history[0].compile_count(), 1);
EXPECT_FALSE(activity_history[0].used_persistent_cache());
listener_->ClearListenerHistory();
// Create another DeviceCompiler object pointing to the same persistent cache
// directory. It should load the executable instead of building it.
auto xla_device_compiler_2 =
CreateXlaDeviceCompiler(/*enable_persistence=*/true);
core::ScopedUnref xla_device_compiler_ref_2(xla_device_compiler_2);
auto profiler = new DeviceCompilationProfiler();
core::ScopedUnref profiler_ref(profiler);
const XlaCompiler::CompilationResult* compilation_result_2 = nullptr;
xla::LocalExecutable* xla_executable_2 = nullptr;
TF_EXPECT_OK(xla_device_compiler_2->CompileIfNeeded(
options, fn, args, XlaCompiler::CompileOptions{},
DeviceCompileMode::kStrict, profiler, &compilation_result_2,
&xla_executable_2));
EXPECT_TRUE(compilation_result_2 != nullptr);
EXPECT_TRUE(xla_executable_2 != nullptr);
activity_history = listener_->GetListenerHistory();
EXPECT_EQ(activity_history.size(), 1);
EXPECT_EQ(activity_history[0].cluster_name(), fn.name());
EXPECT_EQ(activity_history[0].compile_count(), 1);
// Verify that the executable was loaded instead of built.
EXPECT_TRUE(activity_history[0].used_persistent_cache());
}
TEST_F(DeviceCompilerTest, CompileFailedToLoadFromPersistentCache) {
auto xla_device_compiler =
CreateXlaDeviceCompiler(/*enable_persistence=*/true);
core::ScopedUnref xla_device_compiler_ref(xla_device_compiler);
NameAttrList fn;
fn.set_name("foo");
auto args = SampleArgsForAddXY();
XlaCompiler::Options options = GetDefaultXlaOptions();
const XlaCompiler::CompilationResult* compilation_result = nullptr;
xla::LocalExecutable* xla_executable = nullptr;
// Persist an executable.
TF_EXPECT_OK(xla_device_compiler->CompileIfNeeded(
options, fn, args, XlaCompiler::CompileOptions{},
DeviceCompileMode::kStrict, profiler_, &compilation_result,
&xla_executable));
// Corrupt the file which contains the serialized executable.
std::vector<std::string> files;
TF_ASSERT_OK(Env::Default()->GetChildren(testing::TmpDir(), &files));
std::string const* serialized_executable_filename = nullptr;
for (const auto& file : files) {
if (absl::StartsWith(file, "xla__")) {
serialized_executable_filename = &file;
break;
}
}
EXPECT_TRUE(serialized_executable_filename != nullptr);
std::string serialized_executable_filepath =
io::JoinPath(testing::TmpDir(), *serialized_executable_filename);
std::unique_ptr<WritableFile> serialized_executable_file;
TF_ASSERT_OK(Env::Default()->NewWritableFile(serialized_executable_filepath,
&serialized_executable_file));
TF_ASSERT_OK(serialized_executable_file->Append("Garbage."));
TF_ASSERT_OK(serialized_executable_file->Close());
// Create another DeviceCompiler object pointing to the same persistent cache
// directory. It should error out while loading the executable from the
// corrupt file.
auto xla_device_compiler_2 =
CreateXlaDeviceCompiler(/*enable_persistence=*/true);
core::ScopedUnref xla_device_compiler_ref_2(xla_device_compiler_2);
const XlaCompiler::CompilationResult* compilation_result_2 = nullptr;
xla::LocalExecutable* xla_executable_2 = nullptr;
EXPECT_FALSE(xla_device_compiler_2
->CompileIfNeeded(options, fn, args,
XlaCompiler::CompileOptions{},
DeviceCompileMode::kStrict, profiler_,
&compilation_result_2, &xla_executable_2)
.ok());
EXPECT_TRUE(compilation_result_2 == nullptr);
EXPECT_TRUE(xla_executable_2 == nullptr);
}
TEST_F(DeviceCompilerTest, CompileStrictPersistentCacheFailedToPersist) {
auto xla_compiler_client =
std::make_unique<XlaDeviceCompilerClient>(GetLocalClient());
auto xla_persistor = std::make_unique<MockXlaDeviceExecutablePersistor>();
auto xla_device_compiler = new XlaDeviceCompiler(
std::move(xla_persistor), std::move(xla_compiler_client));
core::ScopedUnref xla_device_compiler_ref(xla_device_compiler);
NameAttrList fn;
fn.set_name("foo");
auto args = SampleArgsForAddXY();
XlaCompiler::Options options = GetDefaultXlaOptions();
const XlaCompiler::CompilationResult* compilation_result = nullptr;
xla::LocalExecutable* xla_executable = nullptr;
auto persistor = absl::down_cast<MockXlaDeviceExecutablePersistor*>(
xla_device_compiler->persistor());
TF_ASSERT_OK_AND_ASSIGN(auto signature, Signature::Build(fn, args));
EXPECT_CALL(*persistor,
TryToPersistExecutable(Signature::Hash()(signature),
signature.HumanString(), _, _, _, _))
.WillOnce(Return(absl::FailedPreconditionError("Random error.")));
EXPECT_THAT(xla_device_compiler->CompileIfNeeded(
options, fn, args, XlaCompiler::CompileOptions{},
DeviceCompileMode::kStrict, profiler_, &compilation_result,
&xla_executable),
absl_testing::StatusIs(error::FAILED_PRECONDITION,
::testing::HasSubstr("Random error.")));
EXPECT_TRUE(compilation_result == nullptr);
EXPECT_TRUE(xla_executable == nullptr);
}
class DeviceCompilerTestWithDump : public DeviceCompilerTest {
protected:
explicit DeviceCompilerTestWithDump() {
dump_dir_ = tsl::io::JoinPath(
testing::TmpDir(),
absl::StrCat("dump_test_", absl::ToUnixNanos(absl::Now())));
CHECK_OK(Env::Default()->RecursivelyCreateDir(dump_dir_));
setenv("TF_GRAPH_TO_HLO_COMPILER_DUMP_DIR", dump_dir_.c_str(), 1);
}
~DeviceCompilerTestWithDump() override {
unsetenv("TF_GRAPH_TO_HLO_COMPILER_DUMP_DIR");
}
std::string dump_dir_;
};
TEST_F(DeviceCompilerTestWithDump, CompileStrictDebugInformationDumpWorks) {
// We create a new XlaDeviceCompiler here so that we ensure that the cached
// result is not used and `CompileStrict` is called.
auto xla_device_compiler =
CreateXlaDeviceCompiler(/*enable_persistence=*/false);
ASSERT_THAT(xla_device_compiler, NotNull());
core::ScopedUnref xla_device_compiler_ref(xla_device_compiler);
// Now we run the compilation. We only care that `CompileStrict` has been
// actually called, and we don't care about whether it succeeded.
const XlaCompiler::CompilationResult* compilation_result = nullptr;
xla::LocalExecutable* xla_executable = nullptr;
XlaCompiler::Options options = GetDefaultXlaOptions();
XlaCompiler::CompileOptions compile_options;
NameAttrList fn;
fn.set_name("foo");
xla_device_compiler
->CompileIfNeeded(options, fn, SampleArgsForAddXY(), compile_options,
DeviceCompileMode::kStrict, profiler_,
&compilation_result, &xla_executable)
.IgnoreError();
// Check the directory structure.
std::vector<std::string> dump_dir_files;
EXPECT_THAT(Env::Default()->GetChildren(dump_dir_, &dump_dir_files), IsOk());
EXPECT_THAT(dump_dir_files, SizeIs(1));
absl::string_view dump_subdir = dump_dir_files[0];
std::vector<std::string> dump_subdir_files;
EXPECT_THAT(
Env::Default()->GetChildren(tsl::io::JoinPath(dump_dir_, dump_subdir),
&dump_subdir_files),
IsOk());
// `options.pb` and one pb for the compile call arguments.
EXPECT_THAT(dump_subdir_files,
UnorderedElementsAre("options.pb", EndsWith(".pb")));
// First check the options file.
std::string options_file =
tsl::io::JoinPath(dump_dir_, dump_subdir, "options.pb");
std::string options_contents;
EXPECT_THAT(ReadFileToString(Env::Default(), options_file, &options_contents),
IsOk());
TfGraphToHloCompilerOptions options_proto;
EXPECT_THAT(options_proto.ParseFromString(options_contents), IsTrue());
EXPECT_THAT(absl::StrCat(tsl::DeterministicProtoHash64(options_proto)),
StrEq(dump_subdir));
EXPECT_THAT(options_proto.device_type(),
StrEq(options.device_type.type_string()));
EXPECT_THAT(options_proto.flib_def(),
EqualsProto(options.flib_def->ToProto()));
EXPECT_THAT(options_proto.graph_def_version(), Eq(options.graph_def_version));
// Now check the compile call arguments file.
absl::string_view compile_call_args_file_name =
(dump_subdir_files[0] == "options.pb" ? dump_subdir_files[1]
: dump_subdir_files[0]);
std::string compile_call_args_file =
tsl::io::JoinPath(dump_dir_, dump_subdir, compile_call_args_file_name);
std::string compile_call_args_contents;
EXPECT_THAT(ReadFileToString(Env::Default(), compile_call_args_file,
&compile_call_args_contents),
IsOk());
TfGraphToHloCompilerCompileCallArgs compile_call_args_proto;
EXPECT_THAT(
compile_call_args_proto.ParseFromString(compile_call_args_contents),
IsTrue());
EXPECT_THAT(compile_call_args_proto.compile_options(),
StrEq(compile_options.DebugString()));
EXPECT_THAT(compile_call_args_proto.function(), EqualsProto(fn));
std::vector<XlaArgument> xla_arguments_from_proto;
absl::c_transform(compile_call_args_proto.xla_args(),
std::back_inserter(xla_arguments_from_proto),
[](const tf2xla::XlaArgumentProto& proto) -> XlaArgument {
auto arg = XlaArgument::FromProto(proto);
CHECK_OK(arg); // Crash OK
return arg.value();
});
EXPECT_THAT(xla_arguments_from_proto, ElementsAreArray(SampleArgsForAddXY()));
}
TEST_F(OpsTestBase, CompileSingleOpSuccess) {
TF_EXPECT_OK(NodeDefBuilder("identity_op", "Identity")
.Input(FakeInput(DT_FLOAT))
.Attr("T", DT_FLOAT)
.Finalize(node_def()));
TF_EXPECT_OK(InitOp());
AddInputFromArray<float>(TensorShape({1, 2}), {6.9, 4.2});
TF_EXPECT_OK(RunOpKernel());
auto xla_device_compiler = CreateXlaDeviceCompiler();
core::ScopedUnref xla_device_compiler_ref(xla_device_compiler);
auto profiler = new DeviceCompilationProfiler();
core::ScopedUnref profiler_ref(profiler);
const XlaCompiler::CompilationResult* compilation_result = nullptr;
xla::LocalExecutable* xla_executable = nullptr;
XlaOpRegistry::RegisterCompilationKernels();
auto flib_def = std::make_unique<FunctionLibraryDefinition>(
OpRegistry::Global(), FunctionDefLibrary());
XlaCompiler::Options options;
options.device_type = DeviceType(DEVICE_GPU_XLA_JIT);
options.client = GetLocalClient();
options.flib_def = flib_def.get();
std::vector<XlaCompiler::Argument> args(1);
args[0].kind = XlaCompiler::Argument::kConstant;
args[0].type = DT_FLOAT;
args[0].shape = TensorShape({1, 2});
args[0].constant_value = GetInput(0);
args[0].initialized = true;
NameAttrList fn;
fn.set_name("foo");
TF_EXPECT_OK(xla_device_compiler->CompileSingleOpIfNeeded(
options, args, XlaCompiler::CompileOptions{}, context_.get(), profiler,
&compilation_result, &xla_executable));
EXPECT_TRUE(compilation_result != nullptr);
EXPECT_TRUE(xla_executable != nullptr);
}
TEST_F(DeviceCompilerTest, Finalize) {
XlaCompiler::Options options = GetDefaultXlaOptions();
NameAttrList fn;
fn.set_name("foo");
const XlaCompiler::CompilationResult* compilation_result = nullptr;
xla::LocalExecutable* xla_executable = nullptr;
TF_EXPECT_OK(xla_device_compiler_->CompileIfNeeded(
options, fn, SampleArgsForAddXY(), XlaCompiler::CompileOptions{},
DeviceCompileMode::kStrict, profiler_, &compilation_result,
&xla_executable));
ASSERT_TRUE(compilation_result != nullptr);
// Cast to `ResourceBase` to verify that the `Finalize` implementation
// overrides the base class's.
static_cast<ResourceBase*>(xla_device_compiler_)->Finalize();
TF_EXPECT_OK(xla_device_compiler_->CompileIfNeeded(
options, fn, SampleArgsForAddXY(), XlaCompiler::CompileOptions{},
DeviceCompileMode::kStrict, profiler_, &compilation_result,
&xla_executable));
ASSERT_TRUE(compilation_result != nullptr);
}
} // namespace
} // namespace tensorflow
tensorflow/compiler/jit/device_context_test.cc
+96
View File
@@ -0,0 +1,96 @@
/* Copyright 2023 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 <memory>
#include <string_view>
#include <utility>
#include <vector>
#include <gtest/gtest.h>
#include "tensorflow/compiler/jit/flags.h"
#include "tensorflow/compiler/tf2xla/xla_op_registry.h"
#include "xla/tsl/lib/core/status_test_util.h"
#include "tensorflow/core/framework/tensor_testutil.h"
namespace tensorflow {
namespace {
static bool Initialized = [] {
auto& rollout_config = GetXlaOpsCommonFlags()->tf_xla_use_device_api;
rollout_config.enabled_for_xla_launch_ = true;
rollout_config.enabled_for_compile_on_demand_ = true;
tensorflow::GetXlaDeviceFlags()->tf_xla_enable_xla_devices = true;
return true;
}();
class DeviceContextTest : public ::testing::Test {
public:
void SetDevice(const std::string& device_type) {
auto& rollout_config = GetXlaOpsCommonFlags()->tf_xla_use_device_api;
rollout_config.AllowForDeviceInXlaLaunch(DeviceType(device_type));
rollout_config.AllowForDeviceInXlaCompileOnDemand(DeviceType(device_type));
auto device_factory = DeviceFactory::GetFactory(device_type);
SessionOptions options;
std::vector<std::unique_ptr<Device>> devices;
absl::Status s = device_factory->CreateDevices(
options, "/job:worker/replica:0/task:0", &devices);
device_ = std::move(devices[0]);
tensorflow::AllocatorAttributes host_alloc_attr;
host_alloc_attr.set_on_host(true);
host_allocator_ = device_->GetAllocator(host_alloc_attr);
tensorflow::AllocatorAttributes device_alloc_attr;
device_alloc_attr.set_on_host(false);
device_allocator_ = device_->GetAllocator(device_alloc_attr);
tensorflow::DeviceContext* device_context;
auto status = device_->TryGetDeviceContext(&device_context);
TF_EXPECT_OK(status);
device_context_.reset(device_context);
}
std::unique_ptr<Device> device_;
tensorflow::core::RefCountPtr<DeviceContext> device_context_;
tensorflow::Allocator* host_allocator_;
tensorflow::Allocator* device_allocator_;
};
#if GOOGLE_CUDA || TENSORFLOW_USE_ROCM
TEST_F(DeviceContextTest, TestXlaGpuRoundTripTransferWithDeviceApi) {
SetDevice(DEVICE_XLA_GPU);
tensorflow::Tensor origin_cpu_tensor(host_allocator_, tensorflow::DT_FLOAT,
tensorflow::TensorShape({2, 2}));
tensorflow::test::FillValues<float>(&origin_cpu_tensor, {1.2, 2.3, 3.4, 4.5});
tensorflow::Tensor device_tensor(device_allocator_, tensorflow::DT_FLOAT,
tensorflow::TensorShape({2, 2}));
tensorflow::Tensor dest_cpu_tensor(host_allocator_, tensorflow::DT_FLOAT,
tensorflow::TensorShape({2, 2}));
TF_ASSERT_OK(device_context_->CopyCPUTensorToDeviceSync(
&origin_cpu_tensor, device_.get(), &device_tensor));
TF_ASSERT_OK(device_context_->CopyDeviceTensorToCPUSync(
&device_tensor, "", device_.get(), &dest_cpu_tensor));
LOG(INFO) << "H2D - D2H roundtrip completes. tensor: "
<< dest_cpu_tensor.DebugString(4);
tensorflow::test::ExpectClose(origin_cpu_tensor, dest_cpu_tensor);
}
#endif
} // namespace
} // namespace tensorflow
tensorflow/compiler/jit/device_executable_persistor.cc
+37
View File
@@ -0,0 +1,37 @@
/* Copyright 2024 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/compiler/jit/device_executable_persistor.h"
#include <string>
#include "absl/strings/str_cat.h"
namespace tensorflow {
std::string XlaSerializedCacheKeyToFileName(const XlaSerializedCacheKey& key) {
static constexpr char kXlaSerializedCacheKeySeparator[] = "__";
return absl::StrCat(
key.prefix(), key.prefix().empty() ? "" : kXlaSerializedCacheKeySeparator,
key.signature_fingerprint(), kXlaSerializedCacheKeySeparator,
key.cluster_fingerprint(), kXlaSerializedCacheKeySeparator,
key.device_type(),
key.compiled_using_pjrt()
? absl::StrCat(kXlaSerializedCacheKeySeparator, "pjrt")
: "",
".pb");
}
} // namespace tensorflow
tensorflow/compiler/jit/device_executable_persistor.h
+404
View File
@@ -0,0 +1,404 @@
/* Copyright 2022 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_COMPILER_JIT_DEVICE_EXECUTABLE_PERSISTOR_H_
#define TENSORFLOW_COMPILER_JIT_DEVICE_EXECUTABLE_PERSISTOR_H_
#include <optional>
#include <string>
#include <utility>
#include "absl/log/log.h"
#include "absl/status/status.h"
#include "absl/strings/str_cat.h"
#include "tensorflow/compiler/jit/device_compiler_client.h"
#include "tensorflow/compiler/jit/xla_compilation_cache.pb.h"
#include "tensorflow/compiler/jit/xla_device_compiler_client.h"
#include "tensorflow/compiler/tf2xla/xla_compiler.h"
#include "xla/pjrt/pjrt_client.h"
#include "xla/service/hlo.pb.h"
#include "xla/util.h"
#include "tensorflow/core/framework/device.h"
#include "tensorflow/core/lib/strings/proto_serialization.h"
#include "tensorflow/core/platform/path.h"
#include "tensorflow/core/platform/status.h"
#include "tensorflow/core/platform/statusor.h"
#include "tensorflow/core/platform/types.h"
#include "tsl/platform/errors.h"
#include "tsl/platform/statusor.h"
namespace tensorflow {
// Returns the persisted compilation cache file name for the given key.
std::string XlaSerializedCacheKeyToFileName(const XlaSerializedCacheKey& key);
// Offers a way to persist and/or load compiled `ExecutableType`s along with the
// corresponding HLO (`CompilationResult`) to/from `persistent_cache_directory`
// (if one was provided during construction) on disk using `ClientType`.
template <typename ExecutableType, typename ClientType>
class DeviceExecutablePersistor {
public:
// Configuration for setting up persistence (directory, filename prefix, etc).
struct Config {
Config() = default;
explicit Config(absl::string_view persistent_cache_directory,
bool disable_strict_signature_checks,
absl::string_view persistence_prefix,
bool persistent_cache_directory_read_only)
: persistent_cache_directory(persistent_cache_directory),
disable_strict_signature_checks(disable_strict_signature_checks),
persistence_prefix(persistence_prefix),
persistent_cache_directory_read_only(
persistent_cache_directory_read_only) {}
explicit Config(absl::string_view persistent_cache_directory,
bool disable_strict_signature_checks,
absl::string_view persistence_prefix)
: persistent_cache_directory(persistent_cache_directory),
disable_strict_signature_checks(disable_strict_signature_checks),
persistence_prefix(persistence_prefix) {}
// If non-empty, JIT-compiled executables are saved to and loaded from the
// specified file system directory path.
std::string persistent_cache_directory;
// Disable strict signature checks for entries loaded into the cache from
// external sources.
bool disable_strict_signature_checks = false;
// The cache persistence prefix to use if serializing/deserialzing entries.
std::string persistence_prefix;
// Cache is read-only if set to true.
bool persistent_cache_directory_read_only = false;
};
DeviceExecutablePersistor(const Config& config,
const DeviceType& device_type);
virtual ~DeviceExecutablePersistor() = default;
// Returns std::nullopt if persistence is not enabled (i.e.
// `persistent_cache_directory_` is empty) or if the serialized entry is not
// found on disk. Otherwise, loads and returns the serialized executable
// (or returns a status).
// TODO(b/255826209): Take in Signature instead of hash and string once cache
// is refactored.
std::optional<StatusOr<std::unique_ptr<ExecutableType>>> TryToLoadExecutable(
uint64_t signature_hash, const std::string& signature_str,
const XlaCompiler::Options& options,
const XlaCompiler::CompilationResult& compilation_result,
DeviceCompilerClient<ExecutableType, ClientType>* client) const;
// Tries to serialize an already built `executable` and persist it on disk. If
// unable to do so, tries to build a serialized executable using the AOT
// pipeline and persists that to disk.
// TODO(b/255826209): Take in Signature instead hash and string once cache
// is refactored.
virtual absl::Status TryToPersistExecutable(
uint64_t signature_hash, const std::string& signature_str,
const XlaCompiler::Options& options,
const XlaCompiler::CompilationResult& compilation_result,
const ExecutableType& executable,
DeviceCompilerClient<ExecutableType, ClientType>* client) const;
const DeviceType& device_type() const { return device_type_; }
const std::string& persistence_prefix() const { return persistence_prefix_; }
const std::string& persistent_cache_directory() const {
return persistent_cache_directory_;
}
private:
// Returns a cache key proto that identifies an entry in the compilation
// cache.
XlaSerializedCacheKey BuildSerializedCacheKey(
uint64_t signature_hash, const xla::HloModuleProto& hlo_module) const;
XlaSerializedCacheKey BuildSerializedCacheKey(
uint64_t signature_hash, const xla::HloModuleProto& hlo_module,
bool compiled_using_pjrt) const;
// Serializes the signature and its corresponding entry to a proto message.
absl::StatusOr<XlaSerializedCacheEntry> SerializeEntry(
uint64_t signature_hash, const XlaCompiler::Options& options,
const XlaCompiler::CompilationResult& compilation_result,
const ExecutableType& executable,
DeviceCompilerClient<ExecutableType, ClientType>* compiler_client) const;
// Saves the cache entry in the file directory supplied during the
// construction of this class. Overwrites existing entries.
absl::Status SaveSerializedEntry(const XlaSerializedCacheEntry& entry) const;
// Tries to read a cache entry given a `key` by searching the file directory
// supplied during the construction of this class. Returns std::nullopt if no
// cache entry is found.
absl::StatusOr<std::optional<XlaSerializedCacheEntry>>
TryToReadSerializedEntry(const XlaSerializedCacheKey& key) const;
// Checks if the loaded `entry` matches the expected `key` and `hlo_module`.
absl::Status VerifyLoadedCacheEntry(
const XlaSerializedCacheKey& key, const xla::HloModuleProto& hlo_module,
const XlaSerializedCacheEntry& entry) const;
std::string GetFilePath(const XlaSerializedCacheKey& key) const;
const DeviceType device_type_;
const bool disable_strict_signature_checks_;
const std::string persistence_prefix_;
// If non-empty, JIT-compiled executables are saved to and loaded from the
// specified file system directory path.
const std::string persistent_cache_directory_;
// Cache is read-only if set to true.
const bool persistent_cache_directory_read_only_;
DeviceExecutablePersistor(const DeviceExecutablePersistor&) = delete;
void operator=(const DeviceExecutablePersistor&) = delete;
};
template <typename ExecutableType, typename ClientType>
DeviceExecutablePersistor<ExecutableType, ClientType>::
DeviceExecutablePersistor(const Config& config,
const DeviceType& device_type)
: device_type_(device_type),
disable_strict_signature_checks_(config.disable_strict_signature_checks),
persistence_prefix_(config.persistence_prefix),
persistent_cache_directory_(config.persistent_cache_directory),
persistent_cache_directory_read_only_(
config.persistent_cache_directory_read_only) {}
template <typename ExecutableType, typename ClientType>
std::string DeviceExecutablePersistor<ExecutableType, ClientType>::GetFilePath(
const XlaSerializedCacheKey& key) const {
const std::string file_name = XlaSerializedCacheKeyToFileName(key);
return io::JoinPath(persistent_cache_directory_, file_name);
}
template <typename ExecutableType, typename ClientType>
XlaSerializedCacheKey
DeviceExecutablePersistor<ExecutableType, ClientType>::BuildSerializedCacheKey(
uint64_t signature_hash, const xla::HloModuleProto& hlo_module,
bool compiled_using_pjrt) const {
XlaSerializedCacheKey key;
key.set_signature_fingerprint(signature_hash);
key.set_cluster_fingerprint(DeterministicProtoHash64(hlo_module));
key.set_device_type(device_type().type_string());
key.set_prefix(persistence_prefix());
key.set_compiled_using_pjrt(compiled_using_pjrt);
return key;
}
template <typename ExecutableType, typename ClientType>
XlaSerializedCacheKey
DeviceExecutablePersistor<ExecutableType, ClientType>::BuildSerializedCacheKey(
uint64_t signature_hash, const xla::HloModuleProto& hlo_module) const {
return BuildSerializedCacheKey(signature_hash, hlo_module, false);
}
// This template specialization sets compiled_using_prjt to true in the cache
// key when the template arguments are PjRtLoadedExecutable and PjRtClient.
template <>
inline XlaSerializedCacheKey
DeviceExecutablePersistor<xla::PjRtLoadedExecutable, xla::PjRtClient>::
BuildSerializedCacheKey(uint64_t signature_hash,
const xla::HloModuleProto& hlo_module) const {
return BuildSerializedCacheKey(signature_hash, hlo_module, true);
}
template <typename ExecutableType, typename ClientType>
absl::StatusOr<std::optional<XlaSerializedCacheEntry>>
DeviceExecutablePersistor<ExecutableType, ClientType>::TryToReadSerializedEntry(
const XlaSerializedCacheKey& key) const {
Env* env = Env::Default();
const std::string file_path = GetFilePath(key);
if (!env->FileExists(file_path).ok()) {
return absl::StatusOr<std::optional<XlaSerializedCacheEntry>>(std::nullopt);
}
XlaSerializedCacheEntry entry;
TF_RETURN_IF_ERROR(ReadTextOrBinaryProto(env, file_path, &entry));
return std::optional<XlaSerializedCacheEntry>(entry);
}
template <typename ExecutableType, typename ClientType>
absl::Status
DeviceExecutablePersistor<ExecutableType, ClientType>::VerifyLoadedCacheEntry(
const XlaSerializedCacheKey& key, const xla::HloModuleProto& hlo_module,
const XlaSerializedCacheEntry& entry) const {
XLA_SCOPED_LOGGING_TIMER(absl::StrCat("Verifying loaded cache entry: ",
hlo_module.entry_computation_name()));
if (!AreSerializedProtosEqual(key, entry.key())) {
VLOG(2) << "Serialized cache key does not match:\n"
<< "got:\n"
<< entry.key().DebugString() << "\nexpected:\n"
<< key.DebugString() << "\n";
return absl::InvalidArgumentError("Serialized cache key does not match.");
}
// Perform a stricter (slower) check of the snapshot to verify that they
// match exactly.
if (!disable_strict_signature_checks_) {
if (!AreSerializedProtosEqual(hlo_module, entry.hlo_module())) {
VLOG(2) << "HLOs do not match:\n"
<< "got:\n"
<< hlo_module.DebugString() << "\nexpected:\n"
<< entry.hlo_module().DebugString() << "\n";
return absl::InvalidArgumentError("Serialized HLO does not match.");
}
}
if (entry.executable().empty()) {
return absl::InvalidArgumentError("No binary found in serialized entry.");
}
return absl::OkStatus();
}
template <typename ExecutableType, typename ClientType>
absl::Status
DeviceExecutablePersistor<ExecutableType, ClientType>::SaveSerializedEntry(
const XlaSerializedCacheEntry& entry) const {
Env* env = Env::Default();
TF_RETURN_IF_ERROR(env->RecursivelyCreateDir(persistent_cache_directory_));
// The cache on the filesystem can be read while we're writing out the proto.
// To prevent reads of partially-written files, we write the proto to a temp
// file, then move it into place once we're done writing. And we warn the
// user if these moves are not known to be atomic.
bool has_atomic_move = false;
env->HasAtomicMove(persistent_cache_directory_, &has_atomic_move)
.IgnoreError();
if (!has_atomic_move) {
LOG_EVERY_POW_2(WARNING)
<< "Filesystem for XLA persistent cache at "
<< persistent_cache_directory_
<< " does not support atomic moves. Therefore the persistent cache is "
"racy if you have multiple XLA compilations occurring "
"simultaneously! You have been warned. :)";
}
// Write to temp location, then when that completes, atomically move into the
// final location.
std::string temp_path =
io::JoinPath(persistent_cache_directory_,
XlaSerializedCacheKeyToFileName(entry.key()));
if (!env->CreateUniqueFileName(&temp_path, ".tmp")) {
return absl::UnavailableError(absl::StrCat(
"Could not create a unique file inside ", persistent_cache_directory_));
}
TF_RETURN_IF_ERROR(WriteBinaryProto(env, temp_path, entry));
return env->RenameFile(temp_path, GetFilePath(entry.key()));
}
template <typename ExecutableType, typename ClientType>
absl::StatusOr<XlaSerializedCacheEntry>
DeviceExecutablePersistor<ExecutableType, ClientType>::SerializeEntry(
uint64_t signature_hash, const XlaCompiler::Options& options,
const XlaCompiler::CompilationResult& compilation_result,
const ExecutableType& executable,
DeviceCompilerClient<ExecutableType, ClientType>* compiler_client) const {
XlaSerializedCacheEntry serialized_entry;
const xla::HloModuleProto& hlo_module =
compilation_result.computation->proto();
*serialized_entry.mutable_key() =
BuildSerializedCacheKey(signature_hash, hlo_module);
*serialized_entry.mutable_hlo_module() = hlo_module;
// XLA compiler supports exporting executables as an AOT compilation result
// to avoid running potentially expensive compilation pipeline twice.
// Check if XLA compiler can export available executable.
if (auto serialized_executable =
compiler_client->SerializeExecutable(executable);
serialized_executable.ok()) {
serialized_entry.set_executable(std::move(*serialized_executable));
return serialized_entry;
} else if (serialized_executable.status().code() == error::UNIMPLEMENTED) {
VLOG(1) << "Executable export is not implemented";
} else {
return serialized_executable.status();
}
TF_ASSIGN_OR_RETURN(
auto serialized_executable,
compiler_client->BuildSerializedExecutable(options, compilation_result));
serialized_entry.set_executable(std::move(serialized_executable));
return serialized_entry;
}
template <typename ExecutableType, typename ClientType>
std::optional<StatusOr<std::unique_ptr<ExecutableType>>>
DeviceExecutablePersistor<ExecutableType, ClientType>::TryToLoadExecutable(
uint64_t signature_hash, const std::string& signature_str,
const XlaCompiler::Options& options,
const XlaCompiler::CompilationResult& compilation_result,
DeviceCompilerClient<ExecutableType, ClientType>* compiler_client) const {
if (persistent_cache_directory_.empty()) {
return std::nullopt;
}
const xla::HloModuleProto& hlo_module =
compilation_result.computation->proto();
XlaSerializedCacheKey cache_key =
BuildSerializedCacheKey(signature_hash, hlo_module);
std::optional<XlaSerializedCacheEntry> serialized_entry;
{
XLA_SCOPED_LOGGING_TIMER(
absl::StrCat("Try loading serialized cache entry:", signature_str));
TF_ASSIGN_OR_RETURN(serialized_entry, TryToReadSerializedEntry(cache_key));
}
if (!serialized_entry.has_value()) {
return std::nullopt;
}
TF_RETURN_IF_ERROR(
VerifyLoadedCacheEntry(cache_key, hlo_module, *serialized_entry));
VLOG(1) << "Loading cached entry for: " << signature_str;
return compiler_client->LoadExecutable(options, compilation_result,
serialized_entry->executable());
}
template <typename ExecutableType, typename ClientType>
absl::Status
DeviceExecutablePersistor<ExecutableType, ClientType>::TryToPersistExecutable(
uint64_t signature_hash, const std::string& signature_str,
const XlaCompiler::Options& options,
const XlaCompiler::CompilationResult& compilation_result,
const ExecutableType& executable,
DeviceCompilerClient<ExecutableType, ClientType>* client) const {
if (persistent_cache_directory_.empty() ||
persistent_cache_directory_read_only_) {
VLOG(1) << "Not persisting executable. No `persistent_cache_directory` "
"provided or cache is read-only.";
return absl::OkStatus();
}
XLA_SCOPED_LOGGING_TIMER(
absl::StrCat("Serializing and saving cache entry: ", signature_str));
TF_ASSIGN_OR_RETURN(XlaSerializedCacheEntry serialized_entry,
SerializeEntry(signature_hash, options,
compilation_result, executable, client));
TF_RETURN_IF_ERROR(SaveSerializedEntry(std::move(serialized_entry)));
VLOG(2) << "XlaSerializedCacheEntry saved for signature: [" << signature_str
<< "] with signature hash: " << signature_hash;
return absl::OkStatus();
}
} // namespace tensorflow
#endif // TENSORFLOW_COMPILER_JIT_DEVICE_EXECUTABLE_PERSISTOR_H_
tensorflow/compiler/jit/device_executable_persistor_test.cc
+637
View File
@@ -0,0 +1,637 @@
/* Copyright 2022 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/compiler/jit/device_executable_persistor.h"
#include <stdlib.h>
#include <memory>
#include <string>
#include <utility>
#include <vector>
#include <gmock/gmock.h>
#include <gtest/gtest.h>
#include "tensorflow/cc/framework/scope.h"
#include "tensorflow/cc/ops/function_ops.h"
#include "tensorflow/cc/ops/math_ops.h"
#include "tensorflow/compiler/jit/device_compiler_client.h"
#include "tensorflow/compiler/jit/pjrt_device_compiler_client.h"
#include "tensorflow/compiler/jit/xla_compilation_cache.pb.h"
#include "tensorflow/compiler/jit/xla_device_compiler_client.h"
#include "xla/client/client_library.h"
#include "xla/client/executable_build_options.h"
#include "xla/client/local_client.h"
#include "xla/pjrt/pjrt_client.h"
#include "xla/pjrt/plugin/xla_cpu/cpu_client_options.h"
#include "xla/pjrt/plugin/xla_cpu/xla_cpu_pjrt_client.h"
#include "tensorflow/core/lib/core/status_test_util.h"
#include "tensorflow/core/platform/errors.h"
#include "tensorflow/core/platform/status_matchers.h"
#include "tensorflow/core/platform/statusor.h"
#include "tensorflow/core/tfrt/common/create_pjrt_client_util.h"
#include "tensorflow/core/tfrt/common/pjrt_util.h"
namespace tensorflow {
namespace {
using ::testing::_;
using ::testing::ByMove;
using ::testing::Return;
using XlaDeviceExecutablePersistor =
DeviceExecutablePersistor<xla::LocalExecutable, xla::LocalClient>;
using PjRtDeviceExecutablePersistor =
DeviceExecutablePersistor<xla::PjRtLoadedExecutable, xla::PjRtClient>;
class DeviceExecutionPersistorTest : public ::testing::Test {
protected:
void SetUp() override {
xla_compiler_client_ = std::make_unique<XlaDeviceCompilerClient>(
xla::ClientLibrary::LocalClientOrDie());
TF_ASSERT_OK(CreatePjRtCompilerClient());
XlaOpRegistry::RegisterCompilationKernels();
flib_def_ = std::make_unique<FunctionLibraryDefinition>(
OpRegistry::Global(), FunctionDefLibrary());
cache_dir_ = testing::TmpDir();
TF_ASSERT_OK_AND_ASSIGN(compilation_result_add_,
BuildSampleCompilationResult());
}
absl::StatusOr<std::unique_ptr<xla::LocalExecutable>>
BuildSampleExecutable() {
return xla_compiler_client_->BuildExecutable(DefaultXlaOptions(),
compilation_result_add_);
}
absl::StatusOr<std::unique_ptr<xla::PjRtLoadedExecutable>>
BuildSamplePjRtExecutable() {
return pjrt_compiler_client_->BuildExecutable(DefaultPjRtOptions(),
compilation_result_add_);
}
absl::StatusOr<XlaCompiler::CompilationResult> BuildSampleCompilationResult(
bool mul = false) {
std::unique_ptr<Graph> graph(new Graph(OpRegistry::Global()));
Scope scope = Scope::NewRootScope().ExitOnError();
auto a = ops::_Arg(scope.WithOpName("A"), DT_INT32, 0);
auto b = ops::_Arg(scope.WithOpName("B"), DT_INT32, 1);
if (mul) {
auto c = ops::Mul(scope.WithOpName("C"), a, b);
auto d = ops::_Retval(scope.WithOpName("D"), c, 0);
TF_RETURN_IF_ERROR(scope.ToGraph(graph.get()));
} else {
auto c = ops::Add(scope.WithOpName("C"), a, b);
auto d = ops::_Retval(scope.WithOpName("D"), c, 0);
TF_RETURN_IF_ERROR(scope.ToGraph(graph.get()));
}
// Builds a description of the arguments.
std::vector<XlaCompiler::Argument> args(2);
args[0].kind = XlaCompiler::Argument::kParameter;
args[0].type = DT_INT32;
args[0].shape = TensorShape({2});
args[1].kind = XlaCompiler::Argument::kParameter;
args[1].type = DT_INT32;
args[1].shape = TensorShape({2});
// Compiles the graph.
XlaCompiler compiler(DefaultXlaOptions());
XlaCompiler::CompilationResult compilation_result;
TF_RETURN_IF_ERROR(compiler.CompileGraph(XlaCompiler::CompileOptions(),
"graph", std::move(graph), args,
&compilation_result));
return compilation_result;
}
XlaCompiler::Options DefaultXlaOptions() {
XlaCompiler::Options options;
options.device_type = DeviceType(DEVICE_CPU_XLA_JIT);
options.client = xla_compiler_client_->client();
options.flib_def = flib_def_.get();
return options;
}
XlaCompiler::Options DefaultPjRtOptions() {
XlaCompiler::Options options;
options.device_type = DeviceType(DEVICE_CPU_XLA_JIT);
options.client = nullptr;
options.flib_def = flib_def_.get();
return options;
}
absl::Status CreatePjRtCompilerClient() {
// Create PjRtClient manually while GetOrCreatePjRtClient() is WIP.
xla::CpuClientOptions options;
options.asynchronous = true;
options.cpu_device_count = 1;
TF_RETURN_IF_ERROR(SetPjRtClientInTFGlobalResourceManager(
DEVICE_CPU_XLA_JIT, xla::GetXlaPjrtCpuClient(options).value()));
TF_ASSIGN_OR_RETURN(auto pjrt_client,
GetOrCreatePjRtClient(DeviceType(DEVICE_CPU_XLA_JIT)));
pjrt_compiler_client_ =
std::make_unique<PjRtDeviceCompilerClient>(pjrt_client);
return absl::OkStatus();
}
std::unique_ptr<FunctionLibraryDefinition> flib_def_;
std::unique_ptr<XlaDeviceCompilerClient> xla_compiler_client_;
std::unique_ptr<PjRtDeviceCompilerClient> pjrt_compiler_client_;
XlaCompiler::CompilationResult compilation_result_add_;
std::string serialized_xla_executable_ = "serialized_xla_executable";
std::string serialized_pjrt_executable_ = "serialized_pjrt_executable";
std::string cache_dir_;
};
// Using a mock to make testing different branches and triggering errors easier.
// Currently the `XlaDeviceCompilerClient`'s load/serialize functions don't work
// with the current test setup.
// TODO(b/255826209): Look into using a real object for most tests.
class MockXlaCompilerClient : public XlaDeviceCompilerClient {
public:
MockXlaCompilerClient() : XlaDeviceCompilerClient(nullptr) {}
MOCK_METHOD(absl::StatusOr<std::string>, SerializeExecutable,
(const xla::LocalExecutable& executable), (override));
MOCK_METHOD(absl::StatusOr<std::string>, BuildSerializedExecutable,
(const XlaCompiler::Options& options,
const XlaCompiler::CompilationResult& result),
(override));
MOCK_METHOD(absl::StatusOr<std::unique_ptr<xla::LocalExecutable>>,
LoadExecutable,
(const XlaCompiler::Options& options,
const XlaCompiler::CompilationResult& result,
const std::string& serialized_executable),
(override));
};
class MockPjRtCompilerClient : public PjRtDeviceCompilerClient {
public:
MockPjRtCompilerClient() : PjRtDeviceCompilerClient(nullptr) {}
MOCK_METHOD(absl::StatusOr<std::string>, SerializeExecutable,
(const xla::PjRtLoadedExecutable& executable), (override));
MOCK_METHOD(absl::StatusOr<std::string>, BuildSerializedExecutable,
(const XlaCompiler::Options& options,
const XlaCompiler::CompilationResult& result),
(override));
MOCK_METHOD(absl::StatusOr<std::unique_ptr<xla::PjRtLoadedExecutable>>,
LoadExecutable,
(const XlaCompiler::Options& options,
const XlaCompiler::CompilationResult& result,
const std::string& serialized_executable),
(override));
};
std::string GetFilePath(XlaSerializedCacheKey key,
const std::string& persistent_cache_dir) {
static constexpr char kXlaSerializedCacheKeySeparator[] = "__";
std::string file_name = absl::StrCat(
key.prefix(), key.prefix().empty() ? "" : kXlaSerializedCacheKeySeparator,
key.signature_fingerprint(), kXlaSerializedCacheKeySeparator,
key.cluster_fingerprint(), kXlaSerializedCacheKeySeparator,
key.device_type(),
key.compiled_using_pjrt()
? absl::StrCat(kXlaSerializedCacheKeySeparator, "pjrt")
: "",
".pb");
return io::JoinPath(persistent_cache_dir, file_name);
}
absl::StatusOr<XlaSerializedCacheEntry> ReadCacheEntryFromFile(
XlaSerializedCacheKey key, const std::string& persistent_cache_dir) {
std::string file_path = GetFilePath(key, persistent_cache_dir);
XlaSerializedCacheEntry entry;
TF_RETURN_IF_ERROR(ReadTextOrBinaryProto(Env::Default(), file_path, &entry));
return entry;
}
XlaSerializedCacheKey CreateCacheKey(
uint64_t signature_hash,
const XlaCompiler::CompilationResult& compilation_result,
const DeviceType& device_type, const std::string& persistence_prefix,
bool compiled_using_pjrt = false) {
XlaSerializedCacheKey key;
key.set_signature_fingerprint(signature_hash);
key.set_cluster_fingerprint(
DeterministicProtoHash64(compilation_result.computation->proto()));
key.set_device_type(device_type.type_string());
key.set_prefix(persistence_prefix);
key.set_compiled_using_pjrt(compiled_using_pjrt);
return key;
}
TEST_F(DeviceExecutionPersistorTest, PersistCacheDirNotSet) {
XlaDeviceExecutablePersistor::Config config(
/*persistent_cache_directory=*/"",
/*disable_strict_signature_checks=*/false,
/*persistence_prefix=*/"xla");
XlaDeviceExecutablePersistor persistor(config,
DefaultXlaOptions().device_type);
MockXlaCompilerClient mock_client;
TF_ASSERT_OK_AND_ASSIGN(auto executable, BuildSampleExecutable());
TF_EXPECT_OK(persistor.TryToPersistExecutable(
/*signature_hash=*/123, "signature_string", DefaultXlaOptions(),
compilation_result_add_, *executable, &mock_client));
auto key =
CreateCacheKey(/*signature_hash=*/123, compilation_result_add_,
persistor.device_type(), persistor.persistence_prefix());
auto entry = ReadCacheEntryFromFile(key, "");
EXPECT_FALSE(entry.ok());
}
TEST_F(DeviceExecutionPersistorTest, PersistCacheDirReadOnly) {
XlaDeviceExecutablePersistor::Config config(
/*persistent_cache_directory=*/"cache_dir_",
/*disable_strict_signature_checks=*/false,
/*persistence_prefix=*/"xla",
/*persistent_cache_directory_read_only=*/true);
XlaDeviceExecutablePersistor persistor(config,
DefaultXlaOptions().device_type);
MockXlaCompilerClient mock_client;
TF_ASSERT_OK_AND_ASSIGN(auto executable, BuildSampleExecutable());
TF_EXPECT_OK(persistor.TryToPersistExecutable(
/*signature_hash=*/123, "signature_string", DefaultXlaOptions(),
compilation_result_add_, *executable, &mock_client));
auto key =
CreateCacheKey(/*signature_hash=*/123, compilation_result_add_,
persistor.device_type(), persistor.persistence_prefix());
auto entry = ReadCacheEntryFromFile(key, "");
EXPECT_FALSE(entry.ok());
}
TEST_F(DeviceExecutionPersistorTest, PersistSerializeAlreadyBuiltExecutable) {
XlaDeviceExecutablePersistor::Config config(
/*persistent_cache_directory=*/cache_dir_,
/*disable_strict_signature_checks=*/false,
/*persistence_prefix=*/"xla");
XlaDeviceExecutablePersistor persistor(config,
DefaultXlaOptions().device_type);
MockXlaCompilerClient mock_client;
EXPECT_CALL(mock_client, SerializeExecutable(_))
.WillOnce(
Return(absl::StatusOr<std::string>(serialized_xla_executable_)));
TF_ASSERT_OK_AND_ASSIGN(auto executable, BuildSampleExecutable());
TF_EXPECT_OK(persistor.TryToPersistExecutable(
/*signature_hash=*/123, "signature_string", DefaultXlaOptions(),
compilation_result_add_, *executable, &mock_client));
auto key =
CreateCacheKey(/*signature_hash=*/123, compilation_result_add_,
persistor.device_type(), persistor.persistence_prefix());
TF_ASSERT_OK_AND_ASSIGN(auto entry, ReadCacheEntryFromFile(key, cache_dir_));
EXPECT_EQ(entry.executable(), serialized_xla_executable_);
}
TEST_F(DeviceExecutionPersistorTest, PersistBuildSerializedExecutable) {
XlaDeviceExecutablePersistor::Config config(
/*persistent_cache_directory=*/cache_dir_,
/*disable_strict_signature_checks=*/false,
/*persistence_prefix=*/"xla");
XlaDeviceExecutablePersistor persistor(config,
DefaultXlaOptions().device_type);
MockXlaCompilerClient mock_client;
EXPECT_CALL(mock_client, SerializeExecutable(_))
.WillOnce(Return(absl::UnimplementedError("Unimplemented.")));
EXPECT_CALL(mock_client, BuildSerializedExecutable(_, _))
.WillOnce(Return(serialized_xla_executable_));
TF_ASSERT_OK_AND_ASSIGN(auto executable, BuildSampleExecutable());
TF_EXPECT_OK(persistor.TryToPersistExecutable(
/*signature_hash=*/123, "signature_string", DefaultXlaOptions(),
compilation_result_add_, *executable, &mock_client));
auto key =
CreateCacheKey(/*signature_hash=*/123, compilation_result_add_,
persistor.device_type(), persistor.persistence_prefix());
TF_ASSERT_OK_AND_ASSIGN(auto entry, ReadCacheEntryFromFile(key, cache_dir_));
EXPECT_EQ(entry.executable(), serialized_xla_executable_);
}
TEST_F(DeviceExecutionPersistorTest, PersistSerializeExecutableError) {
XlaDeviceExecutablePersistor::Config config(
/*persistent_cache_directory=*/cache_dir_,
/*disable_strict_signature_checks=*/false,
/*persistence_prefix=*/"xla");
XlaDeviceExecutablePersistor persistor(config,
DefaultXlaOptions().device_type);
MockXlaCompilerClient mock_client;
EXPECT_CALL(mock_client, SerializeExecutable(_))
.WillOnce(Return(absl::InvalidArgumentError("InvalidArgument.")));
TF_ASSERT_OK_AND_ASSIGN(auto executable, BuildSampleExecutable());
EXPECT_THAT(
persistor.TryToPersistExecutable(
/*signature_hash=*/123, "signature_string", DefaultXlaOptions(),
compilation_result_add_, *executable, &mock_client),
absl_testing::StatusIs(error::INVALID_ARGUMENT));
}
TEST_F(DeviceExecutionPersistorTest, PersistExecutableEmpty) {
XlaDeviceExecutablePersistor::Config config(
/*persistent_cache_directory=*/cache_dir_,
/*disable_strict_signature_checks=*/false,
/*persistence_prefix=*/"xla");
XlaDeviceExecutablePersistor persistor(config,
DefaultXlaOptions().device_type);
MockXlaCompilerClient mock_client;
xla::LocalExecutable empty_executable(
nullptr, nullptr,
GetExecutableBuildOptions(DefaultXlaOptions(), compilation_result_add_,
0));
EXPECT_CALL(mock_client, SerializeExecutable(_))
.WillOnce(Return(absl::FailedPreconditionError("Failed precondition.")));
EXPECT_THAT(
persistor.TryToPersistExecutable(
/*signature_hash=*/123, "signature_string", DefaultXlaOptions(),
compilation_result_add_, empty_executable, &mock_client),
absl_testing::StatusIs(error::FAILED_PRECONDITION));
}
TEST_F(DeviceExecutionPersistorTest, LoadCacheDirNotSet) {
XlaDeviceExecutablePersistor::Config config(
/*persistent_cache_directory=*/"",
/*disable_strict_signature_checks=*/false,
/*persistence_prefix=*/"xla");
XlaDeviceExecutablePersistor persistor(config,
DefaultXlaOptions().device_type);
MockXlaCompilerClient mock_client;
auto executable = persistor.TryToLoadExecutable(
123, "signature_string", DefaultXlaOptions(), compilation_result_add_,
&mock_client);
EXPECT_FALSE(executable.has_value());
}
TEST_F(DeviceExecutionPersistorTest, LoadSuccess) {
XlaDeviceExecutablePersistor::Config config(
/*persistent_cache_directory=*/cache_dir_,
/*disable_strict_signature_checks=*/false,
/*persistence_prefix=*/"xla");
XlaDeviceExecutablePersistor persistor(config,
DefaultXlaOptions().device_type);
MockXlaCompilerClient mock_client;
TF_ASSERT_OK_AND_ASSIGN(auto executable, BuildSampleExecutable());
EXPECT_CALL(mock_client, LoadExecutable(_, _, serialized_xla_executable_))
.WillOnce(Return(ByMove(std::move(executable))));
auto loaded_executable = persistor.TryToLoadExecutable(
/*signature_hash=*/123, "signature_string", DefaultXlaOptions(),
compilation_result_add_, &mock_client);
EXPECT_TRUE(loaded_executable.has_value());
EXPECT_TRUE(loaded_executable.value().ok());
EXPECT_TRUE((*loaded_executable.value())->executable() != nullptr);
}
TEST_F(DeviceExecutionPersistorTest, LoadFileDoesntExist) {
XlaDeviceExecutablePersistor::Config config(
/*persistent_cache_directory=*/cache_dir_,
/*disable_strict_signature_checks=*/false,
/*persistence_prefix=*/"xla");
XlaDeviceExecutablePersistor persistor(config,
DefaultXlaOptions().device_type);
MockXlaCompilerClient mock_client;
// Try to load an executable for a different signature hash (which hasn't been
// persisted).
auto loaded_executable = persistor.TryToLoadExecutable(
/*signature_hash=*/12345, "different_signature", DefaultXlaOptions(),
compilation_result_add_, &mock_client);
EXPECT_FALSE(loaded_executable.has_value());
}
TEST_F(DeviceExecutionPersistorTest, LoadSerializedKeyMismatch) {
XlaDeviceExecutablePersistor::Config config(
/*persistent_cache_directory=*/cache_dir_,
/*disable_strict_signature_checks=*/false,
/*persistence_prefix=*/"xla");
XlaDeviceExecutablePersistor persistor(config,
DefaultXlaOptions().device_type);
auto key1 =
CreateCacheKey(/*signature_hash=*/123, compilation_result_add_,
persistor.device_type(), persistor.persistence_prefix());
auto key2 =
CreateCacheKey(/*signature_hash=*/456, compilation_result_add_,
persistor.device_type(), persistor.persistence_prefix());
// File for key2 contains the same content as key1.
TF_ASSERT_OK(Env::Default()->CopyFile(
GetFilePath(key1, persistor.persistent_cache_directory()),
GetFilePath(key2, persistor.persistent_cache_directory())));
MockXlaCompilerClient mock_client;
// Try to load an executable from file corresponding to key2 (whose file
// content corresponds to key1).
auto loaded_executable = persistor.TryToLoadExecutable(
/*signature_hash=*/456, "different_signature", DefaultXlaOptions(),
compilation_result_add_, &mock_client);
EXPECT_TRUE(loaded_executable.has_value());
EXPECT_FALSE(loaded_executable->ok());
EXPECT_THAT(loaded_executable.value(),
absl_testing::StatusIs(error::INVALID_ARGUMENT));
}
TEST_F(DeviceExecutionPersistorTest, LoadSerializedHloMismatch) {
XlaDeviceExecutablePersistor::Config config(
/*persistent_cache_directory=*/cache_dir_,
/*disable_strict_signature_checks=*/false,
/*persistence_prefix=*/"xla");
XlaDeviceExecutablePersistor persistor(config,
DefaultXlaOptions().device_type);
TF_ASSERT_OK_AND_ASSIGN(auto compilation_result_mul,
BuildSampleCompilationResult(true));
auto key1 =
CreateCacheKey(/*signature_hash=*/123, compilation_result_add_,
persistor.device_type(), persistor.persistence_prefix());
auto key2 =
CreateCacheKey(/*signature_hash=*/123, compilation_result_mul,
persistor.device_type(), persistor.persistence_prefix());
// Read serialized entry corresponding to key1.
XlaSerializedCacheEntry entry;
TF_ASSERT_OK(ReadTextOrBinaryProto(
Env::Default(), GetFilePath(key1, persistor.persistent_cache_directory()),
&entry));
// Change the entry's key to key2.
*entry.mutable_key() = key2;
// Write the modified entry to file corresponding to key2.
TF_ASSERT_OK(WriteBinaryProto(
Env::Default(), GetFilePath(key2, persistor.persistent_cache_directory()),
entry));
MockXlaCompilerClient mock_client;
// Try to load executable corresponding to key2 (whose file contains HLO
// corresponding to key1).
auto loaded_executable = persistor.TryToLoadExecutable(
/*signature_hash=*/123, "signature", DefaultXlaOptions(),
compilation_result_mul, &mock_client);
EXPECT_TRUE(loaded_executable.has_value());
EXPECT_FALSE(loaded_executable->ok());
EXPECT_THAT(loaded_executable.value(),
absl_testing::StatusIs(error::INVALID_ARGUMENT));
}
TEST_F(DeviceExecutionPersistorTest, LoadStrictChecksDisabled) {
XlaDeviceExecutablePersistor::Config config(
/*persistent_cache_directory=*/cache_dir_,
/*disable_strict_signature_checks=*/true,
/*persistence_prefix=*/"xla");
XlaDeviceExecutablePersistor persistor(config,
DefaultXlaOptions().device_type);
TF_ASSERT_OK_AND_ASSIGN(auto compilation_result_mul,
BuildSampleCompilationResult(true));
auto key1 =
CreateCacheKey(/*signature_hash=*/123, compilation_result_add_,
persistor.device_type(), persistor.persistence_prefix());
auto key2 =
CreateCacheKey(/*signature_hash=*/123, compilation_result_mul,
persistor.device_type(), persistor.persistence_prefix());
// Read serialized entry corresponding to key1.
XlaSerializedCacheEntry entry;
TF_ASSERT_OK(ReadTextOrBinaryProto(
Env::Default(), GetFilePath(key1, persistor.persistent_cache_directory()),
&entry));
// Change the entry's key to key2.
*entry.mutable_key() = key2;
// Write the modified entry to file corresponding to key2.
TF_ASSERT_OK(WriteBinaryProto(
Env::Default(), GetFilePath(key2, persistor.persistent_cache_directory()),
entry));
MockXlaCompilerClient mock_client;
TF_ASSERT_OK_AND_ASSIGN(auto executable, BuildSampleExecutable());
EXPECT_CALL(mock_client, LoadExecutable(_, _, serialized_xla_executable_))
.WillOnce(Return(ByMove(std::move(executable))));
auto loaded_executable =
persistor.TryToLoadExecutable(123, "signature", DefaultXlaOptions(),
compilation_result_mul, &mock_client);
EXPECT_TRUE(loaded_executable.has_value());
EXPECT_TRUE(loaded_executable->ok());
}
TEST_F(DeviceExecutionPersistorTest, LoadSerializedExecutableEmpty) {
XlaDeviceExecutablePersistor::Config config(
/*persistent_cache_directory=*/cache_dir_,
/*disable_strict_signature_checks=*/false,
/*persistence_prefix=*/"xla");
XlaDeviceExecutablePersistor persistor(config,
DefaultXlaOptions().device_type);
auto key =
CreateCacheKey(/*signature_hash=*/123, compilation_result_add_,
persistor.device_type(), persistor.persistence_prefix());
// Read serialized entry.
XlaSerializedCacheEntry entry;
TF_ASSERT_OK(ReadTextOrBinaryProto(
Env::Default(), GetFilePath(key, persistor.persistent_cache_directory()),
&entry));
entry.clear_executable();
// Write entry to another file.
TF_ASSERT_OK(WriteBinaryProto(
Env::Default(), GetFilePath(key, persistor.persistent_cache_directory()),
entry));
MockXlaCompilerClient mock_client;
auto loaded_executable = persistor.TryToLoadExecutable(
/*signature_hash=*/123, "signature", DefaultXlaOptions(),
compilation_result_add_, &mock_client);
EXPECT_TRUE(loaded_executable.has_value());
EXPECT_FALSE(loaded_executable->ok());
EXPECT_THAT(loaded_executable.value(),
absl_testing::StatusIs(error::INVALID_ARGUMENT));
}
TEST_F(DeviceExecutionPersistorTest, PersistPjRtAndXlaExecutables) {
// Persist PJRT executable.
PjRtDeviceExecutablePersistor::Config pjrt_config(
/*persistent_cache_directory=*/cache_dir_,
/*disable_strict_signature_checks=*/false,
/*persistence_prefix=*/"xla");
PjRtDeviceExecutablePersistor pjrt_persistor(
pjrt_config, DefaultPjRtOptions().device_type);
MockPjRtCompilerClient mock_pjrt_client;
EXPECT_CALL(mock_pjrt_client, SerializeExecutable(_))
.WillOnce(Return(serialized_pjrt_executable_));
TF_ASSERT_OK_AND_ASSIGN(auto pjrt_executable, BuildSamplePjRtExecutable());
TF_EXPECT_OK(pjrt_persistor.TryToPersistExecutable(
/*signature_hash=*/123, "signature_string", DefaultPjRtOptions(),
compilation_result_add_, *pjrt_executable, &mock_pjrt_client));
// Persist XLA executable.
XlaDeviceExecutablePersistor::Config xla_config(
/*persistent_cache_directory=*/cache_dir_,
/*disable_strict_signature_checks=*/false,
/*persistence_prefix=*/"xla");
XlaDeviceExecutablePersistor xla_persistor(xla_config,
DefaultXlaOptions().device_type);
MockXlaCompilerClient mock_xla_client;
EXPECT_CALL(mock_xla_client, SerializeExecutable(_))
.WillOnce(Return(serialized_xla_executable_));
TF_ASSERT_OK_AND_ASSIGN(auto xla_executable, BuildSampleExecutable());
TF_EXPECT_OK(xla_persistor.TryToPersistExecutable(
/*signature_hash=*/123, "signature_string", DefaultXlaOptions(),
compilation_result_add_, *xla_executable, &mock_xla_client));
// Read and verify serialized PJRT executable.
auto pjrt_key = CreateCacheKey(
/*signature_hash=*/123, compilation_result_add_,
pjrt_persistor.device_type(), pjrt_persistor.persistence_prefix(), true);
TF_ASSERT_OK_AND_ASSIGN(auto entry,
ReadCacheEntryFromFile(pjrt_key, cache_dir_));
EXPECT_EQ(entry.executable(), serialized_pjrt_executable_);
// Read and verify serialized XLA executable.
auto xla_key = CreateCacheKey(/*signature_hash=*/123, compilation_result_add_,
pjrt_persistor.device_type(),
pjrt_persistor.persistence_prefix());
TF_ASSERT_OK_AND_ASSIGN(entry, ReadCacheEntryFromFile(xla_key, cache_dir_));
EXPECT_EQ(entry.executable(), serialized_xla_executable_);
}
} // namespace
} // namespace tensorflow
tensorflow/compiler/jit/device_util.cc
+236
View File
@@ -0,0 +1,236 @@
/* Copyright 2019 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/compiler/jit/device_util.h"
#include "absl/algorithm/container.h"
#include "absl/container/flat_hash_set.h"
#include "xla/status_macros.h"
namespace tensorflow {
namespace jit {
void DeviceSet::Insert(DeviceId device_id) {
int word_index = device_id.id() / kWordSize;
int bit_index = device_id.id() % kWordSize;
const int storage_size = storage_.size();
if (word_index >= storage_size) {
storage_.resize(word_index + 1, 0);
}
storage_[word_index] |= (1ull << bit_index);
}
void DeviceSet::UnionWith(const DeviceSet& other) {
if (other.storage_.size() > storage_.size()) {
storage_.resize(other.storage_.size(), 0);
}
for (int i = 0, end = other.storage_.size(); i < end; i++) {
storage_[i] |= other.storage_[i];
}
}
bool DeviceSet::IsEmpty() const {
return absl::c_all_of(storage_, [&](uint64_t val) { return val == 0; });
}
absl::StatusOr<DeviceId> DeviceInfoCache::GetIdFor(absl::string_view name) {
TF_RET_CHECK(!name.empty());
auto it = name_to_id_.find(name);
if (it != name_to_id_.end()) {
return it->second;
}
int new_id = names_.size();
names_.push_back(std::string(name));
id_to_device_type_.push_back(std::make_unique<DeviceType>(""));
DeviceType* device_type = id_to_device_type_.back().get();
TF_RETURN_IF_ERROR(DeviceNameToDeviceType(names_.back(), device_type));
is_cpu_.push_back(device_type->type_string() == DEVICE_CPU);
is_gpu_.push_back(device_type->type_string() == DEVICE_GPU);
name_to_id_.emplace(std::string(name), DeviceId(new_id));
const XlaOpRegistry::DeviceRegistration* compilation_device;
if (!XlaOpRegistry::GetCompilationDevice(device_type->type(),
&compilation_device)) {
compilation_device = nullptr;
}
id_to_compilation_device_.push_back(compilation_device);
return DeviceId(new_id);
}
std::string DeviceInfoCache::DebugString(const DeviceSet& device_set) const {
std::vector<std::string> names;
device_set.ForEach([&](DeviceId device_id) {
names.push_back(std::string(GetNameFor(device_id)));
return true;
});
return absl::StrCat("[", absl::StrJoin(names, ","), "]");
}
} // namespace jit
absl::Status DeviceNameToDeviceType(const std::string& device,
DeviceType* device_type) {
DeviceNameUtils::ParsedName parsed;
if (!DeviceNameUtils::ParseFullName(device, &parsed)) {
return absl::InternalError(
absl::StrCat("Malformed assigned device '", device, "'"));
}
*device_type = DeviceType(parsed.type);
return absl::OkStatus();
}
absl::StatusOr<std::optional<jit::DeviceId>> PickDeviceForXlaImpl(
const jit::DeviceInfoCache& device_info_cache,
const jit::DeviceSet& devices, bool allow_mixing_unknown_and_cpu,
bool failure_to_pick_is_error) {
#define FAILED_TO_PICK_DEVICE(failing_status) \
do { \
if (failure_to_pick_is_error) { \
return failing_status; \
} else { \
return {std::nullopt}; \
} \
} while (false)
std::optional<jit::DeviceId> maybe_gpu_device;
std::optional<jit::DeviceId> maybe_cpu_device;
std::optional<jit::DeviceId> maybe_unknown_device;
bool multiple_cpu_devices = false;
bool multiple_gpu_devices = false;
bool multiple_unknown_devices = false;
// Returns 'true' if d0 and d1 are conflicting devices. If they are
// compatible, update d1 with a more specific one.
// TODO(sanjoy): Cache DeviceNameUtils::ParsedName inside device_info_cache.
const auto is_multiple_devices =
[&](const jit::DeviceId& d0, std::optional<jit::DeviceId>* d1) -> bool {
const absl::string_view name0 = device_info_cache.GetNameFor(d0);
const absl::string_view name1 = device_info_cache.GetNameFor(d1->value());
DeviceNameUtils::ParsedName parsed0, parsed1;
if (!DeviceNameUtils::ParseFullName(name0, &parsed0) ||
!DeviceNameUtils::ParseFullName(name1, &parsed1) ||
!DeviceNameUtils::AreCompatibleDevNames(parsed0, parsed1)) {
return true;
}
if (DeviceNameUtils::IsSpecification(parsed0, parsed1)) {
return false;
}
if (DeviceNameUtils::IsSpecification(parsed1, parsed0)) {
*d1 = d0;
return false;
}
return true;
};
devices.ForEach([&](jit::DeviceId device) {
if (device_info_cache.IsGpu(device)) {
if (maybe_gpu_device) {
multiple_gpu_devices = is_multiple_devices(device, &maybe_gpu_device);
if (multiple_gpu_devices) return false;
} else {
maybe_gpu_device = device;
}
} else if (device_info_cache.IsCpu(device)) {
if (maybe_cpu_device) {
multiple_cpu_devices = is_multiple_devices(device, &maybe_cpu_device);
if (multiple_cpu_devices) return false;
} else {
maybe_cpu_device = device;
}
} else {
if (maybe_unknown_device) {
multiple_unknown_devices = true;
return false;
}
maybe_unknown_device = device;
}
return true;
});
if (multiple_cpu_devices) {
FAILED_TO_PICK_DEVICE(absl::InternalError(absl::StrCat(
"Multiple CPU devices ", device_info_cache.DebugString(devices))));
}
if (multiple_gpu_devices) {
FAILED_TO_PICK_DEVICE(absl::InternalError(absl::StrCat(
"Multiple GPU devices ", device_info_cache.DebugString(devices))));
}
if (multiple_unknown_devices) {
FAILED_TO_PICK_DEVICE(absl::InternalError(absl::StrCat(
"Multiple unknown devices ", device_info_cache.DebugString(devices))));
}
if (maybe_unknown_device && maybe_gpu_device) {
FAILED_TO_PICK_DEVICE(absl::InternalError(
absl::StrCat("Found both unknown and GPU devices: ",
device_info_cache.GetNameFor(*maybe_unknown_device), ", ",
device_info_cache.GetNameFor(*maybe_gpu_device))));
}
if (!allow_mixing_unknown_and_cpu) {
if (maybe_unknown_device && maybe_cpu_device) {
FAILED_TO_PICK_DEVICE(absl::InternalError(
absl::StrCat("Found both unknown and CPU devices: ",
device_info_cache.GetNameFor(*maybe_unknown_device),
", ", device_info_cache.GetNameFor(*maybe_cpu_device))));
}
}
if (maybe_gpu_device) {
return {*maybe_gpu_device};
} else if (maybe_unknown_device) {
return {*maybe_unknown_device};
} else if (maybe_cpu_device) {
return {*maybe_cpu_device};
}
FAILED_TO_PICK_DEVICE(absl::InternalError("Empty device set!"));
#undef FAILED_TO_PICK_DEVICE
}
absl::StatusOr<jit::DeviceId> PickDeviceForXla(
const jit::DeviceInfoCache& device_info_cache,
const jit::DeviceSet& devices, bool allow_mixing_unknown_and_cpu) {
TF_ASSIGN_OR_RETURN(std::optional<jit::DeviceId> device_id,
PickDeviceForXlaImpl(device_info_cache, devices,
allow_mixing_unknown_and_cpu,
/*failure_to_pick_is_error=*/true));
return *device_id;
}
absl::StatusOr<std::optional<jit::DeviceId>> MaybePickDeviceForXla(
const jit::DeviceInfoCache& device_info_cache,
const jit::DeviceSet& devices, bool allow_mixing_unknown_and_cpu) {
return PickDeviceForXlaImpl(device_info_cache, devices,
allow_mixing_unknown_and_cpu,
/*failure_to_pick_is_error=*/false);
}
} // namespace tensorflow
tensorflow/compiler/jit/device_util.h
+203
View File
@@ -0,0 +1,203 @@
/* Copyright 2019 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_COMPILER_JIT_DEVICE_UTIL_H_
#define TENSORFLOW_COMPILER_JIT_DEVICE_UTIL_H_
#include <functional>
#include <memory>
#include "absl/container/flat_hash_map.h"
#include "absl/numeric/bits.h"
#include "absl/status/statusor.h"
#include "absl/strings/string_view.h"
#include "absl/types/span.h"
#include "tensorflow/compiler/tf2xla/xla_op_registry.h"
#include "xla/status_macros.h"
#include "tensorflow/core/framework/types.h"
namespace tensorflow {
namespace jit {
class DeviceInfoCache;
class DeviceSet;
// Instances of DeviceId represent TensorFlow devices as integers.
//
// This helps avoid having to manipulate device names as strings when
// auto-clustering.
class DeviceId {
public:
DeviceId(DeviceId&&) = default;
DeviceId(const DeviceId&) = default;
DeviceId& operator=(const DeviceId&) = default;
bool operator==(const DeviceId& other) const { return id() == other.id(); }
bool operator!=(const DeviceId& other) const { return !(*this == other); }
private:
int id_;
explicit DeviceId(int id) : id_(id) {}
int id() const { return id_; }
friend class DeviceInfoCache;
friend class DeviceSet;
};
// A set of DeviceIds, represented as a bitmap.
class DeviceSet {
public:
void Insert(DeviceId device_id);
void UnionWith(const DeviceSet& other);
bool IsEmpty() const;
// Calls `func` on each DeviceId in the set. Stops iterating early if `func`
// return false.
//
// TODO(sanjoy): Change this to take a typed std::function if that's
// performance neutral.
template <typename FnTy>
void ForEach(FnTy func) const {
// This is really a poor man's iterator, we should consider writing a proper
// iterator if this ends up being used widely.
for (int word_index = 0, end = storage_.size(); word_index < end;
word_index++) {
uint64_t word = storage_[word_index];
while (word != 0) {
uint64_t only_lowest_bit_set = word & -word;
// The number of trailing zeros in a non-zero word is the index of the
// least significant 1.
int bit_index = absl::countr_zero(word);
if (!func(DeviceId(word_index * kWordSize + bit_index))) {
return;
}
word ^= only_lowest_bit_set;
}
}
}
private:
absl::InlinedVector<uint64_t, 1> storage_;
const int kWordSize = 64;
};
// Caches some miscellaneous information about TF devices. Thread compatible.
class DeviceInfoCache {
public:
bool IsGpu(DeviceId device) const { return is_gpu_[device.id()]; }
bool IsCpu(DeviceId device) const { return is_cpu_[device.id()]; }
absl::string_view GetNameFor(DeviceId device) const {
return names_[device.id()];
}
absl::StatusOr<DeviceId> GetIdFor(absl::string_view name);
using DeviceRegistration = const XlaOpRegistry::DeviceRegistration;
DeviceRegistration* GetCompilationDevice(DeviceId device) const {
return id_to_compilation_device_[device.id()];
}
absl::StatusOr<DeviceRegistration*> GetCompilationDevice(
absl::string_view name) {
TF_ASSIGN_OR_RETURN(DeviceId device_id, GetIdFor(name));
return GetCompilationDevice(device_id);
}
const DeviceType& GetDeviceTypeFor(DeviceId device) const {
return *id_to_device_type_[device.id()];
}
using DeviceTypeConstRef = std::reference_wrapper<const DeviceType>;
absl::StatusOr<DeviceTypeConstRef> GetDeviceTypeFor(
absl::string_view device_name) {
TF_ASSIGN_OR_RETURN(DeviceId device_id, GetIdFor(device_name));
return std::cref(*id_to_device_type_[device_id.id()]);
}
std::string DebugString(const DeviceSet& device_set) const;
private:
absl::flat_hash_map<std::string, DeviceId> name_to_id_;
// These fields are populated for a device in GetIdFor, *before* we give out a
// DeviceId.
std::vector<const XlaOpRegistry::DeviceRegistration*>
id_to_compilation_device_;
std::vector<std::unique_ptr<DeviceType>> id_to_device_type_;
std::vector<std::string> names_;
std::vector<bool> is_cpu_;
std::vector<bool> is_gpu_;
};
} // namespace jit
// Returns the DeviceType corresponding to 'device'.
absl::Status DeviceNameToDeviceType(const std::string& device,
DeviceType* device_type);
// Picks the device for which XLA should compile a cluster that contains
// operations placed in devices in `devices`. For instance a cluster that
// contains operations solely placed on the CPU will be compiled into a CPU
// executable by XLA, whereas a cluster that contains operations placed on the
// CPU and also operations placed on the GPU will be compiled into a GPU
// executable.
//
// Returns a non-OK Status if no unambiguous choice of device exists.
//
// We choose the device using the following rules:
//
// - It is an error for `device_names` to contain more than one device of the
// same type.
// - GPU is preferred over CPU.
// - If `allow_mixing_unknown_and_cpu` is true then unknown devices are
// preferred over CPU.
// - XLA devices count as "unrecognized devices".
//
// This set of rules above implicitly assume that XLA:GPU can compile all
// operations in the cluster that XLA:CPU can compile, and if
// `allow_mixing_unknown_and_cpu` then the unrecognized device can also compile
// all operations in the cluster that XLA:CPU can compile.
//
// We provide the `allow_mixing_unknown_and_cpu` knob so that we can do both of
// the following things:
//
// - Let MarkForCompilationPass not inject CPU-placed operations into clusters
// that will run on unknown devices (because the unknown XLA backend may not
// support every operation supported by CPU).
// - Let BuildXlaOpsPass successfully infer a compilation device for a cluster
// that contains nodes placed on both the CPU and on unknown devices. In this
// case it is the responsibility of the optimization pass that injected the
// CPU nodes into the cluster to ensure that these nodes can be compiled by
// the unknown XLA backend.
absl::StatusOr<jit::DeviceId> PickDeviceForXla(
const jit::DeviceInfoCache& device_info_cache,
const jit::DeviceSet& devices, bool allow_mixing_unknown_and_cpu);
// This is like `PickDeviceForXla` except that it returns nullopt (instead of a
// non-OK Status) if no unambiguous choice of device exists.
//
// We return a failing Status for errors unrelated to the device choice
// algorithm itself.
absl::StatusOr<std::optional<jit::DeviceId>> MaybePickDeviceForXla(
const jit::DeviceInfoCache& device_info_cache,
const jit::DeviceSet& devices, bool allow_mixing_unknown_and_cpu);
} // namespace tensorflow
#endif // TENSORFLOW_COMPILER_JIT_DEVICE_UTIL_H_
tensorflow/compiler/jit/device_util_test.cc
+144
View File
@@ -0,0 +1,144 @@
/* Copyright 2019 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/compiler/jit/device_util.h"
#include "tensorflow/core/lib/core/status_test_util.h"
#include "tensorflow/core/platform/test.h"
namespace tensorflow {
namespace {
absl::Status PickDeviceHelper(bool allow_mixing_unknown_and_cpu,
absl::Span<const absl::string_view> device_names,
std::string* result) {
jit::DeviceInfoCache cache;
jit::DeviceSet device_set;
for (absl::string_view name : device_names) {
TF_ASSIGN_OR_RETURN(jit::DeviceId device_id, cache.GetIdFor(name));
device_set.Insert(device_id);
}
TF_ASSIGN_OR_RETURN(
jit::DeviceId result_id,
PickDeviceForXla(cache, device_set, allow_mixing_unknown_and_cpu));
*result = std::string(cache.GetNameFor(result_id));
return absl::OkStatus();
}
void CheckPickDeviceResult(absl::string_view expected_result,
bool allow_mixing_unknown_and_cpu,
absl::Span<const absl::string_view> inputs) {
std::string result;
TF_ASSERT_OK(PickDeviceHelper(allow_mixing_unknown_and_cpu, inputs, &result))
<< "inputs = [" << absl::StrJoin(inputs, ", ")
<< "], allow_mixing_unknown_and_cpu=" << allow_mixing_unknown_and_cpu
<< ", expected_result=" << expected_result;
EXPECT_EQ(result, expected_result);
}
void CheckPickDeviceHasError(bool allow_mixing_unknown_and_cpu,
absl::Span<const absl::string_view> inputs) {
std::string result;
EXPECT_FALSE(
PickDeviceHelper(allow_mixing_unknown_and_cpu, inputs, &result).ok());
}
const char* kCPU0 = "/job:localhost/replica:0/task:0/device:CPU:0";
const char* kGPU0 = "/job:localhost/replica:0/task:0/device:GPU:0";
const char* kXPU0 = "/job:localhost/replica:0/task:0/device:XPU:0";
const char* kYPU0 = "/job:localhost/replica:0/task:0/device:YPU:0";
const char* kCPU1 = "/job:localhost/replica:0/task:0/device:CPU:1";
const char* kGPU1 = "/job:localhost/replica:0/task:0/device:GPU:1";
const char* kXPU1 = "/job:localhost/replica:0/task:0/device:XPU:1";
const char* kCPU0Partial = "/device:CPU:0";
const char* kGPU0Partial = "/device:GPU:0";
const char* kXPU0Partial = "/device:XPU:0";
TEST(PickDeviceForXla, UniqueDevice) {
CheckPickDeviceResult(kGPU0, false, {kGPU0, kGPU0});
}
TEST(PickDeviceForXla, MoreSpecificDevice) {
CheckPickDeviceResult(kCPU0, false, {kCPU0, kCPU0Partial});
CheckPickDeviceResult(kGPU0, false, {kGPU0, kGPU0Partial});
// Unknown devices do not support merging of full and partial specifications.
CheckPickDeviceHasError(false, {kXPU1, kXPU0Partial});
}
TEST(PickDeviceForXla, DeviceOrder) {
CheckPickDeviceResult(kGPU0, false, {kGPU0, kCPU0});
CheckPickDeviceResult(kGPU0, false, {kCPU0, kGPU0});
CheckPickDeviceResult(kXPU0, true, {kXPU0, kCPU0});
}
TEST(PickDeviceForXla, MultipleUnknownDevices) {
CheckPickDeviceHasError(false, {kXPU0, kYPU0});
}
TEST(PickDeviceForXla, GpuAndUnknown) {
CheckPickDeviceHasError(false, {kGPU0, kXPU1});
}
TEST(PickDeviceForXla, UnknownAndCpu) {
CheckPickDeviceHasError(false, {kXPU0, kCPU1});
}
TEST(PickDeviceForXla, MultipleDevicesOfSameType) {
CheckPickDeviceHasError(true, {kCPU0, kCPU1});
CheckPickDeviceHasError(false, {kCPU0, kCPU1});
CheckPickDeviceHasError(false, {kGPU0, kGPU1});
CheckPickDeviceHasError(false, {kXPU0, kXPU1});
CheckPickDeviceHasError(false, {kCPU0, kCPU1, kGPU0});
}
void SimpleRoundTripTestForDeviceSet(int num_devices) {
jit::DeviceSet device_set;
jit::DeviceInfoCache device_info_cache;
std::vector<std::string> expected_devices, actual_devices;
for (int i = 0; i < num_devices; i++) {
std::string device_name =
absl::StrCat("/job:localhost/replica:0/task:0/device:XPU:", i);
TF_ASSERT_OK_AND_ASSIGN(jit::DeviceId device_id,
device_info_cache.GetIdFor(device_name));
device_set.Insert(device_id);
expected_devices.push_back(device_name);
}
device_set.ForEach([&](jit::DeviceId device_id) {
actual_devices.push_back(
std::string(device_info_cache.GetNameFor(device_id)));
return true;
});
EXPECT_EQ(expected_devices, actual_devices);
}
TEST(DeviceSetTest, SimpleRoundTrip_One) { SimpleRoundTripTestForDeviceSet(1); }
TEST(DeviceSetTest, SimpleRoundTrip_Small) {
SimpleRoundTripTestForDeviceSet(8);
}
TEST(DeviceSetTest, SimpleRoundTrip_Large) {
SimpleRoundTripTestForDeviceSet(800);
}
} // namespace
} // namespace tensorflow
tensorflow/compiler/jit/encapsulate_subgraphs_pass.cc
+1334
View File
File diff suppressed because it is too large Load Diff
tensorflow/compiler/jit/encapsulate_subgraphs_pass.h
+108
View File
@@ -0,0 +1,108 @@
/* Copyright 2017 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.
==============================================================================*/
// An optimization pass that groups nodes marked with a common
// kXlaClusterAttr into functions, and replaces the original nodes by
// calls. The calls are annotated with kXlaCompiledKernelAttr.
#ifndef TENSORFLOW_COMPILER_JIT_ENCAPSULATE_SUBGRAPHS_PASS_H_
#define TENSORFLOW_COMPILER_JIT_ENCAPSULATE_SUBGRAPHS_PASS_H_
#include "tensorflow/core/common_runtime/optimization_registry.h"
#include "tensorflow/core/framework/function.h"
#include "tensorflow/core/graph/graph.h"
#include "tensorflow/core/lib/core/status.h"
namespace tensorflow {
// EncapsulateSubgraphs pass takes all the nodes with the same cluster ID
// (derived from kXlaClusterAttr=ID (kXlaClusterAttr) attribute), puts them into
// a TF function, and replaces the subgraph in the main graph with a call to
// that TF function annotated with kXlaCompiledKernelAttr (_XlaCompiledKernel).
class EncapsulateSubgraphsPass : public GraphOptimizationPass {
public:
absl::Status Run(const GraphOptimizationPassOptions& options) override;
};
// A rewriting function to apply to each subgraph during encapsulation.
// 'arg_source_tensors' are the tensors corresponding to the arguments in the
// original source graph (*not* 'graph').
//
// 'graph' is the subgraph. The rewriting may renumber the inputs and outputs;
// 'input_permutation' is a mapping from old argument numbers to new argument
// numbers, whereas 'output_permutation' is the same for outputs. Both
// 'input_permutation' and 'output_permutation' are initialized to the identity
// permutation. 'nodedef' is the NodeDef for the call to the function under
// construction, provided to allow additional attributes to be set.
// The rewrite may also change the NodeDef's operator name, and that
// name will be used as the name of the generated function.
typedef std::function<absl::Status(
const std::vector<OutputTensor>& arg_source_tensors,
std::unique_ptr<Graph>* graph, std::vector<int>* input_permutation,
std::vector<int>* output_permutation, NodeDef* node_def)>
RewriteSubgraphFn;
// Transformation that finds subgraphs whose nodes are marked with
// 'group_attribute', splits those subgraphs into functions, and replaces
// the originals with function calls.
//
// 'group_attribute' must be a string valued-attribute that names the new
// functions to introduce.
//
// If 'rewrite_subgraph_fn' is set, it is applied to each subgraph before
// function conversion.
//
// If 'reuse_existing_functions' is set, use an existing function with the
// same name, if any.
//
// TODO(phawkins): currently, some information in control edges
// is not preserved. Suppose you have A and B in the main
// graph, C and D in a subgraph. B and C have control deps from A, D has control
// dep from B. Originally D must run after C, post-transformation this
// dependency is lost.
absl::Status EncapsulateSubgraphsInFunctions(
std::string group_attribute, const Graph& graph_in,
const RewriteSubgraphFn& rewrite_subgraph_fn, bool reuse_existing_functions,
std::unique_ptr<Graph>* graph_out, FunctionLibraryDefinition* library);
// The attribute that marks function calls produced by the encapsulate
// subgraphs pass and that should in turn be compiled via XlaLaunch operators.
extern const char* const kXlaCompiledKernelAttr;
// Does `node` have the kXlaCompiledKernelAttr attribute?
bool IsXlaCompiledKernel(const Node& node);
// Functions produced by the EncapsulateSubgraphs pass have their arguments in
// the order:
// 1) compile-time constant arguments, in host memory,
// 2) other arguments, in device memory.
// 3) resource variable arguments, in host memory. Note that only the resource
// Tensor itself is in host memory; the underlying value may be in device
// memory.
// The functions are annotated with the following attributes that describe how
// many constant and resource arguments there are:
// Name of the attribute containing the number of constant arguments.
extern const char* const kXlaNumConstantArgsAttr;
// Name of the attribute containing the number of resource variable arguments.
extern const char* const kXlaNumResourceArgsAttr;
// Name of the attribute defining whether the cluster has reference variables.
extern const char* const kXlaHasReferenceVarsAttr;
} // namespace tensorflow
#endif // TENSORFLOW_COMPILER_JIT_ENCAPSULATE_SUBGRAPHS_PASS_H_
tensorflow/compiler/jit/encapsulate_subgraphs_pass_test.cc
+2788
View File
File diff suppressed because it is too large Load Diff
tensorflow/compiler/jit/encapsulate_util.cc
+418
View File
@@ -0,0 +1,418 @@
/* 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/compiler/jit/encapsulate_util.h"
#include <algorithm>
#include <iterator>
#include "absl/container/flat_hash_map.h"
#include "absl/container/flat_hash_set.h"
#include "absl/strings/str_cat.h"
#include "absl/types/optional.h"
#include "tensorflow/compiler/jit/shape_inference.h"
#include "tensorflow/compiler/tf2xla/tf2xla_util.h"
#include "tensorflow/core/framework/node_def_util.h"
#include "tensorflow/core/graph/node_builder.h"
#include "tensorflow/core/protobuf/error_codes.pb.h"
using tsl::StatusOr;
namespace tensorflow {
namespace {
// Returns string attribute value for the node if the attribute is present,
// otherwise returns empty optional value.
std::optional<std::string> GetStringAttr(const Node& n,
const std::string& attr_name) {
auto attr = n.attrs().Find(attr_name);
if (!attr) {
return std::nullopt;
} else {
return attr->s();
}
}
// Adds a value to the node's list attribute.
template <typename T>
absl::Status AppendToListAttr(Node* n, const std::string& attr_name,
const std::string& value) {
std::vector<T> attr_value;
absl::Status s = GetNodeAttr(n->attrs(), attr_name, &attr_value);
if (!s.ok() && s.code() != error::NOT_FOUND) {
return s;
}
n->ClearAttr(attr_name);
attr_value.push_back(value);
n->AddAttr(attr_name, attr_value);
return absl::OkStatus();
}
// Replaces attribute value.
template <typename T>
void ReplaceAttr(Node* n, const std::string& attr_name, const T& value) {
n->ClearAttr(attr_name);
n->AddAttr(attr_name, value);
}
// Step 1 for `PreprocessEdgesBetweenOutsideCompilations`. See comments of
// `PreprocessEdgesBetweenOutsideCompilations` for details.
absl::Status PreprocessControlEdgesBetweenOutsideCompilations(
Graph* g, const std::string& outside_compilation_attr_name) {
// Gather edges to remove. We should not remove the edge while iterating.
std::vector<const Edge*> edges_to_remove;
for (const Edge* e : g->edges()) {
if (!e->IsControlEdge()) {
continue;
}
auto src_outside_compilation =
GetStringAttr(*e->src(), outside_compilation_attr_name);
auto dst_outside_compilation =
GetStringAttr(*e->dst(), outside_compilation_attr_name);
if (src_outside_compilation && dst_outside_compilation) {
if (*src_outside_compilation != *dst_outside_compilation) {
// Case 1a: outside compilation to outside compilation control edge.
edges_to_remove.push_back(e);
TF_RETURN_IF_ERROR(AppendToListAttr<std::string>(
e->dst(), kXlaControlDependenciesWithinXlaClusterAttrName,
e->src()->name()));
}
} else if (src_outside_compilation && !dst_outside_compilation) {
// Case 1b: outside compilation to its XLA computation control edge.
ReplaceAttr(e->src(), kXlaConnectedToXlaComputationAttrName, true);
} else if (!src_outside_compilation && dst_outside_compilation) {
// Case 1b: XLA computation to outside compilation in it control edge.
ReplaceAttr(e->dst(), kXlaConnectedFromXlaComputationAttrName, true);
}
}
for (auto e : edges_to_remove) {
g->RemoveEdge(e);
}
return absl::OkStatus();
}
// Step 2 for `PreprocessEdgesBetweenOutsideCompilations`. See comments of
// `PreprocessEdgesBetweenOutsideCompilations` for details.
absl::Status PreprocessDataEdgesBetweenOutsideCompilations(
Graph* g, const std::string& outside_compilation_attr_name) {
// Gather edges between outside compilation and host computation. Notice that
// we do not store `Edge*` directly because we remove some nodes while adding
// Identity nodes, and those Edge pointers might be invalidated.
struct EdgeInfo {
int dst_input, dst_node_id;
};
std::vector<EdgeInfo> edges;
for (const Edge* e : g->edges()) {
if (e->IsControlEdge()) {
continue;
}
auto src_outside_compilation =
GetStringAttr(*e->src(), outside_compilation_attr_name);
auto dst_outside_compilation =
GetStringAttr(*e->dst(), outside_compilation_attr_name);
if (src_outside_compilation && dst_outside_compilation &&
*src_outside_compilation != *dst_outside_compilation) {
edges.push_back(EdgeInfo{e->dst_input(), e->dst()->id()});
VLOG(4) << "Oc -> oc edge: " << e->DebugString();
}
}
// Remove the edge from host to outside compilation. Add a placeholder as
// outside compilation node input.
std::map<std::pair<std::string, int>, Node*> placeholders;
for (int i = 0, end = edges.size(); i < end; i++) {
Node* dst = g->FindNodeId(edges[i].dst_node_id);
const Edge* e;
TF_RETURN_IF_ERROR(dst->input_edge(edges[i].dst_input, &e));
Node* src = e->src();
int src_output = e->src_output(), dst_input = e->dst_input();
g->RemoveEdge(e);
// Find or create placeholder node.
std::string new_name =
absl::StrCat(src->name(), "_oc_to_oc_placeholder_", src_output);
auto placeholder_index = std::make_pair(src->name(), src_output);
auto iter = placeholders.find(placeholder_index);
Node* placeholder_node;
if (iter == placeholders.end()) {
NodeDefBuilder placeholder_builder(new_name, "Placeholder");
placeholder_builder.Attr("dtype", src->output_type(src_output));
std::string outside_compilation_attr;
TF_RETURN_IF_ERROR(GetNodeAttr(dst->attrs(),
outside_compilation_attr_name,
&outside_compilation_attr));
placeholder_builder.Attr(outside_compilation_attr_name,
outside_compilation_attr);
placeholder_builder.Attr(kOutsideCompilationOriginalNodeAttrName,
src->name());
placeholder_builder.Attr(kOutsideCompilationSrcOutputAttrName,
src_output);
NodeDef placeholder_def;
TF_RETURN_IF_ERROR(placeholder_builder.Finalize(&placeholder_def));
TF_ASSIGN_OR_RETURN(placeholder_node, g->AddNode(placeholder_def));
placeholders[placeholder_index] = placeholder_node;
} else {
placeholder_node = iter->second;
}
g->AddEdge(placeholder_node, 0, dst, dst_input);
// Replace `e->dst()` because its input node changed.
NodeDef new_def = dst->def();
*new_def.mutable_input(dst_input) = placeholder_node->name();
TF_ASSIGN_OR_RETURN(Node * dst_replace_node, ReplaceNode(g, dst, new_def));
// Other edge in `edges` might have `e->dst()` as src or dst
// node. Before removing `e->dst()`, replace those edges with
// corresponding edges for `dst_replace_node`.
for (int j = i + 1, end = edges.size(); j < end; j++) {
if (edges[j].dst_node_id == edges[i].dst_node_id) {
edges[j].dst_node_id = dst_replace_node->id();
}
}
}
return absl::OkStatus();
}
// Step 1 for `PostprocessEdgesBetweenOutsideCompilations`. See comments of
// `PostprocessEdgesBetweenOutsideCompilations` for details.
absl::Status PostprocessDataEdgesBetweenOutsideCompilations(
Graph* g, const std::string& outside_compilation_attr_name) {
// Gather all outside compilation to outside compilation nodes.
std::vector<Node*> placeholder_nodes;
for (Node* n : g->nodes()) {
if (n->type_string() == "Placeholder" &&
HasNodeAttr(n->def(), kOutsideCompilationOriginalNodeAttrName)) {
placeholder_nodes.push_back(n);
}
}
// Remove the placeholder nodes, and reconnect original edge.
auto node_name_index = g->BuildNodeNameIndex();
for (auto n : placeholder_nodes) {
std::string node_name;
int node_src_output;
TF_RETURN_IF_ERROR(GetNodeAttr(
n->attrs(), kOutsideCompilationOriginalNodeAttrName, &node_name));
TF_RETURN_IF_ERROR(GetNodeAttr(
n->attrs(), kOutsideCompilationSrcOutputAttrName, &node_src_output));
auto iter = node_name_index.find(node_name);
if (iter == node_name_index.end()) {
return absl::InternalError(absl::StrCat(
"Cannot find original node for oc -> host placeholder node ",
node_name));
}
// Change all usage node to use the original node instead.
Node* original_node = iter->second;
std::vector<const Edge*> control_edges;
std::vector<OutEdgeInfo> data_edges;
for (auto e : n->out_edges()) {
if (e->IsControlEdge()) {
control_edges.push_back(e);
} else {
data_edges.push_back({e->dst(), e->src_output(), e->dst_input()});
}
}
for (const Edge* e : control_edges) {
g->AddControlEdge(original_node, e->dst());
g->RemoveEdge(e);
}
for (int i = 0, end = data_edges.size(); i < end; i++) {
Node* dst = data_edges[i].dst;
NodeDef new_def = dst->def();
int dst_input = data_edges[i].dst_input;
*new_def.mutable_input(dst_input) =
absl::StrCat(original_node->name(), ":", node_src_output);
TF_ASSIGN_OR_RETURN(Node * replace_node, ReplaceNode(g, dst, new_def));
const Edge* edge_to_replace = nullptr;
TF_RETURN_IF_ERROR(replace_node->input_edge(dst_input, &edge_to_replace));
g->RemoveEdge(edge_to_replace);
g->AddEdge(original_node, node_src_output, replace_node, dst_input);
// Other edges might have `dst` as dst node. Update those edges with
// `replace_node`.
for (int j = i + 1, end = data_edges.size(); j < end; j++) {
if (data_edges[j].dst == dst) {
data_edges[j].dst = replace_node;
}
}
// Other placeholder node might have `dst` as original node. Update
// `node_name_index` with `replace_node`.
node_name_index[replace_node->name()] = replace_node;
}
// Remove placeholder node.
g->RemoveNode(n);
}
return absl::OkStatus();
}
// Step 2 for `PostprocessEdgesBetweenOutsideCompilations`. See comments of
// `PostprocessEdgesBetweenOutsideCompilations` for details.
absl::Status PostprocessControlEdgesBetweenOutsideCompilations(
Graph* g, const std::string& outside_compilation_attr_name) {
auto node_name_index = g->BuildNodeNameIndex();
// Reconnect outside compilation to outside compilation control edge.
for (Node* n : g->nodes()) {
std::vector<std::string> control_deps;
absl::Status s =
GetNodeAttr(n->attrs(), kXlaControlDependenciesWithinXlaClusterAttrName,
&control_deps);
if (!s.ok()) {
if (s.code() != error::NOT_FOUND) {
return s;
} else {
continue;
}
} else {
n->ClearAttr(kXlaControlDependenciesWithinXlaClusterAttrName);
for (const std::string& control_input : control_deps) {
auto iter = node_name_index.find(control_input);
if (iter == node_name_index.end()) {
return absl::InternalError(
absl::StrCat("Cannot find original node for ", control_input));
}
g->AddControlEdge(iter->second, n);
}
}
}
return absl::OkStatus();
}
} // namespace
const char kXlaInferredShapesAttrName[] = "_xla_inferred_shapes";
const char kXlaConnectedToXlaComputationAttrName[] =
"_xla_connected_to_xla_computation";
const char kXlaConnectedFromXlaComputationAttrName[] =
"_xla_connected_from_xla_computation";
const char kOutsideCompilationOriginalNodeAttrName[] =
"_xla_oc_to_oc_node_name";
const char kOutsideCompilationSrcOutputAttrName[] = "_xla_oc_to_oc_src_output";
const char kXlaControlDependenciesWithinXlaClusterAttrName[] =
"_xla_control_dependencies_within_xla_cluster";
const char kXlaIsLiftedArgAttrName[] = "_xla_is_lifted_arg";
const char kXlaLiftedArgOutsideCompilationAttrName[] = "_xla_lifted_arg_oc";
const char kXlaOutsideCompilationInputsAttrName[] = "_xla_oc_inputs";
const char kXlaIsPlaceholderForArg[] = "_xla_is_placeholder_for_arg";
absl::Status PerformStaticShapeInferenceBeforeEncapsulation(Graph* g) {
// Perform shape inference.
std::map<int, InferredShape> arg_shapes;
GraphShapeInfo shape_info;
TF_RETURN_IF_ERROR(
InferShapes(g, arg_shapes, /*fnlib_def=*/nullptr, &shape_info));
// Add attribute for output shapes.
auto node_name_index = g->BuildNodeNameIndex();
for (auto iter : shape_info) {
std::vector<PartialTensorShape> output_shapes;
std::transform(iter.second.begin(), iter.second.end(),
std::back_inserter(output_shapes),
[](const InferredShape& inferred_shape) {
return inferred_shape.shape;
});
Node* n = node_name_index[iter.first];
n->AddAttr(kXlaInferredShapesAttrName, output_shapes);
}
return absl::OkStatus();
}
absl::StatusOr<
std::unique_ptr<absl::flat_hash_map<std::string, std::vector<std::string>>>>
OutsideCompilationClusterDependencies(
const Graph* g, const std::string& outside_compilation_attr_name) {
auto cluster_deps = std::make_unique<
absl::flat_hash_map<std::string, absl::flat_hash_set<std::string>>>();
for (const Edge* e : g->edges()) {
auto src_outside_compilation =
GetStringAttr(*e->src(), outside_compilation_attr_name);
auto dst_outside_compilation =
GetStringAttr(*e->dst(), outside_compilation_attr_name);
if (src_outside_compilation && dst_outside_compilation &&
*src_outside_compilation != *dst_outside_compilation) {
auto dst_deps_it = cluster_deps->find(*dst_outside_compilation);
if (dst_deps_it == cluster_deps->end()) {
cluster_deps->insert(std::make_pair(
*dst_outside_compilation,
absl::flat_hash_set<std::string>({*src_outside_compilation})));
} else {
dst_deps_it->second.insert(*src_outside_compilation);
}
}
}
auto cluster_deps_ordered = std::make_unique<
absl::flat_hash_map<std::string, std::vector<std::string>>>();
for (auto it = cluster_deps->begin(); it != cluster_deps->end(); it++) {
std::vector<std::string> ordered_deps(it->second.begin(), it->second.end());
std::sort(ordered_deps.begin(), ordered_deps.end());
cluster_deps_ordered->insert(std::make_pair(it->first, ordered_deps));
}
return std::move(cluster_deps_ordered);
}
absl::Status PreprocessEdgesBetweenOutsideCompilations(
Graph* g, const std::string& outside_compilation_attr_name) {
// Remove edges from source node to outside compilation nodes, and edges
// from outside compilation nodes to sink node.
std::vector<const Edge*> edges_to_remove;
for (const Edge* e : g->source_node()->out_edges()) {
if (HasNodeAttr(e->dst()->def(), outside_compilation_attr_name)) {
edges_to_remove.push_back(e);
}
}
for (const Edge* e : g->sink_node()->in_edges()) {
if (HasNodeAttr(e->src()->def(), outside_compilation_attr_name)) {
edges_to_remove.push_back(e);
}
}
for (auto e : edges_to_remove) {
g->RemoveEdge(e);
}
TF_RETURN_IF_ERROR(PreprocessControlEdgesBetweenOutsideCompilations(
g, outside_compilation_attr_name));
TF_RETURN_IF_ERROR(PreprocessDataEdgesBetweenOutsideCompilations(
g, outside_compilation_attr_name));
return absl::OkStatus();
}
absl::Status PostprocessEdgesBetweenOutsideCompilations(
Graph* g, const std::string& outside_compilation_attr_name) {
TF_RETURN_IF_ERROR(PostprocessDataEdgesBetweenOutsideCompilations(
g, outside_compilation_attr_name));
TF_RETURN_IF_ERROR(PostprocessControlEdgesBetweenOutsideCompilations(
g, outside_compilation_attr_name));
return absl::OkStatus();
}
} // namespace tensorflow
tensorflow/compiler/jit/encapsulate_util.h
+155
View File
@@ -0,0 +1,155 @@
/* 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.
==============================================================================*/
// This file contains some utility functions for encapsulating XLA computation
// in host graph and encapsulating outside compilation in XLA computation.
#ifndef TENSORFLOW_COMPILER_JIT_ENCAPSULATE_UTIL_H_
#define TENSORFLOW_COMPILER_JIT_ENCAPSULATE_UTIL_H_
#include "absl/container/flat_hash_map.h"
#include "tensorflow/core/graph/graph.h"
namespace tensorflow {
// Attribute marking output tensor shapes inferred by XLA. Attribute value is
// a list of PartialTensorShape objects.
extern const char kXlaInferredShapesAttrName[];
// Infers output shapes for all nodes in graph `g`. The output shapes will be
// stored in node attribute `kXlaInferredShapesAttrName`.
//
// We have to perform shape inference before encapsulation because after
// encapsulation, some nodes will be encapsulated into function call, and shape
// inference does not handle function call at the moment.
absl::Status PerformStaticShapeInferenceBeforeEncapsulation(Graph* g);
// Attribute indicating that some ops in this node's XLA computation has control
// dependency on this node. Attribute value will always be "true".
extern const char kXlaConnectedToXlaComputationAttrName[];
// Attribute indicating that this node has control dependency on some ops in
// this node's XLA computation. Attribute value will always be "true".
extern const char kXlaConnectedFromXlaComputationAttrName[];
// Attribute indicating that this is an Placeholder node added to act as a
// temporary input node for an outside compilation node. Attribute value will be
// string (original input node name).
extern const char kOutsideCompilationOriginalNodeAttrName[];
// Attribute indicating that this is an Placeholder node added to act as a
// temporary input node for an outside compilation node. Attribute value will be
// int (src_output for original edge).
extern const char kOutsideCompilationSrcOutputAttrName[];
// Attribute indicating that this node has control dependencies on some other
// nodes within the same XLA cluster. Attribute value will be a list of string
// (node names).
extern const char kXlaControlDependenciesWithinXlaClusterAttrName[];
// Attribute indicating that this node is an outside compilation node which is
// lifted out of If/While/function node. Attribute value will always be boolean
// value "true".
extern const char kXlaIsLiftedArgAttrName[];
// Attribute indicating that this node is a Placeholder node for an _Arg node
// lifted out of If/While/function node. Attribute value will be a string, which
// is the outside compilation cluster name sending the lifted arg node to host.
extern const char kXlaLiftedArgOutsideCompilationAttrName[];
// Attribute indicating that this is an IdentityN node receiving inputs for a
// outside compilation Placeholder node (the original outside compilation node
// is moved out of TPU computation, and we left a Placeholder node there).
// Attribute value will be a string, which is the outside compilation cluster
// name for the outside compilation Placeholder node.
extern const char kXlaOutsideCompilationInputsAttrName[];
// Attribute indicating that this is a Placeholder node for an _Arg node used in
// outside compilation. We should not move this node out of XLA computation.
// Attribute value will always be boolean value "true".
extern const char kXlaIsPlaceholderForArg[];
// Information for XLA computation.
struct XlaClusterInfo {
// Add an explicitly-defined default constructor for this class.
//
// The compiler may delete the default constructor here because
// host_compute_core is a const member whose type (std::map) doesn't
// necessarily have a user provided constructor -- while libc++ and
// libstdc++ 4.8 provide a user defined default constructor, libstdc++ at
// least >= 7.3 does not. See also c++11 [class.ctor] p5.
//
// TODO(klimek): In c++17 we'll be able to initialize host_compute_core
// without losing aggregate initialization, which allows us to get rid of
// the constructor definitions again.
XlaClusterInfo() {}
XlaClusterInfo(const std::string& cluster_name,
const NameAttrList& func_name_attrs, Node* node,
const std::map<std::string, int>& host_compute_core)
: cluster_name(cluster_name),
func_name_attrs(func_name_attrs),
node(node),
host_compute_core(host_compute_core) {}
// XLA cluster name. It might be different from `func_name`.
const std::string cluster_name;
// Name and attributes of XLA computation function.
const NameAttrList func_name_attrs;
// The XLA computation node in the graph.
Node* node;
// A mapping from outside compilation cluster name to its device assignment.
const std::map<std::string, int> host_compute_core;
};
// Finds dependencies between outside compilation clusters, including both data
// dependencies and control dependencies. cluster_deps maps the name name of an
// outside compilation cluster to a set of names of outside compilation clusters
// that it depends on.
absl::StatusOr<
std::unique_ptr<absl::flat_hash_map<std::string, std::vector<std::string>>>>
OutsideCompilationClusterDependencies(
const Graph* g, const std::string& outside_compilation_attr_name);
// Preprocesses edges within the same XLA cluster. It will perform the following
// operations in order:
//
// 0. Remove edges from source node to outside compilation nodes, and edges
// from outside compilation nodes to sink node.
// 1a. For edges between different outside compilation clusters, remove the edge
// and add attr "kXlaControlDependenciesWithinXlaClusterAttrName = src node
// name" to dst node.
// 1b. For control edges between outside compilation and its XLA computation,
// add attr "kXlaConnected{From, To}XlaComputationAttrName = true" to the
// outside compilation node.
// 2. For data edges between different outside compilations, remove the edge
// and create a Placeholder node as dst node's input.
absl::Status PreprocessEdgesBetweenOutsideCompilations(
Graph* g, const std::string& outside_compilation_attr_name);
// Postprocesses edges within the same XLA cluster. This function reverts what
// `PreprocessEdgesBetweenOutsideCompilations` did. It will perform the
// following operations in order:
//
// 1. Remove Placeholder nodes between different outside compilations (created
// in `PreprocessEdgesBetweenOutsideCompilations` step 2).
// 2a. Reconnect control edges between different outside compilations (marked by
// `PreprocessEdgesBetweenOutsideCompilations` step 1a).
// Notice that control edges marked by
// `PreprocessEdgesBetweenOutsideCompilations` step 1b are not handled here.
// They are handled in `RewriteOutsideCompilationSubgraphFn`.
absl::Status PostprocessEdgesBetweenOutsideCompilations(
Graph* g, const std::string& outside_compilation_attr_name);
} // namespace tensorflow
#endif // TENSORFLOW_COMPILER_JIT_ENCAPSULATE_UTIL_H_

Some files were not shown because too many files have changed in this diff Show More