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 kKeyValueSortSymbolName; 75 extern const char* const kTopKF32SymbolName; 76 extern const char* const kAllReduceSymbolName; 77 extern const char* const kCollectivePermuteSymbolName; 78 extern const char* const kReplicaIdSymbolName; 79 extern const char* const kTracingStartSymbolName; 80 extern const char* const kTracingEndSymbolName; 81 extern const char* const kAllToAllSymbolName; 82 83 // All symbol names for XLA CPU runtime functions need to start with this 84 // prefix. 85 extern const char* const kXlaCpuRuntimeSymbolNamePrefix; 86 87 // Returns the infeed manager used by the CPU runtime for the CPU device 88 // `device_ordinal`. Note the device ordinal does not name a CPU 89 XfeedManager* GetXfeedManager(int device_ordinal); 90 91 } // namespace runtime 92 } // namespace cpu 93 } // namespace xla 94 95 extern "C" { 96 97 extern xla::int64 __xla_cpu_runtime_TracingStart( 98 const void* /* xla::ExecutableRunOptions* */ run_options_ptr, 99 const char* name); 100 extern void __xla_cpu_runtime_TracingEnd( 101 const void* /* xla::ExecutableRunOptions* */ run_options_ptr, 102 xla::int64 id); 103 104 // Some things common to all of the runtime entry points below: 105 // 106 // * The shape pointer and shape_length reflect values that can be deserialized 107 // via llvm_ir::DecodeSelfDescribingShapeConstant. This is the way we pass 108 // reified type information from the generated program to the runtime, which 109 // helps check the type safety and contract for the emitted-code/runtime 110 // communication. 111 // 112 // * run_options is used to look up the device ordinal for the stream executor 113 // we're executing under. If it is null the device ordinal is assumed to be 114 // 0 (this behavior helps in writing tests). 115 116 // Note: in the runtime entry points below, the shape pointer and shape_length 117 // reflect values that can be deserialized via 118 // llvm_ir::DecodeSelfDescribingShapeConstant. This is the way we pass reified 119 // type information from the generated program to the runtime, which helps check 120 // the type safety and contract for the emitted-code/runtime communication. 121 122 // Blocks until the next infeed buffer is ready to be dequeued, then 123 // returns it. Fails catastrophically if the next enqueued buffer is 124 // not of the correct length in bytes. Checking the shape rather than 125 // the length would be more exact, but the length check is chosen as a 126 // tradeoff between error checking and speed/simplicity. 127 extern void* __xla_cpu_runtime_AcquireInfeedBufferForDequeue( 128 const xla::ExecutableRunOptions* run_options, xla::int32 buffer_length, 129 const void* shape, xla::int32 shape_length); 130 131 // Relinquishes the next infeed buffer that was returned by 132 // __xla_cpu_runtime_AcquireInfeedBufferForDequeue. Once this call 133 // completes the data at buffer_ptr may no longer be 134 // accessed. buffer_length must match the length passed to the call to 135 // __xla_cpu_runtime_AcquireInfeedBufferForDequeue that returned 136 // buffer_ptr. This function must be called before the next buffer is 137 // acquired, i.e., there may only be one outstanding infeed buffer in 138 // use by the runtime. TODO(b/31340454) investigate whether or not it 139 // is worth supporting zero-copy infeed where the buffer is retained 140 // by the compiled code until it has been used. If zero-copy infeed is 141 // implemented we will add support for multiple outstanding buffers 142 // that can be returned out of order. 143 extern void __xla_cpu_runtime_ReleaseInfeedBufferAfterDequeue( 144 const xla::ExecutableRunOptions* run_options, xla::int32 buffer_length, 145 void* buffer_ptr, const void* shape_ptr, xla::int32 shape_length); 146 147 // Blocks until the next outfeed buffer is available to be populated, then 148 // returns it. 149 extern void* __xla_cpu_runtime_AcquireOutfeedBufferForPopulation( 150 const xla::ExecutableRunOptions* run_options, xla::int32 buffer_length, 151 const void* shape_ptr, xla::int32 shape_length); 152 153 // Relinquishes the outfeed buffer after it has been populated. 154 // buffer_ptr must have been previously returned by 155 // __xla_cpu_runtime_AcquireOutfeedBufferForPopulation. 156 // Once this call completes, buffer_ptr may no longer be accessed. 157 // buffer_length must match the length passed to the call to 158 // __xla_cpu_runtime_AcquireInfeedBufferForDequeue that returned 159 // buffer_ptr. This function must be called before the next buffer is 160 // acquired, i.e., there may only be one outstanding outfeed buffer in 161 // use by the runtime. 162 extern void __xla_cpu_runtime_ReleaseOutfeedBufferAfterPopulation( 163 const xla::ExecutableRunOptions* run_options, xla::int32 buffer_length, 164 void* buffer_ptr, const void* shape_ptr, xla::int32 shape_length); 165 166 // Perform all reduce on a CPU. 167 // 168 // participating_replicas: array of replica IDs participating in the reduction, 169 // cf. GetParticipatingReplicas. 170 // channel_id_present, op_id: whether op_id is a channel ID or a module ID. 171 // reduction_kind: operator used for a reduction, cf. ReductionKind. 172 // shape_ptr: shape of all input/output buffers. 173 extern void __xla_cpu_runtime_AllReduce( 174 const xla::ExecutableRunOptions* run_options, 175 const void* replica_groups_str, xla::int32 replica_groups_str_size, 176 xla::int32 channel_id_present, xla::int64 op_id, xla::int32 reduction_kind, 177 const void* shape_ptr, xla::int32 shape_length, xla::int32 num_buffers, 178 void** input_buffers, void** output_buffers); 179 180 extern void __xla_cpu_runtime_CollectivePermute( 181 const xla::ExecutableRunOptions* run_options, xla::int32 channel_id_present, 182 xla::int64 op_id, xla::int32 byte_size, void* input_buffer, 183 void* output_buffer, const void* source_target_pairs, 184 xla::int32 source_target_pairs_size); 185 186 extern void __xla_cpu_runtime_AllToAll( 187 const xla::ExecutableRunOptions* run_options, xla::int32 channel_id_present, 188 xla::int64 op_id, const void* replica_groups_str, 189 xla::int32 replica_groups_str_size, xla::int32 num_buffers, 190 xla::int64 buffer_size, void** source_buffers, void** destination_buffers); 191 192 // Write the replica ID into the output buffer. 193 extern void __xla_cpu_runtime_ReplicaId( 194 const xla::ExecutableRunOptions* run_options, void* output_buffer); 195 196 } // extern "C" 197 198 #endif // TENSORFLOW_COMPILER_XLA_SERVICE_CPU_CPU_RUNTIME_H_ 199