wehub/tensorflow--tensorflow
1
0
Fork 0
Code Issues Pull Requests Actions 3 Deployments Models Agent Notes Packages Projects Releases Wiki Activity
Files
main
tensorflow--tensorflow/tensorflow/compiler/jit/flags.cc
T
wehub-resource-sync 8a852e4b4e
cffconvert / validate (push) Has been skipped
Details
License Check / license-check (push) Failing after 2s
Details
chore: import upstream snapshot with attribution
2026-07-13 12:14:16 +08:00

558 lines
26 KiB
C++
Raw Permalink Blame History

/* Copyright 2018 The TensorFlow Authors. All Rights Reserved.
Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at
http://www.apache.org/licenses/LICENSE-2.0
Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.
==============================================================================*/
#include "tensorflow/compiler/jit/flags.h"
#include <limits>
#include <mutex> // NOLINT
#include <optional>
#include <vector>
#include "absl/base/call_once.h"
#include "absl/strings/numbers.h"
#include "absl/strings/str_split.h"
#include "absl/strings/strip.h"
#include "tensorflow/compiler/mlir/tensorflow/utils/dump_graph.h"
#include "xla/parse_flags_from_env.h"
#include "tensorflow/core/platform/macros.h"
#include "tensorflow/core/tpu/kernels/sparse_core_xla_flags_defaults.h"
#include "tensorflow/core/util/command_line_flags.h"
namespace tensorflow {
namespace {
BuildXlaOpsPassFlags* build_ops_flags;
MarkForCompilationPassFlags* mark_for_compilation_flags;
XlaDeviceFlags* device_flags;
XlaSparseCoreFlags* sparse_core_flags;
XlaOpsCommonFlags* ops_flags;
XlaCallModuleFlags* call_module_flags;
MlirCommonFlags* mlir_flags;
JitRtFlags* jitrt_flags;
std::vector<Flag>* jitrt_flag_list;
std::vector<Flag>* flag_list;
absl::once_flag flags_init;
bool SetterForXlaAutoJitFlag(const std::string& value) {
int32_t opt_level;
// We need to use the mark_for_compilation_flags directly here instead of
// going via GetMarkForCompilationPassFlags() to avoid infinite recursion. The
// latter will try to setup and parse flags, which would bring us back to this
// setter.
if (absl::SimpleAtoi(value, &opt_level)) {
mark_for_compilation_flags->xla_auto_jit_flag
.optimization_level_single_gpu = opt_level;
mark_for_compilation_flags->xla_auto_jit_flag.optimization_level_general =
opt_level;
return true;
}
if (value == "fusible") {
mark_for_compilation_flags->xla_auto_jit_flag
.optimization_level_single_gpu = 1;
mark_for_compilation_flags->xla_auto_jit_flag.optimization_level_general =
1;
mark_for_compilation_flags->tf_xla_ops_to_cluster = "FUSIBLE";
return true;
}
absl::string_view value_sv(value);
if (!absl::ConsumePrefix(&value_sv, "single-gpu(") ||
!absl::ConsumeSuffix(&value_sv, ")") ||
!absl::SimpleAtoi(value_sv, &opt_level)) {
return false;
}
mark_for_compilation_flags->xla_auto_jit_flag.optimization_level_single_gpu =
opt_level;
return true;
}
bool SetterForXlaCallModuleDisabledChecks(const std::string& value) {
auto directives = absl::StrSplit(value, ',', absl::SkipEmpty());
call_module_flags->disabled_checks.insert(directives.begin(),
directives.end());
return true;
}
void AppendMarkForCompilationPassFlagsInternal(std::vector<Flag>* flag_list) {
std::vector<Flag> new_flags = {
Flag("tf_xla_auto_jit", SetterForXlaAutoJitFlag, "0",
"Control compilation of operators into XLA computations on CPU and "
"GPU devices. 0 = use ConfigProto setting; -1 = off; 1 = on for "
"things very likely to be improved; 2 = on for everything; "
"(experimental) fusible = only for Tensorflow operations that XLA "
"knows how to fuse. "
"If set to single-gpu(<N>) then this resolves to <N> for single-GPU "
"graphs (graphs that have at least one node placed on a GPU and no "
"more than one GPU is in use through the entire graph) and 0 "
"otherwise. Experimental."),
Flag("tf_xla_min_cluster_size",
&mark_for_compilation_flags->tf_xla_min_cluster_size,
"Minimum number of operators in an XLA compilation. Ignored for "
"operators placed on an XLA device or operators explicitly marked "
"for compilation."),
Flag("tf_xla_max_cluster_size",
&mark_for_compilation_flags->tf_xla_max_cluster_size,
"Maximum number of operators in an XLA compilation."),
Flag(
"tf_xla_ops_to_cluster",
&mark_for_compilation_flags->tf_xla_ops_to_cluster,
"(experimental) "
"Limit the operations clustered by XLA to these operations. "
"If multiple, separate them with commas. Shortcuts: "
" PW: All point-wise operations."
" RED: All reduction operations."
" MISC: Mixed operations."
" PWRED: TF operations that get converted to PW+RED operation in XLA."
" REDUCEWINDOW: TF operations like MaxPool/AvgPool that get "
"converted to ReduceWindow in XLA."
" REDUCEWINDOWPW: Operation that get converted to ReduceWindow + PW "
"(LRN, LRNGrad)."
" BN: TF FusedBatchNorm* operations."
" FUSIBLE: All TF operations that XLA can fuse (All the above). "
"You can also put any TF operation name, e.g. 'FUSIBLE,MatMul'."),
Flag("tf_xla_cluster_exclude_ops",
&mark_for_compilation_flags->tf_xla_cluster_exclude_ops,
"(experimental) "
"Exclude the operations from auto-clustering. "
"If multiple, separate them with commas."
" Where, Some_other_ops"),
Flag("tf_xla_clustering_debug",
&mark_for_compilation_flags->tf_xla_clustering_debug,
"Dump graphs during XLA compilation."),
Flag("tf_xla_cpu_global_jit",
&mark_for_compilation_flags->tf_xla_cpu_global_jit,
"Enables global JIT compilation for CPU via SessionOptions."),
Flag("tf_xla_clustering_fuel",
&mark_for_compilation_flags->tf_xla_clustering_fuel,
"Places an artificial limit on the number of ops marked as "
"eligible for clustering."),
Flag("tf_xla_disable_deadness_safety_checks_for_debugging",
&mark_for_compilation_flags
->tf_xla_disable_deadness_safety_checks_for_debugging,
"Disable deadness related safety checks when clustering (this is "
"unsound)."),
Flag("tf_xla_disable_resource_variable_safety_checks_for_debugging",
&mark_for_compilation_flags
->tf_xla_disable_resource_variable_safety_checks_for_debugging,
"Disable resource variables related safety checks when clustering "
"(this is unsound)."),
Flag("tf_xla_deterministic_cluster_names",
&mark_for_compilation_flags->tf_xla_deterministic_cluster_names,
"Causes the function names assigned by auto clustering to be "
"deterministic from run to run."),
Flag("tf_xla_persistent_cache_directory",
&mark_for_compilation_flags->tf_xla_persistent_cache_directory,
"If non-empty, JIT-compiled executables are saved to and loaded "
"from the specified file system directory path. Empty by default."),
Flag("tf_xla_persistent_cache_device_types",
&mark_for_compilation_flags->tf_xla_persistent_cache_device_types,
"If non-empty, the persistent cache will only be used for the "
"specified devices (comma separated). Each device type should be "
"able to be converted to `DeviceType`."),
Flag("tf_xla_persistent_cache_read_only",
&mark_for_compilation_flags->tf_xla_persistent_cache_read_only,
"If true, the persistent cache will be read-only."),
Flag("tf_xla_disable_strict_signature_checks",
&mark_for_compilation_flags->tf_xla_disable_strict_signature_checks,
"If true, entires loaded into the XLA compile cache will not have "
"their signatures checked strictly. Defaults to false."),
Flag("tf_xla_persistent_cache_prefix",
&mark_for_compilation_flags->tf_xla_persistent_cache_prefix,
"Specifies the persistance cache prefix. Default is "
"\"xla_compile_cache\""),
Flag("tf_xla_sparse_core_disable_table_stacking",
&sparse_core_flags->tf_xla_sparse_core_disable_table_stacking,
"Disable table stacking for all the tables passed to the SparseCore"
"mid level API."),
Flag("tf_xla_sparse_core_minibatch_max_division_level",
&sparse_core_flags->tf_xla_sparse_core_minibatch_max_division_level,
"Max level of division to split input data into minibatches."),
Flag("tf_xla_sparse_core_stacking_mem_limit_bytes",
&sparse_core_flags->tf_xla_sparse_core_stacking_mem_limit_bytes,
"If non-zero, limits the size of the activations for a given table"
"to be below these many bytes."),
Flag("tf_xla_sparse_core_stacking_table_shard_limit_bytes",
&sparse_core_flags
->tf_xla_sparse_core_stacking_table_shard_limit_bytes,
"If non-zero, limits the size of any table shard to be below these"
"many bytes.")};
flag_list->insert(flag_list->end(), new_flags.begin(), new_flags.end());
}
void AllocateAndParseJitRtFlags() {
jitrt_flags = new JitRtFlags;
jitrt_flags->always_specialize = false;
jitrt_flags->cost_driven_async_parallel_for = false;
jitrt_flags->enable_crash_reproducer = false;
jitrt_flags->log_query_of_death = false;
jitrt_flags->vectorize = false;
jitrt_flag_list = new std::vector<Flag>({
Flag("always_specialize", &jitrt_flags->always_specialize, ""),
Flag("cost_driven_async_parallel_for",
&jitrt_flags->cost_driven_async_parallel_for, ""),
Flag("enable_crash_reproducer", &jitrt_flags->enable_crash_reproducer,
""),
Flag("log_query_of_death", &jitrt_flags->log_query_of_death, ""),
Flag("vectorize", &jitrt_flags->vectorize, ""),
});
xla::ParseFlagsFromEnvAndDieIfUnknown("TF_JITRT_FLAGS", *jitrt_flag_list);
}
void AllocateAndParseFlags() {
build_ops_flags = new BuildXlaOpsPassFlags;
build_ops_flags->tf_xla_enable_lazy_compilation = true;
build_ops_flags->tf_xla_print_cluster_outputs = false;
build_ops_flags->tf_xla_check_cluster_input_numerics = false;
build_ops_flags->tf_xla_check_cluster_output_numerics = false;
build_ops_flags->tf_xla_disable_constant_folding = false;
build_ops_flags->tf_xla_disable_full_embedding_pipelining = false;
build_ops_flags->tf_xla_disable_full_embedding_pipelining_with_summaries =
true;
build_ops_flags->tf_xla_embedding_parallel_iterations = 0;
mark_for_compilation_flags = new MarkForCompilationPassFlags;
mark_for_compilation_flags->xla_auto_jit_flag.optimization_level_single_gpu =
0;
mark_for_compilation_flags->xla_auto_jit_flag.optimization_level_general = 0;
mark_for_compilation_flags->tf_xla_min_cluster_size = 4;
mark_for_compilation_flags->tf_xla_max_cluster_size =
std::numeric_limits<int32_t>::max();
mark_for_compilation_flags->tf_xla_clustering_debug = false;
mark_for_compilation_flags->tf_xla_cpu_global_jit = false;
mark_for_compilation_flags->tf_xla_clustering_fuel =
std::numeric_limits<int64_t>::max();
mark_for_compilation_flags
->tf_xla_disable_deadness_safety_checks_for_debugging = false;
mark_for_compilation_flags
->tf_xla_disable_resource_variable_safety_checks_for_debugging = false;
mark_for_compilation_flags->tf_xla_deterministic_cluster_names = false;
mark_for_compilation_flags->tf_xla_persistent_cache_directory = "";
mark_for_compilation_flags->tf_xla_persistent_cache_device_types = "";
mark_for_compilation_flags->tf_xla_persistent_cache_read_only = false;
mark_for_compilation_flags->tf_xla_disable_strict_signature_checks = false;
mark_for_compilation_flags->tf_xla_persistent_cache_prefix =
"xla_compile_cache";
device_flags = new XlaDeviceFlags;
device_flags->tf_xla_compile_on_demand = false;
device_flags->tf_xla_enable_xla_devices = false;
sparse_core_flags = new XlaSparseCoreFlags;
sparse_core_flags->tf_xla_sparse_core_minibatch_max_division_level =
kDefaultSparseCoreMinibatchMaxDivisionLevel;
sparse_core_flags->tf_xla_sparse_core_disable_table_stacking =
kDefaultDisableTableStacking;
sparse_core_flags->tf_xla_sparse_core_stacking_mem_limit_bytes =
kDefaultXlaSparseCoreStackingMemLimit;
sparse_core_flags->tf_xla_sparse_core_stacking_table_shard_limit_bytes =
kDefaultXlaSparseCoreStackingTableShardLimit;
ops_flags = new XlaOpsCommonFlags;
ops_flags->tf_xla_always_defer_compilation = false;
ops_flags->tf_xla_async_compilation = false;
ops_flags->tf_xla_use_device_api.enabled_for_xla_launch_ = true;
ops_flags->tf_xla_use_device_api.enabled_for_compile_on_demand_ = true;
ops_flags->tf_xla_use_device_api.enabled_for_compile_and_run_ = true;
ops_flags->tf_xla_use_device_api.enabled_for_all_ = false;
ops_flags->tf_xla_use_device_api.enabled_for_gpu_ = true;
call_module_flags = new XlaCallModuleFlags;
// The `enable_mlir_bridge` flag allows the user to explicitly request that
// their program is (or isn't) compiled using the MLIR-based TF-to-XLA bridge.
//
// The `enable_mlir_bridge_is_explicit` variable tracks whether or not the
// user has made an explicit request. That is, if this variable is set to
// true, the program honors the user's request as per `enable_mlir_bridge`; if
// it's set to false, the default behavior is used (which may run either
// bridge, on a per-graph basis).
bool enable_mlir_bridge = false;
bool enable_mlir_bridge_is_explicit = false;
bool enable_mlir_merge_control_flow_pass = true;
bool enable_mlir_convert_control_to_data_outputs_pass = false;
bool enable_mlir_composite_tpuexecute_side_effects = false;
bool enable_mlir_strict_clusters = false;
bool enable_mlir_multiple_local_cpu_devices = false;
bool enable_mlir_debug_info_serialization = true;
// Dump graphs in TFG dialect.
bool use_tfg_graph_dumper = false;
bool enable_tpu_variable_runtime_reformatting_pass = true;
bool enable_serialize_mlir_to_compressed_bytecode = false;
flag_list = new std::vector<Flag>(
{Flag("tf_xla_enable_lazy_compilation",
&build_ops_flags->tf_xla_enable_lazy_compilation, ""),
Flag("tf_xla_print_cluster_outputs",
&build_ops_flags->tf_xla_print_cluster_outputs,
"If true then insert Print nodes to print out values produced by "
"XLA clusters."),
Flag("tf_xla_check_cluster_input_numerics",
&build_ops_flags->tf_xla_check_cluster_input_numerics,
"If true then insert CheckNumerics nodes to check all cluster "
"inputs."),
Flag("tf_xla_check_cluster_output_numerics",
&build_ops_flags->tf_xla_check_cluster_output_numerics,
"If true then insert CheckNumerics nodes to check all cluster "
"outputs."),
Flag("tf_xla_disable_constant_folding",
&build_ops_flags->tf_xla_disable_constant_folding,
"If true then disables constant folding on TF graph before XLA "
"compilation."),
Flag("tf_xla_disable_full_embedding_pipelining",
&build_ops_flags->tf_xla_disable_full_embedding_pipelining,
"If true then disables full embedding pipelining and instead use "
"strict SparseCore / TensorCore sequencing."),
Flag("tf_xla_disable_full_embedding_pipelining_with_summaries",
&build_ops_flags
->tf_xla_disable_full_embedding_pipelining_with_summaries,
"If true then disables full embedding pipelining when summary ops "
"are detected."),
Flag("tf_xla_embedding_parallel_iterations",
&build_ops_flags->tf_xla_embedding_parallel_iterations,
"If >0 then use this many parallel iterations in "
"embedding_pipelining and embedding_sequency. By default, use the "
"parallel_iterations on the original model WhileOp."),
Flag("tf_xla_compile_on_demand", &device_flags->tf_xla_compile_on_demand,
"Switch a device into 'on-demand' mode, where instead of "
"autoclustering ops are compiled one by one just-in-time."),
Flag("tf_xla_enable_xla_devices",
&device_flags->tf_xla_enable_xla_devices,
"Generate XLA_* devices, where placing a computation on such a "
"device"
"forces compilation by XLA. Deprecated."),
Flag("tf_xla_always_defer_compilation",
&ops_flags->tf_xla_always_defer_compilation, ""),
Flag("tf_xla_async_compilation", &ops_flags->tf_xla_async_compilation,
"When lazy compilation is enabled, asynchronous compilation starts "
"the cluster compilation in the background, and the fallback path "
"is executed until the compilation has finished."),
Flag("tf_xla_use_device_api_for_xla_launch",
&ops_flags->tf_xla_use_device_api.enabled_for_xla_launch_,
"If true, uses Device API (PjRt) for single device compilation and "
"execution of functions marked for JIT compilation i.e. "
"jit_compile=True. Defaults to false."),
Flag("tf_xla_use_device_api_for_compile_on_demand",
&ops_flags->tf_xla_use_device_api.enabled_for_compile_on_demand_,
"If true, uses Device API (PjRt) for compiling and executing ops "
"one by one in 'on-demand' mode. Defaults to false."),
Flag("tf_xla_use_device_api_for_auto_jit",
&ops_flags->tf_xla_use_device_api.enabled_for_compile_and_run_,
"If true, uses Device API (PjRt) for compilation and execution "
"when auto-clustering is enabled. Defaults to false."),
Flag("tf_xla_use_device_api",
&ops_flags->tf_xla_use_device_api.enabled_for_all_,
"If true, uses Device API (PjRt) for compilation and execution "
"of ops one-by-one in 'on-demand' mode, for functions marked for "
"JIT compilation, or when auto-clustering is enabled. Defaults to "
"false."),
Flag("tf_xla_enable_device_api_for_gpu",
&ops_flags->tf_xla_use_device_api.enabled_for_gpu_,
"If true, uses Device API (PjRt) for TF GPU device. This is a "
"helper flag so that individual tests can turn on PjRt for GPU "
"specifically."),
Flag("tf_xla_call_module_disabled_checks",
SetterForXlaCallModuleDisabledChecks, "",
"A comma-sepated list of directives specifying the safety checks "
"to be skipped when compiling XlaCallModuleOp. See the op "
"documentation for the recognized values."),
Flag("tf_mlir_enable_mlir_bridge", &enable_mlir_bridge,
"Enables experimental MLIR-Based TensorFlow Compiler Bridge.",
&enable_mlir_bridge_is_explicit),
Flag("tf_mlir_enable_merge_control_flow_pass",
&enable_mlir_merge_control_flow_pass,
"Enables MergeControlFlow pass for MLIR-Based TensorFlow Compiler "
"Bridge."),
Flag("tf_mlir_enable_convert_control_to_data_outputs_pass",
&enable_mlir_convert_control_to_data_outputs_pass,
"Enables `tf-executor-convert-control-to-data-outputs` pass for "
"MLIR-Based TensorFlow Compiler Bridge."),
Flag("tf_mlir_composite_tpuexecute_side_effects",
&enable_mlir_composite_tpuexecute_side_effects,
"Enables certain TPUExecute ops to run in parallel if they only "
"operate on resources that live on composite devices."),
Flag("tf_mlir_enable_strict_clusters", &enable_mlir_strict_clusters,
"Do not allow clusters that have cyclic control dependencies."),
Flag("tf_mlir_enable_multiple_local_cpu_devices",
&enable_mlir_multiple_local_cpu_devices,
"Enable multiple local CPU devices. CPU ops which are outside "
"compiled inside the tpu cluster will also be replicated across "
"multiple cpu devices."),
Flag("tf_mlir_enable_debug_info_serialization",
&enable_mlir_debug_info_serialization,
"Enable debug info serialization in MLIR."),
Flag("tf_dump_graphs_in_tfg", &use_tfg_graph_dumper,
"When tf_dump_graphs_in_tfg is true, graphs after transformations "
"are dumped in MLIR TFG dialect and not in GraphDef"),
Flag("tf_mlir_enable_tpu_variable_runtime_reformatting_pass",
&enable_tpu_variable_runtime_reformatting_pass,
"Enables TPUVariableRuntimeReformatting pass for MLIR-Based "
"TensorFlow Compiler Bridge. This enables weight update sharding "
"and creates TPUReshardVariables ops."),
Flag("tf_serialize_mlir_to_compressed_bytecode",
&enable_serialize_mlir_to_compressed_bytecode,
"If true, serialize MLIR to compressed bytecode.")});
AppendMarkForCompilationPassFlagsInternal(flag_list);
xla::ParseFlagsFromEnvAndDieIfUnknown("TF_XLA_FLAGS", *flag_list);
mlir_flags = new MlirCommonFlags;
if (!enable_mlir_bridge_is_explicit) {
mlir_flags->tf_mlir_enable_mlir_bridge =
ConfigProto::Experimental::MLIR_BRIDGE_ROLLOUT_UNSPECIFIED;
} else if (enable_mlir_bridge) {
mlir_flags->tf_mlir_enable_mlir_bridge =
ConfigProto::Experimental::MLIR_BRIDGE_ROLLOUT_ENABLED;
} else {
mlir_flags->tf_mlir_enable_mlir_bridge =
ConfigProto::Experimental::MLIR_BRIDGE_ROLLOUT_DISABLED;
}
mlir_flags->tf_mlir_enable_merge_control_flow_pass =
enable_mlir_merge_control_flow_pass;
mlir_flags->tf_mlir_enable_convert_control_to_data_outputs_pass =
enable_mlir_convert_control_to_data_outputs_pass;
mlir_flags->tf_mlir_enable_composite_tpuexecute_side_effects =
enable_mlir_composite_tpuexecute_side_effects;
mlir_flags->tf_mlir_enable_strict_clusters = enable_mlir_strict_clusters;
mlir_flags->tf_mlir_enable_tpu_variable_runtime_reformatting_pass =
enable_tpu_variable_runtime_reformatting_pass;
mlir_flags->tf_mlir_enable_multiple_local_cpu_devices =
enable_mlir_multiple_local_cpu_devices;
mlir_flags->tf_mlir_enable_debug_info_serialization =
enable_mlir_debug_info_serialization;
mlir_flags->tf_serialize_mlir_to_compressed_bytecode =
enable_serialize_mlir_to_compressed_bytecode;
if (use_tfg_graph_dumper) {
UseMlirForGraphDump(MlirDumpConfig{}.elide_large_attributes().emit_dialect(
MlirDumpConfig::Dialect::kTFG));
}
AllocateAndParseJitRtFlags();
}
void ResetFlags() {
delete build_ops_flags;
delete mark_for_compilation_flags;
delete device_flags;
delete ops_flags;
delete mlir_flags;
delete flag_list;
delete jitrt_flags;
delete jitrt_flag_list;
AllocateAndParseFlags();
}
} // namespace
bool SetXlaAutoJitFlagFromFlagString(const std::string& value) {
absl::call_once(flags_init, &AllocateAndParseFlags);
return SetterForXlaAutoJitFlag(value);
}
BuildXlaOpsPassFlags* GetBuildXlaOpsPassFlags() {
absl::call_once(flags_init, &AllocateAndParseFlags);
return build_ops_flags;
}
MarkForCompilationPassFlags* GetMarkForCompilationPassFlags() {
absl::call_once(flags_init, &AllocateAndParseFlags);
return mark_for_compilation_flags;
}
XlaSparseCoreFlags* GetXlaSparseCoreFlags() {
absl::call_once(flags_init, &AllocateAndParseFlags);
return sparse_core_flags;
}
XlaDeviceFlags* GetXlaDeviceFlags() {
absl::call_once(flags_init, &AllocateAndParseFlags);
return device_flags;
}
XlaOpsCommonFlags* GetXlaOpsCommonFlags() {
absl::call_once(flags_init, &AllocateAndParseFlags);
return ops_flags;
}
XlaCallModuleFlags* GetXlaCallModuleFlags() {
absl::call_once(flags_init, &AllocateAndParseFlags);
return call_module_flags;
}
MlirCommonFlags* GetMlirCommonFlags() {
absl::call_once(flags_init, &AllocateAndParseFlags);
return mlir_flags;
}
void ResetJitCompilerFlags() { ResetFlags(); }
const JitRtFlags& GetJitRtFlags() {
absl::call_once(flags_init, &AllocateAndParseFlags);
return *jitrt_flags;
}
ConfigProto::Experimental::MlirBridgeRollout GetMlirBridgeRolloutState(
std::optional<const ConfigProto> config_proto) {
// TF1 graphs that do not override Sessions's ConfigProto and TF2 graphs
// can enable/disable the graph via tf_mlir_enable_mlir_bridge.
auto tf_mlir_enable_mlir_bridge =
GetMlirCommonFlags()->tf_mlir_enable_mlir_bridge;
if (tf_mlir_enable_mlir_bridge !=
ConfigProto::Experimental::MLIR_BRIDGE_ROLLOUT_UNSPECIFIED) {
return tf_mlir_enable_mlir_bridge;
}
// If a ConfigProto was not passed in, we can assume the caller is
// checking if TF2 graph should have the bridge enabled / disabled. In that
// case, we have already checked tf_mlir_enable_mlir_bridge so it is safe to
// return UNSPECIFIED here.
if (!config_proto.has_value()) {
return ConfigProto::Experimental::MLIR_BRIDGE_ROLLOUT_UNSPECIFIED;
}
// TF1 graphs that do override Session's ConfigProto and set
// ConfigProto's enable_mlir_bridge or mlir_bridge_rollout fields will not
// update tf_mlir_enable_mlir_bridge so check their values.
// ConfigProto's enable_mlir_bridge defaults to false so only respect it
// when it is true.
if (config_proto.value().experimental().enable_mlir_bridge()) {
return ConfigProto::Experimental::MLIR_BRIDGE_ROLLOUT_ENABLED;
}
return config_proto.value().experimental().mlir_bridge_rollout();
}
void AppendMarkForCompilationPassFlags(std::vector<Flag>* flag_list) {
absl::call_once(flags_init, &AllocateAndParseFlags);
AppendMarkForCompilationPassFlagsInternal(flag_list);
}
static std::atomic<bool> xla_compilation_disabled(false);
void DisableXlaCompilation() { xla_compilation_disabled = true; }
void EnableXlaCompilation() { xla_compilation_disabled = false; }
bool FailOnXlaCompilation() { return xla_compilation_disabled; }
} // namespace tensorflow
Reference in New Issue View Git Blame Copy Permalink