XlaCompiler::Options options;
options.client = client;
- options.device_type = cache->device_type();
+ options.device_type = &cache->device_type();
options.flib_def = ctx->function_library()->GetFunctionLibraryDefinition();
options.graph_def_version = ctx->function_library()->graph_def_version();
options.allow_cpu_custom_calls = (platform_id_ == se::host::kHostPlatformId);
core::ScopedUnref cache_ref(cache);
XlaCompiler::Options options;
- options.device_type = metadata.jit_device_type();
+ DeviceType device_type = metadata.jit_device_type();
+ options.device_type = &device_type;
options.client = metadata.client();
options.flib_def =
new FunctionLibraryDefinition(OpRegistry::Global(), FunctionDefLibrary{});
// Compile the graph into an XLA computation.
XlaCompiler::Options compiler_options;
compiler_options.client = client;
- compiler_options.device_type = DeviceType(DEVICE_CPU_XLA_JIT);
+ DeviceType device_type(DEVICE_CPU_XLA_JIT);
+ compiler_options.device_type = &device_type;
compiler_options.flib_def = &graph->flib_def();
compiler_options.graph_def_version = graph->versions().producer();
compiler_options.allow_cpu_custom_calls = true;
: options_(options),
initialization_status_(Status::OK()),
next_step_id_(1),
- device_(new XlaCompilationDevice(SessionOptions(), options_.device_type)),
+ device_(
+ new XlaCompilationDevice(SessionOptions(), *options_.device_type)),
device_mgr_({device_}) {
- CHECK(!options_.device_type.type_string().empty());
+ // We no longer need the device_type.
+ options_.device_type = nullptr;
+
if (options_.populate_resource_manager) {
initialization_status_ =
(*options_.populate_resource_manager)(device_->resource_manager());
return CompileGraph(options, name, std::move(graph), args, result);
}
-namespace {
-
-// Check that the ops of all non-functional nodes have been registered.
-string ValidateFunctionDef(const FunctionDef* fdef,
- const FunctionLibraryDefinition& flib_def) {
- std::vector<string> invalid_ops;
- for (const NodeDef& node : fdef->node_def()) {
- const string& op = node.op();
- if (op == FunctionLibraryDefinition::kGradientOp || flib_def.Find(op)) {
- continue;
- }
- const OpDef* op_def;
- if (!OpRegistry::Global()->LookUpOpDef(op, &op_def).ok()) {
- invalid_ops.push_back(op);
- }
- }
- return tensorflow::str_util::Join(invalid_ops, ", ");
-}
-
-// Check that the graph doesn't have any nodes incompatible with given
-// device_type.
-Status ValidateGraph(const Graph* graph,
- const FunctionLibraryDefinition& flib_def,
- const DeviceType& device_type, const string& name) {
- std::vector<string> invalid_ops;
- for (const Node* node : graph->nodes()) {
- if (node->type_string() == FunctionLibraryDefinition::kGradientOp) {
- continue;
- }
- const FunctionDef* fdef = flib_def.Find(node->def().op());
- if (fdef) {
- string error_msg = ValidateFunctionDef(fdef, flib_def);
- if (!error_msg.empty()) {
- invalid_ops.push_back(
- strings::StrCat(node->def().op(), ":{", error_msg, "}"));
- }
- continue;
- }
- if (!FindKernelDef(device_type, node->def(), nullptr, nullptr).ok()) {
- invalid_ops.push_back(node->def().op());
- }
- }
- if (!invalid_ops.empty()) {
- return errors::InvalidArgument(strings::StrCat(
- "Detected unsupported operations when trying to compile graph ", name,
- " on ", device_type.type_string(), ":",
- tensorflow::str_util::Join(invalid_ops, ", ")));
- }
- return Status::OK();
-}
-
-} // namespace
-
Status XlaCompiler::CompileGraph(const XlaCompiler::CompileOptions& options,
string const& name,
std::unique_ptr<Graph> graph,
FunctionalizeControlFlow(flib_runtime_->GetFunctionLibraryDefinition(),
graph.get(), local_flib_def_.get()));
- // Detect ops incompatible with the device_type.
- // FunctionalizeControlFlow may remove some unsupported ops.
- TF_RETURN_IF_ERROR(ValidateGraph(graph.get(), *options_.flib_def,
- options_.device_type, name));
-
xla::XlaBuilder builder(name);
XlaContext* context = new XlaContext(
this, &builder, options_.allow_cpu_custom_calls,
#include "tensorflow/compiler/tf2xla/host_compute_metadata.pb.h"
#include "tensorflow/compiler/tf2xla/xla_compilation_device.h"
-#include "tensorflow/compiler/tf2xla/xla_op_registry.h"
#include "tensorflow/compiler/xla/client/local_client.h"
#include "tensorflow/core/common_runtime/device.h"
#include "tensorflow/core/common_runtime/device_mgr.h"
typedef std::function<TensorShape(const TensorShape&, DataType)>
ShapeRepresentationFn;
struct Options {
- // Name of the compilation device to use. It must be set by the caller.
- // The default empty value is invalid.
- DeviceType device_type = DeviceType("");
+ // Name of the compilation device to use. Needs to be live only during
+ // XlaCompiler's constructor.
+ const DeviceType* device_type = nullptr;
xla::Client* client = nullptr;
class XlaCompilerTest : public ::testing::Test {
protected:
+ XlaCompilerTest() : cpu_device_type_(DEVICE_CPU_XLA_JIT) {}
+
void SetUp() override {
client_ = xla::ClientLibrary::LocalClientOrDie();
XlaCompiler::Options DefaultOptions() {
XlaCompiler::Options options;
- options.device_type = DeviceType(DEVICE_CPU_XLA_JIT);
+ options.device_type = &cpu_device_type_;
options.client = client_;
options.flib_def = flib_def_.get();
return options;
return compiler->local_flib_def_.get();
}
+ DeviceType cpu_device_type_;
xla::Client* client_;
std::unique_ptr<FunctionLibraryDefinition> flib_def_;
};
EXPECT_TRUE(xla::LiteralTestUtil::Equal(*expected_literal, *actual_literal));
}
-// Tests a graph which has a function with an invalid op.
-TEST_F(XlaCompilerTest, FunctionWithInvalidOp) {
- XlaCompiler compiler(DefaultOptions());
-
- FunctionDefLibrary flib;
- FunctionDef fn = FillFn();
- NodeDef* node = fn.add_node_def();
- node->set_name("Invalid");
- node->set_op("InvalidOp"); /* unsupported op */
- node = fn.add_node_def();
- node->set_name("Switch");
- node->set_op("Switch"); /* control flow node */
- *flib.add_function() = fn;
-
- TF_ASSERT_OK(flib_def_->AddFunctionDef(fn));
-
- std::unique_ptr<Graph> graph(new Graph(OpRegistry::Global()));
-
- Scope scope = Scope::NewRootScope().ExitOnError();
- auto value = ops::Const<int32>(scope.WithOpName("value"), 1, {});
- auto shape = ops::Const<int32>(scope.WithOpName("shape"), {5}, {1});
- TF_ASSERT_OK(scope.graph()->AddFunctionLibrary(flib));
-
- NodeDef def;
- TF_ASSERT_OK(NodeDefBuilder("fill_fn", "FillFn", flib_def_.get())
- .Input(value.name(), 0, DT_INT32)
- .Input(shape.name(), 1, DT_INT32)
- .Finalize(&def));
- Status status;
- Node* fill = scope.graph()->AddNode(def, &status);
- TF_ASSERT_OK(status);
- TF_ASSERT_OK(scope.DoShapeInference(fill));
- scope.graph()->AddEdge(value.node(), 0, fill, 0);
- scope.graph()->AddEdge(shape.node(), 0, fill, 1);
-
- auto retval = ops::_Retval(scope.WithOpName("retval"), Output(fill), 0);
-
- TF_ASSERT_OK(scope.ToGraph(graph.get()));
-
- std::vector<XlaCompiler::Argument> args;
- XlaCompiler::CompilationResult result;
- status = compiler.CompileGraph(XlaCompiler::CompileOptions(), "fill",
- std::move(graph), args, &result);
- ASSERT_FALSE(status.ok());
- EXPECT_TRUE(
- str_util::StrContains(status.error_message(), "FillFn:{InvalidOp}"))
- << status.error_message();
-}
-
} // namespace
} // namespace tensorflow
projects / platform / upstream / tensorflow.git / commitdiff