1 /* Copyright 2015 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 #define EIGEN_USE_THREADS 17 18 #include <vector> 19 #include "tensorflow/core/framework/register_types.h" 20 #include "tensorflow/core/kernels/concat_lib.h" 21 #include "tensorflow/core/util/work_sharder.h" 22 23 namespace tensorflow { 24 25 // ElementCopier must be a struct with a single Copy function, which is passed 26 // the output pointer, input pointer, input index, and number of elements to 27 // copy from input to output. 28 template <typename T, typename ElementCopier> 29 void ConcatCPUImpl( 30 DeviceBase* d, 31 const std::vector<std::unique_ptr<typename TTypes<T, 2>::ConstMatrix>>& 32 inputs, 33 int64 cost_per_unit, ElementCopier copier, 34 typename TTypes<T, 2>::Matrix* output) { 35 size_t num_inputs = inputs.size(); 36 37 std::vector<ptrdiff_t> sizes; 38 sizes.reserve(num_inputs); 39 int64 row_size = 0; 40 for (const auto& input : inputs) { 41 sizes.push_back(input->dimension(1)); 42 row_size += sizes.back(); 43 } 44 45 auto worker_threads = d->tensorflow_cpu_worker_threads(); 46 int num_threads = std::min(4, worker_threads->num_threads); 47 // strings define a different amount of work (generally much more) compared 48 // with standard POD, so we parallelize differently. 49 if (!std::is_same<T, string>::value) { 50 num_threads = 51 static_cast<int>(std::min<int64>(num_threads, output->size() / 4096)); 52 } 53 // Single threaded mode. 54 // TODO(dga): Deduplicate this code w.r.t. sharded code below. 55 if (num_threads == 0) { 56 T* out = &(*output)(0, 0); 57 std::vector<const T*> inp; 58 inp.reserve(num_inputs); 59 for (const auto& input : inputs) { 60 inp.push_back(&(*input)(0, 0)); 61 } 62 const int64 dim0 = output->dimension(0); 63 for (int64 i = 0; i < dim0; ++i) { 64 for (int64 j = 0; j < num_inputs; ++j) { 65 auto size = sizes[j]; 66 copier.Copy(out, inp[j], j, size); 67 out += size; 68 inp[j] += size; 69 } 70 } 71 return; 72 } 73 74 // Sharded mode. 75 auto work = [&row_size, &sizes, &inputs, &output, &copier, &num_inputs]( 76 int64 start, int64 end) { 77 int64 skipped_rows = start / row_size; 78 T* out = output->data() + skipped_rows * row_size; 79 T* out_start = output->data() + start; 80 T* out_end = output->data() + end; 81 82 // Handle partial row at start 83 if (out < out_start) { 84 for (size_t j = 0; j < num_inputs; ++j) { 85 ptrdiff_t size = sizes[j]; 86 ptrdiff_t offset = out_start - out; 87 if (size <= offset) { 88 out += size; 89 continue; 90 } 91 const T* inp = &(*inputs[j])(skipped_rows, 0); 92 if (offset > 0) { 93 out += offset; 94 inp += offset; 95 size -= offset; 96 } 97 size = std::min(size, out_end - out); 98 if (size <= 0) break; 99 copier.Copy(out, inp, j, size); 100 out += size; 101 } 102 ++skipped_rows; 103 } 104 if (out == out_end) return; 105 CHECK(out >= out_start); 106 CHECK(out < out_end); 107 108 // Copy remaining data. 109 std::vector<const T*> inp; 110 inp.reserve(num_inputs); 111 for (const auto& input : inputs) { 112 inp.push_back(&(*input)(skipped_rows, 0)); 113 } 114 const int64 dim0 = output->dimension(0); 115 for (int64 i = skipped_rows; i < dim0; ++i) { 116 for (int64 j = 0; j < num_inputs; ++j) { 117 ptrdiff_t size = std::min(sizes[j], out_end - out); 118 copier.Copy(out, inp[j], j, size); 119 out += size; 120 inp[j] += size; 121 if (out == out_end) return; 122 } 123 } 124 }; 125 Shard(worker_threads->num_threads, worker_threads->workers, output->size(), 126 cost_per_unit, work); 127 } 128 129 #ifdef TENSORFLOW_USE_SYCL 130 template <typename T, typename ElementCopier> 131 void ConcatSYCLImpl( 132 const Eigen::SyclDevice& d, 133 const std::vector<std::unique_ptr<typename TTypes<T, 2>::ConstMatrix>>& 134 inputs, 135 int64 cost_per_unit, ElementCopier copier, 136 typename TTypes<T, 2>::Matrix* output) { 137 size_t num_inputs = inputs.size(); 138 139 std::vector<ptrdiff_t> sizes; 140 sizes.reserve(num_inputs); 141 int64 row_size = 0; 142 for (const auto& input : inputs) { 143 sizes.push_back(input->dimension(1)); 144 row_size += sizes.back(); 145 } 146 147 T* out = &(*output)(0, 0); 148 std::vector<const T*> inp; 149 inp.reserve(num_inputs); 150 for (const auto& input : inputs) { 151 inp.push_back(&(*input)(0, 0)); 152 } 153 const int64 dim0 = output->dimension(0); 154 for (int64 i = 0; i < dim0; ++i) { 155 for (int64 j = 0; j < num_inputs; ++j) { 156 auto size = sizes[j]; 157 d.memcpy(out, inp[j], size * sizeof(T)); 158 out += size; 159 inp[j] += size; 160 } 161 } 162 } 163 #endif // TENSORFLOW_USE_SYCL 164 } // namespace tensorflow 165
