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 #include "tensorflow/compiler/xla/service/cpu/cpu_runtime.h" 17 18 #include <functional> 19 20 #include "tensorflow/compiler/xla/service/llvm_ir/llvm_util.h" 21 #include "tensorflow/core/platform/logging.h" 22 #include "tensorflow/core/platform/macros.h" 23 #include "tensorflow/core/platform/types.h" 24 25 namespace xla { 26 namespace cpu { 27 namespace runtime { 28 29 XfeedManager* GetXfeedManager() { 30 static XfeedManager* manager = new XfeedManager; 31 return manager; 32 } 33 34 extern const char* const kEigenMatMulF32SymbolName = 35 "__xla_cpu_runtime_EigenMatMulF32"; 36 extern const char* const kEigenMatMulF64SymbolName = 37 "__xla_cpu_runtime_EigenMatMulF64"; 38 extern const char* const kEigenConvF16SymbolName = 39 "__xla_cpu_runtime_EigenConvF16"; 40 extern const char* const kEigenConvF32SymbolName = 41 "__xla_cpu_runtime_EigenConvF32"; 42 extern const char* const kEigenFftSymbolName = "__xla_cpu_runtime_EigenFft"; 43 extern const char* const kEigenSingleThreadedMatMulF32SymbolName = 44 "__xla_cpu_runtime_EigenSingleThreadedMatMulF32"; 45 extern const char* const kEigenSingleThreadedMatMulF64SymbolName = 46 "__xla_cpu_runtime_EigenSingleThreadedMatMulF64"; 47 extern const char* const kEigenSingleThreadedConvF16SymbolName = 48 "__xla_cpu_runtime_EigenSingleThreadedConvF16"; 49 extern const char* const kEigenSingleThreadedConvF32SymbolName = 50 "__xla_cpu_runtime_EigenSingleThreadedConvF32"; 51 extern const char* const kAcquireInfeedBufferForDequeueSymbolName = 52 "__xla_cpu_runtime_AcquireInfeedBufferForDequeue"; 53 extern const char* const kReleaseInfeedBufferAfterDequeueSymbolName = 54 "__xla_cpu_runtime_ReleaseInfeedBufferAfterDequeue"; 55 extern const char* const kAcquireOutfeedBufferForPopulationSymbolName = 56 "__xla_cpu_runtime_AcquireOutfeedBufferForPopulation"; 57 extern const char* const kReleaseOutfeedBufferAfterPopulationSymbolName = 58 "__xla_cpu_runtime_ReleaseOutfeedBufferAfterPopulation"; 59 extern const char* const kParallelForkJoinSymbolName = 60 "__xla_cpu_runtime_ParallelForkJoin"; 61 62 extern const char* const kXlaCpuRuntimeSymbolNamePrefix = "__xla_cpu_runtime_"; 63 } // namespace runtime 64 } // namespace cpu 65 } // namespace xla 66 67 namespace { 68 69 tensorflow::string ShapeString(const void* shape_ptr, xla::int32 shape_length) { 70 xla::StatusOr<xla::Shape> shape = 71 xla::llvm_ir::DecodeSelfDescribingShapeConstant(shape_ptr, shape_length); 72 if (shape.ok()) { 73 return xla::ShapeUtil::HumanStringWithLayout(shape.ValueOrDie()); 74 } 75 return "<invalid shape>"; 76 } 77 78 } // namespace 79 80 void* __xla_cpu_runtime_AcquireInfeedBufferForDequeue(xla::int32 buffer_length, 81 const void* shape, 82 xla::int32 shape_length) { 83 if (VLOG_IS_ON(2)) { 84 LOG(INFO) << "AcquireInfeedBufferForDequeue: " 85 << ShapeString(shape, shape_length); 86 } 87 xla::cpu::runtime::XfeedManager* xfeed = xla::cpu::runtime::GetXfeedManager(); 88 // Wait until there's a buffer to dequeue. 89 xla::cpu::runtime::XfeedBuffer* buffer = 90 xfeed->infeed()->BlockingDequeueBuffer(); 91 CHECK_EQ(buffer->length(), buffer_length) 92 << "XLA program infeed request buffer size " << buffer_length 93 << " did not match the runtime's infed buffer length " << buffer->length() 94 << "; program reports desired shape: " 95 << ShapeString(shape, shape_length); 96 return buffer->data(); 97 } 98 99 void __xla_cpu_runtime_ReleaseInfeedBufferAfterDequeue( 100 xla::int32 buffer_length, void* buffer_ptr, const void* shape_ptr, 101 xla::int32 shape_length) { 102 if (VLOG_IS_ON(2)) { 103 LOG(INFO) << "ReleaseInfeedBufferAfterDeque: " 104 << ShapeString(shape_ptr, shape_length); 105 } 106 xla::cpu::runtime::XfeedManager* xfeed = xla::cpu::runtime::GetXfeedManager(); 107 xla::StatusOr<xla::Shape> shape = 108 xla::llvm_ir::DecodeSelfDescribingShapeConstant(shape_ptr, shape_length); 109 xfeed->infeed()->ReleaseCurrentBuffer(buffer_length, buffer_ptr, 110 std::move(shape)); 111 } 112 113 void* __xla_cpu_runtime_AcquireOutfeedBufferForPopulation( 114 xla::int32 buffer_length, const void* shape_ptr, xla::int32 shape_length) { 115 if (VLOG_IS_ON(2)) { 116 LOG(INFO) << "AcquireOutfeedBufferForPopulation: " 117 << ShapeString(shape_ptr, shape_length); 118 } 119 xla::cpu::runtime::XfeedManager* xfeed = xla::cpu::runtime::GetXfeedManager(); 120 // Wait until there's a buffer to dequeue. 121 xla::cpu::runtime::XfeedBuffer* buffer = 122 xfeed->outfeed()->BlockingDequeueBuffer(); 123 CHECK_EQ(buffer->length(), buffer_length) 124 << "XLA program outfeed request buffer size " << buffer_length 125 << " did not match the runtime's outfeed buffer length " 126 << buffer->length() << "; program reports outfed shape: " 127 << ShapeString(shape_ptr, shape_length); 128 return buffer->data(); 129 } 130 131 void __xla_cpu_runtime_ReleaseOutfeedBufferAfterPopulation( 132 xla::int32 buffer_length, void* buffer_ptr, const void* shape_ptr, 133 xla::int32 shape_length) { 134 if (VLOG_IS_ON(2)) { 135 LOG(INFO) << "ReleaseOutfeedBufferAfterPopulation: " 136 << ShapeString(shape_ptr, shape_length); 137 } 138 xla::cpu::runtime::XfeedManager* xfeed = xla::cpu::runtime::GetXfeedManager(); 139 xla::StatusOr<xla::Shape> shape = 140 xla::llvm_ir::DecodeSelfDescribingShapeConstant(shape_ptr, shape_length); 141 xfeed->outfeed()->ReleaseCurrentBuffer(buffer_length, buffer_ptr, 142 std::move(shape)); 143 } 144
