[wip] wip on AI.MODELRUN optimizations

src/backends/tensorflow.c

@@ -371,7 +371,9 @@ RAI_Model *RAI_ModelCreateTF(RAI_Backend backend, const char* devicestr, RAI_Mod
   ret->backend = backend;
   ret->devicestr = RedisModule_Strdup(devicestr);
   ret->inputs = inputs_;
+  ret->ninputs = ninputs;
   ret->outputs = outputs_;
+  ret->noutputs = ninputs;
   ret->opts = opts;
   ret->refCount = 1;
 

src/config.h

@@ -27,4 +27,7 @@ typedef enum {
 #define RAI_COPY_RUN_OUTPUT
 #define RAI_PRINT_BACKEND_ERRORS
 
+#define MODELRUN_BATCH_INITIAL_CAPACITY 10
+#define MODELRUN_PARAM_INITIAL_CAPACITY 10
+
 #endif /* SRC_CONFIG_H_ */

src/model.c

@@ -304,38 +304,39 @@ void RAI_ModelFree(RAI_Model* model, RAI_Error* err) {
 }
 
 RAI_ModelRunCtx* RAI_ModelRunCtxCreate(RAI_Model* model) {
-#define BATCH_INITIAL_SIZE 10
   RAI_ModelRunCtx* mctx = RedisModule_Calloc(1, sizeof(*mctx));
   mctx->model = RAI_ModelGetShallowCopy(model);
-  mctx->batches = array_new(RAI_ModelCtxBatch, BATCH_INITIAL_SIZE);
-#undef BATCH_INITIAL_SIZE
+  mctx->nbatches=0;
+  mctx->batches = array_new(RAI_ModelCtxBatch, MODELRUN_BATCH_INITIAL_CAPACITY);
   return mctx;
 }
 
-static int Model_RunCtxAddParam(RAI_ModelRunCtx* mctx, RAI_ModelCtxParam** paramArr,
+static int Model_RunCtxAddParam(RAI_ModelRunCtx* mctx, RAI_ModelCtxParam** paramArr, 
                                 const char* name, RAI_Tensor* tensor) {
 
   RAI_ModelCtxParam param = {
       .name = name,
       .tensor = tensor ? RAI_TensorGetShallowCopy(tensor): NULL,
   };
   *paramArr = array_append(*paramArr, param);
-  return 1;
+  return REDISMODULE_OK;
 }
 
 int RAI_ModelRunCtxAddInput(RAI_ModelRunCtx* mctx, size_t id, const char* inputName, RAI_Tensor* inputTensor) {
   if (id >= RAI_ModelRunCtxNumBatches(mctx)) {
     // TODO error
-    return 0;
+    return REDISMODULE_ERR;
   }
+  mctx->batches[id].ninputs++;
   return Model_RunCtxAddParam(mctx, &mctx->batches[id].inputs, inputName, inputTensor);
 }
 
 int RAI_ModelRunCtxAddOutput(RAI_ModelRunCtx* mctx, size_t id, const char* outputName) {
   if (id >= RAI_ModelRunCtxNumBatches(mctx)) {
     // TODO error
-    return 0;
+    return REDISMODULE_ERR;
   }
+  mctx->batches[id].noutputs++;
   return Model_RunCtxAddParam(mctx, &mctx->batches[id].outputs, outputName, NULL);
 }
 
@@ -344,40 +345,42 @@ size_t RAI_ModelRunCtxNumInputs(RAI_ModelRunCtx* mctx) {
     return 0;
   }
   // Here we assume batch is well-formed (i.e. number of outputs is equal in all batches)
-  return array_len(mctx->batches[0].inputs);
+  return mctx->batches[0].ninputs;
 }
 
 size_t RAI_ModelRunCtxNumOutputs(RAI_ModelRunCtx* mctx) {
-  if (RAI_ModelRunCtxNumBatches(mctx) == 0) {
+  if (RAI_ModelRunCtxNumBatches(mctx)) {
     return 0;
   }
   // Here we assume batch is well-formed (i.e. number of outputs is equal in all batches)
-  return array_len(mctx->batches[0].outputs);
+  return mctx->batches[0].noutputs;
 }
 
 int RAI_ModelRunCtxAddBatch(RAI_ModelRunCtx* mctx) {
-#define PARAM_INITIAL_SIZE 10
   RAI_ModelCtxBatch batch = {
-    .inputs = array_new(RAI_ModelCtxParam, PARAM_INITIAL_SIZE),
-    .outputs = array_new(RAI_ModelCtxParam, PARAM_INITIAL_SIZE)
+    .inputs = array_new(RAI_ModelCtxParam, MODELRUN_PARAM_INITIAL_CAPACITY),
+    .ninputs = 0,
+    .outputs = array_new(RAI_ModelCtxParam, MODELRUN_PARAM_INITIAL_CAPACITY),
+    .noutputs = 0
   };
-#undef PARAM_INITIAL_SIZE
   array_append(mctx->batches, batch);
+  mctx->nbatches++;
   return array_len(mctx->batches)-1;
 }
 
 size_t RAI_ModelRunCtxNumBatches(RAI_ModelRunCtx* mctx) {
-  return array_len(mctx->batches);
+  return mctx->nbatches;
 }
 
 void RAI_ModelRunCtxCopyBatch(RAI_ModelRunCtx* dest, size_t id_dest, RAI_ModelRunCtx* src, size_t id_src) {
-  size_t ninputs = array_len(src->batches[id_src].inputs);
+  const size_t ninputs = src->batches[id_src].ninputs;
+  const size_t noutputs = src->batches[id_src].noutputs;
+
   for (size_t i=0; i<ninputs; i++) {
     RAI_ModelCtxParam param = src->batches[id_src].inputs[i];
     RAI_ModelRunCtxAddInput(dest, id_dest, param.name, param.tensor);
   }
 
-  size_t noutputs = array_len(src->batches[id_src].outputs);
   for (size_t i=0; i<noutputs; i++) {
     RAI_ModelCtxParam param = src->batches[id_src].outputs[i];
     RAI_ModelRunCtxAddOutput(dest, id_dest, param.name);
@@ -423,7 +426,7 @@ void RAI_ModelRunCtxFree(RAI_ModelRunCtx* mctx) {
 }
 
 int RAI_ModelRun(RAI_ModelRunCtx* mctx, RAI_Error* err) {
-  int ret;
+  int ret = REDISMODULE_ERR;
 
   switch (mctx->model->backend) {
     case RAI_BACKEND_TENSORFLOW:

src/model_struct.h

@@ -35,13 +35,16 @@ typedef struct RAI_ModelCtxParam {
 
 typedef struct RAI_ModelCtxBatch {
   RAI_ModelCtxParam* inputs;
+  size_t ninputs;
   RAI_ModelCtxParam* outputs;
+  size_t noutputs;
 } RAI_ModelCtxBatch;
 
 typedef struct RAI_ModelRunCtx {
   size_t ctxtype;
   RAI_Model* model;
   RAI_ModelCtxBatch* batches;
+  size_t nbatches;
 } RAI_ModelRunCtx;
 
 #endif /* SRC_MODEL_STRUCT_H_ */