scalar mul and sum tensor

build/Makefile

@@ -7,7 +7,7 @@ CFLAGS = -Wall -Wextra -std=c++11
 NVCCFLAGS = -std=c++11
 
 # Directories
-SRCDIR = ../csrc
+SRCDIR = ../norch/csrc
 BUILDDIR = ../build
 TARGET = ../build/libtensor.so
 

norch/csrc/cpu.cpp

@@ -0,0 +1,63 @@
+#include "tensor.h"
+#include <stdio.h>
+#include <stdlib.h>
+#include <string.h>
+#include <math.h>
+
+void add_tensor_cpu(Tensor* tensor1, Tensor* tensor2, float* result_data) {
+    
+    for (int i = 0; i < tensor1->size; i++) {
+        result_data[i] = tensor1->data[i] + tensor2->data[i];
+    }
+}
+
+void sub_tensor_cpu(Tensor* tensor1, Tensor* tensor2, float* result_data) {
+    
+    for (int i = 0; i < tensor1->size; i++) {
+        result_data[i] = tensor1->data[i] - tensor2->data[i];
+    }
+}
+
+void elementwise_mul_tensor_cpu(Tensor* tensor1, Tensor* tensor2, float* result_data) {
+    
+    for (int i = 0; i < tensor1->size; i++) {
+        result_data[i] = tensor1->data[i] * tensor2->data[i];
+    }
+}
+
+void scalar_mul_tensor_cpu(Tensor* tensor, float scalar, float* result_data) {
+    
+    for (int i = 0; i < tensor->size; i++) {
+        result_data[i] = scalar * tensor->data[i];
+    }
+}
+
+void matmul_tensor_cpu(Tensor* tensor1, Tensor* tensor2, float* result_data) {
+    
+    for (int i = 0; i < tensor1->shape[0]; i++) {
+        for (int j = 0; j < tensor2->shape[1]; j++) {
+            float sum = 0.0;
+            for (int k = 0; k < tensor1->shape[1]; k++) {
+                sum += tensor1->data[i * tensor1->shape[1] + k] * tensor2->data[k * tensor2->shape[1] + j];
+            }
+            result_data[i * tensor2->shape[1] + j] = sum;
+        }
+    }
+}
+
+void pow_tensor_cpu(Tensor* tensor, float power, float* result_data) {
+    
+    for (int i = 0; i < tensor->size; i++) {
+        result_data[i] = powf(tensor->data[i], power);
+    }
+}
+
+void sum_tensor_cpu(Tensor* tensor, float* result_data) {
+    float sum = 0.0;
+
+    for (int i = 0; i < tensor->size; i++) {
+        sum += tensor->data[i];
+    }
+
+    *result_data = sum;
+}

norch/csrc/cpu.h

@@ -0,0 +1,14 @@
+#ifndef CPU_H
+#define CPU_H
+
+#include "tensor.h"
+
+void add_tensor_cpu(Tensor* tensor1, Tensor* tensor2, float* result_data);
+void sum_tensor_cpu(Tensor* tensor1, float* result_data);
+void sub_tensor_cpu(Tensor* tensor1, Tensor* tensor2, float* result_data);
+void elementwise_mul_tensor_cpu(Tensor* tensor1, Tensor* tensor2, float* result_data);
+void matmul_tensor_cpu(Tensor* tensor1, Tensor* tensor2, float* result_data);
+void pow_tensor_cpu(Tensor* tensor, float power, float* result_data);
+void scalar_mul_tensor_cpu(Tensor* tensor, float scalar, float* result_data);
+
+#endif /* CPU_H */

norch/csrc/cuda.cu

