42 #ifndef STOKHOS_CUDA_CRS_PRODUCT_TENSOR_HPP 43 #define STOKHOS_CUDA_CRS_PRODUCT_TENSOR_HPP 47 #include "Kokkos_Core.hpp" 56 #include "Teuchos_TestForException.hpp" 58 #include "cuda_profiler_api.h" 70 template<
typename TensorScalar,
71 typename MatrixScalar,
72 typename VectorScalar >
75 MatrixScalar, Kokkos::Cuda >,
76 Kokkos::View<VectorScalar**, Kokkos::LayoutLeft, Kokkos::Cuda>,
77 Kokkos::View<VectorScalar**, Kokkos::LayoutLeft, Kokkos::Cuda> >
86 typedef Kokkos::View< VectorScalar**, Kokkos::LayoutLeft, Kokkos::Cuda >
vector_type;
93 class MultiplyKernel {
105 : m_A( A ), m_x(
x ), m_y(
y ), BlockSize(block_size) {}
114 volatile VectorScalar *
const sh_x =
115 kokkos_impl_cuda_shared_memory<VectorScalar>();
116 volatile MatrixScalar *
const sh_A = sh_x + BlockSize*dim;
117 volatile VectorScalar *
const sh_y = sh_A + BlockSize*dim;
118 #if !HAVE_CUDA_SHUFFLE 119 volatile VectorScalar *
const sh_t = sh_y + dim;
122 const size_type nid = blockDim.x * blockDim.y;
123 const size_type tid = threadIdx.x + blockDim.x * threadIdx.y;
126 for (
size_type i = tid; i < dim; i += nid ) {
133 const size_type iBlockEntryEnd = m_A.
graph.row_map[ blockIdx.x + 1 ];
134 for (
size_type iBlockEntry=iBlockEntryBeg; iBlockEntry<iBlockEntryEnd;
135 iBlockEntry += BlockSize) {
137 (iBlockEntryEnd-iBlockEntry < BlockSize) ?
138 iBlockEntryEnd-iBlockEntry : BlockSize;
145 for (
size_type col = 0; col < block_size; ++col ) {
147 const size_type iBlockColumn = m_A.
graph.entries( iBlockEntry + col );
148 const VectorScalar *
const x = & m_x( 0, iBlockColumn );
149 const MatrixScalar *
const A = & m_A.
values( 0, iBlockEntry + col );
152 for (
size_type i = tid; i < dim; i += nid ) {
153 sh_x[col + i * BlockSize] =
x[i];
154 sh_A[col + i * BlockSize] = A[i];
162 for (
size_type i = threadIdx.y; i < dim; i += blockDim.y ) {
170 for (
size_type l = lBeg+threadIdx.x; l < lEnd; l += blockDim.x ) {
174 const TensorScalar v = m_A.
block.value( l );
175 const size_type j = ( kj & 0x0ffff ) * BlockSize ;
176 const size_type k = ( kj >> 16 ) * BlockSize ;
178 for (
size_type col = 0; col < block_size; ++col ) {
179 y += v * ( sh_A[col+
j] * sh_x[col+k] +
180 sh_A[col+k] * sh_x[col+
j] );
186 #if HAVE_CUDA_SHUFFLE 187 if (blockDim.x >= 2)
y +=
shfl_down(
y, 1, blockDim.x);
188 if (blockDim.x >= 4)
y +=
shfl_down(
y, 2, blockDim.x);
189 if (blockDim.x >= 8)
y +=
shfl_down(
y, 4, blockDim.x);
190 if (blockDim.x >= 16)
y +=
shfl_down(
y, 8, blockDim.x);
191 if (blockDim.x >= 32)
y +=
shfl_down(
y, 16, blockDim.x);
192 if ( threadIdx.x == 0 ) sh_y[i] +=
y;
195 if (threadIdx.x+16 < blockDim.x) sh_t[tid] += sh_t[tid+16];
196 if (threadIdx.x+ 8 < blockDim.x) sh_t[tid] += sh_t[tid+ 8];
197 if (threadIdx.x+ 4 < blockDim.x) sh_t[tid] += sh_t[tid+ 4];
198 if (threadIdx.x+ 2 < blockDim.x) sh_t[tid] += sh_t[tid+ 2];
199 if (threadIdx.x+ 1 < blockDim.x) sh_t[tid] += sh_t[tid+ 1];
200 if (threadIdx.x == 0) sh_y[i] += sh_t[tid];
211 for (
size_type i = tid; i < dim; i += nid ) {
212 m_y( i, blockIdx.x ) = sh_y[ i ];
221 class MultiplyKernel {
224 const matrix_type m_A;
225 const vector_type m_x;
226 const vector_type m_y;
227 const size_type BlockSize;
229 MultiplyKernel(
const matrix_type & A,
230 const vector_type &
x,
231 const vector_type &
y,
232 const size_type block_size )
233 : m_A( A ), m_x(
x ), m_y(
y ), BlockSize(block_size) {}
236 void operator()(
void)
const 239 const size_type dim = m_A.block.dimension();
241 volatile VectorScalar *
const sh_x =
242 kokkos_impl_cuda_shared_memory<VectorScalar>();
243 volatile VectorScalar *
const sh_y = sh_x + BlockSize*dim;
244 #if !HAVE_CUDA_SHUFFLE 245 volatile VectorScalar *
const sh_t = sh_y + dim;
248 const size_type nid = blockDim.x * blockDim.y;
249 const size_type tid = threadIdx.x + blockDim.x * threadIdx.y;
252 for ( size_type i = tid; i < dim; i += nid ) {
258 const size_type iBlockEntryBeg = m_A.graph.row_map[ blockIdx.x ];
259 const size_type iBlockEntryEnd = m_A.graph.row_map[ blockIdx.x + 1 ];
260 for (size_type iBlockEntry=iBlockEntryBeg; iBlockEntry<iBlockEntryEnd;
261 iBlockEntry += BlockSize) {
262 const size_type block_size =
263 (iBlockEntryEnd-iBlockEntry < BlockSize) ?
264 iBlockEntryEnd-iBlockEntry : BlockSize;
271 for ( size_type col = 0; col < block_size; ++col ) {
273 const size_type iBlockColumn = m_A.graph.entries( iBlockEntry + col );
274 const VectorScalar *
const x = & m_x( 0, iBlockColumn );
277 for ( size_type i = tid; i < dim; i += nid ) {
278 sh_x[col + i * BlockSize] =
x[i];
286 for ( size_type i = threadIdx.y; i < dim; i += blockDim.y ) {
290 const size_type lBeg = m_A.block.entry_begin( i );
291 const size_type lEnd = m_A.block.entry_end( i );
294 for ( size_type l = lBeg+threadIdx.x; l < lEnd; l += blockDim.x ) {
297 const size_type kj = m_A.block.coord( l );
298 const TensorScalar v = m_A.block.value( l );
299 const size_type
j = ( kj & 0x0ffff ) ;
300 const size_type k = ( kj >> 16 ) ;
302 for ( size_type col = 0; col < block_size; ++col ) {
303 const size_type bCol = iBlockEntry + col;
304 #if (__CUDA_ARCH__ >= 350) 305 y += v * ( __ldg(&m_A.values(
j,bCol)) * sh_x[col+k*BlockSize] +
306 __ldg(&m_A.values(k,bCol)) * sh_x[col+
j*BlockSize] );
308 y += v * ( m_A.values(
j,bCol) * sh_x[col+k*BlockSize] +
309 m_A.values(k,bCol) * sh_x[col+
j*BlockSize] );
316 #if HAVE_CUDA_SHUFFLE 317 if (blockDim.x >= 2)
y +=
shfl_down(
y, 1, blockDim.x);
318 if (blockDim.x >= 4)
y +=
shfl_down(
y, 2, blockDim.x);
319 if (blockDim.x >= 8)
y +=
shfl_down(
y, 4, blockDim.x);
320 if (blockDim.x >= 16)
y +=
shfl_down(
y, 8, blockDim.x);
321 if (blockDim.x >= 32)
y +=
shfl_down(
y, 16, blockDim.x);
322 if ( threadIdx.x == 0 ) sh_y[i] +=
y;
325 if (threadIdx.x+16 < blockDim.x) sh_t[tid] += sh_t[tid+16];
326 if (threadIdx.x+ 8 < blockDim.x) sh_t[tid] += sh_t[tid+ 8];
327 if (threadIdx.x+ 4 < blockDim.x) sh_t[tid] += sh_t[tid+ 4];
328 if (threadIdx.x+ 2 < blockDim.x) sh_t[tid] += sh_t[tid+ 2];
329 if (threadIdx.x+ 1 < blockDim.x) sh_t[tid] += sh_t[tid+ 1];
330 if (threadIdx.x == 0) sh_y[i] += sh_t[tid];
341 for ( size_type i = tid; i < dim; i += nid ) {
342 m_y( i, blockIdx.x ) = sh_y[ i ];
351 struct TensorReadEntry {
362 const size_type tensor_align = tensor_dimension;
363 const size_type avg_tensor_entries_per_row = A.
block.avg_entries_per_row();
389 Kokkos::Impl::cuda_parallel_launch_local_memory<MultiplyKernel>);
391 (warp_size*regs_per_thread + reg_bank_size-1) & ~(reg_bank_size-1);
393 (max_regs_per_sm/regs_per_warp) & ~(warp_granularity-1);
395 (max_regs_per_block/regs_per_warp) & ~(warp_granularity-1);
404 avg_tensor_entries_per_row >= 88 ? 32 : 16;
405 const size_type rows_per_warp = warp_size / threads_per_row;
407 const size_type vec_scalar_size =
sizeof(VectorScalar);
409 const size_type mat_scalar_size =
sizeof(MatrixScalar);
412 #define USE_FIXED_BLOCKSIZE 0 414 #if USE_FIXED_BLOCKSIZE 417 size_type nw = warps_per_sm / num_blocks;
418 while (nw > 1 && num_blocks*nw % warp_granularity) --nw;
420 const size_type sh_per_block = shcap / num_blocks;
422 device_prop.
has_shuffle ? 0 : vec_scalar_size*warp_size*num_warp;
424 size_type bs = ((sh_per_block - sr) / tensor_align - vec_scalar_size) /
427 size_type bs = ((sh_per_block - sr) / tensor_align - vec_scalar_size) /
428 (vec_scalar_size+mat_scalar_size);
430 if (bs % 2 == 0) --bs;
436 ( (vec_scalar_size*block_size+vec_scalar_size)*tensor_align + sr + sh_granularity-1 ) & ~(sh_granularity-1);
439 ( ((vec_scalar_size+mat_scalar_size)*block_size+vec_scalar_size)*tensor_align + sr + sh_granularity-1 ) & ~(sh_granularity-1);
454 const size_type half_nnz_per_row = fem_nnz_per_row / 2 + 1;
456 half_nnz_per_row % 2 ? half_nnz_per_row + 1 : half_nnz_per_row;
457 Teuchos::Array<TensorReadEntry> reads_per_thread;
458 for (
size_type bs = block_size_min; bs<=block_size_max; bs+=2) {
462 device_prop.
has_shuffle ? 0 : vec_scalar_size*warp_size*warps_per_block;
465 (vec_scalar_size*bs+vec_scalar_size)*tensor_align+sr;
468 ((vec_scalar_size+mat_scalar_size)*bs+vec_scalar_size)*tensor_align+sr;
470 shmem = (shmem + sh_granularity-1) & ~(sh_granularity-1);
471 if (shmem <= max_shmem_per_block) {
473 size_type tensor_reads = (fem_nnz_per_row+bs-1) / bs;
476 min_warps_per_block),
477 max_warps_per_block);
478 while (num_warp > 1 && num_blocks*num_warp % warp_granularity)
480 TensorReadEntry entry;
481 entry.block_size = bs;
483 entry.num_blocks = num_blocks;
484 entry.num_warp = num_warp;
488 entry.reads = (
static_cast<double>(tensor_reads) /
489 static_cast<double>(factor*num_blocks*num_warp));
490 reads_per_thread.push_back(entry);
493 TEUCHOS_TEST_FOR_EXCEPTION(
494 reads_per_thread.size() == 0, std::logic_error,
495 "Stochastic problem dimension is too large to fit in shared memory");
497 double min_reads = reads_per_thread[0].reads;
498 for (
int i=1; i<reads_per_thread.size(); ++i) {
499 if (reads_per_thread[i].reads < min_reads) {
501 min_reads = reads_per_thread[i].reads;
505 const size_type block_size = reads_per_thread[idx].block_size;
506 const size_type shmem = reads_per_thread[idx].shmem;
507 const size_type num_blocks = reads_per_thread[idx].num_blocks;
508 const size_type num_warp = reads_per_thread[idx].num_warp;
513 const dim3 dBlock( threads_per_row , rows_per_warp*num_warp , 1 );
514 const dim3 dGrid( row_count, 1, 1 );
517 std::cout <<
"block_size = " << block_size
518 <<
" tensor reads = " << (fem_nnz_per_row+block_size-1)/block_size
519 <<
" regs_per_thread = " << regs_per_thread
520 <<
" num blocks = " << num_blocks
521 <<
" num warps = " << num_warp
522 <<
" num rows = " << tensor_dimension
523 <<
" rows/warp = " << tensor_dimension / (num_warp*rows_per_warp)
524 <<
" avg entries/row = " << avg_tensor_entries_per_row
530 Kokkos::Impl::cuda_parallel_launch_local_memory<<< dGrid, dBlock, shmem >>>
531 ( MultiplyKernel( A,
x,
y, block_size ) );
__device__ void operator()(void) const
KOKKOS_INLINE_FUNCTION Scalar shfl_down(const Scalar &val, const int &delta, const int &width)
Kokkos::View< VectorScalar **, Kokkos::LayoutLeft, Kokkos::Cuda > vector_type
BlockCrsMatrix< tensor_type, MatrixScalar, execution_space > matrix_type
size_type max_blocks_per_sm
static void apply(const matrix_type &A, const vector_type &x, const vector_type &y)
size_type warp_granularity
const size_type BlockSize
size_type max_regs_per_sm
CrsProductTensor< TensorScalar, execution_space > tensor_type
size_type max_shmem_per_block
KOKKOS_INLINE_FUNCTION PCE< Storage > min(const typename PCE< Storage >::value_type &a, const PCE< Storage > &b)
Sparse product tensor with replicated entries to provide subsets with a given coordinate.
size_type shared_memory_granularity
const IndexType const IndexType const IndexType const IndexType const ValueType const ValueType * x
KOKKOS_INLINE_FUNCTION PCE< Storage > max(const typename PCE< Storage >::value_type &a, const PCE< Storage > &b)
Top-level namespace for Stokhos classes and functions.
size_type shared_memory_capacity
size_type max_threads_per_block
size_type max_warps_per_sm
size_type max_regs_per_block
MultiplyKernel(const matrix_type &A, const vector_type &x, const vector_type &y, const size_type block_size)
CRS matrix of dense blocks.
execution_space::size_type size_type
size_type get_kernel_registers(Kernel kernel)
const IndexType const IndexType const IndexType const IndexType const ValueType const ValueType ValueType * y
Kokkos::Cuda execution_space