diff --git a/README.md b/README.md index d6eb1bb..e7042b9 100644 --- a/README.md +++ b/README.md @@ -29,7 +29,7 @@ x1 = norch.Tensor([[1, 2], [3, 4]], requires_grad=True).to("cuda") x2 = norch.Tensor([[4, 3], - [2, 1]], requires_grad=True).to("cuda) + [2, 1]], requires_grad=True).to("cuda") x3 = x1 @ x2 result = x3.sum() diff --git a/build/cpu.o b/build/cpu.o index ecb6dd6..6f33bbc 100644 Binary files a/build/cpu.o and b/build/cpu.o differ diff --git a/build/cuda.cu.o b/build/cuda.cu.o index 427e573..17d7461 100644 Binary files a/build/cuda.cu.o and b/build/cuda.cu.o differ diff --git a/build/distributed.o b/build/distributed.o index caf2c5f..5e37508 100644 Binary files a/build/distributed.o and b/build/distributed.o differ diff --git a/build/tensor.o b/build/tensor.o index d48dad3..09cfc55 100644 Binary files a/build/tensor.o and b/build/tensor.o differ diff --git a/norch/__pycache__/__init__.cpython-38.pyc b/norch/__pycache__/__init__.cpython-38.pyc index 9ef6331..873ae18 100644 Binary files a/norch/__pycache__/__init__.cpython-38.pyc and b/norch/__pycache__/__init__.cpython-38.pyc differ diff --git a/norch/__pycache__/tensor.cpython-38.pyc b/norch/__pycache__/tensor.cpython-38.pyc index 4b1058a..3662c85 100644 Binary files a/norch/__pycache__/tensor.cpython-38.pyc and b/norch/__pycache__/tensor.cpython-38.pyc differ diff --git a/norch/csrc/cuda.cu b/norch/csrc/cuda.cu index 701b63a..268d8a2 100644 --- a/norch/csrc/cuda.cu +++ b/norch/csrc/cuda.cu @@ -42,6 +42,10 @@ __host__ void cuda_to_cpu(Tensor* tensor) { strcpy(tensor->device, device_str); } +__host__ void free_cuda(float* data) { + cudaFree(data); +} + __global__ void add_tensor_cuda_kernel(float* data1, float* data2, 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 702dea1..f6395f0 100644 --- a/norch/csrc/cuda.h +++ b/norch/csrc/cuda.h @@ -4,6 +4,7 @@ __host__ void cpu_to_cuda(Tensor* tensor, int device_id); __host__ void cuda_to_cpu(Tensor* tensor); + __host__ void free_cuda(float* data); __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); diff --git a/norch/csrc/tensor.cpp b/norch/csrc/tensor.cpp index 74d8178..e188e7e 100644 --- a/norch/csrc/tensor.cpp +++ b/norch/csrc/tensor.cpp @@ -16,28 +16,27 @@ extern "C" { 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->device = strdup(device); tensor->strides = (int*)malloc(ndim * sizeof(int)); - if (tensor->strides == NULL) { + tensor->shape = (int*)malloc(ndim * sizeof(int)); + tensor->data = (float*)malloc(tensor->size * sizeof(float)); + + if (tensor->device == NULL || tensor->strides == NULL || tensor->shape == NULL || tensor->data == NULL) { fprintf(stderr, "Memory allocation failed\n"); exit(1); } + + memcpy(tensor->shape, shape, tensor->ndim * sizeof(int)); + strcpy(tensor->device, device); + memcpy(tensor->data, data, tensor->size * sizeof(float)); + int stride = 1; for (int i = ndim - 1; i >= 0; i--) { tensor->strides[i] = stride; @@ -47,6 +46,43 @@ extern "C" { return tensor; } + void delete_tensor(Tensor* tensor) { + if (tensor != NULL) { + + if (tensor->strides != NULL) { + free(tensor->strides); + tensor->strides = NULL; + } + + if (tensor->shape != NULL) { + free(tensor->shape); + tensor->shape = NULL; + } + + if (tensor->data != NULL) { + if (strcmp(tensor->device, "cpu") == 0) { + free(tensor->data); + } else { + free_cuda(tensor->data); + } + tensor->data = NULL; + } + + if (tensor->device != NULL) { + free(tensor->device); + tensor->device = NULL; + } + + if (tensor->strides != NULL) { + free(tensor->strides); + tensor->strides = NULL; + } + + free(tensor); + } + } + + float get_item(Tensor* tensor, int* indices) { int index = 0; for (int i = 0; i < tensor->ndim; i++) { diff --git a/norch/csrc/tensor.h b/norch/csrc/tensor.h index 5669d4f..2465125 100644 --- a/norch/csrc/tensor.h +++ b/norch/csrc/tensor.h @@ -5,8 +5,6 @@ typedef struct { float* data; int* strides; int* shape; - int* strides_cuda; - int* shape_cuda; int ndim; int size; char* device; @@ -14,6 +12,7 @@ typedef struct { extern "C" { Tensor* create_tensor(float* data, int* shape, int ndim, char* device); + void delete_tensor(Tensor* tensor); float get_item(Tensor* tensor, int* indices); Tensor* add_tensor(Tensor* tensor1, Tensor* tensor2); Tensor* sum_tensor(Tensor* tensor, int axis, bool keepdims); diff --git a/norch/libtensor.so b/norch/libtensor.so index 68eb2f2..d7435de 100755 Binary files a/norch/libtensor.so and b/norch/libtensor.so differ diff --git a/norch/optim/optimizers/sgd.py b/norch/optim/optimizers/sgd.py index 62375f6..616d1b9 100644 --- a/norch/optim/optimizers/sgd.py +++ b/norch/optim/optimizers/sgd.py @@ -26,3 +26,4 @@ class SGD(Optimizer): parameter.detach() velocity.detach() + del parameter diff --git a/norch/tensor.py b/norch/tensor.py index 0a1ecab..22061be 100644 --- a/norch/tensor.py +++ b/norch/tensor.py @@ -81,6 +81,12 @@ class Tensor: flat_data, shape = flatten_recursively(nested_list) return flat_data, shape + def __del__(self): + if self.tensor is not None: + Tensor._C.delete_tensor.argtypes = [ctypes.POINTER(CTensor)] + Tensor._C.delete_tensor.restype = None + Tensor._C.delete_tensor(self.tensor) + def __setattr__(self, name, value): if name == 'grad': for hook in self.hooks: