transpose tensor and matmul autograd

norch/autograd/functions.py

@@ -20,7 +20,6 @@ class ScalarMulBackward:
     def backward(self, gradient):
         return [gradient * self.scalar]
 
-
 class ElementwiseMulBackward:
     def __init__(self, x, y):
         self.input = [x, y]
@@ -28,13 +27,22 @@ class ElementwiseMulBackward:
     def backward(self, gradient):
         return [gradient * self.input[1], gradient * self.input[0]]
     
+class MatmulBackward:
+    def __init__(self, x, y):
+        self.input = [x, y]
+
+    def backward(self, gradient):
+        x, y = self.input
+        return [gradient @ y.T, x.T @ gradient]
+        
+
+    
 class PowBackward:
     def __init__(self, x, power):
         self.input = [x]
         self.power = power
 
     def backward(self, gradient):
-        print(self.input[0], "@@@")
         return [(gradient * self.power) * (self.input[0]) ** (self.power - 1)]
        
 class SumBackward:
@@ -43,6 +51,4 @@ class SumBackward:
 
     def backward(self, gradient):
         # Since sum reduces a tensor to a scalar, gradient is broadcasted to match the original shape.
-        return [float(gradient.tensor.contents.data[0]) * self.input[0].ones_like()]
-
-
+        return [float(gradient.tensor.contents.data[0]) * self.input[0].ones_like()]

norch/csrc/cpu.cpp

@@ -75,3 +75,14 @@ void zeros_like_tensor_cpu(Tensor* tensor, float* result_data) {
         result_data[i] = 0.0;
     }
 }
+
+void transpose_tensor_cpu(Tensor* tensor, float* result_data) {
+    int rows = tensor->shape[0];
+    int cols = tensor->shape[1];
+
+    for (int i = 0; i < rows; i++) {
+        for (int j = 0; j < cols; j++) {
+            result_data[j * rows + i] = tensor->data[i * cols + j];
+        }
+    }
+}

norch/csrc/cpu.h

@@ -12,5 +12,6 @@ void pow_tensor_cpu(Tensor* tensor, float power, float* result_data);
 void scalar_mul_tensor_cpu(Tensor* tensor, float scalar, float* result_data);
 void ones_like_tensor_cpu(Tensor* tensor, float* result_data);
 void zeros_like_tensor_cpu(Tensor* tensor, float* result_data);
+void transpose_tensor_cpu(Tensor* tensor, float* result_data);
 
 #endif /* CPU_H */

norch/csrc/cuda.cu

@@ -188,8 +188,8 @@ __host__ void matmul_tensor_cuda(Tensor* tensor1, Tensor* tensor2, float* result
     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);
+    dim3 number_of_blocks((cols2 + threadsPerBlock.x - 1) / threadsPerBlock.x, (rows1 + threadsPerBlock.y - 1) / threadsPerBlock.y);
+    matmul_tensor_cuda_kernel<<<number_of_blocks, threadsPerBlock>>>(tensor1->data, tensor2->data, result_data, rows1, cols1, cols2);
 
 
     cudaError_t error = cudaGetLastError();
@@ -267,4 +267,32 @@ __host__ void zeros_like_tensor_cuda(Tensor* tensor, float* result_data) {
     cudaDeviceSynchronize();
 }
 
+__global__ void transpose_tensor_cuda_kernel(float* data, float* result_data, int rows, int cols) {
+    int tid_x = blockIdx.x * blockDim.x + threadIdx.x;
+    int tid_y = blockIdx.y * blockDim.y + threadIdx.y;
+
+    if (tid_x < cols && tid_y < rows) {
+        result_data[tid_x * rows + tid_y] = data[tid_y * cols + tid_x];
+    }
+}
+
+__host__ void transpose_tensor_cuda(Tensor* tensor, float* result_data) {
+    
+    int rows = tensor->shape[0];
+    int cols = tensor->shape[1];
+
+    dim3 threadsPerBlock(16, 16);
+    dim3 number_of_blocks((cols + threadsPerBlock.x - 1) / threadsPerBlock.x, (rows + threadsPerBlock.y - 1) / threadsPerBlock.y);
+    transpose_tensor_cuda_kernel<<<number_of_blocks, threadsPerBlock>>>(tensor->data, result_data, rows, cols);
+
+
+    cudaError_t error = cudaGetLastError();
+    if (error != cudaSuccess) {
+        printf("CUDA error: %s\n", cudaGetErrorString(error));
+        exit(-1);
+    }
+
+    cudaDeviceSynchronize();
+}
+
 

