std::unique_ptr<Backend> execute_backend)
: Service(options, std::move(execute_backend)) {}
+namespace {
+
+// Retrieves the parameter metadata for the given computation and parameter
+// number.
+//
+// If the parameter number is invalid for this computation, nullopt is
+// returned. When the return value has_value(), nullptr will never be
+// the held value.
+tensorflow::gtl::optional<const OpMetadata*> ParameterMetadata(
+ const XlaComputation& computation, int parameter_number) {
+ for (const HloComputationProto& comp : computation.proto().computations()) {
+ if (comp.id() == computation.proto().entry_computation_id()) {
+ for (const HloInstructionProto& instr : comp.instructions()) {
+ if (instr.opcode() == HloOpcodeString(HloOpcode::kParameter) &&
+ instr.parameter_number() == parameter_number) {
+ if (!instr.has_metadata()) {
+ return tensorflow::gtl::nullopt;
+ }
+ return &instr.metadata();
+ }
+ }
+ }
+ }
+ return tensorflow::gtl::nullopt;
+}
+
+ExecutionOptions CreateExecutionOptions(
+ const ExecutableBuildOptions& build_options,
+ const ProgramShape* program_shape) {
+ ExecutionOptions execution_options = CreateDefaultExecutionOptions();
+ if (build_options.hlo_profile().has_value()) {
+ execution_options.mutable_debug_options()->set_xla_hlo_profile(
+ *build_options.hlo_profile());
+ }
+ if (build_options.generate_hlo_graph().has_value()) {
+ execution_options.mutable_debug_options()->set_xla_generate_hlo_graph(
+ build_options.generate_hlo_graph().value());
+ }
+ if (build_options.dump_optimized_hlo_proto_to().has_value()) {
+ execution_options.mutable_debug_options()
+ ->set_xla_dump_optimized_hlo_proto_to(
+ build_options.dump_optimized_hlo_proto_to().value());
+ }
+ if (build_options.dump_per_pass_hlo_proto_to().has_value()) {
+ execution_options.mutable_debug_options()
+ ->set_xla_dump_per_pass_hlo_proto_to(
+ build_options.dump_per_pass_hlo_proto_to().value());
+ }
+ if (build_options.result_layout() != nullptr) {
+ *execution_options.mutable_shape_with_output_layout() =
+ *build_options.result_layout();
+ } else {
+ *execution_options.mutable_shape_with_output_layout() =
+ program_shape->result();
+ LayoutUtil::SetToDefaultLayout(
+ execution_options.mutable_shape_with_output_layout());
+ }
+ return execution_options;
+}
+
+} // namespace
+
StatusOr<std::unique_ptr<Executable>> LocalService::CompileExecutable(
const ComputationHandle& computation,
const tensorflow::gtl::ArraySlice<const Shape*> argument_layouts,
*build_options.result_layout(), program_shape->result()));
}
- ExecutionOptions execution_options = CreateDefaultExecutionOptions();
- if (build_options.hlo_profile().has_value()) {
- execution_options.mutable_debug_options()->set_xla_hlo_profile(
- *build_options.hlo_profile());
- }
- if (build_options.generate_hlo_graph().has_value()) {
- execution_options.mutable_debug_options()->set_xla_generate_hlo_graph(
- build_options.generate_hlo_graph().value());
- }
- if (build_options.dump_optimized_hlo_proto_to().has_value()) {
- execution_options.mutable_debug_options()
- ->set_xla_dump_optimized_hlo_proto_to(
- build_options.dump_optimized_hlo_proto_to().value());
- }
- if (build_options.dump_per_pass_hlo_proto_to().has_value()) {
- execution_options.mutable_debug_options()
- ->set_xla_dump_per_pass_hlo_proto_to(
- build_options.dump_per_pass_hlo_proto_to().value());
- }
- if (build_options.result_layout() != nullptr) {
- *execution_options.mutable_shape_with_output_layout() =
- *build_options.result_layout();
- } else {
- *execution_options.mutable_shape_with_output_layout() =
- program_shape->result();
- LayoutUtil::SetToDefaultLayout(
- execution_options.mutable_shape_with_output_layout());
- }
+ ExecutionOptions execution_options =
+ CreateExecutionOptions(build_options, program_shape.get());
TF_ASSIGN_OR_RETURN(std::unique_ptr<HloModuleConfig> module_config,
CreateModuleConfig(*program_shape, argument_layouts,
&execution_options, user_computation));
build_options.device_allocator());
}
+StatusOr<std::unique_ptr<Executable>> LocalService::CompileExecutable(
+ const XlaComputation& computation,
+ const tensorflow::gtl::ArraySlice<const Shape*> argument_layouts,
+ const ExecutableBuildOptions& build_options) {
+ const HloModuleProto& proto = computation.proto();
+ TF_RET_CHECK(proto.has_program_shape());
+ const ProgramShape& program_shape = proto.program_shape();
+
+ // Validate incoming layouts.
+ if (argument_layouts.size() != program_shape.parameters_size()) {
+ return InvalidArgument(
+ "Invalid number of arguments for computation: expected %d, got %zu.",
+ program_shape.parameters_size(), argument_layouts.size());
+ }
+
+ for (int i = 0; i < argument_layouts.size(); ++i) {
+ const Shape& argument_shape = *argument_layouts[i];
+ TF_RETURN_IF_ERROR(ShapeUtil::ValidateShape(argument_shape));
+ if (!ShapeUtil::Compatible(argument_shape, program_shape.parameters(i))) {
+ tensorflow::gtl::optional<const OpMetadata*> metadata =
+ ParameterMetadata(computation, /*parameter_number=*/i);
+ auto metadata_string = [&metadata]() -> string {
+ if (!metadata.has_value()) {
+ return "";
+ }
+ CHECK(metadata.value() != nullptr);
+ const OpMetadata& m = *metadata.value();
+ if (!m.source_file().empty()) {
+ return tensorflow::strings::Printf(
+ " (%s:%d)", m.source_file().c_str(), m.source_line());
+ }
+ return "";
+ };
+ return InvalidArgument(
+ "Invalid argument shape for argument %d%s, expected %s, got %s.", i,
+ metadata_string().c_str(),
+ ShapeUtil::HumanString(program_shape.parameters(i)).c_str(),
+ ShapeUtil::HumanString(argument_shape).c_str());
+ }
+ }
+ if (build_options.result_layout() != nullptr) {
+ TF_RETURN_IF_ERROR(ValidateResultShapeWithLayout(
+ *build_options.result_layout(), program_shape.result()));
+ }
+
+ ExecutionOptions execution_options =
+ CreateExecutionOptions(build_options, &program_shape);
+
+ TF_ASSIGN_OR_RETURN(
+ std::unique_ptr<HloModuleConfig> module_config,
+ CreateModuleConfig(program_shape, argument_layouts, &execution_options));
+
+ TF_ASSIGN_OR_RETURN(
+ se::StreamExecutor * executor,
+ execute_backend_->stream_executor(build_options.device_ordinal()));
+
+ return BuildExecutable(proto, std::move(module_config),
+ execute_backend_.get(), executor,
+ build_options.device_allocator());
+}
+
StatusOr<int> LocalService::ReplicaNumberToDeviceOrdinal(int replica_number) {
return backend().computation_placer()->DeviceId(
replica_number, /*computation=*/0, options_.number_of_replicas(),
return tensorflow::Status::OK();
}
+StatusOr<std::unique_ptr<Executable>> Service::BuildExecutable(
+ const HloModuleProto& module_proto,
+ std::unique_ptr<HloModuleConfig> module_config, Backend* backend,
+ se::StreamExecutor* executor, DeviceMemoryAllocator* device_allocator) {
+ VLOG(1) << Printf(
+ "BuildExecutable on service %p with serialized module proto: %s", this,
+ module_proto.name().c_str());
+
+ TF_ASSIGN_OR_RETURN(std::unique_ptr<HloModule> module,
+ HloModule::CreateFromProto(module_proto, *module_config));
+
+ TF_RETURN_IF_ERROR(MaybeDumpHloModule(*module));
+
+ TF_ASSIGN_OR_RETURN(
+ module, backend->compiler()->RunHloPasses(std::move(module), executor,
+ device_allocator));
+
+ TF_ASSIGN_OR_RETURN(std::unique_ptr<Executable> executable,
+ backend->compiler()->RunBackend(
+ std::move(module), executor, device_allocator));
+
+ return std::move(executable);
+}
+
tensorflow::Status Service::ExecuteGraph(const ExecuteGraphRequest* arg,
ExecuteResponse* result) {
VLOG(1) << "running execute-graph request";
std::vector<std::vector<const ShapedBuffer*>> replicated_arguments,
ResolveAndValidateArguments(arg->arguments(), replicas));
- TF_ASSIGN_OR_RETURN(const auto& config,
+ TF_ASSIGN_OR_RETURN(std::unique_ptr<HloModuleConfig> module_config,
CreateModuleConfig(arg->computation().program_shape(),
replicated_arguments.front(),
arg->execution_options()));
- TF_ASSIGN_OR_RETURN(std::unique_ptr<HloModule> module,
- HloModule::CreateFromProto(arg->computation(), *config));
- TF_RETURN_IF_ERROR(MaybeDumpHloModule(*module));
-
- TF_ASSIGN_OR_RETURN(module, execute_backend_->compiler()->RunHloPasses(
- std::move(module),
- execute_backend_->default_stream_executor(),
- /*device_allocator=*/nullptr));
TF_ASSIGN_OR_RETURN(
std::unique_ptr<Executable> executable,
- execute_backend_->compiler()->RunBackend(
- std::move(module), execute_backend_->default_stream_executor(),
- /*device_allocator=*/nullptr));
+ BuildExecutable(arg->computation(), std::move(module_config),
+ execute_backend_.get(),
+ execute_backend_->default_stream_executor(),
+ /*device_allocator=*/nullptr));
TF_ASSIGN_OR_RETURN(
*result->mutable_output(),
projects / platform / upstream / tensorflow.git / commitdiff