1 /* Copyright 2017 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 // This header declares functions which may be called by the generated code on 17 // the CPU. Calls to these functions must be resolved explicitly in the JIT in 18 // xla::cpu::SimpleResolver. It also defines a per-CpuExecutable context 19 // which is used to cache expensive state and resources utilized by the 20 // aforementioned functions. 21 // 22 // Other functions are declared in individual libraries as well, such as 23 // runtime_conv2d and runtime_matmul. As individual libraries, callers for 24 // ahead-of-time compilation can link only the required subset. 25 26 #ifndef TENSORFLOW_COMPILER_XLA_SERVICE_CPU_CPU_RUNTIME_H_ 27 #define TENSORFLOW_COMPILER_XLA_SERVICE_CPU_CPU_RUNTIME_H_ 28 29 #include "tensorflow/compiler/xla/executable_run_options.h" 30 #include "tensorflow/compiler/xla/service/cpu/xfeed_manager.h" 31 #include "tensorflow/compiler/xla/service/hlo_instructions.h" 32 #include "tensorflow/compiler/xla/types.h" 33 34 namespace xla { 35 namespace cpu { 36 namespace runtime { 37 38 // Names of runtime functions. These get resolved from the generated code to the 39 // right symbol at link time in one of two ways: 40 // 1. When using the JIT, the symbol resolver (SimpleResolver in 41 // third_party/tensorflow/compiler/xla/service/cpu/simple_orc_jit.cc) maps 42 // this symbol name to 43 // the actual symbol. 44 // 2. When using ahead-of-time compilation, the linker can resolve the name 45 // because it is a symbol in the cpu_runtime library. 46 extern const char* const kEigenMatMulF16SymbolName; 47 extern const char* const kEigenMatMulF32SymbolName; 48 extern const char* const kEigenMatMulF64SymbolName; 49 extern const char* const kEigenMatMulC64SymbolName; 50 extern const char* const kEigenMatMulC128SymbolName; 51 extern const char* const kEigenMatMulS32SymbolName; 52 extern const char* const kMKLConvF32SymbolName; 53 extern const char* const kMKLMatMulF32SymbolName; 54 extern const char* const kMKLMatMulF64SymbolName; 55 extern const char* const kMKLSingleThreadedMatMulF32SymbolName; 56 extern const char* const kMKLSingleThreadedMatMulF64SymbolName; 57 extern const char* const kEigenConvF16SymbolName; 58 extern const char* const kEigenConvF32SymbolName; 59 extern const char* const kEigenFftSymbolName; 60 extern const char* const kEigenSingleThreadedFftSymbolName; 61 extern const char* const kEigenSingleThreadedMatMulF16SymbolName; 62 extern const char* const kEigenSingleThreadedMatMulF32SymbolName; 63 extern const char* const kEigenSingleThreadedMatMulF64SymbolName; 64 extern const char* const kEigenSingleThreadedMatMulC64SymbolName; 65 extern const char* const kEigenSingleThreadedMatMulC128SymbolName; 66 extern const char* const kEigenSingleThreadedMatMulS32SymbolName; 67 extern const char* const kEigenSingleThreadedConvF16SymbolName; 68 extern const char* const kEigenSingleThreadedConvF32SymbolName; 69 extern const char* const kAcquireInfeedBufferForDequeueSymbolName; 70 extern const char* const kReleaseInfeedBufferAfterDequeueSymbolName; 71 extern const char* const kAcquireOutfeedBufferForPopulationSymbolName; 72 extern const char* const kReleaseOutfeedBufferAfterPopulationSymbolName; 73 extern const char* const kParallelForkJoinSymbolName; 74 extern const char* const kPrintfToStderrSymbolName; 75 extern const char* const kKeyValueSortSymbolName; 76 extern const char* const kTopKF32SymbolName; 77 extern const char* const kAllReduceSymbolName; 78 extern const char* const kCollectivePermuteSymbolName; 79 extern const char* const kReplicaIdSymbolName; 80 extern const char* const kTracingStartSymbolName; 81 extern const char* const kTracingEndSymbolName; 82 extern const char* const kAllToAllSymbolName; 83 84 // All symbol names for XLA CPU runtime functions need to start with this 85 // prefix. 86 extern const char* const kXlaCpuRuntimeSymbolNamePrefix; 87 88 // Returns the infeed manager used by the CPU runtime for the CPU device 89 // `device_ordinal`. Note the device ordinal does not name a CPU 90 XfeedManager* GetXfeedManager(int device_ordinal); 91 92 } // namespace runtime 93 } // namespace cpu 94 } // namespace xla 95 96 extern "C" { 97 98 extern int __xla_cpu_runtime_PrintfToStderr(const char* format, ...); 99 100 extern xla::int64 __xla_cpu_runtime_TracingStart( 101 const void* /* xla::ExecutableRunOptions* */ run_options_ptr, 102 const char* name); 103 extern void __xla_cpu_runtime_TracingEnd( 104 const void* /* xla::ExecutableRunOptions* */ run_options_ptr, 105 xla::int64 id); 106 107 // Some things common to all of the runtime entry points below: 108 // 109 // * The shape pointer and shape_length reflect values that can be deserialized 110 // via llvm_ir::DecodeSelfDescribingShapeConstant. This is the way we pass 111 // reified type information from the generated program to the runtime, which 112 // helps check the type safety and contract for the emitted-code/runtime 113 // communication. 114 // 115 // * run_options is used to look up the device ordinal for the stream executor 116 // we're executing under. If it is null the device ordinal is assumed to be 117 // 0 (this behavior helps in writing tests). 118 119 // Note: in the runtime entry points below, the shape pointer and shape_length 120 // reflect values that can be deserialized via 121 // llvm_ir::DecodeSelfDescribingShapeConstant. This is the way we pass reified 122 // type information from the generated program to the runtime, which helps check 123 // the type safety and contract for the emitted-code/runtime communication. 124 125 // Blocks until the next infeed buffer is ready to be dequeued, then 126 // returns it. Fails catastrophically if the next enqueued buffer is 127 // not of the correct length in bytes. Checking the shape rather than 128 // the length would be more exact, but the length check is chosen as a 129 // tradeoff between error checking and speed/simplicity. 130 extern void* __xla_cpu_runtime_AcquireInfeedBufferForDequeue( 131 const xla::ExecutableRunOptions* run_options, xla::int32 buffer_length, 132 const void* shape, xla::int32 shape_length); 133 134 // Relinquishes the next infeed buffer that was returned by 135 // __xla_cpu_runtime_AcquireInfeedBufferForDequeue. Once this call 136 // completes the data at buffer_ptr may no longer be 137 // accessed. buffer_length must match the length passed to the call to 138 // __xla_cpu_runtime_AcquireInfeedBufferForDequeue that returned 139 // buffer_ptr. This function must be called before the next buffer is 140 // acquired, i.e., there may only be one outstanding infeed buffer in 141 // use by the runtime. TODO(b/31340454) investigate whether or not it 142 // is worth supporting zero-copy infeed where the buffer is retained 143 // by the compiled code until it has been used. If zero-copy infeed is 144 // implemented we will add support for multiple outstanding buffers 145 // that can be returned out of order. 146 extern void __xla_cpu_runtime_ReleaseInfeedBufferAfterDequeue( 147 const xla::ExecutableRunOptions* run_options, xla::int32 buffer_length, 148 void* buffer_ptr, const void* shape_ptr, xla::int32 shape_length); 149 150 // Blocks until the next outfeed buffer is available to be populated, then 151 // returns it. 152 extern void* __xla_cpu_runtime_AcquireOutfeedBufferForPopulation( 153 const xla::ExecutableRunOptions* run_options, xla::int32 buffer_length, 154 const void* shape_ptr, xla::int32 shape_length); 155 156 // Relinquishes the outfeed buffer after it has been populated. 157 // buffer_ptr must have been previously returned by 158 // __xla_cpu_runtime_AcquireOutfeedBufferForPopulation. 159 // Once this call completes, buffer_ptr may no longer be accessed. 160 // buffer_length must match the length passed to the call to 161 // __xla_cpu_runtime_AcquireInfeedBufferForDequeue that returned 162 // buffer_ptr. This function must be called before the next buffer is 163 // acquired, i.e., there may only be one outstanding outfeed buffer in 164 // use by the runtime. 165 extern void __xla_cpu_runtime_ReleaseOutfeedBufferAfterPopulation( 166 const xla::ExecutableRunOptions* run_options, xla::int32 buffer_length, 167 void* buffer_ptr, const void* shape_ptr, xla::int32 shape_length); 168 169 // Perform all reduce on a CPU. 170 // 171 // participating_replicas: array of replica IDs participating in the reduction, 172 // cf. GetParticipatingIDs. 173 // channel_id_present, op_id: whether op_id is a channel ID or a module ID. 174 // reduction_kind: operator used for a reduction, cf. ReductionKind. 175 // shape_ptr: shape of all input/output buffers. 176 extern void __xla_cpu_runtime_AllReduce( 177 const xla::ExecutableRunOptions* run_options, 178 const void* replica_groups_str, xla::int32 replica_groups_str_size, 179 xla::int32 channel_id_present, xla::int64 op_id, xla::int32 reduction_kind, 180 const void* shape_ptr, xla::int32 shape_length, xla::int32 num_buffers, 181 void** input_buffers, void** output_buffers); 182 183 extern void __xla_cpu_runtime_CollectivePermute( 184 const xla::ExecutableRunOptions* run_options, xla::int32 channel_id_present, 185 xla::int64 op_id, xla::int32 byte_size, void* input_buffer, 186 void* output_buffer, const void* source_target_pairs, 187 xla::int32 source_target_pairs_size); 188 189 extern void __xla_cpu_runtime_AllToAll( 190 const xla::ExecutableRunOptions* run_options, xla::int32 channel_id_present, 191 xla::int64 op_id, const void* replica_groups_str, 192 xla::int32 replica_groups_str_size, xla::int32 num_buffers, 193 xla::int64 buffer_size, void** source_buffers, void** destination_buffers); 194 195 // Write the replica ID into the output buffer. 196 extern void __xla_cpu_runtime_ReplicaId( 197 const xla::ExecutableRunOptions* run_options, void* output_buffer); 198 199 } // extern "C" 200 201 #endif // TENSORFLOW_COMPILER_XLA_SERVICE_CPU_CPU_RUNTIME_H_ 202