| // This file is part of Eigen, a lightweight C++ template library |
| // for linear algebra. |
| // |
| // Copyright (C) 2014 Benoit Steiner <benoit.steiner.goog@gmail.com> |
| // |
| // This Source Code Form is subject to the terms of the Mozilla |
| // Public License v. 2.0. If a copy of the MPL was not distributed |
| // with this file, You can obtain one at http://mozilla.org/MPL/2.0/. |
| |
| #ifndef EIGEN_CXX11_TENSOR_TENSOR_CONTRACTION_CUDA_H |
| #define EIGEN_CXX11_TENSOR_TENSOR_CONTRACTION_CUDA_H |
| |
| #if defined(EIGEN_USE_GPU) && defined(__CUDACC__) |
| |
| namespace Eigen { |
| |
| template<typename Scalar, typename Index, typename LhsMapper, |
| typename RhsMapper, typename OutputMapper, bool needs_edge_check> |
| __device__ EIGEN_STRONG_INLINE void |
| EigenContractionKernelInternal(const LhsMapper lhs, const RhsMapper rhs, |
| const OutputMapper output, volatile Scalar* lhs_shmem, volatile Scalar* rhs_shmem, |
| const Index m_size, const Index n_size, const Index k_size) { |
| |
| const Index m_block_idx = blockIdx.x; |
| const Index n_block_idx = blockIdx.y; |
| |
| const Index base_m = 64 * m_block_idx; |
| const Index base_n = 64 * n_block_idx; |
| |
| // declare and initialize 64 registers for output 8x8 block |
| |
| // prefetch registers |
| Scalar lhs_pf0; |
| Scalar lhs_pf1; |
| Scalar lhs_pf2; |
| Scalar lhs_pf3; |
| Scalar lhs_pf4; |
| Scalar lhs_pf5; |
| Scalar lhs_pf6; |
| Scalar lhs_pf7; |
| |
| Scalar rhs_pf0; |
| Scalar rhs_pf1; |
| Scalar rhs_pf2; |
| Scalar rhs_pf3; |
| Scalar rhs_pf4; |
| Scalar rhs_pf5; |
| Scalar rhs_pf6; |
| Scalar rhs_pf7; |
| |
| // shared memory is formatted |
| // (contract idx in block, nocontract idx in block, block idx) |
| // where block idx is column major. This transposition limits the number of |
| // bank conflicts when reading the LHS. The core idea is that since the contracting |
| // index is shared by both sides, then the contracting index should be in threadIdx.x. |
| |
| // On the LHS, we pad each row inside of each block with an extra element. This makes |
| // each block 8 rows of 9 elements, which is 72 elements. This gives no bank conflicts |
| // on writes and very few 2-way conflicts on reads. There is an 8x8 grid of these blocks. |
| |
| // On the RHS we just add 8 padding elements to the end of each block. This gives no bank |
| // conflicts on writes and also none on reads. |
| |
| // storage indices |
| const Index lhs_store_idx_base = threadIdx.y * 72 + threadIdx.x * 9 + threadIdx.z; |
| const Index rhs_store_idx_base = threadIdx.y * 72 + threadIdx.z * 8 + threadIdx.x; |
| |
| const Index lhs_store_idx_0 = lhs_store_idx_base + 576 * 0; |
| const Index lhs_store_idx_1 = lhs_store_idx_base + 576 * 1; |
| const Index lhs_store_idx_2 = lhs_store_idx_base + 576 * 2; |
| const Index lhs_store_idx_3 = lhs_store_idx_base + 576 * 3; |
| const Index lhs_store_idx_4 = lhs_store_idx_base + 576 * 4; |
| const Index lhs_store_idx_5 = lhs_store_idx_base + 576 * 5; |
| const Index lhs_store_idx_6 = lhs_store_idx_base + 576 * 6; |
| const Index lhs_store_idx_7 = lhs_store_idx_base + 576 * 7; |
| |
| const Index rhs_store_idx_0 = rhs_store_idx_base + 576 * 0; |
| const Index rhs_store_idx_1 = rhs_store_idx_base + 576 * 1; |
| const Index rhs_store_idx_2 = rhs_store_idx_base + 576 * 2; |
| const Index rhs_store_idx_3 = rhs_store_idx_base + 576 * 3; |
| const Index rhs_store_idx_4 = rhs_store_idx_base + 576 * 4; |
| const Index rhs_store_idx_5 = rhs_store_idx_base + 576 * 5; |
| const Index rhs_store_idx_6 = rhs_store_idx_base + 576 * 6; |
| const Index rhs_store_idx_7 = rhs_store_idx_base + 576 * 7; |
| |
| // in the loading code, the following variables are important: |
| // threadIdx.x: the vertical position in an 8x8 block |
| // threadIdx.y: the vertical index of the 8x8 block in the grid |
| // threadIdx.z: the horizontal position in an 8x8 block |
| // k: the horizontal index of the 8x8 block in the grid |
| // |
| // The k parameter is implicit (it was the loop counter for a loop that went |
| // from 0 to <8, but now that loop is unrolled in the below code. |
| |
| const Index load_idx_vert = threadIdx.x + 8 * threadIdx.y; |
| const Index lhs_vert = base_m + load_idx_vert; |
| |
| #define prefetchIntoRegisters(base_k) \ |
| { \ |
| lhs_pf0 = Scalar(0); \ |
| lhs_pf1 = Scalar(0); \ |
| lhs_pf2 = Scalar(0); \ |
| lhs_pf3 = Scalar(0); \ |
| lhs_pf4 = Scalar(0); \ |
| lhs_pf5 = Scalar(0); \ |
| lhs_pf6 = Scalar(0); \ |
| lhs_pf7 = Scalar(0); \ |
| \ |
| rhs_pf0 = Scalar(0); \ |
| rhs_pf1 = Scalar(0); \ |
| rhs_pf2 = Scalar(0); \ |
| rhs_pf3 = Scalar(0); \ |
| rhs_pf4 = Scalar(0); \ |
| rhs_pf5 = Scalar(0); \ |
| rhs_pf6 = Scalar(0); \ |
| rhs_pf7 = Scalar(0); \ |
| \ |
| if (!needs_edge_check || lhs_vert < m_size) { \ |
| const Index lhs_horiz_0 = base_k + threadIdx.z + 0 * 8; \ |
| const Index lhs_horiz_1 = base_k + threadIdx.z + 1 * 8; \ |
| const Index lhs_horiz_2 = base_k + threadIdx.z + 2 * 8; \ |
| const Index lhs_horiz_3 = base_k + threadIdx.z + 3 * 8; \ |
| const Index lhs_horiz_4 = base_k + threadIdx.z + 4 * 8; \ |
| const Index lhs_horiz_5 = base_k + threadIdx.z + 5 * 8; \ |
| const Index lhs_horiz_6 = base_k + threadIdx.z + 6 * 8; \ |
| const Index lhs_horiz_7 = base_k + threadIdx.z + 7 * 8; \ |
| \ |
| if (!needs_edge_check || lhs_horiz_7 < k_size) { \ |
| lhs_pf0 = lhs(lhs_vert, lhs_horiz_0); \ |
| lhs_pf1 = lhs(lhs_vert, lhs_horiz_1); \ |
| lhs_pf2 = lhs(lhs_vert, lhs_horiz_2); \ |
| lhs_pf3 = lhs(lhs_vert, lhs_horiz_3); \ |
| lhs_pf4 = lhs(lhs_vert, lhs_horiz_4); \ |
| lhs_pf5 = lhs(lhs_vert, lhs_horiz_5); \ |
| lhs_pf6 = lhs(lhs_vert, lhs_horiz_6); \ |
| lhs_pf7 = lhs(lhs_vert, lhs_horiz_7); \ |
| } else if (lhs_horiz_6 < k_size) { \ |
| lhs_pf0 = lhs(lhs_vert, lhs_horiz_0); \ |
| lhs_pf1 = lhs(lhs_vert, lhs_horiz_1); \ |
| lhs_pf2 = lhs(lhs_vert, lhs_horiz_2); \ |
| lhs_pf3 = lhs(lhs_vert, lhs_horiz_3); \ |
| lhs_pf4 = lhs(lhs_vert, lhs_horiz_4); \ |
| lhs_pf5 = lhs(lhs_vert, lhs_horiz_5); \ |
| lhs_pf6 = lhs(lhs_vert, lhs_horiz_6); \ |
| } else if (lhs_horiz_5 < k_size) { \ |
| lhs_pf0 = lhs(lhs_vert, lhs_horiz_0); \ |
| lhs_pf1 = lhs(lhs_vert, lhs_horiz_1); \ |
| lhs_pf2 = lhs(lhs_vert, lhs_horiz_2); \ |
| lhs_pf3 = lhs(lhs_vert, lhs_horiz_3); \ |
| lhs_pf4 = lhs(lhs_vert, lhs_horiz_4); \ |
| lhs_pf5 = lhs(lhs_vert, lhs_horiz_5); \ |
| } else if (lhs_horiz_4 < k_size) { \ |
| lhs_pf0 = lhs(lhs_vert, lhs_horiz_0); \ |
| lhs_pf1 = lhs(lhs_vert, lhs_horiz_1); \ |
| lhs_pf2 = lhs(lhs_vert, lhs_horiz_2); \ |
| lhs_pf3 = lhs(lhs_vert, lhs_horiz_3); \ |
| lhs_pf4 = lhs(lhs_vert, lhs_horiz_4); \ |
| } else if (lhs_horiz_3 < k_size) { \ |
| lhs_pf0 = lhs(lhs_vert, lhs_horiz_0); \ |
| lhs_pf1 = lhs(lhs_vert, lhs_horiz_1); \ |
| lhs_pf2 = lhs(lhs_vert, lhs_horiz_2); \ |
| lhs_pf3 = lhs(lhs_vert, lhs_horiz_3); \ |
| } else if (lhs_horiz_2 < k_size) { \ |
| lhs_pf0 = lhs(lhs_vert, lhs_horiz_0); \ |
| lhs_pf1 = lhs(lhs_vert, lhs_horiz_1); \ |
| lhs_pf2 = lhs(lhs_vert, lhs_horiz_2); \ |
| } else if (lhs_horiz_1 < k_size) { \ |
| lhs_pf0 = lhs(lhs_vert, lhs_horiz_0); \ |
| lhs_pf1 = lhs(lhs_vert, lhs_horiz_1); \ |
| } else if (lhs_horiz_0 < k_size) { \ |
| lhs_pf0 = lhs(lhs_vert, lhs_horiz_0); \ |
| } \ |
| } \ |
| \ |
| const Index rhs_vert = base_k + load_idx_vert; \ |
| if (!needs_edge_check || rhs_vert < k_size) { \ |
| const Index rhs_horiz_0 = base_n + threadIdx.z + 0 * 8; \ |
| const Index rhs_horiz_1 = base_n + threadIdx.z + 1 * 8; \ |
| const Index rhs_horiz_2 = base_n + threadIdx.z + 2 * 8; \ |
| const Index rhs_horiz_3 = base_n + threadIdx.z + 3 * 8; \ |
| const Index rhs_horiz_4 = base_n + threadIdx.z + 4 * 8; \ |
| const Index rhs_horiz_5 = base_n + threadIdx.z + 5 * 8; \ |
| const Index rhs_horiz_6 = base_n + threadIdx.z + 6 * 8; \ |
| const Index rhs_horiz_7 = base_n + threadIdx.z + 7 * 8; \ |
| \ |
| if (rhs_horiz_7 < n_size) { \ |
| rhs_pf0 = rhs(rhs_vert, rhs_horiz_0); \ |
| rhs_pf1 = rhs(rhs_vert, rhs_horiz_1); \ |
| rhs_pf2 = rhs(rhs_vert, rhs_horiz_2); \ |
| rhs_pf3 = rhs(rhs_vert, rhs_horiz_3); \ |
| rhs_pf4 = rhs(rhs_vert, rhs_horiz_4); \ |
| rhs_pf5 = rhs(rhs_vert, rhs_horiz_5); \ |
| rhs_pf6 = rhs(rhs_vert, rhs_horiz_6); \ |
| rhs_pf7 = rhs(rhs_vert, rhs_horiz_7); \ |
| } else if (rhs_horiz_6 < n_size) { \ |
| rhs_pf0 = rhs(rhs_vert, rhs_horiz_0); \ |
| rhs_pf1 = rhs(rhs_vert, rhs_horiz_1); \ |
| rhs_pf2 = rhs(rhs_vert, rhs_horiz_2); \ |
| rhs_pf3 = rhs(rhs_vert, rhs_horiz_3); \ |
| rhs_pf4 = rhs(rhs_vert, rhs_horiz_4); \ |
| rhs_pf5 = rhs(rhs_vert, rhs_horiz_5); \ |
| rhs_pf6 = rhs(rhs_vert, rhs_horiz_6); \ |
| } else if (rhs_horiz_5 < n_size) { \ |
| rhs_pf0 = rhs(rhs_vert, rhs_horiz_0); \ |
| rhs_pf1 = rhs(rhs_vert, rhs_horiz_1); \ |
| rhs_pf2 = rhs(rhs_vert, rhs_horiz_2); \ |
| rhs_pf3 = rhs(rhs_vert, rhs_horiz_3); \ |
| rhs_pf4 = rhs(rhs_vert, rhs_horiz_4); \ |
| rhs_pf5 = rhs(rhs_vert, rhs_horiz_5); \ |
| } else if (rhs_horiz_4 < n_size) { \ |
| rhs_pf0 = rhs(rhs_vert, rhs_horiz_0); \ |
| rhs_pf1 = rhs(rhs_vert, rhs_horiz_1); \ |
| rhs_pf2 = rhs(rhs_vert, rhs_horiz_2); \ |
| rhs_pf3 = rhs(rhs_vert, rhs_horiz_3); \ |
| rhs_pf4 = rhs(rhs_vert, rhs_horiz_4); \ |
| } else if (rhs_horiz_3 < n_size) { \ |
| rhs_pf0 = rhs(rhs_vert, rhs_horiz_0); \ |
| rhs_pf1 = rhs(rhs_vert, rhs_horiz_1); \ |
| rhs_pf2 = rhs(rhs_vert, rhs_horiz_2); \ |
| rhs_pf3 = rhs(rhs_vert, rhs_horiz_3); \ |
| } else if (rhs_horiz_2 < n_size) { \ |
| rhs_pf0 = rhs(rhs_vert, rhs_horiz_0); \ |
| rhs_pf1 = rhs(rhs_vert, rhs_horiz_1); \ |
| rhs_pf2 = rhs(rhs_vert, rhs_horiz_2); \ |
| } else if (rhs_horiz_1 < n_size) { \ |
| rhs_pf0 = rhs(rhs_vert, rhs_horiz_0); \ |
| rhs_pf1 = rhs(rhs_vert, rhs_horiz_1); \ |
| } else if (rhs_horiz_0 < n_size) { \ |
| rhs_pf0 = rhs(rhs_vert, rhs_horiz_0); \ |
| } \ |
| } \ |
| } \ |
| |
| #define writeRegToShmem(_) \ |
| lhs_shmem[lhs_store_idx_0] = lhs_pf0; \ |
| rhs_shmem[rhs_store_idx_0] = rhs_pf0; \ |
| \ |
| lhs_shmem[lhs_store_idx_1] = lhs_pf1; \ |
| rhs_shmem[rhs_store_idx_1] = rhs_pf1; \ |
| \ |
| lhs_shmem[lhs_store_idx_2] = lhs_pf2; \ |
| rhs_shmem[rhs_store_idx_2] = rhs_pf2; \ |
| \ |
| lhs_shmem[lhs_store_idx_3] = lhs_pf3; \ |
| rhs_shmem[rhs_store_idx_3] = rhs_pf3; \ |
| \ |
| lhs_shmem[lhs_store_idx_4] = lhs_pf4; \ |
| rhs_shmem[rhs_store_idx_4] = rhs_pf4; \ |
| \ |
| lhs_shmem[lhs_store_idx_5] = lhs_pf5; \ |
| rhs_shmem[rhs_store_idx_5] = rhs_pf5; \ |
| \ |
| lhs_shmem[lhs_store_idx_6] = lhs_pf6; \ |
| rhs_shmem[rhs_store_idx_6] = rhs_pf6; \ |
| \ |
| lhs_shmem[lhs_store_idx_7] = lhs_pf7; \ |
| rhs_shmem[rhs_store_idx_7] = rhs_pf7; \ |
| |
| // declare and initialize result array |
| #define res(i, j) _res_##i##j |
| #define initResultRow(i) \ |
| Scalar res(i, 0) = Scalar(0); \ |
| Scalar res(i, 1) = Scalar(0); \ |
| Scalar res(i, 2) = Scalar(0); \ |
| Scalar res(i, 3) = Scalar(0); \ |
| Scalar res(i, 4) = Scalar(0); \ |
| Scalar res(i, 5) = Scalar(0); \ |
| Scalar res(i, 6) = Scalar(0); \ |
| Scalar res(i, 7) = Scalar(0); \ |
| |
| initResultRow(0); |
| initResultRow(1); |
| initResultRow(2); |
| initResultRow(3); |
| initResultRow(4); |
| initResultRow(5); |
| initResultRow(6); |
| initResultRow(7); |
| #undef initResultRow |
| |
| for (Index base_k = 0; base_k < k_size; base_k += 64) { |
| // wait for previous iteration to finish with shmem. Despite common sense, |
| // the code is a bit faster with this here then at bottom of loop |
| __syncthreads(); |
| |
| prefetchIntoRegisters(base_k); |
| writeRegToShmem(); |
| |
| #undef prefetchIntoRegisters |
| #undef writeRegToShmem |
| |
| // wait for shared mem packing to be done before starting computation |
| __syncthreads(); |
| |
| // compute 8x8 matrix product by outer product. This involves packing one column |
| // of LHS and one row of RHS into registers (takes 16 registers). |
| |
| #define lcol(i) _lcol##i |
| Scalar lcol(0); |
| Scalar lcol(1); |
| Scalar lcol(2); |
| Scalar lcol(3); |
| Scalar lcol(4); |
| Scalar lcol(5); |
| Scalar lcol(6); |
| Scalar lcol(7); |
| |
| #define rrow(j) _rrow##j |
| Scalar rrow(0); |
| Scalar rrow(1); |
| Scalar rrow(2); |
| Scalar rrow(3); |
| Scalar rrow(4); |
| Scalar rrow(5); |
| Scalar rrow(6); |
| Scalar rrow(7); |
| |
| // Now x corresponds to k, y to m, and z to n |
| const volatile Scalar* lhs_block = &lhs_shmem[threadIdx.x + 9 * threadIdx.y]; |
| const volatile Scalar* rhs_block = &rhs_shmem[threadIdx.x + 8 * threadIdx.z]; |
| |
| #define lhs_element(i, j) lhs_block[72 * ((i) + 8 * (j))] |
| #define rhs_element(i, j) rhs_block[72 * ((i) + 8 * (j))] |
| |
| #define loadData(i, j) \ |
| lcol(0) = lhs_element(0, j); \ |
| rrow(0) = rhs_element(i, 0); \ |
| lcol(1) = lhs_element(1, j); \ |
| rrow(1) = rhs_element(i, 1); \ |
| lcol(2) = lhs_element(2, j); \ |
| rrow(2) = rhs_element(i, 2); \ |
| lcol(3) = lhs_element(3, j); \ |
| rrow(3) = rhs_element(i, 3); \ |
| lcol(4) = lhs_element(4, j); \ |
| rrow(4) = rhs_element(i, 4); \ |
| lcol(5) = lhs_element(5, j); \ |
| rrow(5) = rhs_element(i, 5); \ |
| lcol(6) = lhs_element(6, j); \ |
| rrow(6) = rhs_element(i, 6); \ |
| lcol(7) = lhs_element(7, j); \ |
| rrow(7) = rhs_element(i, 7); \ |
| |
| #define computeCol(j) \ |
| res(0, j) += lcol(0) * rrow(j); \ |
| res(1, j) += lcol(1) * rrow(j); \ |
| res(2, j) += lcol(2) * rrow(j); \ |
| res(3, j) += lcol(3) * rrow(j); \ |
| res(4, j) += lcol(4) * rrow(j); \ |
| res(5, j) += lcol(5) * rrow(j); \ |
| res(6, j) += lcol(6) * rrow(j); \ |
| res(7, j) += lcol(7) * rrow(j); \ |
| |
| #define computePass(i) \ |
| loadData(i, i); \ |
| \ |
| computeCol(0); \ |
| computeCol(1); \ |
| computeCol(2); \ |
| computeCol(3); \ |
| computeCol(4); \ |
| computeCol(5); \ |
| computeCol(6); \ |
| computeCol(7); \ |
| |
| computePass(0); |
| computePass(1); |
| computePass(2); |
| computePass(3); |
| computePass(4); |
| computePass(5); |
| computePass(6); |
| computePass(7); |
| |
| #undef lcol |
| #undef rrow |
| #undef lhs_element |
| #undef rhs_element |
| #undef loadData |
| #undef computeCol |
| #undef computePass |
| } // end loop over k |
| |
| // we've now iterated over all of the large (ie width 64) k blocks and |
| // accumulated results in registers. At this point thread (x, y, z) contains |
| // the sum across all big k blocks of the product of little k block of index (x, y) |
| // with block of index (y, z). To compute the final output, we need to reduce |
| // the 8 threads over y by summation. |
| #define shuffleInc(i, j, mask) res(i, j) += __shfl_xor(res(i, j), mask) |
| |
| #define reduceRow(i, mask) \ |
| shuffleInc(i, 0, mask); \ |
| shuffleInc(i, 1, mask); \ |
| shuffleInc(i, 2, mask); \ |
| shuffleInc(i, 3, mask); \ |
| shuffleInc(i, 4, mask); \ |
| shuffleInc(i, 5, mask); \ |
| shuffleInc(i, 6, mask); \ |
| shuffleInc(i, 7, mask); \ |
| |
| #define reduceMatrix(mask) \ |
| reduceRow(0, mask); \ |
| reduceRow(1, mask); \ |
| reduceRow(2, mask); \ |
| reduceRow(3, mask); \ |
| reduceRow(4, mask); \ |
| reduceRow(5, mask); \ |
| reduceRow(6, mask); \ |
| reduceRow(7, mask); \ |
| |
| // actually perform the reduction, now each thread of index (_, y, z) |
| // contains the correct values in its registers that belong in the output |
| // block |
| reduceMatrix(1); |
| reduceMatrix(2); |
| reduceMatrix(4); |
| |
| #undef shuffleInc |
| #undef reduceRow |
| #undef reduceMatrix |
| |
| // now we need to copy the 64 values into main memory. We can't split work |
| // among threads because all variables are in registers. There's 2 ways |
| // to do this: |
| // (1) have 1 thread do 64 writes from registers into global memory |
| // (2) have 1 thread do 64 writes into shared memory, and then 8 threads |
| // each do 8 writes into global memory. We can just overwrite the shared |
| // memory from the problem we just solved. |
| // (2) is slightly faster than (1) due to less branching and more ILP |
| |
| // TODO: won't yield much gain, but could just use currently unused shared mem |
| // and then we won't have to sync |
| // wait for shared mem to be out of use |
| __syncthreads(); |
| |
| #define writeResultShmem(i, j) \ |
| lhs_shmem[i + 8 * threadIdx.y + 64 * threadIdx.z + 512 * j] = res(i, j); \ |
| |
| #define writeRow(i) \ |
| writeResultShmem(i, 0); \ |
| writeResultShmem(i, 1); \ |
| writeResultShmem(i, 2); \ |
| writeResultShmem(i, 3); \ |
| writeResultShmem(i, 4); \ |
| writeResultShmem(i, 5); \ |
| writeResultShmem(i, 6); \ |
| writeResultShmem(i, 7); \ |
| |
| if (threadIdx.x == 0) { |
| writeRow(0); |
| writeRow(1); |
| writeRow(2); |
| writeRow(3); |
| writeRow(4); |
| writeRow(5); |
| writeRow(6); |
| writeRow(7); |
| } |
| #undef writeResultShmem |
| #undef writeRow |
| |
| const int max_i_write = (min)((int)((m_size - base_m - threadIdx.y + 7) / 8), 8); |
| const int max_j_write = (min)((int)((n_size - base_n - threadIdx.z + 7) / 8), 8); |
| |
| if (threadIdx.x < max_i_write) { |
| if (max_j_write == 8) { |
| Scalar val0 = lhs_shmem[threadIdx.x + 8 * threadIdx.y + 64 * threadIdx.z + 512 * 0]; |
| Scalar val1 = lhs_shmem[threadIdx.x + 8 * threadIdx.y + 64 * threadIdx.z + 512 * 1]; |
| Scalar val2 = lhs_shmem[threadIdx.x + 8 * threadIdx.y + 64 * threadIdx.z + 512 * 2]; |
| Scalar val3 = lhs_shmem[threadIdx.x + 8 * threadIdx.y + 64 * threadIdx.z + 512 * 3]; |
| Scalar val4 = lhs_shmem[threadIdx.x + 8 * threadIdx.y + 64 * threadIdx.z + 512 * 4]; |
| Scalar val5 = lhs_shmem[threadIdx.x + 8 * threadIdx.y + 64 * threadIdx.z + 512 * 5]; |
| Scalar val6 = lhs_shmem[threadIdx.x + 8 * threadIdx.y + 64 * threadIdx.z + 512 * 6]; |
| Scalar val7 = lhs_shmem[threadIdx.x + 8 * threadIdx.y + 64 * threadIdx.z + 512 * 7]; |
| |
| output(base_m + threadIdx.y + 8 * threadIdx.x, base_n + threadIdx.z + 8 * 0) = val0; |
| output(base_m + threadIdx.y + 8 * threadIdx.x, base_n + threadIdx.z + 8 * 1) = val1; |
| output(base_m + threadIdx.y + 8 * threadIdx.x, base_n + threadIdx.z + 8 * 2) = val2; |
| output(base_m + threadIdx.y + 8 * threadIdx.x, base_n + threadIdx.z + 8 * 3) = val3; |
| output(base_m + threadIdx.y + 8 * threadIdx.x, base_n + threadIdx.z + 8 * 4) = val4; |
| output(base_m + threadIdx.y + 8 * threadIdx.x, base_n + threadIdx.z + 8 * 5) = val5; |
| output(base_m + threadIdx.y + 8 * threadIdx.x, base_n + threadIdx.z + 8 * 6) = val6; |
| output(base_m + threadIdx.y + 8 * threadIdx.x, base_n + threadIdx.z + 8 * 7) = val7; |
| } else { |
| #pragma unroll 7 |
| for (int j = 0; j < max_j_write; j++) { |
| Scalar val = lhs_shmem[threadIdx.x + 8 * threadIdx.y + 64 * threadIdx.z + 512 * j]; |
| output(base_m + threadIdx.y + 8 * threadIdx.x, base_n + threadIdx.z + 8 * j) = val; |
| } |
| } |
| } |
| #undef res |
| } |
| |
| |
| template<typename Scalar, typename Index, typename LhsMapper, |
| typename RhsMapper, typename OutputMapper> |
| __global__ void |
| __launch_bounds__(512) |
| EigenContractionKernel(const LhsMapper lhs, const RhsMapper rhs, |
| const OutputMapper output, |
| const Index m_size, const Index n_size, const Index k_size) { |
| __shared__ volatile Scalar lhs_shmem[72 * 64]; |
| __shared__ volatile Scalar rhs_shmem[72 * 64]; |
| |
| const Index m_block_idx = blockIdx.x; |
| const Index n_block_idx = blockIdx.y; |
| |
| const Index base_m = 64 * m_block_idx; |
| const Index base_n = 64 * n_block_idx; |
| |
| if (base_m + 63 < m_size && base_n + 63 < n_size) { |
| EigenContractionKernelInternal<Scalar, Index, LhsMapper, RhsMapper, OutputMapper, false>(lhs, rhs, output, lhs_shmem, rhs_shmem, m_size, n_size, k_size); |
| } else { |
| EigenContractionKernelInternal<Scalar, Index, LhsMapper, RhsMapper, OutputMapper, true>(lhs, rhs, output, lhs_shmem, rhs_shmem, m_size, n_size, k_size); |
| } |
| } |
| |
| |
| |
| template<typename Index, typename LhsMapper, |
| typename RhsMapper, typename OutputMapper, bool needs_edge_check> |
| __device__ EIGEN_STRONG_INLINE void |
| EigenFloatContractionKernelInternal(const LhsMapper lhs, const RhsMapper rhs, |
| const OutputMapper output, float4* lhs_shmem4, float2* rhs_shmem2, |
| const Index m_size, const Index n_size, const Index k_size) { |
| typedef float Scalar; |
| |
| const Index m_block_idx = blockIdx.x; |
| const Index n_block_idx = blockIdx.y; |
| |
| const Index base_m = 64 * m_block_idx; |
| const Index base_n = 64 * n_block_idx; |
| |
| const Index lane = threadIdx.x + 8 * (threadIdx.y % 4); |
| |
| // prefetch registers |
| float4 lhs_pf0; |
| float4 lhs_pf1; |
| |
| float4 rhs_pf0; |
| float4 rhs_pf1; |
| |
| // shared memory is formatted |
| // (contract idx in block, nocontract idx in block, block idx) |
| // where block idx is column major. This transposition limits the number of |
| // bank conflicts when reading the LHS. The core idea is that since the contracting |
| // index is shared by both sides, then the contracting index should be in threadIdx.x. |
| |
| // all of these indices assume float4 loading |
| // this thread loads the float4 starting at this index, and then also loads |
| // another float4 starting 32 columns to to the right |
| const Index horiz_block_idx = threadIdx.z / 2; |
| const Index vert_block_idx = threadIdx.x / 2 + 4 * (threadIdx.y % 2); |
| const Index horiz_idx_in_block = threadIdx.y / 2 + 4 * (threadIdx.z % 2); |
| const Index vert_idx_in_block = threadIdx.x % 2; |
| |
| // there's padding in both the LHS and RHS shared memory layouts. This padding |
| // allows for 0 bank conflicts on all shmem stores and loads. |
| // LHS padding: 1 float4 on each 8x8 block of floats |
| // RHS padding: 1 float2 on each block, and 12 additional float2s between vertical blocks |
| // 3 and 4 |
| |
| // storage indices |
| // lhs index with respect to float4s |
| const Index lhs_store_idx_base = |
| 136 * horiz_block_idx + |
| 17 * vert_block_idx + |
| 8 * vert_idx_in_block + |
| horiz_idx_in_block; |
| |
| // rhs index with respect to floats |
| const Index rhs_store_idx_base = |
| 552 * horiz_block_idx + |
| 66 * vert_block_idx + |
| 32 * (horiz_idx_in_block / 4) + (horiz_idx_in_block % 4) + |
| 16 * vert_idx_in_block + |
| ((vert_block_idx < 4) ? 0 : 24); |
| |
| const Index lhs_store_idx_0 = lhs_store_idx_base + 544 * 0; |
| const Index lhs_store_idx_1 = lhs_store_idx_base + 544 * 1; |
| |
| const Index rhs_store_idx_0 = (rhs_store_idx_base / 2) + ((lane < 16) ? 0 : 4); |
| const Index rhs_store_idx_1 = rhs_store_idx_0 + 2; |
| const Index rhs_store_idx_2 = rhs_store_idx_0 + 1104; |
| const Index rhs_store_idx_3 = rhs_store_idx_1 + 1104; |
| |
| // The below diagrams show which shmem index (with respect to floats) each element |
| // in an 8x8 input block gets packed into: |
| // LHS: |
| // 0 4 8 12 16 20 24 28 |
| // 1 5 9 13 17 21 25 29 |
| // 2 6 10 14 18 22 26 30 |
| // 3 7 11 15 19 23 27 31 |
| // 32 36 40 44 48 52 56 60 |
| // ... (pack as 2 rows of float4 indexed row major, each float4 is vertical) |
| // |
| // RHS: |
| // 0 1 2 3 32 33 34 35 |
| // 4 5 6 7 36 37 38 39 |
| // ... (pack as 2 cols of float4 indexed col major, each float4 is horizontal) |
| |
| // Each thread in a warp loads 2 float4s. This happens in 2 instructions. On each of these |
| // instruction, the warp loads 2 columns (2 cols * 64 elements / col = 128 elements = 32 threads |
| // * 4 elements/thread). For the LHS, we're able to store the loaded float4 directly into |
| // shmem (using a 128 bit store instruction). For the RHS, we need to transpose the data. |
| // This is done with warp shuffles. Furthermore, we only use 64 bit stores for the RHS, because |
| // 64 bits is only 2 columns (which is all we load in a warp), and the padding for the RHS |
| // doesn't meet 64 bit alignment requirements (namely, the 4 consecutive floats that we want |
| // to load on the RHS are 8 byte aligned, not 16 byte aligned, which is required for float4). |
| |
| const Index load_idx_vert = 4 * (threadIdx.x + 8 * (threadIdx.y % 2)); |
| const Index load_idx_horiz = (threadIdx.y / 2) + 4 * threadIdx.z; |
| |
| const Index lhs_vert = base_m + load_idx_vert; |
| const Index rhs_horiz_0 = base_n + load_idx_horiz; |
| const Index rhs_horiz_1 = base_n + load_idx_horiz + 32; |
| |
| #define prefetchIntoRegisters(base_k) \ |
| { \ |
| lhs_pf0 = internal::pset1<float4>(0); \ |
| lhs_pf1 = internal::pset1<float4>(0); \ |
| \ |
| rhs_pf0 = internal::pset1<float4>(0); \ |
| rhs_pf1 = internal::pset1<float4>(0); \ |
| \ |
| const Index lhs_horiz_0 = base_k + load_idx_horiz; \ |
| const Index lhs_horiz_1 = base_k + load_idx_horiz + 32; \ |
| if (!needs_edge_check || lhs_vert + 3 < m_size) { \ |
| if (lhs_horiz_1 < k_size) { \ |
| lhs_pf0 = lhs.loadPacket(lhs_vert, lhs_horiz_0); \ |
| lhs_pf1 = lhs.loadPacket(lhs_vert, lhs_horiz_1); \ |
| } else if (lhs_horiz_0 < k_size) { \ |
| lhs_pf0 = lhs.loadPacket(lhs_vert, lhs_horiz_0); \ |
| } \ |
| } else if (lhs_vert + 2 < m_size) { \ |
| if (lhs_horiz_1 < k_size) { \ |
| lhs_pf0.x = lhs(lhs_vert + 0, lhs_horiz_0); \ |
| lhs_pf0.y = lhs(lhs_vert + 1, lhs_horiz_0); \ |
| lhs_pf0.z = lhs(lhs_vert + 2, lhs_horiz_0); \ |
| \ |
| lhs_pf1.x = lhs(lhs_vert + 0, lhs_horiz_1); \ |
| lhs_pf1.y = lhs(lhs_vert + 1, lhs_horiz_1); \ |
| lhs_pf1.z = lhs(lhs_vert + 2, lhs_horiz_1); \ |
| } else if (lhs_horiz_0 < k_size) { \ |
| lhs_pf0.x = lhs(lhs_vert + 0, lhs_horiz_0); \ |
| lhs_pf0.y = lhs(lhs_vert + 1, lhs_horiz_0); \ |
| lhs_pf0.z = lhs(lhs_vert + 2, lhs_horiz_0); \ |
| } \ |
| } else if (lhs_vert + 1 < m_size) { \ |
| if (lhs_horiz_1 < k_size) { \ |
| lhs_pf0.x = lhs(lhs_vert + 0, lhs_horiz_0); \ |
| lhs_pf0.y = lhs(lhs_vert + 1, lhs_horiz_0); \ |
| \ |
| lhs_pf1.x = lhs(lhs_vert + 0, lhs_horiz_1); \ |
| lhs_pf1.y = lhs(lhs_vert + 1, lhs_horiz_1); \ |
| } else if (lhs_horiz_0 < k_size) { \ |
| lhs_pf0.x = lhs(lhs_vert + 0, lhs_horiz_0); \ |
| lhs_pf0.y = lhs(lhs_vert + 1, lhs_horiz_0); \ |
| } \ |
| } else if (lhs_vert < m_size) { \ |
| if (lhs_horiz_1 < k_size) { \ |
| lhs_pf0.x = lhs(lhs_vert + 0, lhs_horiz_0); \ |
| lhs_pf1.x = lhs(lhs_vert + 0, lhs_horiz_1); \ |
| } else if (lhs_horiz_0 < k_size) { \ |
| lhs_pf0.x = lhs(lhs_vert + 0, lhs_horiz_0); \ |
| } \ |
| } \ |
| \ |
| const Index rhs_vert = base_k + load_idx_vert; \ |
| if (rhs_vert + 3 < k_size) { \ |
| if (!needs_edge_check || rhs_horiz_1 < n_size) { \ |
| rhs_pf0 = rhs.loadPacket(rhs_vert, rhs_horiz_0); \ |
| rhs_pf1 = rhs.loadPacket(rhs_vert, rhs_horiz_1); \ |
| } else if (rhs_horiz_0 < n_size) { \ |
| rhs_pf0 = rhs.loadPacket(rhs_vert, rhs_horiz_0); \ |
| } \ |
| } else if (rhs_vert + 2 < k_size) { \ |
| if (!needs_edge_check || rhs_horiz_1 < n_size) { \ |
| rhs_pf0.x = rhs(rhs_vert + 0, rhs_horiz_0); \ |
| rhs_pf0.y = rhs(rhs_vert + 1, rhs_horiz_0); \ |
| rhs_pf0.z = rhs(rhs_vert + 2, rhs_horiz_0); \ |
| \ |
| rhs_pf1.x = rhs(rhs_vert + 0, rhs_horiz_1); \ |
| rhs_pf1.y = rhs(rhs_vert + 1, rhs_horiz_1); \ |
| rhs_pf1.z = rhs(rhs_vert + 2, rhs_horiz_1); \ |
| } else if (rhs_horiz_0 < n_size) { \ |
| rhs_pf0.x = rhs(rhs_vert + 0, rhs_horiz_0); \ |
| rhs_pf0.y = rhs(rhs_vert + 1, rhs_horiz_0); \ |
| rhs_pf0.z = rhs(rhs_vert + 2, rhs_horiz_0); \ |
| } \ |
| } else if (rhs_vert + 1 < k_size) { \ |
| if (!needs_edge_check || rhs_horiz_1 < n_size) { \ |
| rhs_pf0.x = rhs(rhs_vert + 0, rhs_horiz_0); \ |
| rhs_pf0.y = rhs(rhs_vert + 1, rhs_horiz_0); \ |
| \ |
| rhs_pf1.x = rhs(rhs_vert + 0, rhs_horiz_1); \ |
| rhs_pf1.y = rhs(rhs_vert + 1, rhs_horiz_1); \ |
| } else if (rhs_horiz_0 < n_size) { \ |
| rhs_pf0.x = rhs(rhs_vert + 0, rhs_horiz_0); \ |
| rhs_pf0.y = rhs(rhs_vert + 1, rhs_horiz_0); \ |
| } \ |
| } else if (rhs_vert < k_size) { \ |
| if (!needs_edge_check || rhs_horiz_1 < n_size) { \ |
| rhs_pf0.x = rhs(rhs_vert + 0, rhs_horiz_0); \ |
| rhs_pf1.x = rhs(rhs_vert + 0, rhs_horiz_1); \ |
| } else if (rhs_horiz_0 < n_size) { \ |
| rhs_pf0.x = rhs(rhs_vert + 0, rhs_horiz_0); \ |
| } \ |
| } \ |
| \ |
| float swap_val0 = (lane < 16) ? rhs_pf0.z : rhs_pf0.x; \ |
| float swap_val1 = (lane < 16) ? rhs_pf0.w : rhs_pf0.y; \ |
| float swap_val2 = (lane < 16) ? rhs_pf1.z : rhs_pf1.x; \ |
| float swap_val3 = (lane < 16) ? rhs_pf1.w : rhs_pf1.y; \ |
| \ |
| swap_val0 = __shfl_xor(swap_val0, 16); \ |
| swap_val1 = __shfl_xor(swap_val1, 16); \ |
| swap_val2 = __shfl_xor(swap_val2, 16); \ |
| swap_val3 = __shfl_xor(swap_val3, 16); \ |
| \ |
| if (lane < 16) { \ |
| rhs_pf0.z = swap_val0; \ |
| rhs_pf0.w = swap_val1; \ |
| rhs_pf1.z = swap_val2; \ |
| rhs_pf1.w = swap_val3; \ |
| } else { \ |
| rhs_pf0.x = swap_val0; \ |
| rhs_pf0.y = swap_val1; \ |
| rhs_pf1.x = swap_val2; \ |
| rhs_pf1.y = swap_val3; \ |
| } \ |
| } \ |
| |
| |
| #define writeRegToShmem(_) \ |
| lhs_shmem4[lhs_store_idx_0] = lhs_pf0; \ |
| \ |
| rhs_shmem2[rhs_store_idx_0] = make_float2(rhs_pf0.x, rhs_pf0.z); \ |
| rhs_shmem2[rhs_store_idx_1] = make_float2(rhs_pf0.y, rhs_pf0.w); \ |
| \ |
| lhs_shmem4[lhs_store_idx_1] = lhs_pf1; \ |
| \ |
| rhs_shmem2[rhs_store_idx_2] = make_float2(rhs_pf1.x, rhs_pf1.z); \ |
| rhs_shmem2[rhs_store_idx_3] = make_float2(rhs_pf1.y, rhs_pf1.w); \ |
| |
| // declare and initialize result array |
| #define res(i, j) _res_##i##j |
| #define initResultRow(i) \ |
| Scalar res(i, 0) = Scalar(0); \ |
| Scalar res(i, 1) = Scalar(0); \ |
| Scalar res(i, 2) = Scalar(0); \ |
| Scalar res(i, 3) = Scalar(0); \ |
| Scalar res(i, 4) = Scalar(0); \ |
| Scalar res(i, 5) = Scalar(0); \ |
| Scalar res(i, 6) = Scalar(0); \ |
| Scalar res(i, 7) = Scalar(0); \ |
| |
| initResultRow(0); |
| initResultRow(1); |
| initResultRow(2); |
| initResultRow(3); |
| initResultRow(4); |
| initResultRow(5); |
| initResultRow(6); |
| initResultRow(7); |
| #undef initResultRow |
| |
| for (Index base_k = 0; base_k < k_size; base_k += 64) { |
| // wait for previous iteration to finish with shmem. Despite common sense, |
| // the code is a bit faster with this here then at bottom of loop |
| __syncthreads(); |
| |
| prefetchIntoRegisters(base_k); |
| writeRegToShmem(); |
| |
| #undef prefetchIntoRegisters |
| #undef writeRegoToShmem |
| |
| // wait for shared mem packing to be done before starting computation |
| __syncthreads(); |
| |
| // compute 8x8 matrix product by outer product. This involves packing one column |
| // of LHS and one row of RHS into registers (takes 16 registers). |
| |
| float4 _lcol0; |
| float4 _lcol1; |
| float2 _rrow0; |
| float2 _rrow1; |
| float2 _rrow2; |
| float2 _rrow3; |
| |
| #define lcol0 _lcol0.x |
| #define lcol1 _lcol0.y |
| #define lcol2 _lcol0.z |
| #define lcol3 _lcol0.w |
| #define lcol4 _lcol1.x |
| #define lcol5 _lcol1.y |
| #define lcol6 _lcol1.z |
| #define lcol7 _lcol1.w |
| #define rrow0 _rrow0.x |
| #define rrow1 _rrow0.y |
| #define rrow2 _rrow1.x |
| #define rrow3 _rrow1.y |
| #define rrow4 _rrow2.x |
| #define rrow5 _rrow2.y |
| #define rrow6 _rrow3.x |
| #define rrow7 _rrow3.y |
| |
| // Now x corresponds to k, y to m, and z to n |
| const float4* lhs_block = &lhs_shmem4[threadIdx.x + 8 * (threadIdx.y % 2) + 17 * (threadIdx.y / 2)]; |
| const float2* rhs_block = &rhs_shmem2[2 * threadIdx.x + 16 * (threadIdx.z % 2) + 276 * (threadIdx.z / 2)]; |
| |
| #define lhs_element(i, k) lhs_block[68 * i + 136 * k] |
| #define rhs_element(k, j) rhs_block[33 * k + 1104 * j + ((k < 4) ? 0 : 12)] |
| |
| #define loadData(i) \ |
| _lcol0 = lhs_element(0, i); \ |
| _rrow0 = rhs_element(i, 0); \ |
| _rrow1 = *(&(rhs_element(i, 0)) + 1); \ |
| _lcol1 = lhs_element(1, i); \ |
| _rrow2 = rhs_element(i, 1); \ |
| _rrow3 = *(&(rhs_element(i, 1)) + 1); \ |
| |
| #define computeCol(j) \ |
| res(0, j) += lcol0 * rrow##j; \ |
| res(1, j) += lcol1 * rrow##j; \ |
| res(2, j) += lcol2 * rrow##j; \ |
| res(3, j) += lcol3 * rrow##j; \ |
| res(4, j) += lcol4 * rrow##j; \ |
| res(5, j) += lcol5 * rrow##j; \ |
| res(6, j) += lcol6 * rrow##j; \ |
| res(7, j) += lcol7 * rrow##j; \ |
| |
| #define computePass(i) \ |
| loadData(i); \ |
| \ |
| computeCol(0); \ |
| computeCol(1); \ |
| computeCol(2); \ |
| computeCol(3); \ |
| computeCol(4); \ |
| computeCol(5); \ |
| computeCol(6); \ |
| computeCol(7); \ |
| |
| computePass(0); |
| computePass(1); |
| computePass(2); |
| computePass(3); |
| computePass(4); |
| computePass(5); |
| computePass(6); |
| computePass(7); |
| |
| #undef lcol0 |
| #undef lcol1 |
| #undef lcol2 |
| #undef lcol3 |
| #undef lcol4 |
| #undef lcol5 |
| #undef lcol6 |
| #undef lcol7 |
| #undef rrow0 |
| #undef rrow1 |
| #undef rrow2 |
| #undef rrow3 |
| #undef rrow4 |
| #undef rrow5 |
| #undef rrow6 |
| #undef rrow7 |
| |
| #undef computePass |
| #undef computeCol |
| #undef loadData |
| #undef lhs_element |
| #undef rhs_element |
| |
| } // end loop over k |
| |
| // we've now iterated over all of the large (ie width 64) k blocks and |
| // accumulated results in registers. At this point thread (x, y, z) contains |
| // the sum across all big k blocks of the product of little k block of index (x, y) |
| // with block of index (y, z). To compute the final output, we need to reduce |
| // the 8 threads over y by summation. |
| #define shuffleInc(i, j, mask) res(i, j) += __shfl_xor(res(i, j), mask) |
| |
| #define reduceRow(i, mask) \ |
| shuffleInc(i, 0, mask); \ |
| shuffleInc(i, 1, mask); \ |
| shuffleInc(i, 2, mask); \ |
| shuffleInc(i, 3, mask); \ |
| shuffleInc(i, 4, mask); \ |
| shuffleInc(i, 5, mask); \ |
| shuffleInc(i, 6, mask); \ |
| shuffleInc(i, 7, mask); \ |
| |
| #define reduceMatrix(mask) \ |
| reduceRow(0, mask); \ |
| reduceRow(1, mask); \ |
| reduceRow(2, mask); \ |
| reduceRow(3, mask); \ |
| reduceRow(4, mask); \ |
| reduceRow(5, mask); \ |
| reduceRow(6, mask); \ |
| reduceRow(7, mask); \ |
| |
| // actually perform the reduction, now each thread of index (_, y, z) |
| // contains the correct values in its registers that belong in the output |
| // block |
| reduceMatrix(1); |
| reduceMatrix(2); |
| reduceMatrix(4); |
| |
| #undef shuffleInc |
| #undef reduceRow |
| #undef reduceMatrix |
| |
| // now we need to copy the 64 values into main memory. We can't split work |
| // among threads because all variables are in registers. There's 2 ways |
| // to do this: |
| // (1) have 1 thread do 64 writes from registers into global memory |
| // (2) have 1 thread do 64 writes into shared memory, and then 8 threads |
| // each do 8 writes into global memory. We can just overwrite the shared |
| // memory from the problem we just solved. |
| // (3) Copies the values into new registers using conditional logic. |
| |
| #define makeAssignments(i) \ |
| val0 = res(i, 0); \ |
| val1 = res(i, 1); \ |
| val2 = res(i, 2); \ |
| val3 = res(i, 3); \ |
| val4 = res(i, 4); \ |
| val5 = res(i, 5); \ |
| val6 = res(i, 6); \ |
| val7 = res(i, 7); \ |
| |
| Scalar val0; |
| Scalar val1; |
| Scalar val2; |
| Scalar val3; |
| Scalar val4; |
| Scalar val5; |
| Scalar val6; |
| Scalar val7; |
| |
| switch (threadIdx.x) { |
| case 0: |
| makeAssignments(0); |
| break; |
| case 1: |
| makeAssignments(1); |
| break; |
| case 2: |
| makeAssignments(2); |
| break; |
| case 3: |
| makeAssignments(3); |
| break; |
| case 4: |
| makeAssignments(4); |
| break; |
| case 5: |
| makeAssignments(5); |
| break; |
| case 6: |
| makeAssignments(6); |
| break; |
| case 7: |
| makeAssignments(7); |
| break; |
| } |
| |
| #undef res |
| |
| const Index vert_base = base_m + 4 * threadIdx.y + (threadIdx.x % 4) + 32 * (threadIdx.x / 4); |
| const Index horiz_base = base_n + 4 * threadIdx.z; |
| |
| if (!needs_edge_check || vert_base < m_size) { |
| if (!needs_edge_check || horiz_base + 35 < n_size) { |
| output(vert_base, horiz_base + 0) = val0; |
| output(vert_base, horiz_base + 1) = val1; |
| output(vert_base, horiz_base + 2) = val2; |
| output(vert_base, horiz_base + 3) = val3; |
| output(vert_base, horiz_base + 32) = val4; |
| output(vert_base, horiz_base + 33) = val5; |
| output(vert_base, horiz_base + 34) = val6; |
| output(vert_base, horiz_base + 35) = val7; |
| } else if (horiz_base + 34 < n_size) { |
| output(vert_base, horiz_base + 0) = val0; |
| output(vert_base, horiz_base + 1) = val1; |
| output(vert_base, horiz_base + 2) = val2; |
| output(vert_base, horiz_base + 3) = val3; |
| output(vert_base, horiz_base + 32) = val4; |
| output(vert_base, horiz_base + 33) = val5; |
| output(vert_base, horiz_base + 34) = val6; |
| } else if (horiz_base + 33 < n_size) { |
| output(vert_base, horiz_base + 0) = val0; |
| output(vert_base, horiz_base + 1) = val1; |
| output(vert_base, horiz_base + 2) = val2; |
| output(vert_base, horiz_base + 3) = val3; |
| output(vert_base, horiz_base + 32) = val4; |
| output(vert_base, horiz_base + 33) = val5; |
| } else if (horiz_base + 32 < n_size) { |
| output(vert_base, horiz_base + 0) = val0; |
| output(vert_base, horiz_base + 1) = val1; |
| output(vert_base, horiz_base + 2) = val2; |
| output(vert_base, horiz_base + 3) = val3; |
| output(vert_base, horiz_base + 32) = val4; |
| } else if (horiz_base + 3 < n_size) { |
| output(vert_base, horiz_base + 0) = val0; |
| output(vert_base, horiz_base + 1) = val1; |
| output(vert_base, horiz_base + 2) = val2; |
| output(vert_base, horiz_base + 3) = val3; |
| } else if (horiz_base + 2 < n_size) { |
| output(vert_base, horiz_base + 0) = val0; |
| output(vert_base, horiz_base + 1) = val1; |
| output(vert_base, horiz_base + 2) = val2; |
| } else if (horiz_base + 1 < n_size) { |
| output(vert_base, horiz_base + 0) = val0; |
| output(vert_base, horiz_base + 1) = val1; |
| } else if (horiz_base < n_size) { |
| output(vert_base, horiz_base + 0) = val0; |
| } |
| } |
| } |
| |
| |
| template<typename Index, typename LhsMapper, |
| typename RhsMapper, typename OutputMapper> |
| __global__ void |
| __launch_bounds__(512) |
| EigenFloatContractionKernel(const LhsMapper lhs, const RhsMapper rhs, |
| const OutputMapper output, |
| const Index m_size, const Index n_size, const Index k_size) { |
| __shared__ float4 lhs_shmem[(68 * 64) / 4]; |
| __shared__ float2 rhs_shmem[((66 * 8 + 24) * 8) / 2]; |
| |
| const Index m_block_idx = blockIdx.x; |
| const Index n_block_idx = blockIdx.y; |
| |
| const Index base_m = 64 * m_block_idx; |
| const Index base_n = 64 * n_block_idx; |
| |
| if (base_m + 63 < m_size && base_n + 63 < n_size) { |
| EigenFloatContractionKernelInternal<Index, LhsMapper, RhsMapper, OutputMapper, false>(lhs, rhs, output, lhs_shmem, rhs_shmem, m_size, n_size, k_size); |
| } else { |
| EigenFloatContractionKernelInternal<Index, LhsMapper, RhsMapper, OutputMapper, true>(lhs, rhs, output, lhs_shmem, rhs_shmem, m_size, n_size, k_size); |
| } |
| } |
| |
| |
| template<typename Indices, typename LeftArgType, typename RightArgType> |
| struct TensorEvaluator<const TensorContractionOp<Indices, LeftArgType, RightArgType>, GpuDevice> : |
| public TensorContractionEvaluatorBase<TensorEvaluator<const TensorContractionOp<Indices, LeftArgType, RightArgType>, GpuDevice> > { |
| |
| typedef GpuDevice Device; |
| |
| typedef TensorEvaluator<const TensorContractionOp<Indices, LeftArgType, RightArgType>, Device> Self; |
| typedef TensorContractionEvaluatorBase<Self> Base; |
| |
| typedef TensorContractionOp<Indices, LeftArgType, RightArgType> XprType; |
| typedef typename internal::remove_const<typename XprType::Scalar>::type Scalar; |
| typedef typename XprType::Packet Packet; |
| typedef typename XprType::Index Index; |
| typedef typename XprType::CoeffReturnType CoeffReturnType; |
| typedef typename XprType::PacketReturnType PacketReturnType; |
| |
| typedef array<Index, TensorEvaluator<LeftArgType, Device>::Dimensions::count> left_dim_mapper_t; |
| typedef array<Index, TensorEvaluator<RightArgType, Device>::Dimensions::count> right_dim_mapper_t; |
| |
| typedef array<Index, internal::array_size<Indices>::value> contract_t; |
| typedef array<Index, TensorEvaluator<LeftArgType, Device>::Dimensions::count - internal::array_size<Indices>::value> left_nocontract_t; |
| typedef array<Index, TensorEvaluator<RightArgType, Device>::Dimensions::count - internal::array_size<Indices>::value> right_nocontract_t; |
| |
| static const int NumDims = max_n_1<TensorEvaluator<LeftArgType, Device>::Dimensions::count + TensorEvaluator<RightArgType, Device>::Dimensions::count - 2 * internal::array_size<Indices>::value>::size; |
| |
| typedef DSizes<Index, NumDims> Dimensions; |
| |
| // typedefs needed in evalTo |
| typedef typename internal::remove_const<typename LeftArgType::Scalar>::type LhsScalar; |
| typedef typename internal::remove_const<typename RightArgType::Scalar>::type RhsScalar; |
| |
| typedef TensorEvaluator<LeftArgType, Device> LeftEvaluator; |
| typedef TensorEvaluator<RightArgType, Device> RightEvaluator; |
| |
| typedef typename LeftEvaluator::Dimensions LeftDimensions; |
| typedef typename RightEvaluator::Dimensions RightDimensions; |
| |
| EIGEN_DEVICE_FUNC TensorEvaluator(const XprType& op, const Device& device) : |
| Base(op, device) {} |
| |
| // We need to redefine this method to make nvcc happy |
| EIGEN_STRONG_INLINE bool evalSubExprsIfNeeded(Scalar* data) { |
| this->m_leftImpl.evalSubExprsIfNeeded(NULL); |
| this->m_rightImpl.evalSubExprsIfNeeded(NULL); |
| if (data) { |
| evalTo(data); |
| return false; |
| } else { |
| this->m_result = static_cast<Scalar *>(this->m_device.allocate(this->dimensions().TotalSize() * sizeof(Scalar))); |
| evalTo(this->m_result); |
| return true; |
| } |
| } |
| |
| void evalTo(Scalar* buffer) const { |
| if (this->m_lhs_inner_dim_contiguous) { |
| if (this->m_rhs_inner_dim_contiguous) { |
| if (this->m_rhs_inner_dim_reordered) { |
| evalTyped<true, true, true, Unaligned>(buffer); |
| } |
| else { |
| evalTyped<true, true, false, Unaligned>(buffer); |
| } |
| } |
| else { |
| if (this->m_rhs_inner_dim_reordered) { |
| evalTyped<true, false, true, Unaligned>(buffer); |
| } |
| else { |
| evalTyped<true, false, false, Unaligned>(buffer); |
| } |
| } |
| } |
| else { |
| if (this->m_rhs_inner_dim_contiguous) { |
| if (this->m_rhs_inner_dim_reordered) { |
| evalTyped<false, true, true, Unaligned>(buffer); |
| } |
| else { |
| evalTyped<false, true, false, Unaligned>(buffer); |
| } |
| } |
| else { |
| if (this->m_rhs_inner_dim_reordered) { |
| evalTyped<false, false, true, Unaligned>(buffer); |
| } |
| else { |
| evalTyped<false, false, false, Unaligned>(buffer); |
| } |
| } |
| } |
| } |
| |
| template <bool lhs_inner_dim_contiguous, bool rhs_inner_dim_contiguous, bool rhs_inner_dim_reordered, int Alignment> |
| void evalTyped(Scalar* buffer) const { |
| // columns in left side, rows in right side |
| const Index k = this->m_k_size; |
| |
| // rows in left side |
| const Index m = this->m_i_size; |
| |
| // columns in right side |
| const Index n = this->m_j_size; |
| |
| // zero out the result buffer (which must be of size at least m * n * sizeof(Scalar) |
| this->m_device.memset(buffer, 0, m * n * sizeof(Scalar)); |
| |
| typedef internal::TensorContractionInputMapper<LhsScalar, Index, internal::Lhs, |
| LeftEvaluator, left_nocontract_t, |
| contract_t, 4, |
| lhs_inner_dim_contiguous, |
| false, Unaligned> LhsMapper; |
| |
| typedef internal::TensorContractionInputMapper<RhsScalar, Index, internal::Rhs, |
| RightEvaluator, right_nocontract_t, |
| contract_t, 4, |
| rhs_inner_dim_contiguous, |
| rhs_inner_dim_reordered, Unaligned> RhsMapper; |
| |
| typedef internal::blas_data_mapper<Scalar, Index, ColMajor> OutputMapper; |
| |
| |
| // initialize data mappers |
| LhsMapper lhs(this->m_leftImpl, this->m_left_nocontract_strides, this->m_i_strides, |
| this->m_left_contracting_strides, this->m_k_strides); |
| |
| RhsMapper rhs(this->m_rightImpl, this->m_right_nocontract_strides, this->m_j_strides, |
| this->m_right_contracting_strides, this->m_k_strides); |
| |
| OutputMapper output(buffer, m); |
| |
| const Index m_blocks = (m + 63) / 64; |
| const Index n_blocks = (n + 63) / 64; |
| const dim3 num_blocks(m_blocks, n_blocks, 1); |
| const dim3 block_size(8, 8, 8); |
| |
| cudaDeviceSetSharedMemConfig(cudaSharedMemBankSizeEightByte); |
| if (internal::is_same<LhsScalar, float>::value && |
| internal::is_same<RhsScalar, float>::value) { |
| EigenFloatContractionKernel<Index, LhsMapper, RhsMapper, OutputMapper> |
| <<<num_blocks, block_size, 0, this->m_device.stream()>>>(lhs, rhs, output, m, n, k); |
| } else { |
| EigenContractionKernel<Scalar, Index, LhsMapper, RhsMapper, OutputMapper> |
| <<<num_blocks, block_size, 0, this->m_device.stream()>>>(lhs, rhs, output, m, n, k); |
| } |
| |
| assert(cudaGetLastError() == cudaSuccess); |
| } |
| }; |
| |
| } // end namespace Eigen |
| |
| #endif // EIGEN_USE_GPU and __CUDACC__ |
| |
| #endif // EIGEN_CXX11_TENSOR_TENSOR_CONTRACTION_CUDA_H |
