@@ -93,6 +93,90 @@ __device__ inline int64_t ullitolli(uint64_t u)
9393#define ATOMICADDF32 (pAccumulator, value ) atomicAdd (pAccumulator, (value))
9494#define ATOMICSUBF32 (pAccumulator, value ) atomicAdd (pAccumulator, -(value))
9595
96+ #ifdef USE_NVTENSOR
97+ /* Begin: Reduction using tensor units */
98+
99+ // Implementation based on M.Sc. thesis by Gabin Schieffer at KTH:
100+ // "Accelerating a Molecular Docking Application by Leveraging Modern Heterogeneous Computing Systemx"
101+ // https://www.diva-portal.org/smash/get/diva2:1786161/FULLTEXT01.pdf
102+
103+ /*
104+ * WMMA Extension for single precision matmul using Tensor Cores
105+ * and error correction technique (TCEC)
106+ * https://github.com/wmmae/wmma_extension/blob/main/docs/mma_f32.md
107+ */
108+ #include < wmma_extension/tcec/tcec.hpp>
109+ using tf32 = nvcuda::wmma::precision::tf32;
110+
111+ /*
112+ * Tensor Cores
113+ * https://developer.nvidia.com/blog/programming-tensor-cores-cuda-9
114+ *
115+ * Don't forget to compile specifying the architecture, e.g., sm_86.
116+ * For AutoDock-GPU, this can be done via the TARGETS option.
117+ * make DEVICE=GPU TESTLS=ad NUMWI=64 TARGETS=86 test
118+ * https://stackoverflow.com/a/53634598/1616865
119+ */
120+ #include < mma.h>
121+ using namespace nvcuda ;
122+
123+ #define TILE_SIZE (16 * 16 )
124+
125+ constexpr int rowscols_M = 16 ; // Number of rows (or cols) in the M dimension
126+ constexpr int rowscols_N = 16 ; // Number of rows (or cols) in the N dimension
127+ constexpr int rowscols_K = 16 ; // Number of rows (or cols) in the K dimension
128+
129+ __device__ void reduce_via_tensor_units (float *data_to_be_reduced) {
130+ __syncthreads ();
131+
132+ if (threadIdx .x <= 31 ) { // Only one warp performs reduction
133+ __shared__ __align__ (256 ) float Q_square[TILE_SIZE ]; // storage for 16x16 matrix and 4x4 tiles of I4 matrix after
134+
135+ // Declaring and filling fragments - Those are *not* shared
136+ mtk::wmma::tcec::fragment<wmma::matrix_b, rowscols_M, rowscols_N, rowscols_K, tf32, wmma::col_major> frag_P;
137+ mtk::wmma::tcec::fragment<wmma::accumulator, rowscols_M, rowscols_N, rowscols_K, tf32> frag_V;
138+
139+ mtk::wmma::tcec::fragment<wmma::matrix_a, rowscols_M, rowscols_N, rowscols_K, tf32, wmma::col_major> frag_Q;
140+ mtk::wmma::tcec::fragment<wmma::matrix_b, rowscols_M, rowscols_N, rowscols_K, tf32, wmma::col_major> frag_W;
141+ mtk::wmma::tcec::fragment<wmma::accumulator, rowscols_M, rowscols_N, rowscols_K, tf32> frag_C;
142+
143+ mtk::wmma::tcec::fill_fragment (frag_P, 1 .0f ); // P: only ones
144+ mtk::wmma::tcec::fill_fragment (frag_V, 0 .0f ); // Output: initialize to zeros
145+ mtk::wmma::tcec::fill_fragment (frag_C, 0 .0f ); // Final result
146+
147+ // 1. Accumulate the values: V <- AP + V
148+ for (uint i = 0 ; i < (4 * NUM_OF_THREADS_PER_BLOCK )/TILE_SIZE ; i++){
149+ const unsigned int offset = i * TILE_SIZE ;
150+
151+ mtk::wmma::tcec::fragment<wmma::matrix_a, rowscols_M, rowscols_N, rowscols_K, tf32, wmma::col_major> frag_A;
152+ mtk::wmma::tcec::load_matrix_sync (frag_A, data_to_be_reduced + offset, 16 );
153+ mtk::wmma::tcec::mma_sync (frag_V, frag_A, frag_P, frag_V);
154+ }
155+
156+ // W <- V (required since we need V as a "wmma::matrix_b")
157+ mtk::wmma::tcec::store_matrix_sync (Q_square, frag_V, 16 , wmma::mem_col_major);
158+ mtk::wmma::tcec::load_matrix_sync (frag_W, Q_square, 16 );
159+
160+ // 2. Perform line sum: C <- QW + C (zero)
161+ // a) create a 4x4 tiled matrix containing 4x4 identity matrix in each tile:
162+ // - TENSOR=ON requires NUMWI to be larger than 32, so the following works and neatly gets rid of an additional function:
163+ const unsigned int k = (threadIdx .x <<3 );
164+ const unsigned int kk = 16 - (threadIdx .x >>1 );
165+ for (uint i = 0 ; i < 8 ; i++) Q_square[k + i] = ((i + kk) & 3 ) ? 0 .0f : 1 .0f ;
166+ mtk::wmma::tcec::load_matrix_sync (frag_Q, Q_square, 16 );
167+ // b) perform sum
168+ mtk::wmma::tcec::mma_sync (frag_C, frag_Q, frag_W, frag_C);
169+
170+ // 3. Store result in shared memory
171+ mtk::wmma::tcec::store_matrix_sync (data_to_be_reduced, frag_C, 16 , wmma::mem_col_major);
172+ }
173+
174+ __syncthreads ();
175+ }
176+
177+ /* End: Reduction using tensor units */
178+ #endif
179+
96180#define REDUCEFLOATSUM (value, pAccumulator ) \
97181 if (threadIdx .x == 0 ) \
98182 { \
0 commit comments