aparler/PyNorch
1
0
Fork 0
Code Issues Pull requests Projects Releases Packages Wiki Activity Actions

Fix sum tensor reduce cuda

Browse source
This commit is contained in:
lucasdelimanogueira 2024-05-01 14:44:09 -03:00
parent 592dacbde9
commit 2fd5b75ff0
4 changed files with 24 additions and 23 deletions
Download patch file Download diff file Expand all files Collapse all files

BIN
build/cuda.cu.o
View file

Binary file not shown.

BIN
build/libtensor.so
View file

Binary file not shown.

41
norch/csrc/cuda.cu
View file

@ -76,38 +76,39 @@ __global__ void sum_tensor_cuda_kernel(float* data, float* result_data, int size
}
}
__global__ void aux_sum_block_kernel(float* result_data, int size) {
extern __shared__ float sdata[];
__global__ void sum_tensor_cuda_kernel(float* data, float* result_data) {
unsigned int tid = threadIdx.x;
__shared__ int sdata[THREADS_PER_BLOCK];
// each thread loads one element from global to shared mem
// note use of 1D thread indices (only) in this kernel
int i = blockIdx.x*blockDim.x + threadIdx.x;
sdata[threadIdx.x] = data[i];
sdata[tid] = (tid < size) ? result_data[tid] : 0;
__syncthreads();
// do reduction in shared mem
for (int s=1; s < blockDim.x; s *=2)
{
int index = 2 * s * threadIdx.x;;
for (unsigned int s = blockDim.x / 2; s > 0; s >>= 1) {
if (tid < s) {
sdata[tid] += sdata[tid + s];
if (index < blockDim.x)
{
sdata[index] += sdata[index + s];
}
__syncthreads();
}
if (tid == 0) {
result_data[0] = sdata[0];
}
// write result for this block to global mem
if (threadIdx.x == 0)
atomicAdd(result_data, sdata[0]);
}
__host__ void sum_tensor_cuda(Tensor* tensor, float* result_data, int number_of_blocks) {
sum_tensor_cuda_kernel<<<number_of_blocks, THREADS_PER_BLOCK, THREADS_PER_BLOCK * sizeof(float)>>>(tensor->data, result_data, tensor->size);
__host__ void sum_tensor_cuda(Tensor* tensor, float* result_data) {
int remaining = number_of_blocks;
int levelSize;
while (remaining > 1) {
int threads = (remaining + THREADS_PER_BLOCK - 1) / THREADS_PER_BLOCK;
aux_sum_block_kernel<<<1, threads, threads * sizeof(float)>>>(result_data, remaining);
remaining = (remaining + THREADS_PER_BLOCK - 1) / THREADS_PER_BLOCK;
levelSize = remaining;
}
number_of_blocks = tensor->size / THREADS_PER_BLOCK;
sum_tensor_cuda_kernel<<<number_of_blocks, THREADS_PER_BLOCK>>>(tensor->data, result_data, tensor->size);
cudaError_t error = cudaGetLastError();
if (error != cudaSuccess) {

6
norch/csrc/tensor.cpp
View file

@ -171,9 +171,9 @@ extern "C" {
if (strcmp(tensor->device, "cuda") == 0) {
float* result_data;
int number_of_blocks = (tensor->size + THREADS_PER_BLOCK - 1) / THREADS_PER_BLOCK;
cudaMalloc((void**)&result_data, number_of_blocks * sizeof(float));
sum_tensor_cuda(tensor, result_data, number_of_blocks);
cudaMalloc((void**)&result_data, tensor->size * sizeof(float));
result_data[0] = 0;
sum_tensor_cuda(tensor, result_data);
return create_tensor(result_data, shape, ndim, device);
}
else {