1 /* Copyright 2020 The TensorFlow Authors. All Rights Reserved. 2 3 Licensed under the Apache License, Version 2.0 (the "License"); 4 you may not use this file except in compliance with the License. 5 You may obtain a copy of the License at 6 7 http://www.apache.org/licenses/LICENSE-2.0 8 9 Unless required by applicable law or agreed to in writing, software 10 distributed under the License is distributed on an "AS IS" BASIS, 11 WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 12 See the License for the specific language governing permissions and 13 limitations under the License. 14 ==============================================================================*/ 15 16 #ifndef TENSORFLOW_COMPILER_XLA_SERVICE_GPU_GPU_EXECUTABLE_RUN_OPTIONS_H_ 17 #define TENSORFLOW_COMPILER_XLA_SERVICE_GPU_GPU_EXECUTABLE_RUN_OPTIONS_H_ 18 19 #include <functional> 20 #include <string> 21 #include <vector> 22 23 #include "absl/types/optional.h" 24 #include "tensorflow/compiler/xla/service/global_device_id.h" 25 #include "tensorflow/compiler/xla/statusor.h" 26 #include "tensorflow/compiler/xla/types.h" 27 28 namespace xla { 29 namespace gpu { 30 31 // Key for naming up a particular NCCL clique. This is just a set of unique 32 // device IDs (i.e. GPU IDs). The device IDs must be global within a cluster. 33 class NcclCliqueKey { 34 public: 35 explicit NcclCliqueKey(std::vector<GlobalDeviceId> devices); 36 37 template <typename H> AbslHashValue(H h,const NcclCliqueKey & k)38 friend H AbslHashValue(H h, const NcclCliqueKey& k) { 39 return H::combine(std::move(h), k.devices_); 40 } 41 friend bool operator==(const NcclCliqueKey& a, const NcclCliqueKey& b) { 42 return a.devices_ == b.devices_; 43 } 44 devices()45 const std::vector<GlobalDeviceId>& devices() const { return devices_; } 46 47 std::string ToString() const; 48 49 private: 50 std::vector<GlobalDeviceId> devices_; 51 }; 52 53 using NcclUniqueIdCallback = 54 std::function<StatusOr<std::string>(const NcclCliqueKey&)>; 55 56 // GPU-specific executable options. 57 // We keep these separate from ExecutableRunOptions to avoid adding 58 // dependencies to ExecutableRunOptions. 59 class GpuExecutableRunOptions { 60 public: 61 // Sets a mapping from local device ordinals to global device IDs. 62 // Used only on NVidia GPUs for cross-host NCCL collectives. If set, the 63 // elements of `device_assignment` are interpreted as global device IDs, not 64 // local device ordinals. 65 GpuExecutableRunOptions& set_gpu_global_device_ids( 66 absl::optional<std::vector<GlobalDeviceId>> gpu_global_device_ids); 67 const absl::optional<std::vector<GlobalDeviceId>>& gpu_global_device_ids() 68 const; 69 70 // Callback that returns a ncclUniqueId encoded as a string for a group of 71 // communicating GPU devices. Used only on NVidia GPUs. 72 GpuExecutableRunOptions& set_nccl_unique_id_callback( 73 NcclUniqueIdCallback nccl_unique_id_callback); 74 const NcclUniqueIdCallback& nccl_unique_id_callback() const; 75 76 private: 77 absl::optional<std::vector<GlobalDeviceId>> gpu_global_device_ids_; 78 NcclUniqueIdCallback nccl_unique_id_callback_; 79 }; 80 81 } // namespace gpu 82 } // namespace xla 83 84 #endif // TENSORFLOW_COMPILER_XLA_SERVICE_GPU_GPU_EXECUTABLE_RUN_OPTIONS_H_ 85