We still support CUDA down to 10 so maybe wrap this code if we have t…

…he right compute capability to use it.
lwa-project · Jul 17, 2024 · 0b219b6 · 0b219b6
1 parent 8243297
commit 0b219b6
Showing 1 changed file with 11 additions and 0 deletions.
diff --git a/src/transpose.cu b/src/transpose.cu
@@ -146,6 +146,17 @@ BFstatus transpose(int           ndim,
 	                    sizes[0]*out_strides[0] <
 	                    (long)std::numeric_limits<int>::max());
 #if BF_CUDA_ENABLED
+#if BF_GPU_MIN_ARCH < 40
+	if( ELEMENT_SIZE ==  6 ||
+	    ELEMENT_SIZE ==  8 ||
+	    ELEMENT_SIZE == 16 ) {
+		// TODO: Doing this here might be a bad idea
+		cudaDeviceSetSharedMemConfig(cudaSharedMemBankSizeEightByte);
+	}
+	else {
+		cudaDeviceSetSharedMemConfig(cudaSharedMemBankSizeFourByte);
+	}
+#endif
 	if( can_use_int ) {
 		kernel::transpose
 			<TILE_DIM,BLOCK_ROWS,CONDITIONAL_WRITE,