Eigen-Contrib  5.0.1
 
Loading...
Searching...
No Matches
TensorConvolution.h
1// This file is part of Eigen, a lightweight C++ template library
2// for linear algebra.
3//
4// Copyright (C) 2014 Benoit Steiner <benoit.steiner.goog@gmail.com>
5//
6// This Source Code Form is subject to the terms of the Mozilla
7// Public License v. 2.0. If a copy of the MPL was not distributed
8// with this file, You can obtain one at http://mozilla.org/MPL/2.0/.
9// SPDX-License-Identifier: MPL-2.0
10
11#ifndef EIGEN_TENSOR_TENSOR_CONVOLUTION_H
12#define EIGEN_TENSOR_TENSOR_CONVOLUTION_H
13
14// IWYU pragma: private
15#include "./InternalHeaderCheck.h"
16
17namespace Eigen {
18
19namespace internal {
20
21template <typename Index, typename InputDims, int NumKernelDims, int Layout>
22class IndexMapper {
23 public:
24 IndexMapper(const InputDims& input_dims, const array<Index, NumKernelDims>& kernel_dims,
25 const array<Index, NumKernelDims>& indices) {
26 array<Index, NumDims> dimensions = input_dims;
27 for (int i = 0; i < NumKernelDims; ++i) {
28 const Index index = indices[i];
29 const Index input_dim = input_dims[index];
30 const Index kernel_dim = kernel_dims[i];
31 const Index result_dim = input_dim - kernel_dim + 1;
32 dimensions[index] = result_dim;
33 }
34
35 array<Index, NumDims> inputStrides;
36 array<Index, NumDims> outputStrides;
37 EIGEN_IF_CONSTEXPR (static_cast<int>(Layout) == static_cast<int>(ColMajor)) {
38 inputStrides[0] = 1;
39 outputStrides[0] = 1;
40 for (int i = 1; i < NumDims; ++i) {
41 inputStrides[i] = inputStrides[i - 1] * input_dims[i - 1];
42 outputStrides[i] = outputStrides[i - 1] * dimensions[i - 1];
43 }
44 } else {
45 inputStrides[NumDims - 1] = 1;
46 outputStrides[NumDims - 1] = 1;
47 for (int i = static_cast<int>(NumDims) - 2; i >= 0; --i) {
48 inputStrides[i] = inputStrides[i + 1] * input_dims[i + 1];
49 outputStrides[i] = outputStrides[i + 1] * dimensions[i + 1];
50 }
51 }
52
53 array<Index, NumDims> gpuInputDimensions;
54 array<Index, NumDims> gpuOutputDimensions;
55 array<Index, NumDims> tmp = dimensions;
56 array<Index, NumDims> ordering;
57 constexpr size_t offset = static_cast<int>(Layout) == static_cast<int>(ColMajor) ? 0 : NumDims - NumKernelDims;
58 for (int i = 0; i < NumKernelDims; ++i) {
59 const Index index = i + offset;
60 ordering[index] = indices[i];
61 tmp[indices[i]] = -1;
62 gpuInputDimensions[index] = input_dims[indices[i]];
63 gpuOutputDimensions[index] = dimensions[indices[i]];
64 }
65
66 int written = static_cast<int>(Layout) == static_cast<int>(ColMajor) ? NumKernelDims : 0;
67 for (int i = 0; i < NumDims; ++i) {
68 if (tmp[i] >= 0) {
69 ordering[written] = i;
70 gpuInputDimensions[written] = input_dims[i];
71 gpuOutputDimensions[written] = dimensions[i];
72 ++written;
73 }
74 }
75
76 for (int i = 0; i < NumDims; ++i) {
77 m_inputStrides[i] = inputStrides[ordering[i]];
78 m_outputStrides[i] = outputStrides[ordering[i]];
79 }
80
81 EIGEN_IF_CONSTEXPR (static_cast<int>(Layout) == static_cast<int>(ColMajor)) {
82 for (int i = 0; i < NumDims; ++i) {
83 if (i > NumKernelDims) {
84 m_gpuInputStrides[i] = m_gpuInputStrides[i - 1] * gpuInputDimensions[i - 1];
85 m_gpuOutputStrides[i] = m_gpuOutputStrides[i - 1] * gpuOutputDimensions[i - 1];
86 } else {
87 m_gpuInputStrides[i] = 1;
88 m_gpuOutputStrides[i] = 1;
89 }
90 }
91 } else {
92 for (int i = NumDims - 1; i >= 0; --i) {
93 if (i + 1 < static_cast<int>(offset)) {
94 m_gpuInputStrides[i] = m_gpuInputStrides[i + 1] * gpuInputDimensions[i + 1];
95 m_gpuOutputStrides[i] = m_gpuOutputStrides[i + 1] * gpuOutputDimensions[i + 1];
96 } else {
97 m_gpuInputStrides[i] = 1;
98 m_gpuOutputStrides[i] = 1;
99 }
100 }
101 }
102 }
103
104 EIGEN_STRONG_INLINE EIGEN_DEVICE_FUNC Index mapGpuInputPlaneToTensorInputOffset(Index p) const {
105 Index inputIndex = 0;
106 EIGEN_IF_CONSTEXPR (static_cast<int>(Layout) == static_cast<int>(ColMajor)) {
107 for (int d = NumDims - 1; d > NumKernelDims; --d) {
108 const Index idx = p / m_gpuInputStrides[d];
109 inputIndex += idx * m_inputStrides[d];
110 p -= idx * m_gpuInputStrides[d];
111 }
112 EIGEN_IF_CONSTEXPR (NumKernelDims < NumDims) {
113 inputIndex += p * m_inputStrides[NumKernelDims];
114 }
115 } else {
116 std::ptrdiff_t limit = 0;
117 EIGEN_IF_CONSTEXPR (NumKernelDims < NumDims) {
118 limit = NumDims - NumKernelDims - 1;
119 }
120 for (int d = 0; d < limit; ++d) {
121 const Index idx = p / m_gpuInputStrides[d];
122 inputIndex += idx * m_inputStrides[d];
123 p -= idx * m_gpuInputStrides[d];
124 }
125 inputIndex += p * m_inputStrides[limit];
126 }
127 return inputIndex;
128 }
129
130 EIGEN_STRONG_INLINE EIGEN_DEVICE_FUNC Index mapGpuOutputPlaneToTensorOutputOffset(Index p) const {
131 Index outputIndex = 0;
132 EIGEN_IF_CONSTEXPR (static_cast<int>(Layout) == static_cast<int>(ColMajor)) {
133 for (int d = NumDims - 1; d > NumKernelDims; --d) {
134 const Index idx = p / m_gpuOutputStrides[d];
135 outputIndex += idx * m_outputStrides[d];
136 p -= idx * m_gpuOutputStrides[d];
137 }
138 EIGEN_IF_CONSTEXPR (NumKernelDims < NumDims) {
139 outputIndex += p * m_outputStrides[NumKernelDims];
140 }
141 } else {
142 std::ptrdiff_t limit = 0;
143 EIGEN_IF_CONSTEXPR (NumKernelDims < NumDims) {
144 limit = NumDims - NumKernelDims - 1;
145 }
146 for (int d = 0; d < limit; ++d) {
147 const Index idx = p / m_gpuOutputStrides[d];
148 outputIndex += idx * m_outputStrides[d];
149 p -= idx * m_gpuOutputStrides[d];
150 }
151 outputIndex += p * m_outputStrides[limit];
152 }
153 return outputIndex;
154 }
155
156 EIGEN_STRONG_INLINE EIGEN_DEVICE_FUNC Index mapGpuInputKernelToTensorInputOffset(Index i) const {
157 constexpr size_t offset = static_cast<int>(Layout) == static_cast<int>(ColMajor) ? 0 : NumDims - NumKernelDims;
158 return i * m_inputStrides[offset];
159 }
160
161 EIGEN_STRONG_INLINE EIGEN_DEVICE_FUNC Index mapGpuOutputKernelToTensorOutputOffset(Index i) const {
162 constexpr size_t offset = static_cast<int>(Layout) == static_cast<int>(ColMajor) ? 0 : NumDims - NumKernelDims;
163 return i * m_outputStrides[offset];
164 }
165
166 EIGEN_STRONG_INLINE EIGEN_DEVICE_FUNC Index mapGpuInputKernelToTensorInputOffset(Index i, Index j) const {
167 constexpr size_t offset = static_cast<int>(Layout) == static_cast<int>(ColMajor) ? 0 : NumDims - NumKernelDims;
168 return i * m_inputStrides[offset] + j * m_inputStrides[offset + 1];
169 }
170
171 EIGEN_STRONG_INLINE EIGEN_DEVICE_FUNC Index mapGpuOutputKernelToTensorOutputOffset(Index i, Index j) const {
172 constexpr size_t offset = static_cast<int>(Layout) == static_cast<int>(ColMajor) ? 0 : NumDims - NumKernelDims;
173 return i * m_outputStrides[offset] + j * m_outputStrides[offset + 1];
174 }
175
176 EIGEN_STRONG_INLINE EIGEN_DEVICE_FUNC Index mapGpuInputKernelToTensorInputOffset(Index i, Index j, Index k) const {
177 constexpr size_t offset = static_cast<int>(Layout) == static_cast<int>(ColMajor) ? 0 : NumDims - NumKernelDims;
178 return i * m_inputStrides[offset] + j * m_inputStrides[offset + 1] + k * m_inputStrides[offset + 2];
179 }
180
181 EIGEN_STRONG_INLINE EIGEN_DEVICE_FUNC Index mapGpuOutputKernelToTensorOutputOffset(Index i, Index j, Index k) const {
182 constexpr size_t offset = static_cast<int>(Layout) == static_cast<int>(ColMajor) ? 0 : NumDims - NumKernelDims;
183 return i * m_outputStrides[offset] + j * m_outputStrides[offset + 1] + k * m_outputStrides[offset + 2];
184 }
185
186 private:
187 static constexpr int NumDims = internal::array_size<InputDims>::value;
188 array<Index, NumDims> m_inputStrides;
189 array<Index, NumDims> m_outputStrides;
190 array<Index, NumDims> m_gpuInputStrides;
191 array<Index, NumDims> m_gpuOutputStrides;
192};
193
194template <typename Dimensions, typename InputXprType, typename KernelXprType>
195struct traits<TensorConvolutionOp<Dimensions, InputXprType, KernelXprType> > {
196 // Type promotion to handle the case where the types of the lhs and the rhs are different.
197 typedef typename promote_storage_type<typename InputXprType::Scalar, typename KernelXprType::Scalar>::ret Scalar;
198 typedef typename promote_storage_type<typename traits<InputXprType>::StorageKind,
199 typename traits<KernelXprType>::StorageKind>::ret StorageKind;
200 typedef typename promote_index_type<typename traits<InputXprType>::Index, typename traits<KernelXprType>::Index>::type
201 Index;
202 static constexpr int NumDimensions = traits<InputXprType>::NumDimensions;
203 static constexpr int Layout = traits<InputXprType>::Layout;
204 typedef std::conditional_t<Pointer_type_promotion<typename InputXprType::Scalar, Scalar>::val,
205 typename traits<InputXprType>::PointerType, typename traits<KernelXprType>::PointerType>
206 PointerType;
207
208 enum { Flags = 0 };
209};
210
211template <typename Dimensions, typename InputXprType, typename KernelXprType>
212struct eval<TensorConvolutionOp<Dimensions, InputXprType, KernelXprType>, Eigen::Dense> {
213 typedef const TensorConvolutionOp<Dimensions, InputXprType, KernelXprType>& type;
214};
215
216} // end namespace internal
217
221template <typename Indices, typename InputXprType, typename KernelXprType>
222class TensorConvolutionOp
223 : public TensorBase<TensorConvolutionOp<Indices, InputXprType, KernelXprType>, ReadOnlyAccessors> {
224 public:
225 typedef typename Eigen::internal::traits<TensorConvolutionOp>::Scalar Scalar;
226 typedef typename Eigen::NumTraits<Scalar>::Real RealScalar;
227 typedef typename internal::promote_storage_type<typename InputXprType::CoeffReturnType,
228 typename KernelXprType::CoeffReturnType>::ret CoeffReturnType;
229 typedef typename Eigen::internal::ref_selector<TensorConvolutionOp>::type Nested;
230 typedef typename Eigen::internal::traits<TensorConvolutionOp>::StorageKind StorageKind;
231 typedef typename Eigen::internal::traits<TensorConvolutionOp>::Index Index;
232
233 EIGEN_DEVICE_FUNC EIGEN_STRONG_INLINE TensorConvolutionOp(const InputXprType& input, const KernelXprType& kernel,
234 const Indices& dims)
235 : m_input_xpr(input), m_kernel_xpr(kernel), m_indices(dims) {}
236
237 EIGEN_DEVICE_FUNC EIGEN_STRONG_INLINE const Indices& indices() const { return m_indices; }
238
240 EIGEN_DEVICE_FUNC EIGEN_STRONG_INLINE const internal::remove_all_t<typename InputXprType::Nested>& inputExpression()
241 const {
242 return m_input_xpr;
243 }
244
245 EIGEN_DEVICE_FUNC EIGEN_STRONG_INLINE const internal::remove_all_t<typename KernelXprType::Nested>& kernelExpression()
246 const {
247 return m_kernel_xpr;
248 }
249
250 protected:
251 typename InputXprType::Nested m_input_xpr;
252 typename KernelXprType::Nested m_kernel_xpr;
253 const Indices m_indices;
254};
255
256template <typename Indices, typename InputArgType, typename KernelArgType, typename Device>
257struct TensorEvaluator<const TensorConvolutionOp<Indices, InputArgType, KernelArgType>, Device> {
258 typedef TensorConvolutionOp<Indices, InputArgType, KernelArgType> XprType;
259
260 static constexpr int NumDims =
261 internal::array_size<typename TensorEvaluator<InputArgType, Device>::Dimensions>::value;
262 static constexpr int NumKernelDims = internal::array_size<Indices>::value;
263 typedef typename XprType::Index Index;
264 typedef DSizes<Index, NumDims> Dimensions;
265
266 typedef typename XprType::Scalar Scalar;
267 typedef typename XprType::CoeffReturnType CoeffReturnType;
268 typedef typename PacketType<CoeffReturnType, Device>::type PacketReturnType;
269 static constexpr int PacketSize = PacketType<CoeffReturnType, Device>::size;
270 typedef StorageMemory<Scalar, Device> Storage;
271 typedef typename Storage::Type EvaluatorPointerType;
272
273 static constexpr int Layout = TensorEvaluator<InputArgType, Device>::Layout;
274 enum {
275 IsAligned =
276 int(TensorEvaluator<InputArgType, Device>::IsAligned) & int(TensorEvaluator<KernelArgType, Device>::IsAligned),
277 PacketAccess = int(TensorEvaluator<InputArgType, Device>::PacketAccess) &
278 int(TensorEvaluator<KernelArgType, Device>::PacketAccess),
279 BlockAccess = false,
280 PreferBlockAccess = false,
281 CoordAccess = false, // to be implemented
282 RawAccess = false
283 };
284
285 //===- Tensor block evaluation strategy (see TensorBlock.h) -------------===//
286 typedef internal::TensorBlockNotImplemented TensorBlock;
287 //===--------------------------------------------------------------------===//
288
289 EIGEN_STRONG_INLINE TensorEvaluator(const XprType& op, const Device& device)
290 : m_inputImpl(op.inputExpression(), device),
291 m_kernelImpl(op.kernelExpression(), device),
292 m_kernelArg(op.kernelExpression()),
293 m_kernel(nullptr),
294 m_local_kernel(false),
295 m_device(device) {
296 EIGEN_STATIC_ASSERT((static_cast<int>(TensorEvaluator<InputArgType, Device>::Layout) ==
297 static_cast<int>(TensorEvaluator<KernelArgType, Device>::Layout)),
298 YOU_MADE_A_PROGRAMMING_MISTAKE);
299
300 const typename TensorEvaluator<InputArgType, Device>::Dimensions& input_dims = m_inputImpl.dimensions();
301 const typename TensorEvaluator<KernelArgType, Device>::Dimensions& kernel_dims = m_kernelImpl.dimensions();
302
303 EIGEN_IF_CONSTEXPR (static_cast<int>(Layout) == static_cast<int>(ColMajor)) {
304 m_inputStride[0] = 1;
305 for (int i = 1; i < NumDims; ++i) {
306 m_inputStride[i] = m_inputStride[i - 1] * input_dims[i - 1];
307 }
308 } else {
309 m_inputStride[NumDims - 1] = 1;
310 for (int i = NumDims - 2; i >= 0; --i) {
311 m_inputStride[i] = m_inputStride[i + 1] * input_dims[i + 1];
312 }
313 }
314
315 m_dimensions = m_inputImpl.dimensions();
316 EIGEN_IF_CONSTEXPR (static_cast<int>(Layout) == static_cast<int>(ColMajor)) {
317 for (int i = 0; i < NumKernelDims; ++i) {
318 const Index index = op.indices()[i];
319 const Index input_dim = input_dims[index];
320 const Index kernel_dim = kernel_dims[i];
321 const Index result_dim = input_dim - kernel_dim + 1;
322 m_dimensions[index] = result_dim;
323 if (i > 0) {
324 m_kernelStride[i] = m_kernelStride[i - 1] * kernel_dims[i - 1];
325 } else {
326 m_kernelStride[0] = 1;
327 }
328 m_indexStride[i] = m_inputStride[index];
329 }
330
331 m_outputStride[0] = 1;
332 for (int i = 1; i < NumDims; ++i) {
333 m_outputStride[i] = m_outputStride[i - 1] * m_dimensions[i - 1];
334 }
335 } else {
336 for (int i = NumKernelDims - 1; i >= 0; --i) {
337 const Index index = op.indices()[i];
338 const Index input_dim = input_dims[index];
339 const Index kernel_dim = kernel_dims[i];
340 const Index result_dim = input_dim - kernel_dim + 1;
341 m_dimensions[index] = result_dim;
342 if (i < NumKernelDims - 1) {
343 m_kernelStride[i] = m_kernelStride[i + 1] * kernel_dims[i + 1];
344 } else {
345 m_kernelStride[NumKernelDims - 1] = 1;
346 }
347 m_indexStride[i] = m_inputStride[index];
348 }
349
350 m_outputStride[NumDims - 1] = 1;
351 for (int i = NumDims - 2; i >= 0; --i) {
352 m_outputStride[i] = m_outputStride[i + 1] * m_dimensions[i + 1];
353 }
354 }
355 }
356
357 EIGEN_DEVICE_FUNC EIGEN_STRONG_INLINE const Dimensions& dimensions() const { return m_dimensions; }
358
359 EIGEN_DEVICE_FUNC EIGEN_STRONG_INLINE bool evalSubExprsIfNeeded(Scalar*) {
360 m_inputImpl.evalSubExprsIfNeeded(nullptr);
361 preloadKernel();
362 return true;
363 }
364 EIGEN_STRONG_INLINE void cleanup() {
365 m_inputImpl.cleanup();
366 if (m_local_kernel) {
367 m_device.deallocate((void*)m_kernel);
368 m_local_kernel = false;
369 }
370 m_kernel = nullptr;
371 }
372
373 void evalTo(typename XprType::Scalar* buffer) {
374 evalSubExprsIfNeeded(nullptr);
375 for (int i = 0; i < dimensions().TotalSize(); ++i) {
376 buffer[i] += coeff(i);
377 }
378 cleanup();
379 }
380
381 EIGEN_DEVICE_FUNC EIGEN_STRONG_INLINE CoeffReturnType coeff(Index index) const {
382 CoeffReturnType result = CoeffReturnType(0);
383 convolve(firstInput(index), 0, NumKernelDims - 1, result);
384 return result;
385 }
386
387 template <int LoadMode>
388 EIGEN_DEVICE_FUNC PacketReturnType packet(const Index index) const {
389 Index indices[2] = {index, index + PacketSize - 1};
390 Index startInputs[2] = {0, 0};
391 EIGEN_IF_CONSTEXPR (static_cast<int>(Layout) == static_cast<int>(ColMajor)) {
392 for (int i = NumDims - 1; i > 0; --i) {
393 const Index idx0 = indices[0] / m_outputStride[i];
394 const Index idx1 = indices[1] / m_outputStride[i];
395 startInputs[0] += idx0 * m_inputStride[i];
396 startInputs[1] += idx1 * m_inputStride[i];
397 indices[0] -= idx0 * m_outputStride[i];
398 indices[1] -= idx1 * m_outputStride[i];
399 }
400 } else {
401 for (int i = 0; i < NumDims - 1; ++i) {
402 const Index idx0 = indices[0] / m_outputStride[i];
403 const Index idx1 = indices[1] / m_outputStride[i];
404 startInputs[0] += idx0 * m_inputStride[i];
405 startInputs[1] += idx1 * m_inputStride[i];
406 indices[0] -= idx0 * m_outputStride[i];
407 indices[1] -= idx1 * m_outputStride[i];
408 }
409 }
410 startInputs[0] += indices[0];
411 startInputs[1] += indices[1];
412
413 if (startInputs[1] - startInputs[0] == PacketSize - 1) {
414 PacketReturnType result = internal::pset1<PacketReturnType>(0);
415 convolvePacket(startInputs[0], 0, NumKernelDims - 1, result);
416 return result;
417 } else {
418 EIGEN_ALIGN_TO_BOUNDARY(internal::unpacket_traits<PacketReturnType>::alignment) Scalar data[PacketSize];
419 data[0] = Scalar(0);
420 convolve(startInputs[0], 0, NumKernelDims - 1, data[0]);
421 for (int i = 1; i < PacketSize - 1; ++i) {
422 data[i] = Scalar(0);
423 convolve(firstInput(index + i), 0, NumKernelDims - 1, data[i]);
424 }
425 data[PacketSize - 1] = Scalar(0);
426 convolve(startInputs[1], 0, NumKernelDims - 1, data[PacketSize - 1]);
427 return internal::pload<PacketReturnType>(data);
428 }
429 }
430
431 EIGEN_DEVICE_FUNC EIGEN_STRONG_INLINE TensorOpCost costPerCoeff(bool vectorized) const {
432 const double kernel_size = m_kernelImpl.dimensions().TotalSize();
433 // We ignore the use of fused multiply-add.
434 const double convolve_compute_cost = TensorOpCost::AddCost<Scalar>() + TensorOpCost::MulCost<Scalar>();
435 const double firstIndex_compute_cost =
436 NumDims *
437 (2 * TensorOpCost::AddCost<Index>() + 2 * TensorOpCost::MulCost<Index>() + TensorOpCost::DivCost<Index>());
438 return TensorOpCost(0, 0, firstIndex_compute_cost, vectorized, PacketSize) +
439 kernel_size * (m_inputImpl.costPerCoeff(vectorized) + m_kernelImpl.costPerCoeff(vectorized) +
440 TensorOpCost(0, 0, convolve_compute_cost, vectorized, PacketSize));
441 }
442
443 EIGEN_DEVICE_FUNC EvaluatorPointerType data() const { return nullptr; }
444
445 private:
446 EIGEN_DEVICE_FUNC EIGEN_STRONG_INLINE Index firstInput(Index index) const {
447 Index startInput = 0;
448 EIGEN_IF_CONSTEXPR (static_cast<int>(Layout) == static_cast<int>(ColMajor)) {
449 for (int i = NumDims - 1; i > 0; --i) {
450 const Index idx = index / m_outputStride[i];
451 startInput += idx * m_inputStride[i];
452 index -= idx * m_outputStride[i];
453 }
454 } else {
455 for (int i = 0; i < NumDims - 1; ++i) {
456 const Index idx = index / m_outputStride[i];
457 startInput += idx * m_inputStride[i];
458 index -= idx * m_outputStride[i];
459 }
460 }
461 startInput += index;
462 return startInput;
463 }
464
465 EIGEN_DEVICE_FUNC void convolve(Index firstIndex, Index firstKernel, int DimIndex, CoeffReturnType& accum) const {
466 for (int j = 0; j < m_kernelImpl.dimensions()[DimIndex]; ++j) {
467 const Index input = firstIndex + j * m_indexStride[DimIndex];
468 const Index kernel = firstKernel + j * m_kernelStride[DimIndex];
469 if (DimIndex > 0) {
470 convolve(input, kernel, DimIndex - 1, accum);
471 } else {
472 accum += m_inputImpl.coeff(input) * m_kernel[kernel];
473 }
474 }
475 }
476
477 template <typename Packet>
478 EIGEN_DEVICE_FUNC void convolvePacket(Index firstIndex, Index firstKernel, int DimIndex, Packet& accum) const {
479 for (int j = 0; j < m_kernelImpl.dimensions()[DimIndex]; ++j) {
480 const Index input = firstIndex + j * m_indexStride[DimIndex];
481 const Index kernel = firstKernel + j * m_kernelStride[DimIndex];
482 if (DimIndex > 0) {
483 convolvePacket(input, kernel, DimIndex - 1, accum);
484 } else {
485 accum = internal::pmadd<Packet>(m_inputImpl.template packet<Unaligned>(input),
486 internal::pset1<Packet>(m_kernel[kernel]), accum);
487 }
488 }
489 }
490
491 EIGEN_DEVICE_FUNC EIGEN_STRONG_INLINE void preloadKernel() {
492 // Don't make a local copy of the kernel unless we have to (i.e. it's an
493 // expression that needs to be evaluated)
494 const Scalar* in_place = m_kernelImpl.data();
495 if (in_place) {
496 m_kernel = in_place;
497 m_local_kernel = false;
498 } else {
499 size_t kernel_sz = m_kernelImpl.dimensions().TotalSize() * sizeof(Scalar);
500 Scalar* local = (Scalar*)m_device.allocate_temp(kernel_sz);
501 typedef TensorEvalToOp<const KernelArgType> EvalTo;
502 EvalTo evalToTmp(local, m_kernelArg);
503 const bool Vectorize = internal::IsVectorizable<Device, KernelArgType>::value;
504 internal::TensorExecutor<const EvalTo, Device, Vectorize>::run(evalToTmp, m_device);
505
506 m_kernel = local;
507 m_local_kernel = true;
508 }
509 }
510
511 array<Index, NumDims> m_inputStride;
512 array<Index, NumDims> m_outputStride;
513
514 array<Index, NumKernelDims> m_indexStride;
515 array<Index, NumKernelDims> m_kernelStride;
516 TensorEvaluator<InputArgType, Device> m_inputImpl;
517 TensorEvaluator<KernelArgType, Device> m_kernelImpl;
518 Dimensions m_dimensions;
519
520 KernelArgType m_kernelArg;
521 const Scalar* m_kernel;
522 bool m_local_kernel;
523 const Device EIGEN_DEVICE_REF m_device;
524};
525
526// Use an optimized implementation of the evaluation code for GPUs whenever possible.
527#if defined(EIGEN_USE_GPU) && defined(EIGEN_GPUCC)
528
529template <int StaticKernelSize>
530struct GetKernelSize {
531 EIGEN_DEVICE_FUNC EIGEN_STRONG_INLINE int operator()(const int /*kernelSize*/) const { return StaticKernelSize; }
532};
533template <>
534struct GetKernelSize<Dynamic> {
535 EIGEN_DEVICE_FUNC EIGEN_STRONG_INLINE int operator()(const int kernelSize) const { return kernelSize; }
536};
537
538template <typename InputEvaluator, typename Index, typename InputDims, int StaticKernelSize>
539__global__ EIGEN_HIP_LAUNCH_BOUNDS_1024 void EigenConvolutionKernel1D(
540 InputEvaluator eval, const internal::IndexMapper<Index, InputDims, 1, InputEvaluator::Layout> indexMapper,
541 const float* __restrict kernel, const int numPlanes, const int numX, const int maxX, const int kernelSize,
542 float* buffer) {
543#if defined(EIGEN_HIPCC)
544 HIP_DYNAMIC_SHARED(float, s)
545#else
546 extern __shared__ float s[];
547#endif
548
549 const int first_x = blockIdx.x * maxX;
550 const int last_x = (first_x + maxX < numX ? first_x + maxX : numX) - 1;
551 const int num_x_input = last_x - first_x + GetKernelSize<StaticKernelSize>()(kernelSize);
552 const int num_x_output = last_x - first_x + 1;
553
554 const int first_plane = blockIdx.y * blockDim.y;
555 const int plane_stride = blockDim.y * gridDim.y;
556
557 for (int p = first_plane + threadIdx.y; p < numPlanes; p += plane_stride) {
558 // Load inputs to shared memory
559 const int plane_input_offset = indexMapper.mapGpuInputPlaneToTensorInputOffset(p);
560 const int plane_kernel_offset = threadIdx.y * num_x_input;
561#pragma unroll
562 for (int i = threadIdx.x; i < num_x_input; i += blockDim.x) {
563 const int tensor_index = plane_input_offset + indexMapper.mapGpuInputKernelToTensorInputOffset(i + first_x);
564 s[i + plane_kernel_offset] = eval.coeff(tensor_index);
565 }
566
567 __syncthreads();
568
569 // Compute the convolution
570 const int plane_output_offset = indexMapper.mapGpuOutputPlaneToTensorOutputOffset(p);
571
572#pragma unroll
573 for (int i = threadIdx.x; i < num_x_output; i += blockDim.x) {
574 const int kernel_offset = plane_kernel_offset + i;
575 float result = 0.0f;
576#pragma unroll
577 for (int k = 0; k < GetKernelSize<StaticKernelSize>()(kernelSize); ++k) {
578 result += s[k + kernel_offset] * kernel[k];
579 }
580 const int tensor_index = plane_output_offset + indexMapper.mapGpuOutputKernelToTensorOutputOffset(i + first_x);
581 buffer[tensor_index] = result;
582 }
583 __syncthreads();
584 }
585}
586
587template <typename InputEvaluator, typename Index, typename InputDims, int StaticKernelSizeX, int StaticKernelSizeY>
588__global__ EIGEN_HIP_LAUNCH_BOUNDS_1024 void EigenConvolutionKernel2D(
589 InputEvaluator eval, const internal::IndexMapper<Index, InputDims, 2, InputEvaluator::Layout> indexMapper,
590 const float* __restrict kernel, const int numPlanes, const int numX, const int maxX, const int numY, const int maxY,
591 const int kernelSizeX, const int kernelSizeY, float* buffer) {
592#if defined(EIGEN_HIPCC)
593 HIP_DYNAMIC_SHARED(float, s)
594#else
595 extern __shared__ float s[];
596#endif
597
598 const int first_x = blockIdx.x * maxX;
599 const int last_x = (first_x + maxX < numX ? first_x + maxX : numX) - 1;
600 const int num_x_input = last_x - first_x + GetKernelSize<StaticKernelSizeX>()(kernelSizeX);
601 const int num_x_output = last_x - first_x + 1;
602
603 const int first_y = blockIdx.y * maxY;
604 const int last_y = (first_y + maxY < numY ? first_y + maxY : numY) - 1;
605 const int num_y_input = last_y - first_y + GetKernelSize<StaticKernelSizeY>()(kernelSizeY);
606 const int num_y_output = last_y - first_y + 1;
607
608 const int first_plane = blockIdx.z * blockDim.z;
609 const int plane_stride = blockDim.z * gridDim.z;
610
611 for (int p = first_plane + threadIdx.z; p < numPlanes; p += plane_stride) {
612 const int plane_input_offset = indexMapper.mapGpuInputPlaneToTensorInputOffset(p);
613 const int plane_kernel_offset = threadIdx.z * num_y_input;
614
615// Load inputs to shared memory
616#pragma unroll
617 for (int j = threadIdx.y; j < num_y_input; j += blockDim.y) {
618 const int input_offset = num_x_input * (j + plane_kernel_offset);
619#pragma unroll
620 for (int i = threadIdx.x; i < num_x_input; i += blockDim.x) {
621 const int tensor_index =
622 plane_input_offset + indexMapper.mapGpuInputKernelToTensorInputOffset(i + first_x, j + first_y);
623 s[i + input_offset] = eval.coeff(tensor_index);
624 }
625 }
626
627 __syncthreads();
628
629 // Convolution
630 const int plane_output_offset = indexMapper.mapGpuOutputPlaneToTensorOutputOffset(p);
631
632#pragma unroll
633 for (int j = threadIdx.y; j < num_y_output; j += blockDim.y) {
634#pragma unroll
635 for (int i = threadIdx.x; i < num_x_output; i += blockDim.x) {
636 float result = 0.0f;
637#pragma unroll
638 for (int l = 0; l < GetKernelSize<StaticKernelSizeY>()(kernelSizeY); ++l) {
639 const int kernel_offset = kernelSizeX * l;
640 const int input_offset = i + num_x_input * (j + l + plane_kernel_offset);
641#pragma unroll
642 for (int k = 0; k < GetKernelSize<StaticKernelSizeX>()(kernelSizeX); ++k) {
643 result += s[k + input_offset] * kernel[k + kernel_offset];
644 }
645 }
646 const int tensor_index =
647 plane_output_offset + indexMapper.mapGpuOutputKernelToTensorOutputOffset(i + first_x, j + first_y);
648 buffer[tensor_index] = result;
649 }
650 }
651
652 __syncthreads();
653 }
654}
655
656template <typename InputEvaluator, typename Index, typename InputDims>
657__global__ EIGEN_HIP_LAUNCH_BOUNDS_1024 void EigenConvolutionKernel3D(
658 InputEvaluator eval, const internal::IndexMapper<Index, InputDims, 3, InputEvaluator::Layout> indexMapper,
659 const float* __restrict kernel, const size_t numPlanes, const size_t numX, const size_t maxX, const size_t numY,
660 const size_t maxY, const size_t numZ, const size_t maxZ, const int kernelSizeX, const int kernelSizeY,
661 const int kernelSizeZ, float* buffer) {
662#if defined(EIGEN_HIPCC)
663 HIP_DYNAMIC_SHARED(float, s)
664#else
665 extern __shared__ float s[];
666#endif
667
668 // Load inputs to shared memory
669 const int first_x = blockIdx.x * static_cast<int>(maxX);
670 const int last_x = (first_x + static_cast<int>(maxX) < static_cast<int>(numX) ? first_x + static_cast<int>(maxX)
671 : static_cast<int>(numX)) -
672 1;
673 const int num_x_input = last_x - first_x + kernelSizeX;
674
675 const int first_y = blockIdx.y * static_cast<int>(maxY);
676 const int last_y = (first_y + static_cast<int>(maxY) < static_cast<int>(numY) ? first_y + static_cast<int>(maxY)
677 : static_cast<int>(numY)) -
678 1;
679 const int num_y_input = last_y - first_y + kernelSizeY;
680
681 const int first_z = blockIdx.z * static_cast<int>(maxZ);
682 const int last_z = (first_z + static_cast<int>(maxZ) < static_cast<int>(numZ) ? first_z + static_cast<int>(maxZ)
683 : static_cast<int>(numZ)) -
684 1;
685 const int num_z_input = last_z - first_z + kernelSizeZ;
686
687 for (int p = 0; p < static_cast<int>(numPlanes); ++p) {
688 const int plane_input_offset = indexMapper.mapGpuInputPlaneToTensorInputOffset(p);
689 const int plane_kernel_offset = 0;
690
691 for (int k = threadIdx.z; k < num_z_input; k += blockDim.z) {
692 for (int j = threadIdx.y; j < num_y_input; j += blockDim.y) {
693 for (int i = threadIdx.x; i < num_x_input; i += blockDim.x) {
694 const int tensor_index = plane_input_offset + indexMapper.mapGpuInputKernelToTensorInputOffset(
695 i + first_x, j + first_y, k + first_z);
696 s[i + num_x_input * (j + num_y_input * (k + plane_kernel_offset))] = eval.coeff(tensor_index);
697 }
698 }
699 }
700
701 __syncthreads();
702
703 // Convolution
704 const int num_z_output = last_z - first_z + 1;
705 const int num_y_output = last_y - first_y + 1;
706 const int num_x_output = last_x - first_x + 1;
707 const int plane_output_offset = indexMapper.mapGpuOutputPlaneToTensorOutputOffset(p);
708
709 for (int k = threadIdx.z; k < num_z_output; k += blockDim.z) {
710 for (int j = threadIdx.y; j < num_y_output; j += blockDim.y) {
711 for (int i = threadIdx.x; i < num_x_output; i += blockDim.x) {
712 float result = 0.0f;
713 for (int n = 0; n < kernelSizeZ; ++n) {
714 for (int m = 0; m < kernelSizeY; ++m) {
715 for (int l = 0; l < kernelSizeX; ++l) {
716 result += s[i + l + num_x_input * (j + m + num_y_input * (k + n + plane_kernel_offset))] *
717 kernel[l + kernelSizeX * (m + kernelSizeY * n)];
718 }
719 }
720 }
721 const int tensor_index = plane_output_offset + indexMapper.mapGpuOutputKernelToTensorOutputOffset(
722 i + first_x, j + first_y, k + first_z);
723 buffer[tensor_index] = result;
724 }
725 }
726 }
727 __syncthreads();
728 }
729}
730
731template <typename Indices, typename InputArgType, typename KernelArgType>
732struct TensorEvaluator<const TensorConvolutionOp<Indices, InputArgType, KernelArgType>, GpuDevice> {
733 typedef TensorConvolutionOp<Indices, InputArgType, KernelArgType> XprType;
734
735 static constexpr int NumDims =
736 internal::array_size<typename TensorEvaluator<InputArgType, GpuDevice>::Dimensions>::value;
737 static constexpr int NumKernelDims = internal::array_size<Indices>::value;
738 typedef typename XprType::Index Index;
739 typedef DSizes<Index, NumDims> Dimensions;
740 typedef typename TensorEvaluator<KernelArgType, GpuDevice>::Dimensions KernelDimensions;
741
742 static constexpr int Layout = TensorEvaluator<InputArgType, GpuDevice>::Layout;
743 enum {
744 IsAligned = int(TensorEvaluator<InputArgType, GpuDevice>::IsAligned) &
745 int(TensorEvaluator<KernelArgType, GpuDevice>::IsAligned),
746 PacketAccess = false,
747 BlockAccess = false,
748 PreferBlockAccess = false,
749 CoordAccess = false, // to be implemented
750 RawAccess = false
751 };
752
753 //===- Tensor block evaluation strategy (see TensorBlock.h) -------------===//
754 typedef internal::TensorBlockNotImplemented TensorBlock;
755 //===--------------------------------------------------------------------===//
756
757 TensorEvaluator(const XprType& op, const GpuDevice& device)
758 : m_inputImpl(op.inputExpression(), device),
759 m_kernelImpl(op.kernelExpression(), device),
760 m_kernelArg(op.kernelExpression()),
761 m_indices(op.indices()),
762 m_buf(nullptr),
763 m_kernel(nullptr),
764 m_local_kernel(false),
765 m_device(device) {
766 EIGEN_STATIC_ASSERT((static_cast<int>(TensorEvaluator<InputArgType, GpuDevice>::Layout) ==
767 static_cast<int>(TensorEvaluator<KernelArgType, GpuDevice>::Layout)),
768 YOU_MADE_A_PROGRAMMING_MISTAKE);
769
770 const typename TensorEvaluator<InputArgType, GpuDevice>::Dimensions& input_dims = m_inputImpl.dimensions();
771 const typename TensorEvaluator<KernelArgType, GpuDevice>::Dimensions& kernel_dims = m_kernelImpl.dimensions();
772
773 m_dimensions = m_inputImpl.dimensions();
774 for (int i = 0; i < NumKernelDims; ++i) {
775 const Index index = op.indices()[i];
776 const Index input_dim = input_dims[index];
777 const Index kernel_dim = kernel_dims[i];
778 const Index result_dim = input_dim - kernel_dim + 1;
779 m_dimensions[index] = result_dim;
780 }
781 }
782
783 typedef typename XprType::CoeffReturnType CoeffReturnType;
784 typedef typename PacketType<CoeffReturnType, GpuDevice>::type PacketReturnType;
785 typedef typename InputArgType::Scalar Scalar;
786 static constexpr int PacketSize = internal::unpacket_traits<PacketReturnType>::size;
787
788 EIGEN_DEVICE_FUNC const Dimensions& dimensions() const { return m_dimensions; }
789
790 EIGEN_DEVICE_FUNC EIGEN_STRONG_INLINE bool evalSubExprsIfNeeded(Scalar* data) {
791 preloadKernel();
792 m_inputImpl.evalSubExprsIfNeeded(nullptr);
793 if (data) {
794 executeEval(data);
795 return false;
796 } else {
797 m_buf = (Scalar*)m_device.allocate(dimensions().TotalSize() * sizeof(Scalar));
798 executeEval(m_buf);
799 return true;
800 }
801 }
802
803 EIGEN_STRONG_INLINE void cleanup() {
804 m_inputImpl.cleanup();
805 if (m_buf) {
806 m_device.deallocate(m_buf);
807 m_buf = nullptr;
808 }
809 if (m_local_kernel) {
810 m_device.deallocate((void*)m_kernel);
811 m_local_kernel = false;
812 }
813 m_kernel = nullptr;
814 }
815
816 EIGEN_STRONG_INLINE void preloadKernel() {
817 // Don't make a local copy of the kernel unless we have to (i.e. it's an
818 // expression that needs to be evaluated)
819 const Scalar* in_place = m_kernelImpl.data();
820 if (in_place) {
821 m_kernel = in_place;
822 m_local_kernel = false;
823 } else {
824 size_t kernel_sz = m_kernelImpl.dimensions().TotalSize() * sizeof(Scalar);
825 Scalar* local = (Scalar*)m_device.allocate(kernel_sz);
826 typedef TensorEvalToOp<const KernelArgType> EvalTo;
827 EvalTo evalToTmp(local, m_kernelArg);
828 const bool PacketAccess = internal::IsVectorizable<GpuDevice, KernelArgType>::value;
829 internal::TensorExecutor<const EvalTo, GpuDevice, PacketAccess>::run(evalToTmp, m_device);
830
831 m_kernel = local;
832 m_local_kernel = true;
833 }
834 }
835
836 static unsigned int ceil(unsigned int num, unsigned int denom) {
837 const unsigned int rounded_toward_zero = num / denom;
838 if (num > rounded_toward_zero * denom) {
839 return rounded_toward_zero + 1;
840 }
841 return rounded_toward_zero;
842 }
843
844 void executeEval(Scalar* data) const {
845 typedef typename TensorEvaluator<InputArgType, GpuDevice>::Dimensions InputDims;
846
847 const int maxSharedMem = m_device.sharedMemPerBlock();
848 const int maxThreadsPerBlock = m_device.maxGpuThreadsPerBlock();
849 const int maxBlocksPerProcessor = m_device.maxGpuThreadsPerMultiProcessor() / maxThreadsPerBlock;
850 const int numMultiProcessors = m_device.getNumGpuMultiProcessors();
851 const int warpSize = 32;
852
853 switch (NumKernelDims) {
854 case 1: {
855 const int kernel_size = m_kernelImpl.dimensions().TotalSize();
856
857 const int numX = dimensions()[m_indices[0]];
858 const int numP = dimensions().TotalSize() / numX;
859 int maxX;
860 dim3 block_size;
861
862 const int single_stride_dim =
863 static_cast<int>(Layout) == static_cast<int>(ColMajor) ? 0 : m_inputImpl.dimensions().rank() - 1;
864 if (m_indices[0] == single_stride_dim) {
865 // Maximize the reuse
866 const int inner_dim = ((maxSharedMem / (sizeof(Scalar)) - kernel_size + 1 + 31) / 32) * 32;
867 maxX = numext::mini<int>(inner_dim, numX);
868 const int maxP = numext::mini<int>(maxSharedMem / ((kernel_size - 1 + maxX) * sizeof(Scalar)), numP);
869 block_size.x = numext::mini(maxThreadsPerBlock, maxX);
870 block_size.y = numext::mini<int>(maxThreadsPerBlock / block_size.x, maxP);
871 } else {
872 // Read as much as possible alongside the inner most dimension, that is the plane
873 const int inner_dim = maxSharedMem / ((warpSize + kernel_size) * sizeof(Scalar));
874 const int maxP = numext::mini<int>(inner_dim, numP);
875 maxX = numext::mini<int>(maxSharedMem / (inner_dim * sizeof(Scalar)) - kernel_size + 1, numX);
876
877 block_size.x = numext::mini(warpSize, maxX);
878 block_size.y = numext::mini<int>(maxThreadsPerBlock / block_size.x, maxP);
879 }
880
881 const int shared_mem = block_size.y * (maxX + kernel_size - 1) * sizeof(Scalar);
882 gpu_assert(shared_mem <= maxSharedMem);
883
884 const int num_x_blocks = ceil(numX, maxX);
885 const int blocksPerProcessor = numext::mini(maxBlocksPerProcessor, maxSharedMem / shared_mem);
886 const int num_y_blocks = ceil(numMultiProcessors * blocksPerProcessor, num_x_blocks);
887
888 dim3 num_blocks(num_x_blocks, numext::mini<int>(num_y_blocks, ceil(numP, block_size.y)));
889
890 const array<Index, 1> indices{m_indices[0]};
891 const array<Index, 1> kernel_dims{m_kernelImpl.dimensions()[0]};
892 internal::IndexMapper<Index, InputDims, 1, Layout> indexMapper(m_inputImpl.dimensions(), kernel_dims, indices);
893 switch (kernel_size) {
894 case 4: {
895 LAUNCH_GPU_KERNEL((EigenConvolutionKernel1D<TensorEvaluator<InputArgType, GpuDevice>, Index, InputDims, 4>),
896 num_blocks, block_size, shared_mem, m_device, m_inputImpl, indexMapper, m_kernel, numP,
897 numX, maxX, 4, data);
898 break;
899 }
900 case 7: {
901 LAUNCH_GPU_KERNEL((EigenConvolutionKernel1D<TensorEvaluator<InputArgType, GpuDevice>, Index, InputDims, 7>),
902 num_blocks, block_size, shared_mem, m_device, m_inputImpl, indexMapper, m_kernel, numP,
903 numX, maxX, 7, data);
904 break;
905 }
906 default: {
907 LAUNCH_GPU_KERNEL(
908 (EigenConvolutionKernel1D<TensorEvaluator<InputArgType, GpuDevice>, Index, InputDims, Dynamic>),
909 num_blocks, block_size, shared_mem, m_device, m_inputImpl, indexMapper, m_kernel, numP, numX, maxX,
910 kernel_size, data);
911 }
912 }
913 break;
914 }
915
916 case 2: {
917 constexpr int idxX = static_cast<int>(Layout) == static_cast<int>(ColMajor) ? 0 : 1;
918 constexpr int idxY = static_cast<int>(Layout) == static_cast<int>(ColMajor) ? 1 : 0;
919 const int kernel_size_x = m_kernelImpl.dimensions()[idxX];
920 const int kernel_size_y = m_kernelImpl.dimensions()[idxY];
921
922 const int numX = dimensions()[m_indices[idxX]];
923 const int numY = dimensions()[m_indices[idxY]];
924 const int numP = dimensions().TotalSize() / (numX * numY);
925
926 const float scaling_factor =
927 sqrtf(static_cast<float>(maxSharedMem) / (sizeof(Scalar) * kernel_size_y * kernel_size_x));
928
929 // Snap maxX to warp size
930 int inner_dim = ((static_cast<int>(scaling_factor * kernel_size_x) - kernel_size_x + 1 + 32) / 32) * 32;
931 const int maxX = numext::mini<int>(inner_dim, numX);
932 const int maxY =
933 numext::mini<int>(maxSharedMem / (sizeof(Scalar) * (maxX + kernel_size_x - 1)) - kernel_size_y + 1, numY);
934 const int maxP = numext::mini<int>(
935 maxSharedMem / ((kernel_size_x - 1 + maxX) * (kernel_size_y - 1 + maxY) * sizeof(Scalar)), numP);
936
937 dim3 block_size;
938 block_size.x = numext::mini(1024, maxX);
939 block_size.y = numext::mini<int>(1024 / block_size.x, maxY);
940 block_size.z = numext::mini<int>(1024 / (block_size.x * block_size.y), maxP);
941
942 const int shared_mem = block_size.z * (maxX + kernel_size_x - 1) * (maxY + kernel_size_y - 1) * sizeof(Scalar);
943 gpu_assert(shared_mem <= maxSharedMem);
944
945 const int num_x_blocks = ceil(numX, maxX);
946 const int num_y_blocks = ceil(numY, maxY);
947 const int blocksPerProcessor = numext::mini(maxBlocksPerProcessor, maxSharedMem / shared_mem);
948 const int num_z_blocks = ceil(numMultiProcessors * blocksPerProcessor, num_x_blocks * num_y_blocks);
949
950 dim3 num_blocks(num_x_blocks, num_y_blocks, numext::mini<int>(num_z_blocks, ceil(numP, block_size.z)));
951
952 const array<Index, 2> indices{m_indices[idxX], m_indices[idxY]};
953 const array<Index, 2> kernel_dims{m_kernelImpl.dimensions()[idxX], m_kernelImpl.dimensions()[idxY]};
954 internal::IndexMapper<Index, InputDims, 2, Layout> indexMapper(m_inputImpl.dimensions(), kernel_dims, indices);
955 switch (kernel_size_x) {
956 case 4: {
957 switch (kernel_size_y) {
958 case 7: {
959 LAUNCH_GPU_KERNEL(
960 (EigenConvolutionKernel2D<TensorEvaluator<InputArgType, GpuDevice>, Index, InputDims, 4, 7>),
961 num_blocks, block_size, shared_mem, m_device, m_inputImpl, indexMapper, m_kernel, numP, numX, maxX,
962 numY, maxY, 4, 7, data);
963 break;
964 }
965 default: {
966 LAUNCH_GPU_KERNEL(
967 (EigenConvolutionKernel2D<TensorEvaluator<InputArgType, GpuDevice>, Index, InputDims, 4, Dynamic>),
968 num_blocks, block_size, shared_mem, m_device, m_inputImpl, indexMapper, m_kernel, numP, numX, maxX,
969 numY, maxY, 4, kernel_size_y, data);
970 break;
971 }
972 }
973 break;
974 }
975 case 7: {
976 switch (kernel_size_y) {
977 case 4: {
978 LAUNCH_GPU_KERNEL(
979 (EigenConvolutionKernel2D<TensorEvaluator<InputArgType, GpuDevice>, Index, InputDims, 7, 4>),
980 num_blocks, block_size, shared_mem, m_device, m_inputImpl, indexMapper, m_kernel, numP, numX, maxX,
981 numY, maxY, 7, 4, data);
982 break;
983 }
984 default: {
985 LAUNCH_GPU_KERNEL(
986 (EigenConvolutionKernel2D<TensorEvaluator<InputArgType, GpuDevice>, Index, InputDims, 7, Dynamic>),
987 num_blocks, block_size, shared_mem, m_device, m_inputImpl, indexMapper, m_kernel, numP, numX, maxX,
988 numY, maxY, 7, kernel_size_y, data);
989 break;
990 }
991 }
992 break;
993 }
994 default: {
995 LAUNCH_GPU_KERNEL((EigenConvolutionKernel2D<TensorEvaluator<InputArgType, GpuDevice>, Index, InputDims,
996 Dynamic, Dynamic>),
997 num_blocks, block_size, shared_mem, m_device, m_inputImpl, indexMapper, m_kernel, numP,
998 numX, maxX, numY, maxY, kernel_size_x, kernel_size_y, data);
999 break;
1000 }
1001 }
1002 break;
1003 }
1004
1005 case 3: {
1006 constexpr int idxX = static_cast<int>(Layout) == static_cast<int>(ColMajor) ? 0 : 2;
1007 constexpr int idxY = static_cast<int>(Layout) == static_cast<int>(ColMajor) ? 1 : 1;
1008 constexpr int idxZ = static_cast<int>(Layout) == static_cast<int>(ColMajor) ? 2 : 0;
1009
1010 const int kernel_size_x = m_kernelImpl.dimensions()[idxX];
1011 const int kernel_size_y = m_kernelImpl.dimensions()[idxY];
1012 const int kernel_size_z = m_kernelImpl.dimensions()[idxZ];
1013
1014 const int numX = dimensions()[m_indices[idxX]];
1015 const int numY = dimensions()[m_indices[idxY]];
1016 const int numZ = dimensions()[m_indices[idxZ]];
1017 const int numP = dimensions().TotalSize() / (numX * numY * numZ);
1018
1019 const int maxX = numext::mini<int>(
1020 128, numext::mini<int>(maxSharedMem / (sizeof(Scalar) * kernel_size_y * kernel_size_z) - kernel_size_x + 1,
1021 numX));
1022 const int maxY = numext::mini<int>(
1023 128, numext::mini<int>(
1024 maxSharedMem / (sizeof(Scalar) * (maxX + kernel_size_x - 1) * kernel_size_z) - kernel_size_y + 1,
1025 numY));
1026 const int maxZ = numext::mini<int>(
1027 128, numext::mini<int>(
1028 maxSharedMem / (sizeof(Scalar) * (maxX + kernel_size_x - 1) * (maxY + kernel_size_y - 1)) -
1029 kernel_size_z + 1,
1030 numZ));
1031
1032 dim3 block_size;
1033 block_size.x = numext::mini(32, maxX);
1034 block_size.y = numext::mini(32, maxY);
1035 block_size.z = numext::mini<int>(1024 / (block_size.x * block_size.y), maxZ);
1036 dim3 num_blocks(ceil(numX, maxX), ceil(numY, maxY), ceil(numZ, maxZ));
1037
1038 const int shared_mem =
1039 (maxX + kernel_size_x - 1) * (maxY + kernel_size_y - 1) * (maxZ + kernel_size_z - 1) * sizeof(Scalar);
1040 gpu_assert(shared_mem <= maxSharedMem);
1041
1042 const array<Index, 3> indices{m_indices[idxX], m_indices[idxY], m_indices[idxZ]};
1043 const array<Index, 3> kernel_dims{m_kernelImpl.dimensions()[idxX], m_kernelImpl.dimensions()[idxY],
1044 m_kernelImpl.dimensions()[idxZ]};
1045 internal::IndexMapper<Index, InputDims, 3, Layout> indexMapper(m_inputImpl.dimensions(), kernel_dims, indices);
1046
1047 LAUNCH_GPU_KERNEL((EigenConvolutionKernel3D<TensorEvaluator<InputArgType, GpuDevice>, Index, InputDims>),
1048 num_blocks, block_size, shared_mem, m_device, m_inputImpl, indexMapper, m_kernel, numP, numX,
1049 maxX, numY, maxY, numZ, maxZ, kernel_size_x, kernel_size_y, kernel_size_z, data);
1050 break;
1051 }
1052
1053 default: {
1054 EIGEN_STATIC_ASSERT((NumKernelDims >= 1 && NumKernelDims <= 3),
1055 THIS_METHOD_IS_ONLY_FOR_OBJECTS_OF_A_SPECIFIC_SIZE);
1056 }
1057 }
1058 }
1059
1060 EIGEN_DEVICE_FUNC EIGEN_STRONG_INLINE CoeffReturnType coeff(Index index) const {
1061 eigen_assert(m_buf);
1062 eigen_assert(index < m_dimensions.TotalSize());
1063 return m_buf[index];
1064 }
1065
1066 template <int LoadMode>
1067 EIGEN_DEVICE_FUNC EIGEN_STRONG_INLINE PacketReturnType packet(const Index index) const {
1068 eigen_assert(m_buf);
1069 eigen_assert(index < m_dimensions.TotalSize());
1070 return internal::ploadt<PacketReturnType, LoadMode>(m_buf + index);
1071 }
1072
1073 EIGEN_DEVICE_FUNC EIGEN_STRONG_INLINE TensorOpCost costPerCoeff(bool vectorized) const {
1074 // TODO(rmlarsen): For now, this is just a copy of the CPU cost
1075 // model.
1076 const double kernel_size = m_kernelImpl.dimensions().TotalSize();
1077 // We ignore the use of fused multiply-add.
1078 const double convolve_compute_cost = TensorOpCost::AddCost<Scalar>() + TensorOpCost::MulCost<Scalar>();
1079 const double firstIndex_compute_cost =
1080 NumDims *
1081 (2 * TensorOpCost::AddCost<Index>() + 2 * TensorOpCost::MulCost<Index>() + TensorOpCost::DivCost<Index>());
1082 return TensorOpCost(0, 0, firstIndex_compute_cost, vectorized, PacketSize) +
1083 kernel_size * (m_inputImpl.costPerCoeff(vectorized) + m_kernelImpl.costPerCoeff(vectorized) +
1084 TensorOpCost(0, 0, convolve_compute_cost, vectorized, PacketSize));
1085 }
1086
1087 private:
1088 TensorEvaluator<InputArgType, GpuDevice> m_inputImpl;
1089 TensorEvaluator<KernelArgType, GpuDevice> m_kernelImpl;
1090 KernelArgType m_kernelArg;
1091 Indices m_indices;
1092 Dimensions m_dimensions;
1093 Scalar* m_buf;
1094 const Scalar* m_kernel;
1095 bool m_local_kernel;
1096
1097 const GpuDevice& m_device;
1098};
1099#endif
1100
1101} // end namespace Eigen
1102
1103#endif // EIGEN_TENSOR_TENSOR_CONVOLUTION_H
Eigen::TensorBase
The tensor base class.
Definition TensorForwardDeclarations.h:69
Eigen::TensorConvolutionOp
Definition TensorConvolution.h:223
Eigen::TensorConvolutionOp::inputExpression
const internal::remove_all_t< typename InputXprType::Nested > & inputExpression() const
Definition TensorConvolution.h:240
Eigen::ColMajor
ColMajor
Eigen
Namespace containing all symbols from the Eigen library.
Eigen::TensorEvaluator
The tensor evaluator class.
Definition TensorEvaluator.h:47