@@ -0,0 +1,226 @@
+#include "tensor.h"
+#include <stdio.h>
+#include <stdlib.h>
+#include <math.h>
+
+#define THREADS_PER_BLOCK 128
+
+__host__ void cpu_to_cuda(Tensor* tensor) {
+    
+    float* data_tmp;
+    cudaMalloc((void **)&data_tmp, tensor->size * sizeof(float));
+    cudaMemcpy(data_tmp, tensor->data, tensor->size * sizeof(float), cudaMemcpyHostToDevice);
+
+    tensor->data = data_tmp;
+
+    const char* device_str = "cuda";
+    tensor->device = (char*)malloc(strlen(device_str) + 1);
+    strcpy(tensor->device, device_str); 
+
+    printf("Successfully sent tensor to: %s\n", tensor->device);
+}
+
+__host__ void cuda_to_cpu(Tensor* tensor) {
+    float* data_tmp = (float*)malloc(tensor->size * sizeof(float));
+
+    cudaMemcpy(data_tmp, tensor->data, tensor->size * sizeof(float), cudaMemcpyDeviceToHost);
+    cudaFree(tensor->data);
+
+    tensor->data = data_tmp;
+
+    const char* device_str = "cpu";
+    tensor->device = (char*)malloc(strlen(device_str) + 1);
+    strcpy(tensor->device, device_str); 
+
+    printf("Successfully sent tensor to: %s\n", tensor->device);
+}
+
+__global__ void add_tensor_cuda_kernel(float* data1, float* data2, float* result_data, int size) {
+    
+    int i = blockIdx.x * blockDim.x + threadIdx.x;
+    if (i < size) {
+        result_data[i] = data1[i] + data2[i];
+    }
+}
+
+__host__ void add_tensor_cuda(Tensor* tensor1, Tensor* tensor2, float* result_data) {
+    
+    int number_of_blocks = (tensor1->size + THREADS_PER_BLOCK - 1) / THREADS_PER_BLOCK;
+    add_tensor_cuda_kernel<<<number_of_blocks, THREADS_PER_BLOCK>>>(tensor1->data, tensor2->data, result_data, tensor1->size);
+
+    cudaError_t error = cudaGetLastError();
+    if (error != cudaSuccess) {
+        printf("CUDA error: %s\n", cudaGetErrorString(error));
+        exit(-1);
+    }
+
+    cudaDeviceSynchronize();
+}
+
+__global__ void sum_tensor_cuda_kernel(float* data, float* result_data, int size) {
+    extern __shared__ float sdata[];
+
+    unsigned int tid = threadIdx.x;
+    unsigned int i = blockIdx.x * blockDim.x + threadIdx.x;
+
+    // Each thread loads one element into shared memory
+    if (i < size) {
+        sdata[tid] = data[i];
+    } else {
+        sdata[tid] = 0.0f;
+    }
+    __syncthreads();
+
+    // Reduce within block using shared memory
+    for (unsigned int s = blockDim.x / 2; s > 0; s >>= 1) {
+        if (tid < s) {
+            sdata[tid] += sdata[tid + s];
+        }
+        __syncthreads();
+    }
+
+    // Write the result for this block to global memory
+    if (tid == 0) {
+        result_data[blockIdx.x] = sdata[0];
+    }
+}
+
+__host__ void sum_tensor_cuda(Tensor* tensor, float* result_data) {
+    
+    int number_of_blocks = (tensor->size + THREADS_PER_BLOCK - 1) / THREADS_PER_BLOCK;
+    sum_tensor_cuda_kernel<<<number_of_blocks, THREADS_PER_BLOCK>>>(tensor->data, result_data, tensor->size);
+
+    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) {
+   
+    int i = blockIdx.x * blockDim.x + threadIdx.x;
+    if (i < size) {
+        result_data[i] = data1[i] - data2[i];
+    }
+}
+
+__host__ void sub_tensor_cuda(Tensor* tensor1, Tensor* tensor2, float* result_data) {
+    
+    int number_of_blocks = (tensor1->size + THREADS_PER_BLOCK - 1) / THREADS_PER_BLOCK;
+    sub_tensor_cuda_kernel<<<number_of_blocks, THREADS_PER_BLOCK>>>(tensor1->data, tensor2->data, result_data, tensor1->size);
+
+    cudaError_t error = cudaGetLastError();
+    if (error != cudaSuccess) {
+        printf("CUDA error: %s\n", cudaGetErrorString(error));
+        exit(-1);
+    }
+
+    cudaDeviceSynchronize();
+}
+
+__global__ void elementwise_mul_tensor_cuda_kernel(float* data1, float* data2, float* result_data, int size) {
+    
+    int i = blockIdx.x * blockDim.x + threadIdx.x;
+    if (i < size) {
+        result_data[i] = data1[i] * data2[i];
+    }
+}
+
+__host__ void elementwise_mul_tensor_cuda(Tensor* tensor1, Tensor* tensor2, float* result_data) {
+    
+    int number_of_blocks = (tensor1->size + THREADS_PER_BLOCK - 1) / THREADS_PER_BLOCK;
+    elementwise_mul_tensor_cuda_kernel<<<number_of_blocks, THREADS_PER_BLOCK>>>(tensor1->data, tensor2->data, result_data, tensor1->size);
+
+    cudaError_t error = cudaGetLastError();
+    if (error != cudaSuccess) {
+        printf("CUDA error: %s\n", cudaGetErrorString(error));
+        exit(-1);
+    }
+
+    cudaDeviceSynchronize();
+}
+
+__global__ void scalar_mul_tensor_cuda_kernel(float* data, float scalar, float* result_data, int size) {
+    
+    int i = blockIdx.x * blockDim.x + threadIdx.x;
+    if (i < size) {
+        result_data[i] = scalar * data[i];
+    }
+}
+
+__host__ void scalar_mul_tensor_cuda(Tensor* tensor, float scalar, float* result_data) {
+    
+    int number_of_blocks = (tensor->size + THREADS_PER_BLOCK - 1) / THREADS_PER_BLOCK;
+    scalar_mul_tensor_cuda_kernel<<<number_of_blocks, THREADS_PER_BLOCK>>>(tensor->data, scalar, result_data, tensor->size);
+
+    cudaError_t error = cudaGetLastError();
+    if (error != cudaSuccess) {
+        printf("CUDA error: %s\n", cudaGetErrorString(error));
+        exit(-1);
+    }
+
+    cudaDeviceSynchronize();
+}
+
+__global__ void matmul_tensor_cuda_kernel(float* data1, float* data2, float* result_data, int rows1, int cols1, int cols2) {    
+
+    int row = blockIdx.y * blockDim.y + threadIdx.y;
+    int col = blockIdx.x * blockDim.x + threadIdx.x;
+
+    if (row < rows1 && col < cols2) {
+        float sum = 0.0;
+        for (int k = 0; k < cols1; k++) {
+            sum += data1[row * cols1 + k] * data2[k * cols2 + col];
+        }
+        result_data[row * cols2 + col] = sum;
+    }
+
+}
+
+__host__ void matmul_tensor_cuda(Tensor* tensor1, Tensor* tensor2, float* result_data) {
+    
+    int rows1 = tensor1->shape[0];
+    int cols1 = tensor1->shape[1];
+    int cols2 = tensor2->shape[1];
+
+    dim3 threadsPerBlock(16, 16);
+    dim3 numBlocks((cols2 + threadsPerBlock.x - 1) / threadsPerBlock.x, (rows1 + threadsPerBlock.y - 1) / threadsPerBlock.y);
+    matmul_tensor_cuda_kernel<<<numBlocks, threadsPerBlock>>>(tensor1->data, tensor2->data, result_data, rows1, cols1, cols2);
+
+
+    cudaError_t error = cudaGetLastError();
+    if (error != cudaSuccess) {
+        printf("CUDA error: %s\n", cudaGetErrorString(error));
+        exit(-1);
+    }
+
+    cudaDeviceSynchronize();
+}
+
+__global__ void pow_tensor_cuda_kernel(float* data, float power, float* result_data, int size) {
+    
+    int i = blockIdx.x * blockDim.x + threadIdx.x;
+    if (i < size) {
+        result_data[i] = powf(data[i], power);
+    }
+}
+
+__host__ void pow_tensor_cuda(Tensor* tensor, float power, float* result_data) {
+    
+    int number_of_blocks = (tensor->size + THREADS_PER_BLOCK - 1) / THREADS_PER_BLOCK;
+    pow_tensor_cuda_kernel<<<number_of_blocks, THREADS_PER_BLOCK>>>(tensor->data, power, result_data, tensor->size);
+
+    cudaError_t error = cudaGetLastError();
+    if (error != cudaSuccess) {
+        printf("CUDA error: %s\n", cudaGetErrorString(error));
+        exit(-1);
+    }
+
+    cudaDeviceSynchronize();
+}
+
+

