diff --git a/norch/csrc/cpu.cpp b/norch/csrc/cpu.cpp index 40eb3b5..63e1615 100644 --- a/norch/csrc/cpu.cpp +++ b/norch/csrc/cpu.cpp @@ -395,6 +395,25 @@ void assign_tensor_cpu(Tensor* tensor, float* result_data) { } } +void make_contiguous_tensor_cpu(Tensor* tensor, float* result_data, int* new_strides) { + + for (int i = 0; i < tensor->size; i++) { + int index = 0; + int offset = i; + for (int j = 0; j < tensor->ndim; j++) { + index += (offset / new_strides[j]) * tensor->strides[j]; + offset %= new_strides[j]; + } + result_data[i] = tensor->data[index]; + } + + // Free old data and update tensor properties + free(tensor->data); + free(tensor->strides); + tensor->data = result_data; + tensor->strides = new_strides; +} + void sin_tensor_cpu(Tensor* tensor, float* result_data) { for (int i = 0; i < tensor->size; i++) { result_data[i] = sinf(tensor->data[i]); diff --git a/norch/csrc/cpu.h b/norch/csrc/cpu.h index 1d9547b..9a27a26 100644 --- a/norch/csrc/cpu.h +++ b/norch/csrc/cpu.h @@ -30,6 +30,7 @@ void transpose_2D_tensor_cpu(Tensor* tensor, float* result_data); void transpose_3D_tensor_cpu(Tensor* tensor, float* result_data); void transpose_axes_cpu(Tensor* tensor, float* result_data, int axis1, int axis2, int* new_shape); void assign_tensor_cpu(Tensor* tensor, float* result_data); +void make_contiguous_tensor_cpu(Tensor* tensor, float* result_data, int* new_strides); void sin_tensor_cpu(Tensor* tensor, float* result_data); void cos_tensor_cpu(Tensor* tensor, float* result_data); void sigmoid_tensor_cpu(Tensor* tensor, float* result_data); diff --git a/norch/csrc/cuda.cu b/norch/csrc/cuda.cu index bda8a7d..064c3d1 100644 --- a/norch/csrc/cuda.cu +++ b/norch/csrc/cuda.cu @@ -1115,6 +1115,48 @@ __host__ void assign_tensor_cuda(Tensor* tensor, float* result_data) { cudaDeviceSynchronize(); } +__global__ void make_contiguous_tensor_cuda_kernel(float* data, float* result_data, int ndim, int size, int* strides, int* new_strides) { + + int i = blockIdx.x * blockDim.x + threadIdx.x; + if (i < size) { + int index = 0; + int offset = i; + for (int j = 0; j < ndim; j++) { + index += (offset / new_strides[j]) * strides[j]; + offset %= new_strides[j]; + } + result_data[i] = data[index]; + } +} + +__host__ void make_contiguous_tensor_cuda(Tensor* tensor, float* result_data, int* new_strides) { + + int* d_strides; + cudaMalloc((void **)&d_strides, tensor->ndim * sizeof(int)); + cudaMemcpy(d_strides, tensor->strides, tensor->ndim * sizeof(int), cudaMemcpyHostToDevice); + + int* d_new_strides; + cudaMalloc((void **)&d_new_strides, tensor->ndim * sizeof(int)); + cudaMemcpy(d_new_strides, new_strides, tensor->ndim * sizeof(int), cudaMemcpyHostToDevice); + + int number_of_blocks = (tensor->size + THREADS_PER_BLOCK - 1) / THREADS_PER_BLOCK; + make_contiguous_tensor_cuda_kernel<<>>(tensor->data, result_data, tensor->ndim, tensor->size, d_strides, d_new_strides); + + cudaError_t error = cudaGetLastError(); + if (error != cudaSuccess) { + printf("CUDA error: %s\n", cudaGetErrorString(error)); + exit(-1); + } + + cudaDeviceSynchronize(); + + // Free old data and update tensor properties + cudaFree(tensor->data); + free(tensor->strides); + tensor->data = result_data; + tensor->strides = new_strides; +} + __global__ void sin_tensor_cuda_kernel(float* data, float* result_data, int size) { int i = blockIdx.x * blockDim.x + threadIdx.x; diff --git a/norch/csrc/cuda.h b/norch/csrc/cuda.h index a684ee2..cc7b385 100644 --- a/norch/csrc/cuda.h +++ b/norch/csrc/cuda.h @@ -85,6 +85,9 @@ __global__ void assign_tensor_cuda_kernel(float* data, float* result_data, int size); __host__ void assign_tensor_cuda(Tensor* tensor, float* result_data); + __global__ void make_contiguous_tensor_cuda_kernel(float* data, float* result_data, int ndim, int size, int* strides, int* new_strides); + __host__ void make_contiguous_tensor_cuda(Tensor* tensor, float* result_data, int* new_strides); + __global__ void sin_tensor_cuda_kernel(float* data, float* result_data, int size); __host__ void sin_tensor_cuda(Tensor* tensor, float* result_data); diff --git a/norch/csrc/tensor.cpp b/norch/csrc/tensor.cpp index 22739b6..d5bfa71 100644 --- a/norch/csrc/tensor.cpp +++ b/norch/csrc/tensor.cpp @@ -1522,17 +1522,10 @@ extern "C" { } void make_contiguous(Tensor* tensor) { - float* new_data = (float*)malloc(tensor->size * sizeof(float)); - if (new_data == NULL) { - // Handle memory allocation failure - return; - } - + int* new_strides = (int*)malloc(tensor->ndim * sizeof(int)); if (new_strides == NULL) { - free(new_data); - // Handle memory allocation failure - return; + fprintf(stderr, "Memory allocation failed\n"); } // Calculate new strides assuming C-contiguous order @@ -1542,21 +1535,17 @@ extern "C" { stride *= tensor->shape[i]; } - // Rearrange data - for (int i = 0; i < tensor->size; i++) { - int index = 0; - int offset = i; - for (int j = 0; j < tensor->ndim; j++) { - index += (offset / new_strides[j]) * tensor->strides[j]; - offset %= new_strides[j]; - } - new_data[i] = tensor->data[index]; - } + if (strcmp(tensor->device, "cuda") == 0) { + float* result_data; + cudaMalloc((void **)&result_data, tensor->size * sizeof(float)); + make_contiguous_tensor_cuda(tensor, result_data, new_strides); - // Free old data and update tensor properties - free(tensor->data); - free(tensor->strides); - tensor->data = new_data; - tensor->strides = new_strides; + } else { + float* result_data = (float*)malloc(tensor->size * sizeof(float)); + if (result_data == NULL) { + fprintf(stderr, "Memory allocation failed\n"); + } + make_contiguous_tensor_cpu(tensor, result_data, new_strides); } } +} diff --git a/norch/libtensor.so b/norch/libtensor.so index 0f5fce8..5f2b381 100755 Binary files a/norch/libtensor.so and b/norch/libtensor.so differ