diff --git a/build/cuda.cu.o b/build/cuda.cu.o index ff4eb61..a60bfa8 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 a24d9f7..f49cc88 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 3fc0459..8102e7c 100644 --- a/norch/csrc/cuda.cu +++ b/norch/csrc/cuda.cu @@ -141,33 +141,24 @@ __global__ void sum_tensor_cuda_kernel(float* data, float* result_data, int size } } -__global__ void sum_tensor_cuda_kernel_axis(float* data, float* result_data, int* shape, int* strides, int ndim, int axis, int axis_size) { +__global__ void sum_tensor_axis_cuda_kernel(float* data, float* result_data, int size, int axis_size) { __shared__ float partial_sum[THREADS_PER_BLOCK]; int tid = threadIdx.x; - int bid = blockIdx.x; - int i = bid * blockDim.x + tid; + int block_offset = blockIdx.x * axis_size; + int i = block_offset + tid; - int total_size = 1; - for (int i = 0; i < ndim; i++) { - total_size *= shape[i]; - } - - 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]; - } + partial_sum[tid] = 0.0f; + while (i < block_offset + axis_size && i < size) { + partial_sum[tid] += data[i]; + i += blockDim.x; } __syncthreads(); // Perform block-wise reduction for (int s = blockDim.x / 2; s > 0; s >>= 1) { - if (tid < s) { + if (tid < s && i < block_offset + axis_size) { partial_sum[tid] += partial_sum[tid + s]; } __syncthreads(); @@ -175,7 +166,7 @@ __global__ void sum_tensor_cuda_kernel_axis(float* data, float* result_data, int // Write block sum to global memory if (tid == 0) { - result_data[bid] = partial_sum[0]; + result_data[blockIdx.x] = partial_sum[0]; } } @@ -203,33 +194,31 @@ __host__ void sum_tensor_cuda(Tensor* tensor, float* result_data, int axis) { } cudaDeviceSynchronize(); - - } else { - if (axis < 0 || axis >= tensor->ndim) { - printf("Invalid axis\n"); - return; - } - int total_size = tensor->size; + } else { int axis_size = tensor->shape[axis]; - int outer_dim_size = total_size / axis_size; + // Allocate memory for temporary storage on the device + float* temp_data; + cudaMalloc(&temp_data, tensor->size * sizeof(float)); - int num_blocks = (outer_dim_size + THREADS_PER_BLOCK - 1) / THREADS_PER_BLOCK; + // Copy tensor data to device + cudaMemcpy(temp_data, tensor->data, tensor->size * sizeof(float), cudaMemcpyHostToDevice); - 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); + int num_blocks = (tensor->size + THREADS_PER_BLOCK - 1) / THREADS_PER_BLOCK; - sum_tensor_cuda_kernel_axis<<>>( - tensor->data, result_data, d_shape, d_strides, tensor->ndim, axis, axis_size - ); + // First-level reduction + sum_tensor_axis_cuda_kernel<<>>(temp_data, result_data, tensor->size, axis_size); - cudaFree(d_shape); - cudaFree(d_strides); + // 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; + } + + // Free allocated memory on the device + cudaFree(temp_data); cudaError_t error = cudaGetLastError(); if (error != cudaSuccess) { diff --git a/norch/csrc/cuda.h b/norch/csrc/cuda.h index 153173b..5e180cf 100644 --- a/norch/csrc/cuda.h +++ b/norch/csrc/cuda.h @@ -14,7 +14,7 @@ __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); - __global__ void sum_tensor_cuda_kernel_axis(float* data, float* result_data, int* shape, int* strides, int ndim, int axis, int axis_size); + __global__ void sum_tensor_axis_cuda_kernel(float* data, float* result_data, int size, 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); diff --git a/norch/csrc/tensor.cpp b/norch/csrc/tensor.cpp index b6d2aae..1f374f8 100644 --- a/norch/csrc/tensor.cpp +++ b/norch/csrc/tensor.cpp @@ -182,6 +182,11 @@ extern "C" { } int ndim; int* shape; + + if (axis > tensor->ndim - 1) { + fprintf(stderr, "Error: axis argument %d must be smaller than tensor dimension %d", axis, tensor->ndim); + } + if (axis == -1) { shape = (int*) malloc(sizeof(int)); diff --git a/norch/libtensor.so b/norch/libtensor.so index 709baff..d5ff2a2 100755 Binary files a/norch/libtensor.so and b/norch/libtensor.so differ diff --git a/norch/tensor.py b/norch/tensor.py index f7f205e..8063cba 100644 --- a/norch/tensor.py +++ b/norch/tensor.py @@ -793,6 +793,9 @@ class Tensor: if axis == None: axis = -1 + if axis > self.ndim - 1: + raise ValueError(f"Error: axis argument {axis} cannot be higher than tensor dimension {self.ndim}") + Tensor._C.sum_tensor.argtypes = [ctypes.POINTER(CTensor), ctypes.c_int, ctypes.c_bool] Tensor._C.sum_tensor.restype = ctypes.POINTER(CTensor) @@ -834,6 +837,9 @@ class Tensor: if axis == None: axis = -1 + if axis > self.ndim - 1: + raise ValueError(f"Error: axis argument {axis} cannot be higher than tensor dimension {self.ndim}") + Tensor._C.max_tensor.argtypes = [ctypes.POINTER(CTensor), ctypes.c_int, ctypes.c_bool] Tensor._C.max_tensor.restype = ctypes.POINTER(CTensor) @@ -874,6 +880,9 @@ class Tensor: if axis == None: axis = -1 + + if axis > self.ndim - 1: + raise ValueError(f"Error: axis argument {axis} must be smaller than tensor dimension {self.ndim}") Tensor._C.min_tensor.argtypes = [ctypes.POINTER(CTensor), ctypes.c_int, ctypes.c_bool] Tensor._C.min_tensor.restype = ctypes.POINTER(CTensor)