norch/csrc/cuda.h

@@ -31,5 +31,9 @@
     __global__ void zeros_like_tensor_cuda_kernel(float* data, float* result_data, int size);
     __host__ void zeros_like_tensor_cuda(Tensor* tensor, float* result_data);
 
+    __global__ void transpose_tensor_cuda_kernel(float* data, float* result_data, int rows, int cols);
+    __host__ void transpose_tensor_cuda(Tensor* tensor, float* result_data);
+
+
 
 #endif /* CUDA_KERNEL_H_ */

norch/csrc/tensor.cpp

@@ -461,7 +461,6 @@ extern "C" {
             stride *= new_shape[i];
         }
     }
-}
 
     Tensor* ones_like_tensor(Tensor* tensor) {
         char* device = (char*)malloc(strlen(tensor->device) + 1);
@@ -535,4 +534,45 @@ extern "C" {
             zeros_like_tensor_cpu(tensor, result_data);
             return create_tensor(result_data, shape, ndim, device);
         }
-    }
+    }
+
+    Tensor* transpose_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 = 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[ndim - 1 - i];
+        }
+
+        int size = tensor->size;
+
+        if (strcmp(tensor->device, "cuda") == 0) {
+
+            float* result_data;
+            cudaMalloc((void **)&result_data, size * sizeof(float));
+            transpose_tensor_cuda(tensor, 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);
+            }
+            transpose_tensor_cpu(tensor, result_data);
+            return create_tensor(result_data, shape, ndim, device);
+        }
+    }
+}

norch/tensor.py

@@ -302,6 +302,10 @@ class Tensor:
         result_data.ndim = 2
         result_data.device = self.device
 
+        result_data.requires_grad = self.requires_grad or other.requires_grad
+        if result_data.requires_grad:
+            result_data.grad_fn = MatmulBackward(self, other)
+
         return result_data
 
     def __pow__(self, power):
@@ -341,4 +345,22 @@ class Tensor:
         if result_data.requires_grad:
             result_data.grad_fn = SumBackward(self)
 
-        return result_data
+        return result_data
+    
+    @property
+    def T(self):
+        if self.ndim != 2:
+            raise ValueError("Transpose requires 2D tensors")
+
+        Tensor._C.transpose_tensor.argtypes = [ctypes.POINTER(CTensor)]
+        Tensor._C.transpose_tensor.restype = ctypes.POINTER(CTensor)
+
+        result_tensor_ptr = Tensor._C.transpose_tensor(self.tensor)
+
+        result_data = Tensor()
+        result_data.tensor = result_tensor_ptr
+        result_data.shape = [self.shape[1], self.shape[0]]
+        result_data.ndim = 2
+        result_data.device = self.device
+
+        return result_data

test.py

@@ -27,12 +27,13 @@ if __name__ == "__main__":
 
     #d = b-c
 
-    a = norch.Tensor([[1, 2, 3], [1, 2, 3], [1, 2, 3]], requires_grad=True)#.to("cuda")
-    b = norch.Tensor([[1, 400, 3], [1, 2, 3], [1, 2, 3]], requires_grad=True)
+    a = norch.Tensor([[1, 2], [1, 2], [1, 2]], requires_grad=True)#.to("cuda")
+    b = norch.Tensor([[1, 400, 3], [1, 2, 3]], requires_grad=True)
 
-    c = (a ** 3)
+    c = (b @ a) * 5
     d = c.sum()
     d.backward()
+
     print(a.grad)
     
     """#print(a)