norch/csrc/cuda.h

@@ -0,0 +1,28 @@
+#ifndef CUDA_KERNEL_H_
+#define CUDA_KERNEL_H_
+
+    __host__ void cpu_to_cuda(Tensor* tensor);
+    __host__ void cuda_to_cpu(Tensor* tensor);
+    
+    __global__ void add_tensor_cuda_kernel(float* data1, float* data2, float* result_data, int size);
+    __host__ void add_tensor_cuda(Tensor* tensor1, Tensor* tensor2, float* result_data);
+    
+    __global__ void sum_tensor_cuda_kernel(float* data, float* result_data, int size);
+    __host__ void sum_tensor_cuda(Tensor* tensor1, float* result_data);
+
+    __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);
+    
+    __global__ void elementwise_mul_tensor_cuda_kernel(float* data1, float* data2, float* result_data, int size);
+    __host__ void elementwise_mul_tensor_cuda(Tensor* tensor1, Tensor* tensor2, float* result_data);
+
+    __global__ void scalar_mul_tensor_cuda_kernel(float* data, float scalar, float* result_data, int size);
+    __host__ void scalar_mul_tensor_cuda(Tensor* tensor, float scalar, float* result_data);
+
+    __global__ void matmul_tensor_cuda_kernel(float* data1, float* data2, float* result_data, int rows1, int cols1, int cols2);
+    __host__ void matmul_tensor_cuda(Tensor* tensor1, Tensor* tensor2, float* result_data);
+
+    __global__ void pow_tensor_cuda_kernel(float* data, float power, float* result_data, int size);
+    __host__ void pow_tensor_cuda(Tensor* tensor, float power, float* result_data);
+
+#endif /* CUDA_KERNEL_H_ */

