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

Merge pull request #30 from lucasdelimanogueira/tmp

Browse source 
Tmp
This commit is contained in:
lucasdelimanogueira 2024-05-02 15:19:14 -03:00 • committed by GitHub
commit 9da59793c5
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
1 changed files with 8 additions and 34 deletions
Download patch file Download diff file Expand all files Collapse all files

42
norch/csrc/cuda.cu
View file

@ -4,8 +4,8 @@
#include <math.h>
#define THREADS_PER_BLOCK 128
#define THREADS_PER_BLOCK_SUM 1024
#define TILE_SIZE 32
#define SHMEM_SIZE THREADS_PER_BLOCK_SUM * sizeof(float)
__host__ void cpu_to_cuda(Tensor* tensor) {
@ -61,7 +61,7 @@ __host__ void add_tensor_cuda(Tensor* tensor1, Tensor* tensor2, float* result_da
__global__ void sum_tensor_cuda_kernel(float* data, float* result_data, int size) {
__shared__ float partial_sum[THREADS_PER_BLOCK_SUM];
__shared__ float partial_sum[SHMEM_SIZE];
int tid = threadIdx.x;
int i = blockIdx.x * blockDim.x + threadIdx.x;
@ -84,46 +84,18 @@ __global__ void sum_tensor_cuda_kernel(float* data, float* result_data, int size
}
}
__global__ void aux_final_sum_kernel(float* result_data, int size) {
__shared__ float partial_sum[THREADS_PER_BLOCK_SUM];
int tid = threadIdx.x;
int i = blockIdx.x * blockDim.x + threadIdx.x;
partial_sum[tid] = (i < size) ? result_data[i] : 0;
__syncthreads();
// Perform final reduction
for (int s = blockDim.x / 2; s > 0; s >>= 1) {
if (tid < s) {
partial_sum[tid] += partial_sum[tid + s];
}
__syncthreads();
}
// Write final result to global memory
if (tid == 0 && blockIdx.x == 0) {
result_data[0] = partial_sum[0];
}
}
__host__ void sum_tensor_cuda(Tensor* tensor, float* result_data) {
cudaMemcpy(result_data, tensor->data, tensor->size * sizeof(float), cudaMemcpyHostToDevice);
int num_blocks = (tensor->size + THREADS_PER_BLOCK_SUM - 1) / THREADS_PER_BLOCK_SUM;
int num_blocks = (tensor->size + THREADS_PER_BLOCK - 1) / THREADS_PER_BLOCK;
// First-level reduction
sum_tensor_cuda_kernel<<<num_blocks, THREADS_PER_BLOCK_SUM>>>(tensor->data, result_data, tensor->size);
sum_tensor_cuda_kernel<<<num_blocks, THREADS_PER_BLOCK>>>(tensor->data, result_data, tensor->size);
// If necessary, perform multiple levels of reduction
while (num_blocks > 1) {
int num_blocks_next = (num_blocks + THREADS_PER_BLOCK_SUM - 1) / THREADS_PER_BLOCK_SUM;
aux_final_sum_kernel<<<num_blocks_next, THREADS_PER_BLOCK_SUM>>>(result_data, num_blocks);
int num_blocks_next = (num_blocks + THREADS_PER_BLOCK - 1) / THREADS_PER_BLOCK;
sum_tensor_cuda_kernel<<<num_blocks_next, THREADS_PER_BLOCK>>>(result_data, result_data, num_blocks);
num_blocks = num_blocks_next;
}
@ -137,6 +109,8 @@ __host__ void sum_tensor_cuda(Tensor* tensor, float* result_data) {
}
__global__ void sub_tensor_cuda_kernel(float* data1, float* data2, float* result_data, int size) {
int i = blockIdx.x * blockDim.x + threadIdx.x;