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_C_API_CONVERSIONS_H_ 17 #define TENSORFLOW_STREAM_EXECUTOR_TPU_C_API_CONVERSIONS_H_ 18 19 #include "absl/container/inlined_vector.h" 20 #include "tensorflow/compiler/xla/executable_run_options.h" 21 #include "tensorflow/compiler/xla/literal.h" 22 #include "tensorflow/compiler/xla/service/hlo_module.h" 23 #include "tensorflow/compiler/xla/service/hlo_module_config.h" 24 #include "tensorflow/compiler/xla/service/maybe_owning_device_memory.h" 25 #include "tensorflow/compiler/xla/service/service_executable_run_options.h" 26 #include "tensorflow/compiler/xla/service/shaped_buffer.h" 27 #include "tensorflow/compiler/xla/shape.h" 28 #include "tensorflow/compiler/xla/shape_util.h" 29 #include "tensorflow/stream_executor/device_memory.h" 30 #include "tensorflow/stream_executor/device_memory_allocator.h" 31 #include "tensorflow/stream_executor/tpu/c_api_decl.h" 32 #include "tensorflow/stream_executor/tpu/tpu_executor_c_api.h" 33 34 // APIs for converting between internal and external versions of 35 // XLA/StreamExecutor data structures. 36 namespace ApiConverter { 37 38 // se::DeviceMemoryBase 39 SE_DeviceMemoryBase ToC(const stream_executor::DeviceMemoryBase& base); 40 void ToC(const stream_executor::DeviceMemoryBase& base, 41 SE_DeviceMemoryBase* se_base); 42 stream_executor::DeviceMemoryBase FromC(const SE_DeviceMemoryBase& se_base); 43 void Free(SE_DeviceMemoryBase*); 44 45 // xla::Shape 46 xla::Shape FromC(const XLA_Shape* c_shape); 47 void ToC(const xla::Shape& xla_shape, XLA_Shape* c_shape); 48 void Free(XLA_Shape* c_shape); 49 50 // xla::Layout 51 xla::Layout FromC(const XLA_Layout* c_layout); 52 void ToC(const xla::Layout& xla_layout, XLA_Layout* c_layout); 53 void Free(XLA_Layout* c_layout); 54 55 // xla::Tile 56 xla::Tile FromC(const XLA_Tile* c_tile); 57 void ToC(const xla::Tile& xla_tile, XLA_Tile* c_tile); 58 void Free(XLA_Tile* c_tile); 59 60 // xla::ShapeIndex 61 XLA_ShapeIndex ToC(const xla::ShapeIndex& xla_shape); 62 xla::ShapeIndex FromC(XLA_ShapeIndex* c_shape); 63 void Free(XLA_ShapeIndex*); 64 65 // Literal 66 void ToC(const xla::LiteralSlice& literal, XLA_Literal* c_literal); 67 xla::MutableBorrowingLiteral FromC(XLA_Literal* c_literal); 68 void Free(XLA_Literal* c_literal); 69 70 // ShapedBuffer 71 void ToC(const xla::ShapedBuffer& buffer, XLA_ShapedBuffer* c_device_buffer); 72 xla::ShapedBuffer FromC(XLA_ShapedBuffer* c_buffer); 73 void Free(XLA_ShapedBuffer* c_buffer); 74 75 // se::DeviceMemoryBase 76 SE_DeviceMemoryBase ToC(const stream_executor::DeviceMemoryBase& base); 77 stream_executor::DeviceMemoryBase FromC(const SE_DeviceMemoryBase& se_base); 78 void Free(SE_DeviceMemoryBase*); 79 80 // Literal 81 void ToC(const xla::LiteralSlice& literal, XLA_Literal* c_literal); 82 xla::MutableBorrowingLiteral FromC(XLA_Literal* c_literal); 83 void Free(XLA_Literal* c_literal); 84 85 // ShapedBuffer 86 void ToC(const xla::ShapedBuffer& buffer, XLA_ShapedBuffer* c_device_buffer); 87 xla::ShapedBuffer FromC(XLA_ShapedBuffer* c_buffer); 88 void Free(XLA_ShapedBuffer* c_buffer); 89 90 xla::MaybeOwningDeviceMemory FromC( 91 SE_MaybeOwningDeviceMemory* se_mem, 92 stream_executor::DeviceMemoryAllocator* allocator); 93 94 // DeviceMemoryAllocator 95 SE_DeviceMemoryAllocator ToC(stream_executor::DeviceMemoryAllocator* allocator); 96 97 // OwningDeviceMemory 98 SE_MaybeOwningDeviceMemory ToC(stream_executor::OwningDeviceMemory* mem); 99 // mem.HasOwnership() may be true if the buffer is aliased and shouldn't be 100 // released. 'aliased' should be true in this case. 'aliased' has no effect if 101 // 'mem' is unowned. 102 SE_MaybeOwningDeviceMemory ToC(xla::MaybeOwningDeviceMemory& mem, bool aliased); 103 104 // HloModule 105 XLA_HloModule ToC(const xla::HloModule& module); 106 xla::StatusOr<std::unique_ptr<xla::HloModule>> FromC( 107 const XLA_HloModule& c_module); 108 void Free(XLA_HloModule* c_module); 109 110 // HloModuleConfig 111 XLA_HloModuleConfig ToC(const xla::HloModuleConfig& config); 112 xla::HloModuleConfig FromC(const XLA_HloModuleConfig& c_config); 113 void Free(XLA_HloModuleConfig* c_config); 114 115 // Helper for managing stack based C -> C++ conversions. 116 template <class CType> 117 struct StackHelper { StackHelperStackHelper118 explicit StackHelper() {} 119 120 template <class CppType> StackHelperStackHelper121 explicit StackHelper(const CppType& t) { 122 ::ApiConverter::ToC(t, &value); 123 } ~StackHelperStackHelper124 ~StackHelper() { ::ApiConverter::Free(&value); } 125 126 template <class CppType> AsCppStackHelper127 CppType AsCpp() const { 128 return ::ApiConverter::FromC(&value); 129 } 130 131 mutable CType value; 132 }; 133 134 } // namespace ApiConverter 135 136 #endif 137