- save index by casting void* to size_t in TensorsNamesToIndices.

- Save the indices of the tensors to persist as values in persistTensors dict, and release TensorsNamesToIndices before DAG run.

Author: alonre24

src/DAG/dag.c

@@ -118,10 +118,8 @@ static int _DAG_PersistTensors(RedisModuleCtx *ctx, RedisAI_RunInfo *rinfo) {
 
     while ((persist_entry = AI_dictNext(persist_iter))) {
         RedisModuleString *persist_key_name = AI_dictGetKey(persist_entry);
-        AI_dictEntry *tensor_entry = AI_dictFind(rinfo->tensorsNamesToIndices, persist_key_name);
-        RedisModule_Assert(tensor_entry);
-        int *index = AI_dictGetVal(tensor_entry);
-        RAI_Tensor *tensor = Dag_GetInternalTensor(rinfo, *index);
+        size_t index = (size_t)AI_dictGetVal(persist_entry);
+        RAI_Tensor *tensor = Dag_GetInternalTensor(rinfo, index);
         tensor = RAI_TensorGetShallowCopy(tensor);
 
         if (_StoreTensorInKeySpace(ctx, tensor, persist_key_name, rinfo->err) == REDISMODULE_ERR) {

src/DAG/dag_builder.c

@@ -31,9 +31,8 @@ int RAI_DAGLoadTensor(RAI_DAGRunCtx *run_info, const char *t_name, RAI_Tensor *t
     RedisModuleString *key_name = RedisModule_CreateString(NULL, t_name, strlen(t_name));
 
     // Add the tensor to the DAG shared tensors and map its name to the relevant index.
-    int *instance = RedisModule_Alloc(sizeof(int));
-    *instance = array_len(rinfo->dagSharedTensors);
-    AI_dictAdd(rinfo->tensorsNamesToIndices, (void *)key_name, (void *)instance);
+    size_t index = array_len(rinfo->dagSharedTensors);
+    AI_dictAdd(rinfo->tensorsNamesToIndices, (void *)key_name, (void *)index);
     RAI_TensorGetShallowCopy(tensor);
     rinfo->dagSharedTensors = array_append(rinfo->dagSharedTensors, (void *)tensor);
     RedisModule_FreeString(NULL, key_name);

src/DAG/dag_execute.c

@@ -17,6 +17,8 @@ int ValidatePersistKeys(RedisAI_RunInfo *rinfo, AI_dict *tensorsNamesToInd,
                 AI_dictReleaseIterator(iter);
                 return REDISMODULE_ERR;
             }
+            size_t index = (size_t)AI_dictGetVal(entry);
+            AI_dictReplace(persistTensorsNames, (void *)persist_key, (void *)index);
         }
         AI_dictReleaseIterator(iter);
     }
@@ -35,29 +37,23 @@ int MapTensorsKeysToIndices(RedisAI_RunInfo *rinfo, AI_dict *tensorsNamesToInd)
                 RAI_SetError(rinfo->err, RAI_EDAGRUN, "ERR INPUT key cannot be found in DAG");
                 return REDISMODULE_ERR;
             }
-            int *ind = AI_dictGetVal(entry);
-            currentOp->inkeys_indices = array_append(currentOp->inkeys_indices, *ind);
+            size_t ind = (size_t)AI_dictGetVal(entry);
+            currentOp->inkeys_indices = array_append(currentOp->inkeys_indices, ind);
         }
 
         for (long long j = 0; j < array_len(currentOp->outkeys); j++) {
             RedisModuleString *key = currentOp->outkeys[j];
-            int *ind = RedisModule_Alloc(sizeof(int));
-            *ind = array_len(rinfo->dagSharedTensors);
+            size_t ind = array_len(rinfo->dagSharedTensors);
 
             // Add a new empty place holder in the array for an output tensor.
-            // If this is MODELSET op, the tensor is already realized.
+            // If this is a TENSORSET op, the tensor is already realized.
             if (currentOp->commandType == REDISAI_DAG_CMD_TENSORSET) {
                 RAI_Tensor *t = RAI_TensorGetShallowCopy(currentOp->outTensor);
                 rinfo->dagSharedTensors = array_append(rinfo->dagSharedTensors, t);
             } else {
                 rinfo->dagSharedTensors = array_append(rinfo->dagSharedTensors, NULL);
             }
-            currentOp->outkeys_indices = array_append(currentOp->outkeys_indices, *ind);
-            AI_dictEntry *entry = AI_dictFind(tensorsNamesToInd, (void *)key);
-            // If this key was already in the dict, remove and free the previous index
-            if (entry) {
-                RedisModule_Free(AI_dictGetVal(entry));
-            }
+            currentOp->outkeys_indices = array_append(currentOp->outkeys_indices, ind);
             AI_dictReplace(tensorsNamesToInd, (void *)key, (void *)ind);
         }
     }

