make contiguous cuda

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]);

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);

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<<<number_of_blocks, THREADS_PER_BLOCK>>>(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;

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);
 

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);
         }
     }
+}