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_OP_EXECUTABLE_H_ 17 #define TENSORFLOW_STREAM_EXECUTOR_TPU_TPU_OP_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 tensorflow { 34 35 // An executable capable of being fed to a TPU device via TpuExecutor. 36 class TpuOpExecutable : public xla::TpuExecutableInterface { 37 public: 38 using HostCommandHandler = std::function<void(uint32, int64_t)>; 39 40 // Constructs an executable that holds a non-owning reference to an 41 // XLA_TpuProgram. 42 explicit TpuOpExecutable(const XLA_TpuProgram* core_program, 43 std::unique_ptr<xla::HloModule> hlo_module, 44 HostCommandHandler host_command_handler = nullptr); 45 ~TpuOpExecutable() 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 xla::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 xla::Shape HostShapeToDeviceShape(const xla::Shape& host_shape) override; 60 61 int64 ShapeSize(const xla::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(TpuOpExecutable); 68 }; 69 70 } // namespace tensorflow 71 72 #endif // TENSORFLOW_STREAM_EXECUTOR_TPU_TPU_OP_EXECUTABLE_H_ 73