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_STREAM_EXECUTOR_TPU_TPU_EXECUTABLE_H_ 17 #define TENSORFLOW_STREAM_EXECUTOR_TPU_TPU_EXECUTABLE_H_ 18 19 #include <functional> 20 #include <memory> 21 22 #include "absl/types/optional.h" 23 #include "absl/types/span.h" 24 #include "tensorflow/compiler/xla/service/hlo_module.h" 25 #include "tensorflow/compiler/xla/service/service_executable_run_options.h" 26 #include "tensorflow/compiler/xla/shape.h" 27 #include "tensorflow/compiler/xla/status.h" 28 #include "tensorflow/compiler/xla/types.h" 29 #include "tensorflow/core/tpu/tpu_ops_c_api.h" 30 #include "tensorflow/stream_executor/device_memory.h" 31 #include "tensorflow/stream_executor/tpu/tpu_executable_interface.h" 32 33 namespace xla { 34 35 // An executable capable of being fed to a TPU device via TpuExecutor. 36 class TpuExecutable : public TpuExecutableInterface { 37 public: 38 using HostCommandHandler = std::function<void(uint32, int64)>; 39 40 // Constructs an executable that holds a non-owning reference to an 41 // XLA_TpuProgram. 42 explicit TpuExecutable(const XLA_TpuProgram* core_program, 43 std::unique_ptr<HloModule> hlo_module, 44 HostCommandHandler host_command_handler = nullptr); 45 ~TpuExecutable() override = default; 46 core_program()47 const XLA_TpuProgram* core_program() const { return core_program_; } 48 49 absl::string_view fingerprint() const override; 50 51 private: 52 Status LoadProgramAndEnqueueToStream( 53 const ServiceExecutableRunOptions& run_options, 54 absl::Span<const stream_executor::DeviceMemoryBase> arguments, 55 stream_executor::DeviceMemoryBase result, 56 absl::optional<stream_executor::DeviceMemoryBase> 57 cross_program_prefetch_addr) override; 58 59 Shape HostShapeToDeviceShape(const Shape& host_shape) override; 60 61 int64 ShapeSize(const Shape& shape) override; 62 63 const XLA_TpuProgram* const core_program_; 64 65 const HostCommandHandler host_command_handler_; 66 67 TF_DISALLOW_COPY_AND_ASSIGN(TpuExecutable); 68 }; 69 70 } // namespace xla 71 72 #endif // TENSORFLOW_STREAM_EXECUTOR_TPU_TPU_EXECUTABLE_H_ 73