diff --git a/src/backends/onnxruntime.c b/src/backends/onnxruntime.c index 60d5fc778..3945df4a5 100644 --- a/src/backends/onnxruntime.c +++ b/src/backends/onnxruntime.c @@ -174,7 +174,7 @@ RAI_Tensor *RAI_TensorCreateFromOrtValue(OrtValue *v, size_t batch_offset, long return NULL; } - ret = RedisModule_Calloc(1, sizeof(*ret)); + ret = RAI_TensorNew(); DLContext ctx = (DLContext){.device_type = kDLCPU, .device_id = 0}; @@ -261,7 +261,6 @@ RAI_Tensor *RAI_TensorCreateFromOrtValue(OrtValue *v, size_t batch_offset, long .manager_ctx = NULL, .deleter = NULL}; - ret->refCount = 1; return ret; } diff --git a/src/backends/tensorflow.c b/src/backends/tensorflow.c index d23fe6a5e..cb42e6d25 100644 --- a/src/backends/tensorflow.c +++ b/src/backends/tensorflow.c @@ -88,7 +88,7 @@ DLDataType RAI_GetDLDataTypeFromTF(TF_DataType dtype) { RAI_Tensor *RAI_TensorCreateFromTFTensor(TF_Tensor *tensor, size_t batch_offset, long long batch_size) { - RAI_Tensor *ret = RedisModule_Calloc(1, sizeof(*ret)); + RAI_Tensor *ret = RAI_TensorNew(); DLContext ctx = (DLContext){.device_type = kDLCPU, .device_id = 0}; @@ -143,7 +143,6 @@ RAI_Tensor *RAI_TensorCreateFromTFTensor(TF_Tensor *tensor, size_t batch_offset, .manager_ctx = NULL, .deleter = NULL}; - ret->refCount = 1; return ret; } diff --git a/src/serialization/RDB/decoder/current/v1/decode_v1.c b/src/serialization/RDB/decoder/current/v1/decode_v1.c index 77a7aa5ba..bfaf3e59d 100644 --- a/src/serialization/RDB/decoder/current/v1/decode_v1.c +++ b/src/serialization/RDB/decoder/current/v1/decode_v1.c @@ -48,7 +48,7 @@ void *RAI_RDBLoadTensor_v1(RedisModuleIO *io) { if (RedisModule_IsIOError(io)) goto cleanup; - RAI_Tensor *ret = RedisModule_Calloc(1, sizeof(*ret)); + RAI_Tensor *ret = RAI_TensorNew(); ret->tensor = (DLManagedTensor){.dl_tensor = (DLTensor){.ctx = ctx, .data = data, .ndim = ndims, @@ -58,7 +58,6 @@ void *RAI_RDBLoadTensor_v1(RedisModuleIO *io) { .byte_offset = byte_offset}, .manager_ctx = NULL, .deleter = NULL}; - ret->refCount = 1; return ret; cleanup: diff --git a/src/serialization/RDB/decoder/previous/v0/decode_v0.c b/src/serialization/RDB/decoder/previous/v0/decode_v0.c index 71facf712..8dfed030b 100644 --- a/src/serialization/RDB/decoder/previous/v0/decode_v0.c +++ b/src/serialization/RDB/decoder/previous/v0/decode_v0.c @@ -41,7 +41,7 @@ void *RAI_RDBLoadTensor_v0(RedisModuleIO *io) { if (RedisModule_IsIOError(io)) goto cleanup; - RAI_Tensor *ret = RedisModule_Calloc(1, sizeof(*ret)); + RAI_Tensor *ret = RAI_TensorNew(); ret->tensor = (DLManagedTensor){.dl_tensor = (DLTensor){.ctx = ctx, .data = data, .ndim = ndims, @@ -51,7 +51,6 @@ void *RAI_RDBLoadTensor_v0(RedisModuleIO *io) { .byte_offset = byte_offset}, .manager_ctx = NULL, .deleter = NULL}; - ret->refCount = 1; return ret; cleanup: diff --git a/src/tensor.c b/src/tensor.c index 7f087501f..079614abc 100644 --- a/src/tensor.c +++ b/src/tensor.c @@ -95,6 +95,12 @@ int Tensor_DataTypeStr(DLDataType dtype, char *dtypestr) { return result; } +RAI_Tensor *RAI_TensorNew(void) { + RAI_Tensor *ret = RedisModule_Calloc(1, sizeof(*ret)); + ret->refCount = 1; + ret->len = LEN_UNKOWN; +} + RAI_Tensor *RAI_TensorCreateWithDLDataType(DLDataType dtype, long long *dims, int ndims, int tensorAllocMode) { @@ -103,7 +109,7 @@ RAI_Tensor *RAI_TensorCreateWithDLDataType(DLDataType dtype, long long *dims, in return NULL; } - RAI_Tensor *ret = RedisModule_Alloc(sizeof(*ret)); + RAI_Tensor *ret = RAI_TensorNew(); int64_t *shape = RedisModule_Alloc(ndims * sizeof(*shape)); int64_t *strides = RedisModule_Alloc(ndims * sizeof(*strides)); @@ -144,7 +150,6 @@ RAI_Tensor *RAI_TensorCreateWithDLDataType(DLDataType dtype, long long *dims, in .manager_ctx = NULL, .deleter = NULL}; - ret->refCount = 1; return ret; } @@ -195,7 +200,7 @@ RAI_Tensor *_TensorCreateWithDLDataTypeAndRString(DLDataType dtype, size_t dtype memcpy(data, blob, nbytes); RAI_HoldString(NULL, rstr); - RAI_Tensor *ret = RedisModule_Alloc(sizeof(*ret)); + RAI_Tensor *ret = RAI_TensorNew(); ret->tensor = (DLManagedTensor){.dl_tensor = (DLTensor){.ctx = ctx, .data = data, .ndim = ndims, @@ -206,7 +211,6 @@ RAI_Tensor *_TensorCreateWithDLDataTypeAndRString(DLDataType dtype, size_t dtype .manager_ctx = rstr, .deleter = RAI_RStringDataTensorDeleter}; - ret->refCount = 1; return ret; } @@ -335,7 +339,7 @@ int RAI_TensorDeepCopy(RAI_Tensor *t, RAI_Tensor **dest) { // Beware: this will take ownership of dltensor RAI_Tensor *RAI_TensorCreateFromDLTensor(DLManagedTensor *dl_tensor) { - RAI_Tensor *ret = RedisModule_Calloc(1, sizeof(*ret)); + RAI_Tensor *ret = RAI_TensorNew(); ret->tensor = (DLManagedTensor){.dl_tensor = (DLTensor){.ctx = dl_tensor->dl_tensor.ctx, @@ -348,7 +352,6 @@ RAI_Tensor *RAI_TensorCreateFromDLTensor(DLManagedTensor *dl_tensor) { .manager_ctx = dl_tensor->manager_ctx, .deleter = dl_tensor->deleter}; - ret->refCount = 1; return ret; } @@ -361,12 +364,15 @@ int RAI_TensorIsDataTypeEqual(RAI_Tensor *t1, RAI_Tensor *t2) { } size_t RAI_TensorLength(RAI_Tensor *t) { - int64_t *shape = t->tensor.dl_tensor.shape; - size_t len = 1; - for (size_t i = 0; i < t->tensor.dl_tensor.ndim; ++i) { - len *= shape[i]; + if (t->len == LEN_UNKOWN) { + int64_t *shape = t->tensor.dl_tensor.shape; + size_t len = 1; + for (size_t i = 0; i < t->tensor.dl_tensor.ndim; ++i) { + len *= shape[i]; + } + t->len = len; } - return len; + return t->len; } size_t RAI_TensorDataSize(RAI_Tensor *t) { return Tensor_DataTypeSize(RAI_TensorDataType(t)); } diff --git a/src/tensor.h b/src/tensor.h index 5768a7fba..45062135c 100644 --- a/src/tensor.h +++ b/src/tensor.h @@ -49,6 +49,14 @@ extern RedisModuleType *RedisAI_TensorType; */ int RAI_TensorInit(RedisModuleCtx *ctx); +/** + * @brief Allocate an empty tensor with no data. + * @note The new tensor ref coutn is 1. + * + * @return RAI_Tensor* - a pointer to the new tensor. + */ +RAI_Tensor *RAI_TensorNew(void); + /** * Allocate the memory and initialise the RAI_Tensor. Creates a tensor based on * the passed 'dataType` string and with the specified number of dimensions diff --git a/src/tensor_struct.h b/src/tensor_struct.h index bb2fd39d2..6203553b6 100644 --- a/src/tensor_struct.h +++ b/src/tensor_struct.h @@ -1,12 +1,12 @@ -#ifndef SRC_TENSOR_STRUCT_H_ -#define SRC_TENSOR_STRUCT_H_ +#pragma once #include "config.h" #include "dlpack/dlpack.h" +#include "limits.h" +#define LEN_UNKOWN ULONG_MAX typedef struct RAI_Tensor { DLManagedTensor tensor; + size_t len; long long refCount; } RAI_Tensor; - -#endif /* SRC_TENSOR_STRUCT_H_ */