norch/csrc/tensor.cpp

@@ -0,0 +1,464 @@
+#include <stdio.h>
+#include <stdlib.h>
+#include <string.h>
+#include <math.h>
+#include <cuda_runtime_api.h>
+#include "tensor.h"
+#include "cuda.h"
+#include "cpu.h"
+
+extern "C" {
+
+    Tensor* create_tensor(float* data, int* shape, int ndim, char* device) {
+
+        printf("Creating tensor\n");
+        Tensor* tensor = (Tensor*)malloc(sizeof(Tensor));
+        if (tensor == NULL) {
+            fprintf(stderr, "Memory allocation failed\n");
+            exit(1);
+        }
+        tensor->data = data;
+        tensor->shape = shape;
+        tensor->ndim = ndim;
+
+        tensor->device = (char*)malloc(strlen(device) + 1);
+        if (device != NULL) {
+            strcpy(tensor->device, device);
+        } else {
+            fprintf(stderr, "Memory allocation failed\n");
+            exit(-1);
+        }
+
+        tensor->size = 1;
+        for (int i = 0; i < ndim; i++) {
+            tensor->size *= shape[i];
+        }
+
+        tensor->strides = (int*)malloc(ndim * sizeof(int));
+        if (tensor->strides == NULL) {
+            fprintf(stderr, "Memory allocation failed\n");
+            exit(1);
+        }
+        int stride = 1;
+        for (int i = ndim - 1; i >= 0; i--) {
+            tensor->strides[i] = stride;
+            stride *= shape[i];
+        }
+
+        printf("Tensor created successfully\n");
+        printf("Tensor information:\n");
+        printf("Number of dimensions: %d\n", tensor->ndim);
+        printf("Number size: %d\n", tensor->size);
+        printf("Device: %s\n", tensor->device);
+
+        printf("Shape: [");
+        for (int i = 0; i < ndim; i++) {
+            printf("%d", tensor->shape[i]);
+            if (i < ndim - 1) {
+                printf(", ");
+            }
+        }
+        printf("]\n");
+
+        /*printf("Data:\n[");
+        for (int i = 0; i < stride; i++) {
+            printf("%.2f", tensor->data[i]);
+            if (i < stride - 1) {
+                printf(", ");
+            }
+        }
+        printf("]\n\n\n");*/
+        
+        return tensor;
+    }
+
+    float get_item(Tensor* tensor, int* indices) {
+        int index = 0;
+        for (int i = 0; i < tensor->ndim; i++) {
+            index += indices[i] * tensor->strides[i];
+        }
+
+        float result;
+        if (strcmp(tensor->device, "cuda") == 0) {
+            cudaMemcpy(&result, tensor->data + index, sizeof(float), cudaMemcpyDeviceToHost);
+        } else {
+            result = tensor->data[index];
+        }
+
+        return result;
+    }
+
+    void to_device(Tensor* tensor, char* target_device) {
+        printf("Sending tensor to device: %s\n", target_device);
+
+        if ((strcmp(target_device, "cuda") == 0) && (strcmp(tensor->device, "cpu") == 0)) {
+            cpu_to_cuda(tensor);
+        }
+
+        else if ((strcmp(target_device, "cpu") == 0) && (strcmp(tensor->device, "cuda") == 0)) {
+            cuda_to_cpu(tensor);
+        }
+    }
+
+    Tensor* add_tensor(Tensor* tensor1, Tensor* tensor2) {
+        if (tensor1->ndim != tensor2->ndim) {
+            fprintf(stderr, "Tensors must have the same number of dimensions %d and %d for addition\n", tensor1->ndim, tensor2->ndim);
+            exit(1);
+        }
+
+        if (strcmp(tensor1->device, tensor2->device) != 0) {
+            fprintf(stderr, "Tensors must be on the same device: %s and %s\n", tensor1->device, tensor2->device);
+            exit(1);
+        }
+
+        char* device = (char*)malloc(strlen(tensor1->device) + 1);
+        if (device != NULL) {
+            strcpy(device, tensor1->device);
+        } else {
+            fprintf(stderr, "Memory allocation failed\n");
+            exit(-1);
+        }
+        int ndim = tensor1->ndim;
+        int* shape = (int*)malloc(ndim * sizeof(int));
+        if (shape == NULL) {
+            fprintf(stderr, "Memory allocation failed\n");
+            exit(1);
+        }
+
+        for (int i = 0; i < ndim; i++) {
+            if (tensor1->shape[i] != tensor2->shape[i]) {
+                fprintf(stderr, "Tensors must have the same shape %d and %d at index %d for addition\n", tensor1->shape[i], tensor2->shape[i], i);
+                exit(1);
+            }
+            shape[i] = tensor1->shape[i];
+        }        
+        if (strcmp(tensor1->device, "cuda") == 0) {
+
+            float* result_data;
+            cudaMalloc((void **)&result_data, tensor1->size * sizeof(float));
+            add_tensor_cuda(tensor1, tensor2, result_data);
+            return create_tensor(result_data, shape, ndim, device);
+        } 
+        else {
+            float* result_data = (float*)malloc(tensor1->size * sizeof(float));
+            if (result_data == NULL) {
+                fprintf(stderr, "Memory allocation failed\n");
+                exit(1);
+            }
+            add_tensor_cpu(tensor1, tensor2, result_data);
+            return create_tensor(result_data, shape, ndim, device);
+        }     
+    }
+
+    Tensor* sum_tensor(Tensor* tensor) {
+
+        char* device = (char*)malloc(strlen(tensor->device) + 1);
+        if (device != NULL) {
+            strcpy(device, tensor->device);
+        } else {
+            fprintf(stderr, "Memory allocation failed\n");
+            exit(-1);
+        }
+        int ndim = 1;
+        int* shape = (int*)malloc(ndim * sizeof(int));
+        if (shape == NULL) {
+            fprintf(stderr, "Memory allocation failed\n");
+            exit(1);
+        }
+
+        shape[0] = 1;
+  
+        if (strcmp(tensor->device, "cuda") == 0) {
+
+            float* result_data;
+            cudaMalloc((void **)&result_data, 1 * sizeof(float));
+            sum_tensor_cuda(tensor, result_data);
+            return create_tensor(result_data, shape, ndim, device);
+        } 
+        else {
+            float* result_data = (float*)malloc(1 * sizeof(float));
+            if (result_data == NULL) {
+                fprintf(stderr, "Memory allocation failed\n");
+                exit(1);
+            }
+            sum_tensor_cpu(tensor, result_data);
+            return create_tensor(result_data, shape, ndim, device);
+        }     
+    }
+
+    Tensor* sub_tensor(Tensor* tensor1, Tensor* tensor2) {
+        if (tensor1->ndim != tensor2->ndim) {
+            fprintf(stderr, "Tensors must have the same number of dimensions %d and %d for subtraction\n", tensor1->ndim, tensor2->ndim);
+            exit(1);
+        }
+
+        if (strcmp(tensor1->device, tensor2->device) != 0) {
+            fprintf(stderr, "Tensors must be on the same device: %s and %s\n", tensor1->device, tensor2->device);
+            exit(1);
+        }
+
+        char* device = (char*)malloc(strlen(tensor1->device) + 1);
+        if (device != NULL) {
+            strcpy(device, tensor1->device);
+        } else {
+            fprintf(stderr, "Memory allocation failed\n");
+            exit(-1);
+        }
+        int ndim = tensor1->ndim;
+        int* shape = (int*)malloc(ndim * sizeof(int));
+        if (shape == NULL) {
+            fprintf(stderr, "Memory allocation failed\n");
+            exit(1);
+        }
+
+        for (int i = 0; i < ndim; i++) {
+            if (tensor1->shape[i] != tensor2->shape[i]) {
+                fprintf(stderr, "Tensors must have the same shape %d and %d at index %d for subtraction\n", tensor1->shape[i], tensor2->shape[i], i);
+                exit(1);
+            }
+            shape[i] = tensor1->shape[i];
+        }
+
+        if (strcmp(tensor1->device, "cuda") == 0) {
+
+            float* result_data;
+            cudaMalloc((void **)&result_data, tensor1->size * sizeof(float));
+            sub_tensor_cuda(tensor1, tensor2, result_data);
+            return create_tensor(result_data, shape, ndim, device);
+        } 
+        else {
+            float* result_data = (float*)malloc(tensor1->size * sizeof(float));
+            if (result_data == NULL) {
+                fprintf(stderr, "Memory allocation failed\n");
+                exit(1);
+            }
+            sub_tensor_cpu(tensor1, tensor2, result_data);
+            return create_tensor(result_data, shape, ndim, device);
+        }
+    }
+
+    Tensor* elementwise_mul_tensor(Tensor* tensor1, Tensor* tensor2) {
+        if (tensor1->ndim != tensor2->ndim) {
+            fprintf(stderr, "Tensors must have the same number of dimensions %d and %d for element-wise multiplication\n", tensor1->ndim, tensor2->ndim);
+            exit(1);
+        }
+
+        if (strcmp(tensor1->device, tensor2->device) != 0) {
+            fprintf(stderr, "Tensors must be on the same device: %s and %s\n", tensor1->device, tensor2->device);
+            exit(1);
+        }
+
+        char* device = (char*)malloc(strlen(tensor1->device) + 1);
+        if (device != NULL) {
+            strcpy(device, tensor1->device);
+        } else {
+            fprintf(stderr, "Memory allocation failed\n");
+            exit(-1);
+        }
+        int ndim = tensor1->ndim;
+        int* shape = (int*)malloc(ndim * sizeof(int));
+        if (shape == NULL) {
+            fprintf(stderr, "Memory allocation failed\n");
+            exit(1);
+        }
+
+        for (int i = 0; i < ndim; i++) {
+            if (tensor1->shape[i] != tensor2->shape[i]) {
+                fprintf(stderr, "Tensors must have the same shape %d and %d at index %d for element-wise multiplication\n", tensor1->shape[i], tensor2->shape[i], i);
+                exit(1);
+            }
+            shape[i] = tensor1->shape[i];
+        }
+
+        if (strcmp(tensor1->device, "cuda") == 0) {
+
+            float* result_data;
+            cudaMalloc((void **)&result_data, tensor1->size * sizeof(float));
+            elementwise_mul_tensor_cuda(tensor1, tensor2, result_data);
+            return create_tensor(result_data, shape, ndim, device);
+        } 
+        else {
+            float* result_data = (float*)malloc(tensor1->size * sizeof(float));
+            if (result_data == NULL) {
+                fprintf(stderr, "Memory allocation failed\n");
+                exit(1);
+            }
+            elementwise_mul_tensor_cpu(tensor1, tensor2, result_data);
+            return create_tensor(result_data, shape, ndim, device);
+        }
+    }
+
+    Tensor* scalar_mul_tensor(Tensor* tensor, float scalar) {
+
+        char* device = (char*)malloc(strlen(tensor->device) + 1);
+        if (device != NULL) {
+            strcpy(device, tensor->device);
+        } else {
+            fprintf(stderr, "Memory allocation failed\n");
+            exit(-1);
+        }
+        int ndim = tensor->ndim;
+        int* shape = (int*)malloc(ndim * sizeof(int));
+        if (shape == NULL) {
+            fprintf(stderr, "Memory allocation failed\n");
+            exit(1);
+        }
+
+        for (int i = 0; i < ndim; i++) {
+            shape[i] = tensor->shape[i];
+        }
+
+        if (strcmp(tensor->device, "cuda") == 0) {
+
+            float* result_data;
+            cudaMalloc((void **)&result_data, tensor->size * sizeof(float));
+            scalar_mul_tensor_cuda(tensor, scalar, result_data);
+            return create_tensor(result_data, shape, ndim, device);
+        } 
+        else {
+            float* result_data = (float*)malloc(tensor->size * sizeof(float));
+            if (result_data == NULL) {
+                fprintf(stderr, "Memory allocation failed\n");
+                exit(1);
+            }
+            scalar_mul_tensor_cpu(tensor, scalar, result_data);
+            return create_tensor(result_data, shape, ndim, device);
+        }
+    }
+
+    Tensor* matmul_tensor(Tensor* tensor1, Tensor* tensor2) {
+        // Check if tensors have compatible shapes for matrix multiplication
+        if (tensor1->shape[1] != tensor2->shape[0]) {
+            fprintf(stderr, "Incompatible shapes for matrix multiplication\n");
+            exit(1);
+        }
+
+        if (strcmp(tensor1->device, tensor2->device) != 0) {
+            fprintf(stderr, "Tensors must be on the same device: %s and %s\n", tensor1->device, tensor2->device);
+            exit(1);
+        }
+
+        char* device = (char*)malloc(strlen(tensor1->device) + 1);
+        if (device != NULL) {
+            strcpy(device, tensor1->device);
+        } else {
+            fprintf(stderr, "Memory allocation failed\n");
+            exit(-1);
+        }
+        int ndim = tensor1->ndim + tensor2->ndim - 2;
+        int* shape = (int*)malloc(ndim * sizeof(int));
+        if (shape == NULL) {
+            fprintf(stderr, "Memory allocation failed\n");
+            exit(1);
+        }
+        for (int i = 0; i < tensor1->ndim - 1; i++) {
+            shape[i] = tensor1->shape[i];
+        }
+        for (int i = tensor1->ndim - 1; i < ndim; i++) {
+            shape[i] = tensor2->shape[i - tensor1->ndim + 2];
+        }
+
+        int size = 1;
+        for (int i = 0; i < ndim; i++) {
+            size *= shape[i];
+        }
+
+        float* result_data = (float*)malloc(size * sizeof(float));
+        if (result_data == NULL) {
+            fprintf(stderr, "Memory allocation failed\n");
+            exit(1);
+        }
+
+        if (strcmp(tensor1->device, "cuda") == 0) {
+
+            float* result_data;
+            cudaMalloc((void **)&result_data, size * sizeof(float));
+            matmul_tensor_cuda(tensor1, tensor2, result_data);
+            return create_tensor(result_data, shape, ndim, device);
+        } 
+        else {
+            float* result_data = (float*)malloc(size * sizeof(float));
+            if (result_data == NULL) {
+                fprintf(stderr, "Memory allocation failed\n");
+                exit(1);
+            }
+            matmul_tensor_cpu(tensor1, tensor2, result_data);
+            return create_tensor(result_data, shape, ndim, device);
+        }
+    }
+
+    Tensor* pow_tensor(Tensor* tensor, float power) {
+        char* device = (char*)malloc(strlen(tensor->device) + 1);
+        if (device != NULL) {
+            strcpy(device, tensor->device);
+        } else {
+            fprintf(stderr, "Memory allocation failed\n");
+            exit(-1);
+        }
+        int ndim = tensor->ndim;
+        int* shape = (int*)malloc(ndim * sizeof(int));
+        if (shape == NULL) {
+            fprintf(stderr, "Memory allocation failed\n");
+            exit(1);
+        }
+
+        for (int i = 0; i < ndim; i++) {
+            shape[i] = tensor->shape[i];
+        }
+
+        if (strcmp(tensor->device, "cuda") == 0) {
+
+            float* result_data;
+            cudaMalloc((void **)&result_data, tensor->size * sizeof(float));
+            pow_tensor_cuda(tensor, power, result_data);
+            return create_tensor(result_data, shape, ndim, device);
+        } 
+        else {
+            float* result_data = (float*)malloc(tensor->size * sizeof(float));
+            if (result_data == NULL) {
+                fprintf(stderr, "Memory allocation failed\n");
+                exit(1);
+            }
+            pow_tensor_cpu(tensor, power, result_data);
+            return create_tensor(result_data, shape, ndim, device);
+        }
+    }
+
+    void reshape_tensor(Tensor* tensor, int* new_shape, int new_ndim) {
+        // Calculate the total number of elements in the new shape
+        int new_size = 1;
+        for (int i = 0; i < new_ndim; i++) {
+            new_size *= new_shape[i];
+        }
+
+        // Check if the total number of elements matches the current tensor's size
+        if (new_size != tensor->size) {
+            fprintf(stderr, "Cannot reshape tensor. Total number of elements in new shape does not match the current size of the tensor.\n");
+            exit(1);
+        }
+
+        // Update the shape
+        tensor->shape = (int*)malloc(new_ndim * sizeof(int));
+        if (tensor->shape == NULL) {
+            fprintf(stderr, "Memory allocation failed\n");
+            exit(1);
+        }
+        for (int i = 0; i < new_ndim; i++) {
+            tensor->shape[i] = new_shape[i];
+        }
+        tensor->ndim = new_ndim;
+
+        // Update the strides
+        tensor->strides = (int*)malloc(new_ndim * sizeof(int));
+        if (tensor->strides == NULL) {
+            fprintf(stderr, "Memory allocation failed\n");
+            exit(1);
+        }
+
+        int stride = 1;
+        for (int i = new_ndim - 1; i >= 0; i--) {
+            tensor->strides[i] = stride;
+            stride *= new_shape[i];
+        }
+    }
+}