src/DAG/dag_parser.c

@@ -59,9 +59,8 @@ static int _ParseDAGLoadArgs(RedisModuleCtx *ctx, RedisModuleString **argv, int
         }
 
         // Add the tensor to the DAG shared tensors and map its name to the relevant index.
-        int *instance = RedisModule_Alloc(sizeof(int));
-        *instance = array_len(*sharedTensors);
-        AI_dictAdd(tensorsToInd, (void *)key_name, (void *)instance);
+        size_t index = array_len(*sharedTensors);
+        AI_dictAdd(tensorsToInd, (void *)key_name, (void *)index);
         *sharedTensors = array_append(*sharedTensors, (void *)RAI_TensorGetShallowCopy(t));
         number_loaded_keys++;
     }
@@ -113,7 +112,7 @@ static int _ParseDAGPersistArgs(RedisModuleString **argv, int argc, AI_dict *per
             RAI_SetError(err, RAI_EDAGRUN, "ERR PERSIST keys must be unique");
             return -1;
         }
-        AI_dictAdd(persistTensorsNames, (void *)argv[argpos], (void *)1);
+        AI_dictAdd(persistTensorsNames, (void *)argv[argpos], NULL);
         number_keys_to_persist++;
     }
     if (number_keys_to_persist != n_keys) {
@@ -313,6 +312,8 @@ int ParseDAGRunCommand(RedisAI_RunInfo *rinfo, RedisModuleCtx *ctx, RedisModuleS
         REDISMODULE_OK) {
         goto cleanup;
     }
+    AI_dictRelease(rinfo->tensorsNamesToIndices);
+    rinfo->tensorsNamesToIndices = NULL;
     res = REDISMODULE_OK;
 
 cleanup:

src/run_info.c

@@ -161,17 +161,10 @@ void RAI_FreeRunInfo(struct RedisAI_RunInfo *rinfo) {
         RAI_TensorFree(rinfo->dagSharedTensors[i]);
     }
     array_free(rinfo->dagSharedTensors);
-    {
-        AI_dictIterator *iter = AI_dictGetSafeIterator(rinfo->tensorsNamesToIndices);
-        AI_dictEntry *entry;
-        while ((entry = AI_dictNext(iter))) {
-            int *val = (int *)AI_dictGetVal(entry);
-            RedisModule_Free(val);
-        }
-        AI_dictReleaseIterator(iter);
-    }
-    AI_dictRelease(rinfo->tensorsNamesToIndices);
     AI_dictRelease(rinfo->persistTensors);
+    if (rinfo->tensorsNamesToIndices) {
+        AI_dictRelease(rinfo->tensorsNamesToIndices);
+    }
 
     if (rinfo->dagOps) {
         for (size_t i = 0; i < array_len(rinfo->dagOps); i++) {

src/run_info.h

@@ -94,7 +94,7 @@ struct RedisAI_RunInfo {
     int single_op_dag;
     int single_device_dag;
     RAI_Tensor **dagSharedTensors;  // Shared array of tensors that dag ops use.
-    AI_dict *persistTensors;        // The keys of the tensors that will be persisted.
+    AI_dict *persistTensors;        // Associates the tensors to persist with their indices .
     AI_dict *tensorsNamesToIndices; // Maps tensor key name to its (maximal) index.
     RAI_DagOp **dagOps;             // all ops in DAG
     RAI_DagOp **dagDeviceOps;       // all ops in DAG for device