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_PJRT_GPU_DEVICE_H_ 17 #define TENSORFLOW_COMPILER_XLA_PJRT_GPU_DEVICE_H_ 18 19 #include <memory> 20 #include <optional> 21 #include <string> 22 #include <utility> 23 24 #include "absl/base/macros.h" 25 #include "tensorflow/compiler/xla/pjrt/distributed/client.h" 26 #include "tensorflow/compiler/xla/pjrt/pjrt_stream_executor_client.h" 27 #include "tensorflow/compiler/xla/statusor.h" 28 #include "tensorflow/core/common_runtime/bfc_allocator.h" 29 30 namespace xla { 31 32 class GpuDevice : public PjRtStreamExecutorDevice { 33 public: 34 GpuDevice(int id, std::unique_ptr<LocalDeviceState> local_device_state, 35 std::string device_kind, std::string device_vendor, int node_id); 36 37 absl::string_view device_vendor(); 38 39 absl::string_view ToString() const override; 40 41 private: 42 std::string device_vendor_; 43 std::string to_string_; 44 }; 45 46 struct GpuAllocatorConfig { 47 enum class Kind { 48 kDefault, // Client picks the best option for the platform. 49 kPlatform, // The platform's default. 50 kBFC, // Allocator using a "Best-Fit with Coalescing" algorithm. Currently 51 // only available for GPU. 52 kCudaAsync, // Use the CUDA async allocator. 53 }; 54 Kind kind = Kind::kDefault; 55 56 // Only used if kind == kBFC. The maximum fraction of available memory to 57 // allocate. 58 double memory_fraction = 0.9; 59 60 // Only used if kind == kBFC. If true, the allocator will immediately allocate 61 // the maximum amount allowed by `memory_fraction`. This reduces 62 // fragmentation, allowing more of the total memory to be used. If false, the 63 // allocator will allocate more memory as allocations are requested. 64 bool preallocate = true; 65 }; 66 67 // distributed_client may be nullptr in non-distributed settings. 68 // distributed_client should be in the connected state before calling this 69 // function. 70 StatusOr<std::unique_ptr<PjRtClient>> GetGpuClient( 71 bool asynchronous, const GpuAllocatorConfig& allocator_config, 72 std::shared_ptr<DistributedRuntimeClient> distributed_client, int node_id, 73 const std::optional<std::set<int>>& allowed_devices = std::nullopt, 74 std::optional<std::string> platform_name = std::nullopt); 75 76 } // namespace xla 77 78 #endif // TENSORFLOW_COMPILER_XLA_PJRT_GPU_DEVICE_H_ 79