norch/csrc/tensor.h

@@ -0,0 +1,29 @@
+#ifndef TENSOR_H
+#define TENSOR_H
+
+typedef struct {
+    float* data;
+    int* strides;
+    int* shape;
+    int* strides_cuda;
+    int* shape_cuda;
+    int ndim;
+    int size;
+    char* device;
+} Tensor;
+
+extern "C" {
+    Tensor* create_tensor(float* data, int* shape, int ndim, char* device);
+    float get_item(Tensor* tensor, int* indices);
+    Tensor* add_tensor(Tensor* tensor1, Tensor* tensor2);
+    Tensor* sum_tensor(Tensor* tensor);
+    Tensor* sub_tensor(Tensor* tensor1, Tensor* tensor2);
+    Tensor* elementwise_mul_tensor(Tensor* tensor1, Tensor* tensor2);
+    Tensor* scalar_mul_tensor(Tensor* tensor, float scalar);
+    void reshape_tensor(Tensor* tensor, int* new_shape, int new_ndim);
+    Tensor* matmul_tensor(Tensor* tensor1, Tensor* tensor2);
+    Tensor* pow_tensor(Tensor* tensor, float power);
+    void to_device(Tensor* tensor, char* device);
+}
+
+#endif /* TENSOR_H */

norch/tensor.py

