wehub/tensorflow--tensorflow
1
0
Fork 0
Code Issues Pull Requests Actions 6 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/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_