From e7e9a0c449803648d5103a39b868ee3abed31bdf Mon Sep 17 00:00:00 2001 From: Qiao Zhang Date: Mon, 14 Dec 2020 15:10:33 -0800 Subject: [PATCH] Move CreateModuleConfig to a new hlo_module_util header. Move CreateExecutionOptions to xla/client/executable_build_options.h. PiperOrigin-RevId: 347478735 Change-Id: Iab07bcd31e89ee8e24f2385d0a515dba03e505ad --- tensorflow/compiler/xla/client/BUILD | 1 + .../xla/client/executable_build_options.cc | 31 +++++ .../xla/client/executable_build_options.h | 6 + tensorflow/compiler/xla/service/BUILD | 17 +++ .../compiler/xla/service/hlo_module_util.cc | 131 ++++++++++++++++++ .../compiler/xla/service/hlo_module_util.h | 44 ++++++ .../compiler/xla/service/local_service.cc | 31 +---- tensorflow/compiler/xla/service/service.cc | 85 +----------- 8 files changed, 238 insertions(+), 108 deletions(-) create mode 100644 tensorflow/compiler/xla/service/hlo_module_util.cc create mode 100644 tensorflow/compiler/xla/service/hlo_module_util.h diff --git a/tensorflow/compiler/xla/client/BUILD b/tensorflow/compiler/xla/client/BUILD index 171afa42351..6baeca85149 100644 --- a/tensorflow/compiler/xla/client/BUILD +++ b/tensorflow/compiler/xla/client/BUILD @@ -95,6 +95,7 @@ cc_library( hdrs = ["executable_build_options.h"], deps = [ "//tensorflow/compiler/xla:debug_options_flags", + "//tensorflow/compiler/xla:execution_options_util", "//tensorflow/compiler/xla:shape_util", "//tensorflow/compiler/xla:util", "//tensorflow/compiler/xla:xla_data_proto_cc", diff --git a/tensorflow/compiler/xla/client/executable_build_options.cc b/tensorflow/compiler/xla/client/executable_build_options.cc index f39a3e79fe5..647232300e2 100644 --- a/tensorflow/compiler/xla/client/executable_build_options.cc +++ b/tensorflow/compiler/xla/client/executable_build_options.cc @@ -17,6 +17,7 @@ limitations under the License. #include "absl/strings/str_format.h" #include "tensorflow/compiler/xla/debug_options_flags.h" +#include "tensorflow/compiler/xla/execution_options_util.h" #include "tensorflow/compiler/xla/shape_util.h" namespace xla { @@ -99,4 +100,34 @@ string ExecutableBuildOptions::ToString() const { device_ordinal_, result_layout, num_replicas_); } +ExecutionOptions CreateExecutionOptions( + const ExecutableBuildOptions& build_options, + const ProgramShape* program_shape) { + ExecutionOptions execution_options = CreateDefaultExecutionOptions(); + if (build_options.has_debug_options()) { + *execution_options.mutable_debug_options() = build_options.debug_options(); + } + if (build_options.result_layout() != nullptr) { + *execution_options.mutable_shape_with_output_layout() = + build_options.result_layout()->ToProto(); + } else { + Shape result_shape(program_shape->result()); + LayoutUtil::SetToDefaultLayout(&result_shape); + *execution_options.mutable_shape_with_output_layout() = + result_shape.ToProto(); + } + execution_options.set_num_replicas(build_options.num_replicas()); + execution_options.set_num_partitions(build_options.num_partitions()); + execution_options.set_use_spmd_partitioning( + build_options.use_spmd_partitioning()); + execution_options.set_deduplicate_hlo(build_options.deduplicate_hlo()); + if (build_options.has_device_assignment()) { + TF_CHECK_OK(build_options.device_assignment().Serialize( + execution_options.mutable_device_assignment())); + } + execution_options.set_alias_passthrough_params( + build_options.alias_passthrough_params()); + return execution_options; +} + } // namespace xla diff --git a/tensorflow/compiler/xla/client/executable_build_options.h b/tensorflow/compiler/xla/client/executable_build_options.h index c55f5750da7..104a49c3ce4 100644 --- a/tensorflow/compiler/xla/client/executable_build_options.h +++ b/tensorflow/compiler/xla/client/executable_build_options.h @@ -141,6 +141,12 @@ class ExecutableBuildOptions { tensorflow::thread::ThreadPool* compile_thread_pool_ = nullptr; }; +// Creates an ExecutionOptions based on a given ExecutableBuildOptions and +// ProgramShape. +ExecutionOptions CreateExecutionOptions( + const ExecutableBuildOptions& build_options, + const ProgramShape* program_shape); + } // namespace xla #endif // TENSORFLOW_COMPILER_XLA_CLIENT_EXECUTABLE_BUILD_OPTIONS_H_ diff --git a/tensorflow/compiler/xla/service/BUILD b/tensorflow/compiler/xla/service/BUILD index 36982700d5e..0a0c31908af 100644 --- a/tensorflow/compiler/xla/service/BUILD +++ b/tensorflow/compiler/xla/service/BUILD @@ -936,6 +936,7 @@ cc_library( ":hlo_evaluator", ":hlo_execution_profile", ":hlo_module_config", + ":hlo_module_util", ":hlo_proto_util", ":platform_util", ":source_map_util", @@ -977,6 +978,7 @@ cc_library( ":hlo", ":hlo_execution_profile", ":hlo_module_config", + ":hlo_module_util", ":platform_util", ":service", ":shaped_buffer", @@ -1528,6 +1530,21 @@ cc_library( ], ) +cc_library( + name = "hlo_module_util", + srcs = ["hlo_module_util.cc"], + hdrs = ["hlo_module_util.h"], + deps = [ + ":compiler", + ":hlo_module_config", + "//tensorflow/compiler/xla:shape_util", + "//tensorflow/compiler/xla:status", + "//tensorflow/compiler/xla:statusor", + "@com_google_absl//absl/types:optional", + "@com_google_absl//absl/types:span", + ], +) + cc_library( name = "hlo_module_group_util", srcs = ["hlo_module_group_util.cc"], diff --git a/tensorflow/compiler/xla/service/hlo_module_util.cc b/tensorflow/compiler/xla/service/hlo_module_util.cc new file mode 100644 index 00000000000..106c50c6e8a --- /dev/null +++ b/tensorflow/compiler/xla/service/hlo_module_util.cc @@ -0,0 +1,131 @@ +/* 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. +==============================================================================*/ + +#include "tensorflow/compiler/xla/service/hlo_module_util.h" + +#include "absl/types/span.h" +#include "tensorflow/compiler/xla/service/compiler.h" +#include "tensorflow/compiler/xla/service/hlo_module_config.h" +#include "tensorflow/compiler/xla/shape.h" +#include "tensorflow/compiler/xla/shape_util.h" +#include "tensorflow/compiler/xla/status.h" +#include "tensorflow/compiler/xla/statusor.h" + +namespace xla { + +namespace { + +Status ValidateResultShape(const Shape& client_shape, + const Shape& result_shape) { + TF_RETURN_IF_ERROR(ShapeUtil::ValidateShapeWithOptionalLayout(client_shape)); + if (!ShapeUtil::Compatible(client_shape, result_shape)) { + return InvalidArgument( + "Shape used to set computation result layout %s is not compatible " + "with result shape %s", + ShapeUtil::HumanStringWithLayout(client_shape), + ShapeUtil::HumanString(result_shape)); + } + return Status::OK(); +} +} // namespace + +StatusOr> CreateModuleConfig( + const ProgramShape& program_shape, + absl::Span argument_shapes, + const ExecutionOptions* execution_options, int default_num_replicas, + absl::optional num_threads, const AotCompilationOptions* aot_options) { + auto config = absl::make_unique(program_shape); + ComputationLayout* computation_layout = + config->mutable_entry_computation_layout(); + const int64 argument_shapes_size = argument_shapes.size(); + if (program_shape.parameters_size() != argument_shapes_size) { + return InvalidArgument("computation takes %d parameters, but %u given", + program_shape.parameters_size(), + argument_shapes.size()); + } + for (int i = 0, end = argument_shapes.size(); i < end; ++i) { + // Verify that shape of arguments matches the shape of the arguments in the + // ProgramShape. + if (!ShapeUtil::Compatible(*argument_shapes[i], + program_shape.parameters(i))) { + return InvalidArgument( + "Argument does not match shape of computation parameter %d: want " + "%s, got %s", + i, ShapeUtil::HumanString(program_shape.parameters(i)), + ShapeUtil::HumanString(*argument_shapes[i])); + } + TF_RETURN_IF_ERROR( + computation_layout->mutable_parameter_layout(i)->CopyLayoutFromShape( + *argument_shapes[i])); + } + if (execution_options != nullptr && + execution_options->has_shape_with_output_layout()) { + const Shape shape_with_output_layout( + execution_options->shape_with_output_layout()); + TF_RETURN_IF_ERROR( + ValidateResultShape(shape_with_output_layout, program_shape.result())); + TF_RETURN_IF_ERROR( + computation_layout->mutable_result_layout()->CopyLayoutFromShape( + shape_with_output_layout)); + } else { + // If the result layout is not set, then choose the default. + computation_layout->mutable_result_layout()->SetToDefaultLayout(); + } + + if (execution_options != nullptr) { + if (execution_options->num_replicas() > 0) { + config->set_replica_count(execution_options->num_replicas()); + } else { + config->set_replica_count(default_num_replicas); + } + if (execution_options->num_partitions() > 0) { + config->set_num_partitions(execution_options->num_partitions()); + } + config->set_use_spmd_partitioning( + execution_options->use_spmd_partitioning()); + config->set_deduplicate_hlo(execution_options->deduplicate_hlo()); + config->set_seed(execution_options->seed()); + config->set_launch_id(execution_options->launch_id()); + config->set_debug_options(execution_options->debug_options()); + } else { + config->set_replica_count(default_num_replicas); + config->set_debug_options(GetDebugOptionsFromFlags()); + } + + if (num_threads.has_value()) { + config->set_intra_op_parallelism_threads(*num_threads); + } + + if (execution_options != nullptr && + execution_options->has_device_assignment()) { + TF_ASSIGN_OR_RETURN( + auto device_assignment, + DeviceAssignment::Deserialize(execution_options->device_assignment())); + config->set_static_device_assignment(*device_assignment); + } + config->set_alias_passthrough_params( + execution_options->alias_passthrough_params()); + + if (aot_options != nullptr && + aot_options->fusion_config_collection() != FusionConfigCollection::kOff) { + config->set_fusion_config_collection( + aot_options->fusion_config_collection()); + *config->mutable_fusion_config() = aot_options->fusion_config(); + } + + return std::move(config); +} + +} // namespace xla diff --git a/tensorflow/compiler/xla/service/hlo_module_util.h b/tensorflow/compiler/xla/service/hlo_module_util.h new file mode 100644 index 00000000000..93d11eae5e6 --- /dev/null +++ b/tensorflow/compiler/xla/service/hlo_module_util.h @@ -0,0 +1,44 @@ +/* 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_XLA_SERVICE_HLO_MODULE_UTIL_H_ +#define TENSORFLOW_COMPILER_XLA_SERVICE_HLO_MODULE_UTIL_H_ + +#include + +#include "absl/types/optional.h" +#include "absl/types/span.h" +#include "tensorflow/compiler/xla/service/compiler.h" +#include "tensorflow/compiler/xla/service/hlo_module_config.h" +#include "tensorflow/compiler/xla/shape.h" +#include "tensorflow/compiler/xla/status.h" +#include "tensorflow/compiler/xla/statusor.h" + +namespace xla { + +// Creates an HloModuleConfig for a given program shape and arguments. +// If execution_options does not set num_replicas, default_num_replicas is used. +// num_threads is optional; if not given, intra_op_parallelism_threads not set. +// aot_options is optional; if not given a default is used. +StatusOr> CreateModuleConfig( + const ProgramShape& program_shape, + absl::Span argument_shapes, + const ExecutionOptions* execution_options, int default_num_replicas, + absl::optional num_threads = absl::nullopt, + const AotCompilationOptions* aot_options = nullptr); + +} // namespace xla + +#endif // TENSORFLOW_COMPILER_XLA_SERVICE_HLO_MODULE_UTIL_H_ diff --git a/tensorflow/compiler/xla/service/local_service.cc b/tensorflow/compiler/xla/service/local_service.cc index ea8c45d3d46..e18767523bb 100644 --- a/tensorflow/compiler/xla/service/local_service.cc +++ b/tensorflow/compiler/xla/service/local_service.cc @@ -32,6 +32,7 @@ limitations under the License. #include "tensorflow/compiler/xla/service/hlo_execution_profile.h" #include "tensorflow/compiler/xla/service/hlo_module.h" #include "tensorflow/compiler/xla/service/hlo_module_config.h" +#include "tensorflow/compiler/xla/service/hlo_module_util.h" #include "tensorflow/compiler/xla/service/platform_util.h" #include "tensorflow/compiler/xla/shape_layout.h" #include "tensorflow/compiler/xla/shape_util.h" @@ -94,36 +95,6 @@ absl::optional ParameterMetadata( return absl::nullopt; } -ExecutionOptions CreateExecutionOptions( - const ExecutableBuildOptions& build_options, - const ProgramShape* program_shape) { - ExecutionOptions execution_options = CreateDefaultExecutionOptions(); - if (build_options.has_debug_options()) { - *execution_options.mutable_debug_options() = build_options.debug_options(); - } - if (build_options.result_layout() != nullptr) { - *execution_options.mutable_shape_with_output_layout() = - build_options.result_layout()->ToProto(); - } else { - Shape result_shape(program_shape->result()); - LayoutUtil::SetToDefaultLayout(&result_shape); - *execution_options.mutable_shape_with_output_layout() = - result_shape.ToProto(); - } - execution_options.set_num_replicas(build_options.num_replicas()); - execution_options.set_num_partitions(build_options.num_partitions()); - execution_options.set_use_spmd_partitioning( - build_options.use_spmd_partitioning()); - execution_options.set_deduplicate_hlo(build_options.deduplicate_hlo()); - if (build_options.has_device_assignment()) { - TF_CHECK_OK(build_options.device_assignment().Serialize( - execution_options.mutable_device_assignment())); - } - execution_options.set_alias_passthrough_params( - build_options.alias_passthrough_params()); - return execution_options; -} - } // namespace StatusOr>> diff --git a/tensorflow/compiler/xla/service/service.cc b/tensorflow/compiler/xla/service/service.cc index cf781b4fcdd..d915f5de4b3 100644 --- a/tensorflow/compiler/xla/service/service.cc +++ b/tensorflow/compiler/xla/service/service.cc @@ -38,6 +38,7 @@ limitations under the License. #include "tensorflow/compiler/xla/service/hlo_instruction.h" #include "tensorflow/compiler/xla/service/hlo_module.h" #include "tensorflow/compiler/xla/service/hlo_module_config.h" +#include "tensorflow/compiler/xla/service/hlo_module_util.h" #include "tensorflow/compiler/xla/service/hlo_proto_util.h" #include "tensorflow/compiler/xla/service/platform_util.h" #include "tensorflow/compiler/xla/service/source_map_util.h" @@ -256,88 +257,16 @@ StatusOr> Service::CreateModuleConfig( absl::Span argument_shapes, const ExecutionOptions* execution_options, const AotCompilationOptions* aot_options) { - auto config = absl::make_unique(program_shape); - ComputationLayout* computation_layout = - config->mutable_entry_computation_layout(); - const int64 argument_shapes_size = argument_shapes.size(); - if (program_shape.parameters_size() != argument_shapes_size) { - return InvalidArgument("computation takes %d parameters, but %u given", - program_shape.parameters_size(), - argument_shapes.size()); - } - for (int i = 0, end = argument_shapes.size(); i < end; ++i) { - // Verify that shape of arguments matches the shape of the arguments in the - // ProgramShape. - if (!ShapeUtil::Compatible(*argument_shapes[i], - program_shape.parameters(i))) { - return InvalidArgument( - "Argument does not match shape of computation parameter %d: want " - "%s, got %s", - i, ShapeUtil::HumanString(program_shape.parameters(i)), - ShapeUtil::HumanString(*argument_shapes[i])); - } - TF_RETURN_IF_ERROR( - computation_layout->mutable_parameter_layout(i)->CopyLayoutFromShape( - *argument_shapes[i])); - } - if (execution_options != nullptr && - execution_options->has_shape_with_output_layout()) { - const Shape shape_with_output_layout( - execution_options->shape_with_output_layout()); - TF_RETURN_IF_ERROR( - ValidateResultShape(shape_with_output_layout, program_shape.result())); - TF_RETURN_IF_ERROR( - computation_layout->mutable_result_layout()->CopyLayoutFromShape( - shape_with_output_layout)); - } else { - // If the result layout is not set, then choose the default. - computation_layout->mutable_result_layout()->SetToDefaultLayout(); - } - - if (execution_options != nullptr) { - if (execution_options->num_replicas() > 0) { - config->set_replica_count(execution_options->num_replicas()); - } else { - config->set_replica_count(options_.number_of_replicas()); - } - if (execution_options->num_partitions() > 0) { - config->set_num_partitions(execution_options->num_partitions()); - } - config->set_use_spmd_partitioning( - execution_options->use_spmd_partitioning()); - config->set_deduplicate_hlo(execution_options->deduplicate_hlo()); - config->set_seed(execution_options->seed()); - config->set_launch_id(execution_options->launch_id()); - config->set_debug_options(execution_options->debug_options()); - } else { - config->set_replica_count(options_.number_of_replicas()); - config->set_debug_options(GetDebugOptionsFromFlags()); - } - + int default_num_replicas = options_.number_of_replicas(); + absl::optional num_threads; if (execute_backend_ != nullptr && execute_backend_->eigen_intra_op_thread_pool() != nullptr) { - config->set_intra_op_parallelism_threads( - execute_backend_->eigen_intra_op_thread_pool()->NumThreads()); + num_threads = execute_backend_->eigen_intra_op_thread_pool()->NumThreads(); } - if (execution_options != nullptr && - execution_options->has_device_assignment()) { - TF_ASSIGN_OR_RETURN( - auto device_assignment, - DeviceAssignment::Deserialize(execution_options->device_assignment())); - config->set_static_device_assignment(*device_assignment); - } - config->set_alias_passthrough_params( - execution_options->alias_passthrough_params()); - - if (aot_options != nullptr && - aot_options->fusion_config_collection() != FusionConfigCollection::kOff) { - config->set_fusion_config_collection( - aot_options->fusion_config_collection()); - *config->mutable_fusion_config() = aot_options->fusion_config(); - } - - return std::move(config); + return xla::CreateModuleConfig(program_shape, argument_shapes, + execution_options, default_num_replicas, + num_threads, aot_options); } StatusOr> Service::CreateModuleConfig(