@@ -163,21 +163,45 @@ class Tensor:
         return result_data
     
     def __mul__(self, other):
-        if self.shape != other.shape:
-            raise ValueError("Tensors must have the same shape for element-wise multiplication")
-        
-        Tensor._C.elementwise_mul_tensor.argtypes = [ctypes.POINTER(CTensor), ctypes.POINTER(CTensor)]
-        Tensor._C.elementwise_mul_tensor.restype = ctypes.POINTER(CTensor)
+        if isinstance(other, (int, float)):
+            result_data = Tensor()
+            result_data.shape = self.shape.copy()
+            result_data.ndim = self.ndim
+            result_data.device = self.device
+            
+            Tensor._C.scalar_mul_tensor.argtypes = [ctypes.POINTER(CTensor), ctypes.c_float]
+            Tensor._C.scalar_mul_tensor.restype = ctypes.POINTER(CTensor)
 
-        result_tensor_ptr = Tensor._C.elementwise_mul_tensor(self.tensor, other.tensor)
+            result_data.tensor = Tensor._C.scalar_mul_tensor(self.tensor, ctypes.c_float(other))
 
-        result_data = Tensor()
-        result_data.tensor = result_tensor_ptr
-        result_data.shape = self.shape.copy()
-        result_data.ndim = self.ndim
-        result_data.device = self.device
+            return result_data
+        elif isinstance(other, Tensor):
+            if self.shape != other.shape:
+                raise ValueError("Tensors must have the same shape for element-wise multiplication")
 
