wehub-resource-sync
382 lines
14 KiB
C++
382 lines
14 KiB
C++
/* 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
|