diff --git a/build/cuda.cu.o b/build/cuda.cu.o index b40f8c4..ff4eb61 100644 Binary files a/build/cuda.cu.o and b/build/cuda.cu.o differ diff --git a/build/tensor.o b/build/tensor.o index eb3c24e..a24d9f7 100644 Binary files a/build/tensor.o and b/build/tensor.o differ diff --git a/norch/csrc/cuda.cu b/norch/csrc/cuda.cu index a2bb8e0..3fc0459 100644 --- a/norch/csrc/cuda.cu +++ b/norch/csrc/cuda.cu @@ -141,28 +141,104 @@ __global__ void sum_tensor_cuda_kernel(float* data, float* result_data, int size } } -__host__ void sum_tensor_cuda(Tensor* tensor, float* result_data) { - cudaMemcpy(result_data, tensor->data, tensor->size * sizeof(float), cudaMemcpyHostToDevice); +__global__ void sum_tensor_cuda_kernel_axis(float* data, float* result_data, int* shape, int* strides, int ndim, int axis, int axis_size) { + __shared__ float partial_sum[THREADS_PER_BLOCK]; - int num_blocks = (tensor->size + THREADS_PER_BLOCK - 1) / THREADS_PER_BLOCK; + int tid = threadIdx.x; + int bid = blockIdx.x; + int i = bid * blockDim.x + tid; - // First-level reduction - sum_tensor_cuda_kernel<<>>(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 - 1) / THREADS_PER_BLOCK; - sum_tensor_cuda_kernel<<>>(result_data, result_data, num_blocks); - num_blocks = num_blocks_next; + int total_size = 1; + for (int i = 0; i < ndim; i++) { + total_size *= shape[i]; } - cudaError_t error = cudaGetLastError(); - if (error != cudaSuccess) { - printf("CUDA error: %s\n", cudaGetErrorString(error)); - exit(-1); + int outer_dim_size = total_size / (axis_size * shape[axis]); + + partial_sum[tid] = 0.0; + if (i < outer_dim_size) { + for (int j = 0; j < axis_size; j++) { + int idx = i * strides[axis] + j * strides[axis + 1]; + partial_sum[tid] += data[idx]; + } } - cudaDeviceSynchronize(); + __syncthreads(); + + // Perform block-wise reduction + for (int s = blockDim.x / 2; s > 0; s >>= 1) { + if (tid < s) { + partial_sum[tid] += partial_sum[tid + s]; + } + __syncthreads(); + } + + // Write block sum to global memory + if (tid == 0) { + result_data[bid] = partial_sum[0]; + } +} + + +__host__ void sum_tensor_cuda(Tensor* tensor, float* result_data, int axis) { + if (axis == -1) { + cudaMemcpy(result_data, tensor->data, tensor->size * sizeof(float), cudaMemcpyHostToDevice); + + int num_blocks = (tensor->size + THREADS_PER_BLOCK - 1) / THREADS_PER_BLOCK; + + // First-level reduction + sum_tensor_cuda_kernel<<>>(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 - 1) / THREADS_PER_BLOCK; + sum_tensor_cuda_kernel<<>>(result_data, result_data, num_blocks); + num_blocks = num_blocks_next; + } + + cudaError_t error = cudaGetLastError(); + if (error != cudaSuccess) { + printf("CUDA error: %s\n", cudaGetErrorString(error)); + exit(-1); + } + + cudaDeviceSynchronize(); + + } else { + if (axis < 0 || axis >= tensor->ndim) { + printf("Invalid axis\n"); + return; + } + + int total_size = tensor->size; + int axis_size = tensor->shape[axis]; + + int outer_dim_size = total_size / axis_size; + + int num_blocks = (outer_dim_size + THREADS_PER_BLOCK - 1) / THREADS_PER_BLOCK; + + int* d_shape; + int* d_strides; + cudaMalloc(&d_shape, tensor->ndim * sizeof(int)); + cudaMalloc(&d_strides, tensor->ndim * sizeof(int)); + cudaMemcpy(d_shape, tensor->shape, tensor->ndim * sizeof(int), cudaMemcpyHostToDevice); + cudaMemcpy(d_strides, tensor->strides, tensor->ndim * sizeof(int), cudaMemcpyHostToDevice); + + sum_tensor_cuda_kernel_axis<<>>( + tensor->data, result_data, d_shape, d_strides, tensor->ndim, axis, axis_size + ); + + cudaFree(d_shape); + cudaFree(d_strides); + + cudaError_t error = cudaGetLastError(); + if (error != cudaSuccess) { + printf("CUDA error: %s\n", cudaGetErrorString(error)); + exit(-1); + } + + cudaDeviceSynchronize(); + } } __global__ void sub_tensor_cuda_kernel(float* data1, float* data2, float* result_data, int size) { diff --git a/norch/csrc/cuda.h b/norch/csrc/cuda.h index 38a6977..153173b 100644 --- a/norch/csrc/cuda.h +++ b/norch/csrc/cuda.h @@ -14,7 +14,8 @@ __host__ void sub_broadcasted_tensor_cuda(Tensor* tensor1, Tensor* tensor2, float* result_data, int* broadcasted_shape, int broadcasted_size); __global__ void sum_tensor_cuda_kernel(float* data, float* result_data); - __host__ void sum_tensor_cuda(Tensor* tensor, float* result_data); + __global__ void sum_tensor_cuda_kernel_axis(float* data, float* result_data, int* shape, int* strides, int ndim, int axis, int axis_size); + __host__ void sum_tensor_cuda(Tensor* tensor, float* result_data, int axis); __global__ void sub_tensor_cuda_kernel(float* data1, float* data2, float* result_data, int size); __host__ void sub_tensor_cuda(Tensor* tensor1, Tensor* tensor2, float* result_data); diff --git a/norch/csrc/tensor.cpp b/norch/csrc/tensor.cpp index 4d4d2fb..b6d2aae 100644 --- a/norch/csrc/tensor.cpp +++ b/norch/csrc/tensor.cpp @@ -207,7 +207,25 @@ extern "C" { float* result_data; cudaMalloc((void**)&result_data, size * sizeof(float)); cudaMemset(result_data, 0, size * sizeof(float)); - sum_tensor_cuda(tensor, result_data); + sum_tensor_cuda(tensor, result_data, axis); + + if (keepdim) { + if (axis == -1){ + ndim = tensor->ndim; + shape = (int*) malloc((tensor->ndim) * sizeof(int)); + for (int i = 0; i < tensor->ndim; i++) { + shape[i] = 1; + } + } else { + shape = (int*) malloc((tensor->ndim) * sizeof(int)); + for (int i = 0; i < tensor->ndim; i++) { + shape[i] = tensor->shape[i]; + } + shape[axis] = 1; + ndim = tensor->ndim; + } + + } return create_tensor(result_data, shape, ndim, device); } else { diff --git a/norch/libtensor.so b/norch/libtensor.so index d9dcf12..709baff 100755 Binary files a/norch/libtensor.so and b/norch/libtensor.so differ