Merge pull request #581 from RedisAI/improve_tensor_len_calculation

tensor len member

Author: DvirDukhan

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

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

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:

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:

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

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

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_ */