-        return result_data
+            Tensor._C.elementwise_mul_tensor.argtypes = [ctypes.POINTER(CTensor), ctypes.POINTER(CTensor)]
+            Tensor._C.elementwise_mul_tensor.restype = ctypes.POINTER(CTensor)
+
+            result_tensor_ptr = Tensor._C.elementwise_mul_tensor(self.tensor, other.tensor)
+
+            result_data = Tensor()
+            result_data.tensor = result_tensor_ptr
+            result_data.shape = self.shape.copy()
+            result_data.ndim = self.ndim
+            result_data.device = self.device
+
+            return result_data
+        else:
+            raise TypeError("Unsupported operand type(s) for *: '{}' and '{}'".format(type(self), type(other)))
+
+    def __rmul__(self, other):
+        return self.__mul__(other)
+
+    def __neg__(self):
+        return self.__mul__(-1)  
+
+    def __pos__(self):
+        return self  
     
     def __matmul__(self, other):
         if self.ndim != 2 or other.ndim != 2:
@@ -214,4 +238,18 @@ class Tensor:
         result_data.ndim = self.ndim
         result_data.device = self.device
 
+        return result_data
+    
+    def sum(self):
+        Tensor._C.sum_tensor.argtypes = [ctypes.POINTER(CTensor)]
+        Tensor._C.sum_tensor.restype = ctypes.POINTER(CTensor)
+
+        result_tensor_ptr = Tensor._C.sum_tensor(self.tensor)
+
+        result_data = Tensor()
+        result_data.tensor = result_tensor_ptr
+        result_data.shape = [1]
+        result_data.ndim = 1
+        result_data.device = self.device
+
         return result_data

test.py

@@ -16,7 +16,7 @@ def matrix_sum(matrix1, matrix2):
 
 if __name__ == "__main__":
     import norch
-    a = norch.Tensor([[1, 2, 3], [1, 2, 3], [1, 2, 3]]).to("cuda")
+    a = norch.Tensor([[1, 2, 3], [1, 2, 3], [1, 2, 3]])#.to("cuda")
     b = norch.Tensor([[1, 2, 3], [1, 2, 3], [1, 2, 3]])
     import time
     import random
@@ -29,7 +29,7 @@ if __name__ == "__main__":
 
     #d = b-c
 
-    b = a @ a
+    b =    a.sum()
     print(b)
     #print(a ** 2)