From 6985a7d565df9cb31972475fb9af0da738fe5bc6 Mon Sep 17 00:00:00 2001
From: "Chayim I. Kirshen" <chayim@gnupower.net>
Date: Tue, 16 Feb 2021 10:01:03 +0200
Subject: [PATCH 01/32] passing arch in to pack.sh

---
 opt/build/onnxruntime/Dockerfile.arm  | 2 +-
 opt/build/onnxruntime/Dockerfile.arm7 | 2 +-
 opt/build/onnxruntime/Dockerfile.x64  | 3 ++-
 opt/build/onnxruntime/pack.sh         | 7 ++++---
 4 files changed, 8 insertions(+), 6 deletions(-)

diff --git a/opt/build/onnxruntime/Dockerfile.arm b/opt/build/onnxruntime/Dockerfile.arm
index 979b3e852..b24c9bf54 100755
--- a/opt/build/onnxruntime/Dockerfile.arm
+++ b/opt/build/onnxruntime/Dockerfile.arm
@@ -48,6 +48,6 @@ RUN set -e ;\
 	./build.sh ${BUILDARGS} --build_shared_lib
 # RUN ./build.sh ${BUILDARGS} --enable_pybind --build_wheel
 
-RUN ./pack.sh ${ONNXRUNTIME_VER}
+RUN ./pack.sh ${ONNXRUNTIME_VER} ${ARCH}
 
 RUN [ "cross-build-end" ]
diff --git a/opt/build/onnxruntime/Dockerfile.arm7 b/opt/build/onnxruntime/Dockerfile.arm7
index 732a9c127..5c20b58e0 100755
--- a/opt/build/onnxruntime/Dockerfile.arm7
+++ b/opt/build/onnxruntime/Dockerfile.arm7
@@ -41,6 +41,6 @@ RUN set -e ;\
 	./build.sh ${BUILDARGS} --build_shared_lib
 # RUN ./build.sh ${BUILDARGS} --enable_pybind --build_wheel
 
-RUN ./pack.sh ${ONNXRUNTIME_VER}
+RUN ./pack.sh ${ONNXRUNTIME_VER} arm7
 
 RUN [ "cross-build-end" ]
diff --git a/opt/build/onnxruntime/Dockerfile.x64 b/opt/build/onnxruntime/Dockerfile.x64
index a16576fb2..a148a06f6 100755
--- a/opt/build/onnxruntime/Dockerfile.x64
+++ b/opt/build/onnxruntime/Dockerfile.x64
@@ -7,6 +7,7 @@ FROM ${OS}
 
 ARG ONNXRUNTIME_REPO=https://github.com/Microsoft/onnxruntime
 ARG ONNXRUNTIME_VER=1.0.0
+ARG ARCH=x64
 
 RUN apt-get -qq update
 RUN apt-get -qq install -y curl wget tar git
@@ -31,4 +32,4 @@ RUN	cd onnxruntime ;\
 	./build.sh ${BUILDARGS} --build_shared_lib
 # RUN ./build.sh ${BUILDARGS} --enable_pybind --build_wheel
 
-RUN ./pack.sh ${ONNXRUNTIME_VER}
+RUN ./pack.sh ${ONNXRUNTIME_VER} ${ARCH}
diff --git a/opt/build/onnxruntime/pack.sh b/opt/build/onnxruntime/pack.sh
index 52eb55cb7..7c689aec4 100755
--- a/opt/build/onnxruntime/pack.sh
+++ b/opt/build/onnxruntime/pack.sh
@@ -2,6 +2,7 @@
 
 set -e
 VER="$1"
+PLATFORM="$2"
 
 mkdir -p pack/include pack/lib
 cp onnxruntime/build/Linux/MinSizeRel/libonnxruntime.so.${VER} pack/lib/
@@ -15,7 +16,7 @@ git rev-parse HEAD > ../pack/GIT_COMMIT_ID
 cd ..
 cp onnxruntime/include/onnxruntime/core/session/onnxruntime_c_api.h pack/include/
 cd pack/lib/
-ln -s libonnxruntime.so.${VER} libonnxruntime.so
+ln -s libonnxruntime.so.${VER} libonnxruntimeo
 cd ../..
-mv pack onnxruntime-linux-arm64-${VER}
-tar czf onnxruntime-linux-arm64-${VER}.tgz onnxruntime-linux-arm64-${VER}/
+mv pack onnxruntime-linux-${PLATFORM}-${VER}
+tar czf onnxruntime-linux-${PLATFORM}-${VER}.tgz onnxruntime-linux-${PLATFORM}-${VER}/

From cf89a897a5061e64c2abaab31aae3d10fedfcbea Mon Sep 17 00:00:00 2001
From: "Chayim I. Kirshen" <chayim@gnupower.net>
Date: Tue, 16 Feb 2021 10:19:33 +0200
Subject: [PATCH 02/32] arm64v7

---
 opt/build/onnxruntime/Dockerfile.arm7 | 3 ++-
 1 file changed, 2 insertions(+), 1 deletion(-)

diff --git a/opt/build/onnxruntime/Dockerfile.arm7 b/opt/build/onnxruntime/Dockerfile.arm7
index 5c20b58e0..4de482563 100755
--- a/opt/build/onnxruntime/Dockerfile.arm7
+++ b/opt/build/onnxruntime/Dockerfile.arm7
@@ -4,6 +4,7 @@ ARG ONNXRUNTIME_REPO=https://github.com/Microsoft/onnxruntime
 ARG ONNXRUNTIME_BRANCH=rel-1.0.0
 ARG ONNXRUNTIME_VER=1.0.0
 ARG ARCH_FLAG="--arm"
+ARG ARCH=arm64v7
 
 RUN [ "cross-build-start" ]
 
@@ -41,6 +42,6 @@ RUN set -e ;\
 	./build.sh ${BUILDARGS} --build_shared_lib
 # RUN ./build.sh ${BUILDARGS} --enable_pybind --build_wheel
 
-RUN ./pack.sh ${ONNXRUNTIME_VER} arm7
+RUN ./pack.sh ${ONNXRUNTIME_VER} ${ARCH}
 
 RUN [ "cross-build-end" ]

From e74a8ab1b76ea4cee6ed5ac244c8f697b65c8b61 Mon Sep 17 00:00:00 2001
From: alonre24 <alonreshef24@gmail.com>
Date: Tue, 16 Feb 2021 18:00:11 +0200
Subject: [PATCH 03/32] POC - set Redis allocator in onnx and use forked onnx
 repo

---
 opt/build/onnxruntime/Dockerfile.x64 |  2 +-
 src/backends/onnxruntime.c           | 81 ++++++++++++++++++++--------
 2 files changed, 60 insertions(+), 23 deletions(-)

diff --git a/opt/build/onnxruntime/Dockerfile.x64 b/opt/build/onnxruntime/Dockerfile.x64
index a16576fb2..96c760457 100755
--- a/opt/build/onnxruntime/Dockerfile.x64
+++ b/opt/build/onnxruntime/Dockerfile.x64
@@ -5,7 +5,7 @@ ARG OS=debian:buster
 #----------------------------------------------------------------------------------------------
 FROM ${OS}
 
-ARG ONNXRUNTIME_REPO=https://github.com/Microsoft/onnxruntime
+ARG ONNXRUNTIME_REPO=https://github.com/RedisAI/onnxruntime
 ARG ONNXRUNTIME_VER=1.0.0
 
 RUN apt-get -qq update
diff --git a/src/backends/onnxruntime.c b/src/backends/onnxruntime.c
index 4ad5af962..5b64ea637 100644
--- a/src/backends/onnxruntime.c
+++ b/src/backends/onnxruntime.c
@@ -7,13 +7,19 @@
 
 #include "onnxruntime_c_api.h"
 
+
+OrtEnv *env = NULL;
+OrtAllocator *global_allocator = NULL;
+
 int RAI_InitBackendORT(int (*get_api_fn)(const char *, void *)) {
     get_api_fn("RedisModule_Alloc", ((void **)&RedisModule_Alloc));
     get_api_fn("RedisModule_Calloc", ((void **)&RedisModule_Calloc));
     get_api_fn("RedisModule_Free", ((void **)&RedisModule_Free));
     get_api_fn("RedisModule_Realloc", ((void **)&RedisModule_Realloc));
     get_api_fn("RedisModule_Strdup", ((void **)&RedisModule_Strdup));
-
+    get_api_fn("RedisModule_Log", ((void **)&RedisModule_Log));
+    get_api_fn("RedisModule_GetThreadSafeContext", ((void **)&RedisModule_GetThreadSafeContext));
+    get_api_fn("RedisModule_FreeThreadSafeContext", ((void **)&RedisModule_FreeThreadSafeContext));
     return REDISMODULE_OK;
 }
 
@@ -85,8 +91,8 @@ OrtValue *RAI_OrtValueFromTensors(RAI_Tensor **ts, size_t count, RAI_Error *erro
         return NULL;
     }
 
-    OrtAllocator *allocator;
-    status = ort->GetAllocatorWithDefaultOptions(&allocator);
+    //OrtAllocator *allocator;
+    //status = ort->GetAllocatorWithDefaultOptions(&allocator);
     if (status != NULL) {
         return NULL;
     }
@@ -114,7 +120,7 @@ OrtValue *RAI_OrtValueFromTensors(RAI_Tensor **ts, size_t count, RAI_Error *erro
 
     if (count > 1) {
         status =
-            ort->CreateTensorAsOrtValue(allocator, batched_shape, t0->tensor.dl_tensor.ndim,
+            ort->CreateTensorAsOrtValue(global_allocator, batched_shape, t0->tensor.dl_tensor.ndim,
                                         RAI_GetOrtDataTypeFromDL(t0->tensor.dl_tensor.dtype), &out);
         if (status != NULL) {
             goto error;
@@ -133,7 +139,7 @@ OrtValue *RAI_OrtValueFromTensors(RAI_Tensor **ts, size_t count, RAI_Error *erro
         }
     } else {
         status = ort->CreateTensorWithDataAsOrtValue(
-            allocator->Info(allocator), t0->tensor.dl_tensor.data, RAI_TensorByteSize(t0),
+          global_allocator->Info(global_allocator), t0->tensor.dl_tensor.data, RAI_TensorByteSize(t0),
             t0->tensor.dl_tensor.shape, t0->tensor.dl_tensor.ndim,
             RAI_GetOrtDataTypeFromDL(t0->tensor.dl_tensor.dtype), &out);
 
@@ -275,7 +281,31 @@ RAI_Tensor *RAI_TensorCreateFromOrtValue(OrtValue *v, size_t batch_offset, long
     return NULL;
 }
 
-OrtEnv *env = NULL;
+const OrtMemoryInfo* myInfo(const OrtAllocator* allocator) {
+    (void)allocator;
+    const OrtApi *ort = OrtGetApiBase()->GetApi(1);
+    OrtMemoryInfo* mem_info;
+    if (ort->CreateCpuMemoryInfo(OrtDeviceAllocator, OrtMemTypeDefault, &mem_info) != NULL) {
+        return NULL;
+    }
+    return mem_info;
+}
+
+void* myAlloc(OrtAllocator *ptr, size_t size) {
+    (void)ptr;
+    RedisModuleCtx *ctx = RedisModule_GetThreadSafeContext(NULL);
+    RedisModule_Log(ctx, "warning", "Allocating %zu\n", size);
+    RedisModule_FreeThreadSafeContext(ctx);
+    return RedisModule_Alloc(size);
+}
+
+void myFree(OrtAllocator *ptr, void* p) {
+    (void)ptr;
+    RedisModuleCtx *ctx = RedisModule_GetThreadSafeContext(NULL);
+    RedisModule_Log(ctx, "warning", "Freeing\n");
+    RedisModule_FreeThreadSafeContext(ctx);
+    return RedisModule_Free(p);
+}
 
 RAI_Model *RAI_ModelCreateORT(RAI_Backend backend, const char *devicestr, RAI_ModelOpts opts,
                               const char *modeldef, size_t modellen, RAI_Error *error) {
@@ -300,6 +330,11 @@ RAI_Model *RAI_ModelCreateORT(RAI_Backend backend, const char *devicestr, RAI_Mo
 
     if (env == NULL) {
         status = ort->CreateEnv(ORT_LOGGING_LEVEL_WARNING, "test", &env);
+        RedisModule_Assert(!status);
+        status = ort->CreateCustomAllocator(myInfo, myAlloc, myFree, 7, &global_allocator);
+        RedisModule_Assert(!status);
+        status = ort->RegisterAllocator(env, global_allocator);
+        RedisModule_Assert(!status);
     }
 
     if (status != NULL || env == NULL) {
@@ -324,8 +359,10 @@ RAI_Model *RAI_ModelCreateORT(RAI_Backend backend, const char *devicestr, RAI_Mo
         ort->ReleaseSessionOptions(session_options);
         goto error;
     }
-    ort->SetIntraOpNumThreads(session_options, (int)opts.backends_intra_op_parallelism);
-    ort->SetInterOpNumThreads(session_options, (int)opts.backends_inter_op_parallelism);
+    RedisModule_Assert(!ort->SetIntraOpNumThreads(session_options, (int)opts.backends_intra_op_parallelism));
+    RedisModule_Assert(!ort->SetInterOpNumThreads(session_options, (int)opts.backends_inter_op_parallelism));
+    RedisModule_Assert(!ort->SetUseEnvAllocators(session_options, "1"));
+    RedisModule_Assert(!ort->DisableCpuMemArena(session_options));
 
     // TODO: we will need to propose a more dynamic way to request a specific provider,
     // e.g. given the name, in ONNXRuntime
@@ -363,13 +400,13 @@ RAI_Model *RAI_ModelCreateORT(RAI_Backend backend, const char *devicestr, RAI_Mo
         goto error;
     }
 
-    OrtAllocator *allocator;
-    status = ort->GetAllocatorWithDefaultOptions(&allocator);
+    //OrtAllocator *allocator;
+    //status = ort->GetAllocatorWithDefaultOptions(&allocator);
 
     inputs_ = array_new(char *, n_input_nodes);
     for (long long i = 0; i < n_input_nodes; i++) {
         char *input_name;
-        status = ort->SessionGetInputName(session, i, allocator, &input_name);
+        status = ort->SessionGetInputName(session, i, global_allocator, &input_name);
         if (status != NULL) {
             goto error;
         }
@@ -379,7 +416,7 @@ RAI_Model *RAI_ModelCreateORT(RAI_Backend backend, const char *devicestr, RAI_Mo
     outputs_ = array_new(char *, n_output_nodes);
     for (long long i = 0; i < n_output_nodes; i++) {
         char *output_name;
-        status = ort->SessionGetOutputName(session, i, allocator, &output_name);
+        status = ort->SessionGetOutputName(session, i, global_allocator, &output_name);
         if (status != NULL) {
             goto error;
         }
@@ -413,14 +450,14 @@ RAI_Model *RAI_ModelCreateORT(RAI_Backend backend, const char *devicestr, RAI_Mo
     if (inputs_) {
         n_input_nodes = array_len(inputs_);
         for (uint32_t i = 0; i < n_input_nodes; i++) {
-            status = ort->AllocatorFree(allocator, inputs_[i]);
+            status = ort->AllocatorFree(global_allocator, inputs_[i]);
         }
         array_free(inputs_);
     }
     if (outputs_) {
         n_output_nodes = array_len(outputs_);
         for (uint32_t i = 0; i < n_output_nodes; i++) {
-            status = ort->AllocatorFree(allocator, outputs_[i]);
+            status = ort->AllocatorFree(global_allocator, outputs_[i]);
         }
         array_free(outputs_);
     }
@@ -433,16 +470,16 @@ void RAI_ModelFreeORT(RAI_Model *model, RAI_Error *error) {
 
     RedisModule_Free(model->data);
     RedisModule_Free(model->devicestr);
-    OrtAllocator *allocator;
+    //OrtAllocator *allocator;
     OrtStatus *status = NULL;
-    status = ort->GetAllocatorWithDefaultOptions(&allocator);
+    //status = ort->GetAllocatorWithDefaultOptions(&allocator);
     for (uint32_t i = 0; i < model->ninputs; i++) {
-        status = ort->AllocatorFree(allocator, model->inputs[i]);
+        status = ort->AllocatorFree(global_allocator, model->inputs[i]);
     }
     array_free(model->inputs);
 
     for (uint32_t i = 0; i < model->noutputs; i++) {
-        status = ort->AllocatorFree(allocator, model->outputs[i]);
+        status = ort->AllocatorFree(global_allocator, model->outputs[i]);
     }
     array_free(model->outputs);
 
@@ -484,8 +521,8 @@ int RAI_ModelRunORT(RAI_ModelRunCtx **mctxs, RAI_Error *error) {
 
     OrtStatus *status = NULL;
 
-    OrtAllocator *allocator;
-    status = ort->GetAllocatorWithDefaultOptions(&allocator);
+    //OrtAllocator *allocator;
+    //status = ort->GetAllocatorWithDefaultOptions(&allocator);
     if (status != NULL) {
         goto error;
     }
@@ -528,7 +565,7 @@ int RAI_ModelRunORT(RAI_ModelRunCtx **mctxs, RAI_Error *error) {
 
         for (size_t i = 0; i < n_input_nodes; i++) {
             char *input_name;
-            status = ort->SessionGetInputName(session, i, allocator, &input_name);
+            status = ort->SessionGetInputName(session, i, global_allocator, &input_name);
             if (status != NULL) {
                 goto error;
             }
@@ -568,7 +605,7 @@ int RAI_ModelRunORT(RAI_ModelRunCtx **mctxs, RAI_Error *error) {
 
         for (size_t i = 0; i < n_output_nodes; i++) {
             char *output_name;
-            status = ort->SessionGetOutputName(session, i, allocator, &output_name);
+            status = ort->SessionGetOutputName(session, i, global_allocator, &output_name);
             if (status != NULL) {
                 goto error;
             }

From 3da382549a6b5c1ba53d5aec03cf3f6bd0e7a709 Mon Sep 17 00:00:00 2001
From: alonre24 <alonreshef24@gmail.com>
Date: Tue, 16 Feb 2021 18:00:55 +0200
Subject: [PATCH 04/32] Update readies

---
 opt/readies | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/opt/readies b/opt/readies
index 03ba961f1..9dec58bd9 160000
--- a/opt/readies
+++ b/opt/readies
@@ -1 +1 @@
-Subproject commit 03ba961f139ac5e2a9ae395502620e4046c4987b
+Subproject commit 9dec58bd9aeb0c8f32b49600f66e83b29ad277f4

From 089464ec1e329987b16e79a97a58908d7fa4b5d3 Mon Sep 17 00:00:00 2001
From: alonre24 <alonreshef24@gmail.com>
Date: Wed, 17 Feb 2021 14:37:19 +0200
Subject: [PATCH 05/32] Fix docker build for onnx

---
 opt/build/onnxruntime/Dockerfile.x64 | 8 +++++++-
 opt/build/onnxruntime/Makefile       | 2 +-
 2 files changed, 8 insertions(+), 2 deletions(-)

diff --git a/opt/build/onnxruntime/Dockerfile.x64 b/opt/build/onnxruntime/Dockerfile.x64
index 9b551490c..65cd9ed72 100755
--- a/opt/build/onnxruntime/Dockerfile.x64
+++ b/opt/build/onnxruntime/Dockerfile.x64
@@ -6,7 +6,7 @@ ARG OS=debian:buster
 FROM ${OS}
 
 ARG ONNXRUNTIME_REPO=https://github.com/RedisAI/onnxruntime
-ARG ONNXRUNTIME_VER=1.0.0
+ARG ONNXRUNTIME_VER=1.6.0
 ARG ARCH=x64
 
 RUN apt-get -qq update
@@ -19,6 +19,12 @@ RUN pip3 install --upgrade pip setuptools wheel
 # RUN pip3 install numpy
 RUN apt-get -q install -y python3-numpy
 
+ENV LANG=en_US.UTF-8
+RUN apt-get install -y locales && \
+    sed -i -e "s/# $LANG.*/$LANG UTF-8/" /etc/locale.gen && \
+    dpkg-reconfigure --frontend=noninteractive locales && \
+    update-locale LANG=$LANG
+
 WORKDIR /build
 
 ADD ./pack.sh /build/
diff --git a/opt/build/onnxruntime/Makefile b/opt/build/onnxruntime/Makefile
index 6e785c455..2ff062829 100755
--- a/opt/build/onnxruntime/Makefile
+++ b/opt/build/onnxruntime/Makefile
@@ -1,7 +1,7 @@
 
 ROOT=.
 
-VERSION ?= 1.0.0
+VERSION ?= 1.6.0
 OSNICK ?= buster
 
 #----------------------------------------------------------------------------------------------

From 16cb1971affc5372a0d4098495e9700ef1c3e7e6 Mon Sep 17 00:00:00 2001
From: Chayim Kirshen <c@kirshen.com>
Date: Wed, 17 Feb 2021 17:24:59 +0200
Subject: [PATCH 06/32] changes to get the tarball from the docker and publish
 to s3

---
 get_deps.sh                    | 20 +++++++++-----------
 opt/build/onnxruntime/Makefile | 10 +++++-----
 2 files changed, 14 insertions(+), 16 deletions(-)

diff --git a/get_deps.sh b/get_deps.sh
index fc8a6d438..f6d8c8a99 100755
--- a/get_deps.sh
+++ b/get_deps.sh
@@ -12,7 +12,7 @@ HERE="$(cd "$(dirname "${BASH_SOURCE[0]}")" >/dev/null 2>&1 && pwd)"
 if [[ $1 == --help || $1 == help ]]; then
 	cat <<-END
 		[ARGVARS...] get_deps.sh [cpu|gpu] [--help|help]
-		
+
 		Argument variables:
 		CPU=1              Get CPU dependencies
 		GPU=1              Get GPU dependencies
@@ -103,7 +103,7 @@ if [[ $WITH_TF != 0 ]]; then
 			fi
 			if [[ $ARCH == x64 ]]; then
 				TF_ARCH=x86_64
-				
+
 				LIBTF_URL_BASE=https://storage.googleapis.com/tensorflow/libtensorflow
 			elif [[ $ARCH == arm64v8 ]]; then
 				TF_ARCH=arm64
@@ -134,7 +134,7 @@ if [[ $WITH_TF != 0 ]]; then
 		mkdir $LIBTENSORFLOW.x
 		tar xf $LIBTF_ARCHIVE --no-same-owner -C $LIBTENSORFLOW.x
 		mv $LIBTENSORFLOW.x $LIBTENSORFLOW
-		
+
 		echo "Done."
 	else
 		echo "TensorFlow is in place."
@@ -152,7 +152,7 @@ if [[ $WITH_TFLITE != 0 ]]; then
 
 	if [[ ! -d $LIBTFLITE ]]; then
 		echo "Installing TensorFlow Lite ..."
-		
+
 		LIBTF_URL_BASE=https://s3.amazonaws.com/redismodules/tensorflow
 		if [[ $OS == linux ]]; then
 			TFLITE_OS="linux"
@@ -183,7 +183,7 @@ if [[ $WITH_TFLITE != 0 ]]; then
 		mkdir $LIBTFLITE.x
 		tar xf $LIBTFLITE_ARCHIVE --no-same-owner -C $LIBTFLITE.x
 		mv $LIBTFLITE.x $LIBTFLITE
-		
+
 		echo "Done."
 	else
 		echo "TensorFlow Lite is in place."
@@ -241,14 +241,14 @@ if [[ $WITH_PT != 0 ]]; then
 
 			[[ ! -f $LIBTORCH_ARCHIVE || $FORCE == 1 ]] && wget -q $LIBTORCH_URL
 		fi
-		
+
 		rm -rf $LIBTORCH.x
 		mkdir $LIBTORCH.x
 
 		tar xf $LIBTORCH_ARCHIVE --no-same-owner -C $LIBTORCH.x
 		mv $LIBTORCH.x/libtorch $LIBTORCH
 		rmdir $LIBTORCH.x
-		
+
 		echo "Done."
 	else
 		echo "libtorch is in place."
@@ -297,15 +297,13 @@ if [[ $WITH_ORT != 0 ]]; then
 			else
 				ORT_BUILD="-gpu"
 			fi
+			ORT_URL_BASE=https://s3.amazonaws.com/redismodules/onnxruntime
 			if [[ $ARCH == x64 ]]; then
 				ORT_ARCH=x64
-				ORT_URL_BASE=https://github.com/microsoft/onnxruntime/releases/download/v${ORT_VERSION}
 			elif [[ $ARCH == arm64v8 ]]; then
 				ORT_ARCH=arm64
-				ORT_URL_BASE=https://s3.amazonaws.com/redismodules/onnxruntime
 			elif [[ $ARCH == arm32v7 ]]; then
 				ORT_ARCH=arm
-				ORT_URL_BASE=https://s3.amazonaws.com/redismodules/onnxruntime
 			fi
 		elif [[ $OS == macos ]]; then
 			ORT_OS=osx
@@ -322,7 +320,7 @@ if [[ $WITH_ORT != 0 ]]; then
 		mkdir $ONNXRUNTIME.x
 		tar xzf ${ORT_ARCHIVE} --no-same-owner --strip-components=1 -C $ONNXRUNTIME.x
 		mv $ONNXRUNTIME.x $ONNXRUNTIME
-		
+
 		echo "Done."
 	else
 		echo "ONNXRuntime is in place."
diff --git a/opt/build/onnxruntime/Makefile b/opt/build/onnxruntime/Makefile
index 2ff062829..5ab04f22e 100755
--- a/opt/build/onnxruntime/Makefile
+++ b/opt/build/onnxruntime/Makefile
@@ -16,7 +16,7 @@ DOCKER_OS.bionic=ubuntu:bionic
 DOCKER_OS.stretch=debian:stretch-slim
 DOCKER_OS.buster=debian:buster-slim
 DOCKER_OS=$(DOCKER_OS.$(OSNICK))
- 
+
 #----------------------------------------------------------------------------------------------
 
 define targets # (1=OP, 2=op)
@@ -40,10 +40,10 @@ IID_$(1)=$(1)_$(VERSION).iid
 CID_$(1)=$(1)_$(VERSION).cid
 
 build_x64:
-	@docker build --iidfile $$(IID_$(1)) -t redisfab/$(STEM)-$(1):$(VERSION) -f Dockerfile.x64 \
+	@docker build --iidfile $$(IID_$(1)) -t redisfab/$(STEM)-x64:$(VERSION) -f Dockerfile.x64 \
 		--build-arg OS=$(DOCKER_OS) $(ROOT)
 	@docker create --cidfile $$(CID_$(1)) `cat $$(IID_$(1))`
-	@docker cp `cat $$(CID_$(1))`:/build/$(STEM)-$(2)-$(VERSION).tgz .
+	@docker cp `cat $$(CID_$(1))`:/build/$(STEM)-x64-$(VERSION).tgz .
 
 .PHONY: build_x64
 endef
@@ -68,7 +68,7 @@ endef
 
 define publish_x64 # (1=arch, 2=tar-arch)
 publish_x64:
-	@aws s3 cp $(STEM)-$(2)-$(VERSION).tar.gz s3://$(S3_URL)/ --acl public-read
+	@aws s3 cp $(STEM)-$(2)-$(VERSION).tgz s3://$(S3_URL)/ --acl public-read
 
 .PHONY: publish_x64
 endef
@@ -101,7 +101,7 @@ endif
 
 publish: $(PUBLISH_TARGETS)
 
-$(eval $(call publish_x64,x64,x86_64))
+$(eval $(call publish_x64,x64,x64))
 $(eval $(call publish_arm,arm64v8,arm64))
 $(eval $(call publish_arm,arm32v7,arm))
 

From e256864fd79aa1d79a72fa2d3d89c54cc5ed3664 Mon Sep 17 00:00:00 2001
From: alonre24 <alonreshef24@gmail.com>
Date: Sun, 21 Feb 2021 15:21:05 +0200
Subject: [PATCH 07/32] testing onnx custom allocator

---
 src/CMakeLists.txt           |   1 +
 src/backends/onnxruntime.c   |   7 ++-
 src/redisai.c                |   5 ++
 src/tensor.h                 |   2 +-
 src/tensor_struct.h          |   3 +-
 tests/module/Makefile        |  12 +++-
 tests/module/OnnxAllocator.c |  83 ++++++++++++++++++++++++++++
 tests/module/onnx_redisai.h  | 103 +++++++++++++++++++++++++++++++++++
 8 files changed, 208 insertions(+), 8 deletions(-)
 create mode 100644 tests/module/OnnxAllocator.c
 create mode 100644 tests/module/onnx_redisai.h

diff --git a/src/CMakeLists.txt b/src/CMakeLists.txt
index 22fceaa99..ac870f0df 100644
--- a/src/CMakeLists.txt
+++ b/src/CMakeLists.txt
@@ -1,3 +1,4 @@
+set(CMAKE_BUILD_TYPE Debug)
 if (CMAKE_BUILD_TYPE STREQUAL Debug)
 	SET(DEBUG_SRC "${CMAKE_CURRENT_SOURCE_DIR}/../opt/readies/cetara/diag/gdb.c")
 endif()
diff --git a/src/backends/onnxruntime.c b/src/backends/onnxruntime.c
index 5b64ea637..dd5b73edf 100644
--- a/src/backends/onnxruntime.c
+++ b/src/backends/onnxruntime.c
@@ -331,9 +331,9 @@ RAI_Model *RAI_ModelCreateORT(RAI_Backend backend, const char *devicestr, RAI_Mo
     if (env == NULL) {
         status = ort->CreateEnv(ORT_LOGGING_LEVEL_WARNING, "test", &env);
         RedisModule_Assert(!status);
-        status = ort->CreateCustomAllocator(myInfo, myAlloc, myFree, 7, &global_allocator);
+        status = ort->CreateCustomDeviceAllocator(7, myAlloc, myFree, myInfo, &global_allocator);
         RedisModule_Assert(!status);
-        status = ort->RegisterAllocator(env, global_allocator);
+        status = ort->RegisterCustomDeviceAllocator(env, global_allocator);
         RedisModule_Assert(!status);
     }
 
@@ -361,7 +361,8 @@ RAI_Model *RAI_ModelCreateORT(RAI_Backend backend, const char *devicestr, RAI_Mo
     }
     RedisModule_Assert(!ort->SetIntraOpNumThreads(session_options, (int)opts.backends_intra_op_parallelism));
     RedisModule_Assert(!ort->SetInterOpNumThreads(session_options, (int)opts.backends_inter_op_parallelism));
-    RedisModule_Assert(!ort->SetUseEnvAllocators(session_options, "1"));
+
+    RedisModule_Assert(!ort->AddSessionConfigEntry(session_options,"session.use_env_allocators", "1"));
     RedisModule_Assert(!ort->DisableCpuMemArena(session_options));
 
     // TODO: we will need to propose a more dynamic way to request a specific provider,
diff --git a/src/redisai.c b/src/redisai.c
index 28d016564..38bd940d1 100644
--- a/src/redisai.c
+++ b/src/redisai.c
@@ -23,6 +23,7 @@
 #include <sys/resource.h>
 #include <sys/time.h>
 #include <unistd.h>
+#include "backends/onnxruntime.h"
 
 #include "rmutil/alloc.h"
 #include "rmutil/args.h"
@@ -1092,6 +1093,10 @@ static int RedisAI_RegisterApi(RedisModuleCtx *ctx) {
     REGISTER_API(DAGRunOpFree, ctx);
     REGISTER_API(DAGFree, ctx);
 
+    // For ORT test module
+    REGISTER_API(LoadDefaultBackend, ctx);
+    //REGISTER_API(ModelCreateORT, ctx);
+
     return REDISMODULE_OK;
 }
 
diff --git a/src/tensor.h b/src/tensor.h
index 45062135c..41b8a24e0 100644
--- a/src/tensor.h
+++ b/src/tensor.h
@@ -11,7 +11,7 @@
 #define SRC_TENSOR_H_
 
 #include "config.h"
-#include "dlpack/dlpack.h"
+//#include "dlpack/dlpack.h"
 #include "err.h"
 #include "redismodule.h"
 #include "redisai.h"
diff --git a/src/tensor_struct.h b/src/tensor_struct.h
index 6203553b6..d5df8288e 100644
--- a/src/tensor_struct.h
+++ b/src/tensor_struct.h
@@ -1,7 +1,8 @@
 #pragma once
 
 #include "config.h"
-#include "dlpack/dlpack.h"
+#include "../deps/linux-x64-cpu/dlpack/include/dlpack/dlpack.h"
+//#include "dlpack/dlpack.h"
 #include "limits.h"
 
 #define LEN_UNKOWN ULONG_MAX
diff --git a/tests/module/Makefile b/tests/module/Makefile
index 2ad5b1984..6673e0f9a 100644
--- a/tests/module/Makefile
+++ b/tests/module/Makefile
@@ -17,7 +17,9 @@ else
 	SHOBJ_LDFLAGS ?= -bundle -undefined dynamic_lookup
 endif
 
-TEST_MODULES = LLAPI.so
+TEST_MODULES = \
+    LLAPI.so \
+    OnnxAllocator.so
 
 .PHONY: all
 
@@ -27,12 +29,16 @@ all: $(TEST_MODULES)
 	$(MAKE) CFLAGS="-m32" LDFLAGS="-melf_i386"
 
 %.o: %.c
-	$(CC) $(DEBUGFLAGS) -I../../src -DREDIS_MODULE_TARGET -DREDISMODULE_EXPERIMENTAL_API $(SHOBJ_CFLAGS) -fPIC -c $< -o $@
+	$(CC) $(DEBUGFLAGS) -I../../src -I../../deps -DREDIS_MODULE_TARGET -DREDISMODULE_EXPERIMENTAL_API $(SHOBJ_CFLAGS) -fPIC -c $< -o $@
 
-%.so: %.o DAG_utils.o
+LLAPI.so: LLAPI.o DAG_utils.o
 	$(CC) -o $@ $^ $(SHOBJ_LDFLAGS) -lc -lm
 	chmod +x LLAPI.so
 
+OnnxAllocator.so: OnnxAllocator.o
+	$(CC) -o $@ $^ $(SHOBJ_LDFLAGS) -lc -lm
+	chmod +x OnnxAllocator.so
+
 .PHONY: clean
 
 clean:
diff --git a/tests/module/OnnxAllocator.c b/tests/module/OnnxAllocator.c
new file mode 100644
index 000000000..cd5bb306c
--- /dev/null
+++ b/tests/module/OnnxAllocator.c
@@ -0,0 +1,83 @@
+#define REDISMODULE_MAIN
+
+typedef unsigned int uint;
+
+#include "err.h"
+#include <string.h>
+#include <stdio.h>
+#include <stdlib.h>
+#include <unistd.h>
+#include "onnx_redisai.h"
+#include "backends/onnxruntime.h"
+#include "config.h"
+#include "err.h"
+#include <limits.h>
+
+size_t _get_memory_usage(RedisModuleCtx *ctx) {
+    RedisModuleCallReply *reply;
+    reply = RedisModule_Call(ctx, "INFO", "c", "memory");
+    if (RedisModule_CallReplyType(reply) == REDISMODULE_REPLY_STRING) {
+        // Retrieve the used_memory field sub string and convert it to a number
+        size_t len;
+        const char *info_str = RedisModule_CallReplyStringPtr(reply, &len);
+        const char *used_memory_str = strchr(info_str, ':') + 1;
+        char *ptr;
+        size_t used_memory = strtoul(used_memory_str, &ptr, 10);
+        printf("used_memory: %zu\n", used_memory);
+        RedisModule_FreeCallReply(reply);
+        return used_memory;
+    }
+    return 0;
+}
+
+int RAI_onnxAllocator_modelSet(RedisModuleCtx *ctx, RedisModuleString **argv, int argc) {
+    REDISMODULE_NOT_USED(argv);
+
+    if(argc > 1) {
+        RedisModule_WrongArity(ctx);
+        return REDISMODULE_OK;
+    }
+    size_t memory_usage_before = _get_memory_usage(ctx);
+    RAI_ModelOpts opts = {0};
+    RAI_Error err = {0};
+    size_t model_len = 26454;
+
+    // Get onnx Model
+    FILE *fp;
+    char model[model_len];
+    size_t index = 0;
+    char model_path[1000];
+    strcat(getcwd(model_path, sizeof(model_path)), "/../../RedisAI/tests/flow/test_data/linear_iris.onnx");
+    fp = fopen(model_path, "r");
+    char c;
+    while((c = getc(fp)) != EOF) {
+        model[index++] = c;
+    }
+    fclose(fp);
+    model_len = index;
+    RedisAI_LoadDefaultBackend(ctx, RAI_BACKEND_ONNXRUNTIME);
+    RedisAI_ModelCreate(RAI_BACKEND_ONNXRUNTIME, "CPU", NULL, opts, 0, NULL, 0, NULL,
+      model, model_len, &err);
+    size_t memory_usage_after = _get_memory_usage(ctx);
+    printf("Redis allocated : %zu\n", memory_usage_after-memory_usage_before);
+    return RedisModule_ReplyWithSimpleString(ctx, "OK");
+}
+
+int RedisModule_OnLoad(RedisModuleCtx *ctx, RedisModuleString **argv, int argc) {
+    REDISMODULE_NOT_USED(argv);
+    REDISMODULE_NOT_USED(argc);
+
+    if(RedisModule_Init(ctx, "RAI_onnxAllocator", 1, REDISMODULE_APIVER_1) ==
+       REDISMODULE_ERR)
+        return REDISMODULE_ERR;
+
+    if(RedisAI_OnnxInitialize(ctx) != REDISMODULE_OK)
+        RedisModule_Log(ctx, "warning",
+          "could not initialize RedisAI api, running without AI support.");
+
+    if(RedisModule_CreateCommand(ctx, "RAI_onnxAllocator.modelSet",
+      RAI_onnxAllocator_modelSet, "", 0, 0, 0) == REDISMODULE_ERR)
+        return REDISMODULE_ERR;
+
+    return REDISMODULE_OK;
+}
\ No newline at end of file
diff --git a/tests/module/onnx_redisai.h b/tests/module/onnx_redisai.h
new file mode 100644
index 000000000..2fbf65ae0
--- /dev/null
+++ b/tests/module/onnx_redisai.h
@@ -0,0 +1,103 @@
+#pragma once
+
+#include <stdbool.h>
+#include "redismodule.h"
+#include "model_struct.h"
+
+#define REDISAI_LLAPI_VERSION 1
+#define MODULE_API_FUNC(x)    (*x)
+
+#ifdef REDISAI_EXTERN
+#define REDISAI_API extern
+#endif
+
+#ifndef REDISAI_API
+#define REDISAI_API
+#endif
+
+#ifndef REDISAI_H_INCLUDE
+typedef struct RAI_Tensor RAI_Tensor;
+typedef struct RAI_Model RAI_Model;
+typedef struct RAI_Script RAI_Script;
+
+typedef struct RAI_Error RAI_Error;
+
+#endif
+
+#define REDISAI_BACKEND_TENSORFLOW  0
+#define REDISAI_BACKEND_TFLITE      1
+#define REDISAI_BACKEND_TORCH       2
+#define REDISAI_BACKEND_ONNXRUNTIME 3
+
+#define REDISAI_DEVICE_CPU 0
+#define REDISAI_DEVICE_GPU 1
+
+typedef enum RedisAI_ErrorCode {
+    RedisAI_ErrorCode_OK = 0,
+    RedisAI_ErrorCode_EMODELIMPORT,
+    RedisAI_ErrorCode_EMODELCONFIGURE,
+    RedisAI_ErrorCode_EMODELCREATE,
+    RedisAI_ErrorCode_EMODELRUN,
+    RedisAI_ErrorCode_EMODELSERIALIZE,
+    RedisAI_ErrorCode_EMODELFREE,
+    RedisAI_ErrorCode_ESCRIPTIMPORT,
+    RedisAI_ErrorCode_ESCRIPTCONFIGURE,
+    RedisAI_ErrorCode_ESCRIPTCREATE,
+    RedisAI_ErrorCode_ESCRIPTRUN,
+    RedisAI_ErrorCode_EUNSUPPORTEDBACKEND,
+    RedisAI_ErrorCode_EBACKENDNOTLOADED,
+    RedisAI_ErrorCode_ESCRIPTFREE,
+    RedisAI_ErrorCode_ETENSORSET,
+    RedisAI_ErrorCode_ETENSORGET,
+    RedisAI_ErrorCode_EDAGBUILDER,
+    RedisAI_ErrorCode_EDAGRUN,
+    RedisAI_ErrorCode_EFINISHCTX
+} RedisAI_ErrorCode;
+
+REDISAI_API int MODULE_API_FUNC(RedisAI_LoadDefaultBackend)(RedisModuleCtx *ctx, int backend);
+REDISAI_API RAI_Model *MODULE_API_FUNC(RedisAI_ModelCreate)(int backend, char *devicestr, char *tag,
+  RAI_ModelOpts opts, size_t ninputs,
+  const char **inputs, size_t noutputs,
+  const char **outputs,
+  const char *modeldef, size_t modellen,
+  RAI_Error *err);
+
+REDISAI_API int MODULE_API_FUNC(RedisAI_GetLLAPIVersion)();
+
+#ifndef __cplusplus
+#define REDISAI_MODULE_INIT_FUNCTION(ctx, name)                                                    \
+    RedisAI_##name = RedisModule_GetSharedAPI(ctx, "RedisAI_" #name);                              \
+    if (!RedisAI_##name) {                                                                         \
+        RedisModule_Log(ctx, "warning", "could not initialize RedisAI_" #name "\r\n");             \
+        return REDISMODULE_ERR;                                                                    \
+    }
+#else
+#define REDISAI_MODULE_INIT_FUNCTION(ctx, name)                                                    \
+    RedisAI_##name = reinterpret_cast<decltype(RedisAI_##name)>(                                   \
+        RedisModule_GetSharedAPI((RedisModuleCtx *)(ctx), "RedisAI_" #name));                      \
+    if (!RedisAI_##name) {                                                                         \
+        RedisModule_Log(ctx, "warning", "could not initialize RedisAI_" #name "\r\n");             \
+        return REDISMODULE_ERR;                                                                    \
+    }
+#endif
+
+static int RedisAI_OnnxInitialize(RedisModuleCtx *ctx) {
+
+    if (!RedisModule_GetSharedAPI) {
+        RedisModule_Log(ctx, "warning",
+                        "redis version is not compatible with module shared api, "
+                        "use redis 5.0.4 or above.");
+        return REDISMODULE_ERR;
+    }
+
+    REDISAI_MODULE_INIT_FUNCTION(ctx, LoadDefaultBackend);
+    REDISAI_MODULE_INIT_FUNCTION(ctx, ModelCreate);
+
+    REDISAI_MODULE_INIT_FUNCTION(ctx, GetLLAPIVersion);
+
+    if (RedisAI_GetLLAPIVersion() < REDISAI_LLAPI_VERSION) {
+        return REDISMODULE_ERR;
+    }
+
+    return REDISMODULE_OK;
+}

From c5a6ef9cc719b688979e04ee9c9bf1a2c260c2ae Mon Sep 17 00:00:00 2001
From: alonre24 <alonreshef24@gmail.com>
Date: Wed, 24 Feb 2021 15:39:48 +0200
Subject: [PATCH 08/32] Introduce an API that allows setting a custom allocator
 for onnxruntime backend. We register RedisModule Alloc and Free functions as
 onnx allocator, so that onnx is now using them in model create, delete and
 run. - The memory usage of onnx (via redis allocator) and the number of
 accesses to the allocator is saved as part of the module info (defined in
 redisai.c), under a new section named "ai_memory_info". These are exported to
 RedisAI as new fields in the RAI_LoadedBackend struct.

---
 src/backends.c                    |  21 +-
 src/backends.h                    |   5 +
 src/backends/onnxruntime.c        | 378 ++++++++++--------------------
 src/backends/onnxruntime.h        |   4 +
 src/redisai.c                     |  10 +-
 tests/flow/tests_onnx.py          |  59 +++++
 tests/module/Makefile             |   3 +
 tests/module/OnnxAllocator.c      |  84 ++++---
 tests/module/onnx_redisai.h       |   5 +
 tests/module/test_data/mul_1.onnx |   3 +
 10 files changed, 282 insertions(+), 290 deletions(-)
 create mode 100644 tests/module/test_data/mul_1.onnx

diff --git a/src/backends.c b/src/backends.c
index e9e3e4d7c..f647fa585 100644
--- a/src/backends.c
+++ b/src/backends.c
@@ -451,10 +451,27 @@ int RAI_LoadBackend_ONNXRuntime(RedisModuleCtx *ctx, const char *path) {
         return REDISMODULE_ERR;
     }
 
-    RAI_backends.onnx = backend;
+    backend.get_memory_info =
+        (unsigned long long (*)(void))(unsigned long)dlsym(handle, "RAI_GetMemoryInfoORT");
+    if (backend.get_memory_info == NULL) {
+        dlclose(handle);
+        RedisModule_Log(ctx, "warning",
+                        "Backend does not export RAI_GetMemoryInfoORT. ONNX backend "
+                        "not loaded from %s",
+                        path);
+    }
+    backend.get_memory_access_num =
+        (unsigned long long (*)(void))(unsigned long)dlsym(handle, "RAI_GetMemoryAccessORT");
+    if (backend.get_memory_access_num == NULL) {
+        dlclose(handle);
+        RedisModule_Log(ctx, "warning",
+                        "Backend does not export RAI_GetMemoryAccessORT. ONNX backend "
+                        "not loaded from %s",
+                        path);
+    }
 
+    RAI_backends.onnx = backend;
     RedisModule_Log(ctx, "notice", "ONNX backend loaded from %s", path);
-
     return REDISMODULE_OK;
 }
 
diff --git a/src/backends.h b/src/backends.h
index 9c2b79231..045526955 100644
--- a/src/backends.h
+++ b/src/backends.h
@@ -77,6 +77,11 @@ typedef struct RAI_LoadedBackend {
     // Returns the backend version.
     const char *(*get_version)(void);
 
+    // Returns the backend's memory usage for INFO report
+    unsigned long long (*get_memory_info)(void);
+
+    // Returns the number of times that Redis accessed backend allocator.
+    unsigned long long (*get_memory_access_num)(void);
 } RAI_LoadedBackend;
 
 typedef struct RAI_LoadedBackends {
diff --git a/src/backends/onnxruntime.c b/src/backends/onnxruntime.c
index dd5b73edf..d287b7699 100644
--- a/src/backends/onnxruntime.c
+++ b/src/backends/onnxruntime.c
@@ -4,12 +4,22 @@
 #include "backends/util.h"
 #include "tensor.h"
 #include "util/arr_rm_alloc.h"
+#include <stdatomic.h>
 
 #include "onnxruntime_c_api.h"
 
+#define ONNX_API(x)                                                                                \
+    if ((status = (x)) != NULL)                                                                    \
+        goto error;
 
 OrtEnv *env = NULL;
 OrtAllocator *global_allocator = NULL;
+unsigned long long OnnxMemory = 0;
+unsigned long long OnnxMemoryAccessCounter = 0;
+
+unsigned long long RAI_GetMemoryInfoORT() { return OnnxMemory; }
+
+unsigned long long RAI_GetMemoryAccessORT() { return OnnxMemoryAccessCounter; }
 
 int RAI_InitBackendORT(int (*get_api_fn)(const char *, void *)) {
     get_api_fn("RedisModule_Alloc", ((void **)&RedisModule_Alloc));
@@ -20,9 +30,40 @@ int RAI_InitBackendORT(int (*get_api_fn)(const char *, void *)) {
     get_api_fn("RedisModule_Log", ((void **)&RedisModule_Log));
     get_api_fn("RedisModule_GetThreadSafeContext", ((void **)&RedisModule_GetThreadSafeContext));
     get_api_fn("RedisModule_FreeThreadSafeContext", ((void **)&RedisModule_FreeThreadSafeContext));
+    get_api_fn("RedisModule_MallocSize", ((void **)&RedisModule_MallocSize));
     return REDISMODULE_OK;
 }
 
+bool setDeviceId(const char *devicestr, OrtSessionOptions *sessionOptions, RAI_Error *error) {
+
+    RAI_Device device;
+    int64_t deviceid;
+    if (!parseDeviceStr(devicestr, &device, &deviceid)) {
+        RAI_SetError(error, RAI_EMODELCREATE, "ERR unsupported device");
+        return false;
+    }
+
+    // TODO: we will need to propose a more dynamic way to request a specific provider,
+    // e.g. given the name, in ONNXRuntime
+#if RAI_ONNXRUNTIME_USE_CUDA
+    if (device == RAI_DEVICE_GPU) {
+        if (deviceid == -1) {
+            // ORT does not like device id as -1
+            deviceid = 0;
+        }
+        OrtSessionOptionsAppendExecutionProvider_CUDA(session_options, deviceid);
+    }
+#else
+    // TODO: Do dynamic device/provider check with GetExecutionProviderType or something on these
+    // lines
+    if (device == RAI_DEVICE_GPU) {
+        RAI_SetError(error, RAI_EMODELCREATE, "ERR GPU requested but ONNX couldn't find CUDA");
+        return false;
+    }
+#endif
+    return true;
+}
+
 ONNXTensorElementDataType RAI_GetOrtDataTypeFromDL(DLDataType dtype) {
     if (dtype.code == kDLFloat) {
         switch (dtype.bits) {
@@ -80,23 +121,12 @@ DLDataType RAI_GetDLDataTypeFromORT(ONNXTensorElementDataType dtype) {
     default:
         return (DLDataType){.bits = 0};
     }
-    return (DLDataType){.bits = 0};
 }
 
 OrtValue *RAI_OrtValueFromTensors(RAI_Tensor **ts, size_t count, RAI_Error *error) {
     OrtStatus *status = NULL;
     const OrtApi *ort = OrtGetApiBase()->GetApi(1);
 
-    if (count == 0) {
-        return NULL;
-    }
-
-    //OrtAllocator *allocator;
-    //status = ort->GetAllocatorWithDefaultOptions(&allocator);
-    if (status != NULL) {
-        return NULL;
-    }
-
     size_t batch_size = 0;
     size_t batch_byte_size = 0;
 
@@ -106,52 +136,36 @@ OrtValue *RAI_OrtValueFromTensors(RAI_Tensor **ts, size_t count, RAI_Error *erro
     }
 
     RAI_Tensor *t0 = ts[0];
-
     const int ndim = t0->tensor.dl_tensor.ndim;
     int64_t batched_shape[ndim];
-
     for (size_t i = 1; i < ndim; i++) {
         batched_shape[i] = t0->tensor.dl_tensor.shape[i];
     }
-
     batched_shape[0] = batch_size;
 
     OrtValue *out;
-
     if (count > 1) {
-        status =
+        ONNX_API(
             ort->CreateTensorAsOrtValue(global_allocator, batched_shape, t0->tensor.dl_tensor.ndim,
-                                        RAI_GetOrtDataTypeFromDL(t0->tensor.dl_tensor.dtype), &out);
-        if (status != NULL) {
-            goto error;
-        }
+                                        RAI_GetOrtDataTypeFromDL(t0->tensor.dl_tensor.dtype), &out))
 
         char *ort_data;
-        status = ort->GetTensorMutableData(out, (void **)&ort_data);
-        if (status != NULL) {
-            goto error;
-        }
-
+        ONNX_API(ort->GetTensorMutableData(out, (void **)&ort_data))
         size_t offset = 0;
         for (size_t i = 0; i < count; i++) {
             memcpy(ort_data + offset, RAI_TensorData(ts[i]), RAI_TensorByteSize(ts[i]));
             offset += RAI_TensorByteSize(ts[i]);
         }
     } else {
-        status = ort->CreateTensorWithDataAsOrtValue(
-          global_allocator->Info(global_allocator), t0->tensor.dl_tensor.data, RAI_TensorByteSize(t0),
-            t0->tensor.dl_tensor.shape, t0->tensor.dl_tensor.ndim,
-            RAI_GetOrtDataTypeFromDL(t0->tensor.dl_tensor.dtype), &out);
-
-        if (status != NULL) {
-            goto error;
-        }
+        ONNX_API(ort->CreateTensorWithDataAsOrtValue(
+            global_allocator->Info(global_allocator), t0->tensor.dl_tensor.data,
+            RAI_TensorByteSize(t0), t0->tensor.dl_tensor.shape, t0->tensor.dl_tensor.ndim,
+            RAI_GetOrtDataTypeFromDL(t0->tensor.dl_tensor.dtype), &out))
     }
-
     return out;
 
 error:
-    RAI_SetError(error, RAI_EMODELCREATE, ort->GetErrorMessage(status));
+    RAI_SetError(error, RAI_EMODELRUN, ort->GetErrorMessage(status));
     ort->ReleaseStatus(status);
     return NULL;
 }
@@ -160,16 +174,12 @@ RAI_Tensor *RAI_TensorCreateFromOrtValue(OrtValue *v, size_t batch_offset, long
                                          RAI_Error *error) {
     OrtStatus *status = NULL;
     const OrtApi *ort = OrtGetApiBase()->GetApi(1);
-
     RAI_Tensor *ret = NULL;
     int64_t *shape = NULL;
     int64_t *strides = NULL;
 
     int is_tensor;
-    status = ort->IsTensor(v, &is_tensor);
-    if (status != NULL)
-        goto error;
-
+    ONNX_API(ort->IsTensor(v, &is_tensor))
     if (!is_tensor) {
         // TODO: if not tensor, flatten the data structure (sequence or map) and store it in a
         // tensor. If return value is string, emit warning.
@@ -177,30 +187,17 @@ RAI_Tensor *RAI_TensorCreateFromOrtValue(OrtValue *v, size_t batch_offset, long
     }
 
     ret = RAI_TensorNew();
-
     DLContext ctx = (DLContext){.device_type = kDLCPU, .device_id = 0};
-
     OrtTensorTypeAndShapeInfo *info;
-    status = ort->GetTensorTypeAndShape(v, &info);
-    if (status != NULL)
-        goto error;
+    ONNX_API(ort->GetTensorTypeAndShape(v, &info))
 
     {
         size_t ndims;
-        status = ort->GetDimensionsCount(info, &ndims);
-        if (status != NULL)
-            goto error;
-
+        ONNX_API(ort->GetDimensionsCount(info, &ndims))
         int64_t dims[ndims];
-        status = ort->GetDimensions(info, dims, ndims);
-        if (status != NULL)
-            goto error;
-
+        ONNX_API(ort->GetDimensions(info, dims, ndims))
         enum ONNXTensorElementDataType ort_dtype;
-        status = ort->GetTensorElementType(info, &ort_dtype);
-        if (status != NULL)
-            goto error;
-
+        ONNX_API(ort->GetTensorElementType(info, &ort_dtype))
         int64_t total_batch_size = dims[0];
         total_batch_size = total_batch_size > 0 ? total_batch_size : 1;
 
@@ -219,23 +216,14 @@ RAI_Tensor *RAI_TensorCreateFromOrtValue(OrtValue *v, size_t batch_offset, long
             strides[i] *= strides[i + 1] * shape[i + 1];
         }
 
-        // size_t sample_bytesize = TF_TensorByteSize(tensor) / total_batch_size;
-
         DLDataType dtype = RAI_GetDLDataTypeFromORT(ort_dtype);
 #ifdef RAI_COPY_RUN_OUTPUT
         char *ort_data;
-        status = ort->GetTensorMutableData(v, (void **)&ort_data);
-        if (status != NULL) {
-            goto error;
-        }
+        ONNX_API(ort->GetTensorMutableData(v, (void **)&ort_data))
         size_t elem_count;
-        status = ort->GetTensorShapeElementCount(info, &elem_count);
-        if (status != NULL) {
-            goto error;
-        }
+        ONNX_API(ort->GetTensorShapeElementCount(info, &elem_count))
 
         const size_t len = dtype.bits * elem_count / 8;
-
         const size_t total_bytesize = len * sizeof(char);
         const size_t sample_bytesize = total_bytesize / total_batch_size;
         const size_t batch_bytesize = sample_bytesize * batch_size;
@@ -281,146 +269,82 @@ RAI_Tensor *RAI_TensorCreateFromOrtValue(OrtValue *v, size_t batch_offset, long
     return NULL;
 }
 
-const OrtMemoryInfo* myInfo(const OrtAllocator* allocator) {
+const OrtMemoryInfo *myInfo(const OrtAllocator *allocator) {
     (void)allocator;
     const OrtApi *ort = OrtGetApiBase()->GetApi(1);
-    OrtMemoryInfo* mem_info;
+    OrtMemoryInfo *mem_info;
     if (ort->CreateCpuMemoryInfo(OrtDeviceAllocator, OrtMemTypeDefault, &mem_info) != NULL) {
         return NULL;
     }
     return mem_info;
 }
 
-void* myAlloc(OrtAllocator *ptr, size_t size) {
+void *myAlloc(OrtAllocator *ptr, size_t size) {
     (void)ptr;
-    RedisModuleCtx *ctx = RedisModule_GetThreadSafeContext(NULL);
-    RedisModule_Log(ctx, "warning", "Allocating %zu\n", size);
-    RedisModule_FreeThreadSafeContext(ctx);
-    return RedisModule_Alloc(size);
+    void *p = RedisModule_Alloc(size);
+    size_t allocated_size = RedisModule_MallocSize(p);
+    atomic_fetch_add(&OnnxMemory, allocated_size);
+    atomic_fetch_add(&OnnxMemoryAccessCounter, 1);
+    return p;
 }
 
-void myFree(OrtAllocator *ptr, void* p) {
+void myFree(OrtAllocator *ptr, void *p) {
     (void)ptr;
-    RedisModuleCtx *ctx = RedisModule_GetThreadSafeContext(NULL);
-    RedisModule_Log(ctx, "warning", "Freeing\n");
-    RedisModule_FreeThreadSafeContext(ctx);
+    size_t allocated_size = RedisModule_MallocSize(p);
+    atomic_fetch_add(&OnnxMemory, -1 * allocated_size);
+    atomic_fetch_add(&OnnxMemoryAccessCounter, 1);
     return RedisModule_Free(p);
 }
 
 RAI_Model *RAI_ModelCreateORT(RAI_Backend backend, const char *devicestr, RAI_ModelOpts opts,
                               const char *modeldef, size_t modellen, RAI_Error *error) {
-    const OrtApi *ort = OrtGetApiBase()->GetApi(1);
 
-    RAI_Device device;
-    int64_t deviceid;
+    const OrtApi *ort = OrtGetApiBase()->GetApi(1);
     char **inputs_ = NULL;
     char **outputs_ = NULL;
-
-    if (!parseDeviceStr(devicestr, &device, &deviceid)) {
-        RAI_SetError(error, RAI_EMODELCREATE, "ERR unsupported device");
-        return NULL;
-    }
-
-    if (deviceid == -1) {
-        // ORT does not like device id as -1
-        deviceid = 0;
-    }
-
+    OrtSessionOptions *session_options = NULL;
+    OrtSession *session = NULL;
     OrtStatus *status = NULL;
 
     if (env == NULL) {
-        status = ort->CreateEnv(ORT_LOGGING_LEVEL_WARNING, "test", &env);
-        RedisModule_Assert(!status);
-        status = ort->CreateCustomDeviceAllocator(7, myAlloc, myFree, myInfo, &global_allocator);
-        RedisModule_Assert(!status);
-        status = ort->RegisterCustomDeviceAllocator(env, global_allocator);
-        RedisModule_Assert(!status);
+        ONNX_API(ort->CreateEnv(ORT_LOGGING_LEVEL_WARNING, "test", &env))
+        ONNX_API(ort->CreateCustomDeviceAllocator(ORT_API_VERSION, myAlloc, myFree, myInfo,
+                                                  &global_allocator))
+        ONNX_API(ort->RegisterCustomDeviceAllocator(env, global_allocator))
     }
 
-    if (status != NULL || env == NULL) {
-        goto error;
-    }
+    ONNX_API(ort->CreateSessionOptions(&session_options))
 
-    // TODO: these options could be configured at the AI.CONFIG level
-    OrtSessionOptions *session_options;
-    status = ort->CreateSessionOptions(&session_options);
-    if (status != NULL) {
-        goto error;
-    }
-
-    // status = ort->SetSessionThreadPoolSize(session_options, 1);
-    if (status != NULL) {
-        ort->ReleaseSessionOptions(session_options);
-        goto error;
-    }
-
-    status = ort->SetSessionGraphOptimizationLevel(session_options, 1);
-    if (status != NULL) {
-        ort->ReleaseSessionOptions(session_options);
-        goto error;
-    }
-    RedisModule_Assert(!ort->SetIntraOpNumThreads(session_options, (int)opts.backends_intra_op_parallelism));
-    RedisModule_Assert(!ort->SetInterOpNumThreads(session_options, (int)opts.backends_inter_op_parallelism));
+    // These are required to ensure that onnx will use the registered REDIS allocator.
+    ONNX_API(ort->AddSessionConfigEntry(session_options, "session.use_env_allocators", "1"))
+    ONNX_API(ort->DisableCpuMemArena(session_options))
 
-    RedisModule_Assert(!ort->AddSessionConfigEntry(session_options,"session.use_env_allocators", "1"));
-    RedisModule_Assert(!ort->DisableCpuMemArena(session_options));
-
-    // TODO: we will need to propose a more dynamic way to request a specific provider,
-    // e.g. given the name, in ONNXRuntime
-#if RAI_ONNXRUNTIME_USE_CUDA
-    if (device == RAI_DEVICE_GPU) {
-        OrtSessionOptionsAppendExecutionProvider_CUDA(session_options, deviceid);
-    }
-#else
-    // TODO: Do dynamic device/provider check with GetExecutionProviderType or something on these
-    // lines
-    if (device == RAI_DEVICE_GPU) {
-        RAI_SetError(error, RAI_EMODELCREATE, "ERR GPU requested but ONNX couldn't find CUDA");
+    // TODO: these options could be configured at the AI.CONFIG level
+    ONNX_API(ort->SetSessionGraphOptimizationLevel(session_options, 1))
+    ONNX_API(ort->SetIntraOpNumThreads(session_options, (int)opts.backends_intra_op_parallelism))
+    ONNX_API(ort->SetInterOpNumThreads(session_options, (int)opts.backends_inter_op_parallelism))
+    if (!setDeviceId(devicestr, session_options, error)) {
         return NULL;
     }
-#endif
 
-    OrtSession *session;
-    status = ort->CreateSessionFromArray(env, modeldef, modellen, session_options, &session);
-
-    ort->ReleaseSessionOptions(session_options);
-
-    if (status != NULL) {
-        goto error;
-    }
+    ONNX_API(ort->CreateSessionFromArray(env, modeldef, modellen, session_options, &session))
 
     size_t n_input_nodes;
-    status = ort->SessionGetInputCount(session, &n_input_nodes);
-    if (status != NULL) {
-        goto error;
-    }
-
+    ONNX_API(ort->SessionGetInputCount(session, &n_input_nodes))
     size_t n_output_nodes;
-    status = ort->SessionGetOutputCount(session, &n_output_nodes);
-    if (status != NULL) {
-        goto error;
-    }
-
-    //OrtAllocator *allocator;
-    //status = ort->GetAllocatorWithDefaultOptions(&allocator);
+    ONNX_API(ort->SessionGetOutputCount(session, &n_output_nodes))
 
     inputs_ = array_new(char *, n_input_nodes);
     for (long long i = 0; i < n_input_nodes; i++) {
         char *input_name;
-        status = ort->SessionGetInputName(session, i, global_allocator, &input_name);
-        if (status != NULL) {
-            goto error;
-        }
+        ONNX_API(ort->SessionGetInputName(session, i, global_allocator, &input_name))
         inputs_ = array_append(inputs_, input_name);
     }
 
     outputs_ = array_new(char *, n_output_nodes);
     for (long long i = 0; i < n_output_nodes; i++) {
         char *output_name;
-        status = ort->SessionGetOutputName(session, i, global_allocator, &output_name);
-        if (status != NULL) {
-            goto error;
-        }
+        ONNX_API(ort->SessionGetOutputName(session, i, global_allocator, &output_name))
         outputs_ = array_append(outputs_, output_name);
     }
 
@@ -448,6 +372,9 @@ RAI_Model *RAI_ModelCreateORT(RAI_Backend backend, const char *devicestr, RAI_Mo
 
 error:
     RAI_SetError(error, RAI_EMODELCREATE, ort->GetErrorMessage(status));
+    if (session_options) {
+        ort->ReleaseSessionOptions(session_options);
+    }
     if (inputs_) {
         n_input_nodes = array_len(inputs_);
         for (uint32_t i = 0; i < n_input_nodes; i++) {
@@ -462,48 +389,52 @@ RAI_Model *RAI_ModelCreateORT(RAI_Backend backend, const char *devicestr, RAI_Mo
         }
         array_free(outputs_);
     }
+    if (session) {
+        ort->ReleaseSession(session);
+    }
     ort->ReleaseStatus(status);
     return NULL;
 }
 
 void RAI_ModelFreeORT(RAI_Model *model, RAI_Error *error) {
     const OrtApi *ort = OrtGetApiBase()->GetApi(1);
+    OrtStatus *status = NULL;
 
-    RedisModule_Free(model->data);
     RedisModule_Free(model->devicestr);
-    //OrtAllocator *allocator;
-    OrtStatus *status = NULL;
-    //status = ort->GetAllocatorWithDefaultOptions(&allocator);
     for (uint32_t i = 0; i < model->ninputs; i++) {
-        status = ort->AllocatorFree(global_allocator, model->inputs[i]);
+        ONNX_API(ort->AllocatorFree(global_allocator, model->inputs[i]))
     }
     array_free(model->inputs);
 
     for (uint32_t i = 0; i < model->noutputs; i++) {
-        status = ort->AllocatorFree(global_allocator, model->outputs[i]);
+        ONNX_API(ort->AllocatorFree(global_allocator, model->outputs[i]))
     }
     array_free(model->outputs);
 
+    RedisModule_Free(model->data);
     ort->ReleaseSession(model->session);
-
     model->model = NULL;
     model->session = NULL;
+    return;
+
+error:
+    RAI_SetError(error, RAI_EMODELFREE, ort->GetErrorMessage(status));
+    ort->ReleaseStatus(status);
 }
 
 int RAI_ModelRunORT(RAI_ModelRunCtx **mctxs, RAI_Error *error) {
     const OrtApi *ort = OrtGetApiBase()->GetApi(1);
 
     OrtSession *session = mctxs[0]->model->session;
-
     if (session == NULL) {
         RAI_SetError(error, RAI_EMODELRUN, "ERR ONNXRuntime session was not allocated");
-        return 1;
+        return REDISMODULE_ERR;
     }
 
     const size_t nbatches = array_len(mctxs);
     if (nbatches == 0) {
         RAI_SetError(error, RAI_EMODELRUN, "ERR No batches to run");
-        return 1;
+        return REDISMODULE_ERR;
     }
 
     size_t batch_sizes[nbatches];
@@ -521,24 +452,11 @@ int RAI_ModelRunORT(RAI_ModelRunCtx **mctxs, RAI_Error *error) {
     }
 
     OrtStatus *status = NULL;
-
-    //OrtAllocator *allocator;
-    //status = ort->GetAllocatorWithDefaultOptions(&allocator);
-    if (status != NULL) {
-        goto error;
-    }
-
     size_t n_input_nodes;
-    status = ort->SessionGetInputCount(session, &n_input_nodes);
-    if (status != NULL) {
-        goto error;
-    }
+    ONNX_API(ort->SessionGetInputCount(session, &n_input_nodes))
 
     size_t n_output_nodes;
-    status = ort->SessionGetOutputCount(session, &n_output_nodes);
-    if (status != NULL) {
-        goto error;
-    }
+    ONNX_API(ort->SessionGetOutputCount(session, &n_output_nodes))
 
     {
         const char *input_names[n_input_nodes];
@@ -554,23 +472,19 @@ int RAI_ModelRunORT(RAI_ModelRunCtx **mctxs, RAI_Error *error) {
             char msg[70];
             sprintf(msg, "ERR Expected %li inputs but got %li", n_input_nodes, ninputs);
             RAI_SetError(error, RAI_EMODELRUN, msg);
-            return 1;
+            return REDISMODULE_ERR;
         }
 
         if (noutputs != n_output_nodes) {
             char msg[70];
             sprintf(msg, "ERR Expected %li outputs but got %li", n_output_nodes, noutputs);
             RAI_SetError(error, RAI_EMODELRUN, msg);
-            return 1;
+            return REDISMODULE_ERR;
         }
 
         for (size_t i = 0; i < n_input_nodes; i++) {
             char *input_name;
-            status = ort->SessionGetInputName(session, i, global_allocator, &input_name);
-            if (status != NULL) {
-                goto error;
-            }
-
+            ONNX_API(ort->SessionGetInputName(session, i, global_allocator, &input_name))
             input_names[i] = input_name;
 
             RAI_Tensor *batched_input_tensors[nbatches];
@@ -581,75 +495,34 @@ int RAI_ModelRunORT(RAI_ModelRunCtx **mctxs, RAI_Error *error) {
             inputs[i] = RAI_OrtValueFromTensors(batched_input_tensors, nbatches, error);
             if (error->code != RAI_OK) {
                 ort->ReleaseStatus(status);
-                return 1;
+                return REDISMODULE_ERR;
             }
-
-// TODO: use this to check input dim, shapes
-#if 0
-    OrtTypeInfo* typeinfo;
-    status = OrtSessionGetInputTypeInfo(session, i, &typeinfo);
-    const OrtTensorTypeAndShapeInfo* tensor_info = OrtCastTypeInfoToTensorInfo(typeinfo);
-    ONNXTensorElementDataType type = OrtGetTensorElementType(tensor_info);
-    // printf("Input %zu : type=%d", i, type);
-
-    size_t num_dims = OrtGetDimensionsCount(tensor_info);
-    // printf("Input %zu : num_dims=%zu", i, num_dims);
-    input_node_dims.resize(num_dims);
-    OrtGetDimensions(tensor_info, (int64_t*)input_node_dims.data(), num_dims);
-    for (size_t j = 0; j < num_dims; j++) {
-      // printf("Input %zu : dim %zu=%jd", i, j, input_node_dims[j]);
-    }
-
-    OrtReleaseTypeInfo(typeinfo);
-#endif
         }
 
         for (size_t i = 0; i < n_output_nodes; i++) {
             char *output_name;
-            status = ort->SessionGetOutputName(session, i, global_allocator, &output_name);
-            if (status != NULL) {
-                goto error;
-            }
-
+            ONNX_API(ort->SessionGetOutputName(session, i, global_allocator, &output_name))
             output_names[i] = output_name;
             outputs[i] = NULL;
         }
 
-        // ORT_API_STATUS(OrtRun, _Inout_ OrtSession* sess,
-        //                _In_ OrtRunOptions* run_options,
-        //                _In_ const char* const* input_names, _In_ const OrtValue* const* input,
-        //                size_t input_len, _In_ const char* const* output_names, size_t
-        //                output_names_len, _Out_ OrtValue** output);
         OrtRunOptions *run_options = NULL;
-        status = ort->Run(session, run_options, input_names, (const OrtValue *const *)inputs,
-                          n_input_nodes, output_names, n_output_nodes, outputs);
-
-        if (status) {
-            goto error;
-        }
+        ONNX_API(ort->Run(session, run_options, input_names, (const OrtValue *const *)inputs,
+                          n_input_nodes, output_names, n_output_nodes, outputs))
 
         for (size_t i = 0; i < n_output_nodes; i++) {
             if (nbatches > 1) {
                 OrtTensorTypeAndShapeInfo *info;
-                status = ort->GetTensorTypeAndShape(outputs[i], &info);
-                if (status != NULL)
-                    goto error;
-
+                ONNX_API(ort->GetTensorTypeAndShape(outputs[i], &info))
                 size_t ndims;
-                status = ort->GetDimensionsCount(info, &ndims);
-                if (status != NULL)
-                    goto error;
-
+                ONNX_API(ort->GetDimensionsCount(info, &ndims))
                 int64_t dims[ndims];
-                status = ort->GetDimensions(info, dims, ndims);
-                if (status != NULL)
-                    goto error;
-
+                ONNX_API(ort->GetDimensions(info, dims, ndims))
                 if (dims[0] != total_batch_size) {
                     RAI_SetError(error, RAI_EMODELRUN,
                                  "ERR Model did not generate the expected batch size");
                     ort->ReleaseStatus(status);
-                    return 1;
+                    return REDISMODULE_ERR;
                 }
 
                 for (size_t b = 0; b < nbatches; b++) {
@@ -657,7 +530,7 @@ int RAI_ModelRunORT(RAI_ModelRunCtx **mctxs, RAI_Error *error) {
                         outputs[i], batch_offsets[b], batch_sizes[b], error);
                     if (error->code != RAI_OK) {
                         ort->ReleaseStatus(status);
-                        return 1;
+                        return REDISMODULE_ERR;
                     }
                     if (output_tensor) {
                         mctxs[b]->outputs[i].tensor = RAI_TensorGetShallowCopy(output_tensor);
@@ -671,7 +544,7 @@ int RAI_ModelRunORT(RAI_ModelRunCtx **mctxs, RAI_Error *error) {
                 RAI_Tensor *output_tensor = RAI_TensorCreateFromOrtValue(outputs[i], 0, -1, error);
                 if (error->code != RAI_OK) {
                     ort->ReleaseStatus(status);
-                    return 1;
+                    return REDISMODULE_ERR;
                 }
                 if (output_tensor) {
                     mctxs[0]->outputs[i].tensor = RAI_TensorGetShallowCopy(output_tensor);
@@ -681,21 +554,18 @@ int RAI_ModelRunORT(RAI_ModelRunCtx **mctxs, RAI_Error *error) {
                            "unsupported)");
                 }
             }
-
             ort->ReleaseValue(outputs[i]);
         }
-
         for (size_t i = 0; i < n_input_nodes; i++) {
             ort->ReleaseValue(inputs[i]);
         }
-
-        return 0;
+        return REDISMODULE_OK;
     }
 
 error:
     RAI_SetError(error, RAI_EMODELRUN, ort->GetErrorMessage(status));
     ort->ReleaseStatus(status);
-    return 1;
+    return REDISMODULE_ERR;
 }
 
 int RAI_ModelSerializeORT(RAI_Model *model, char **buffer, size_t *len, RAI_Error *error) {
@@ -703,7 +573,7 @@ int RAI_ModelSerializeORT(RAI_Model *model, char **buffer, size_t *len, RAI_Erro
     memcpy(*buffer, model->data, model->datalen);
     *len = model->datalen;
 
-    return 0;
+    return REDISMODULE_OK;
 }
 
 const char *RAI_GetBackendVersionORT(void) { return OrtGetApiBase()->GetVersionString(); }
diff --git a/src/backends/onnxruntime.h b/src/backends/onnxruntime.h
index ceb233a49..1f68c0051 100644
--- a/src/backends/onnxruntime.h
+++ b/src/backends/onnxruntime.h
@@ -6,6 +6,10 @@
 #include "model_struct.h"
 #include "err.h"
 
+unsigned long long RAI_GetMemoryInfoORT(void);
+
+unsigned long long RAI_GetMemoryAccessORT(void);
+
 int RAI_InitBackendORT(int (*get_api_fn)(const char *, void *));
 
 RAI_Model *RAI_ModelCreateORT(RAI_Backend backend, const char *devicestr, RAI_ModelOpts opts,
diff --git a/src/redisai.c b/src/redisai.c
index 38bd940d1..5a0669d2d 100644
--- a/src/redisai.c
+++ b/src/redisai.c
@@ -1095,7 +1095,6 @@ static int RedisAI_RegisterApi(RedisModuleCtx *ctx) {
 
     // For ORT test module
     REGISTER_API(LoadDefaultBackend, ctx);
-    //REGISTER_API(ModelCreateORT, ctx);
 
     return REDISMODULE_OK;
 }
@@ -1107,8 +1106,15 @@ void RAI_moduleInfoFunc(RedisModuleInfoCtx *ctx, int for_crash_report) {
     RedisModule_InfoAddFieldLongLong(ctx, "threads_per_queue", perqueueThreadPoolSize);
     RedisModule_InfoAddFieldLongLong(ctx, "inter_op_parallelism", getBackendsInterOpParallelism());
     RedisModule_InfoAddFieldLongLong(ctx, "intra_op_parallelism", getBackendsIntraOpParallelism());
-    struct rusage self_ru, c_ru;
+    RedisModule_InfoAddSection(ctx, "memory_usage");
+    if (RAI_backends.onnx.get_memory_info) {
+        RedisModule_InfoAddFieldULongLong(ctx, "onnxruntime_memory",
+                                          RAI_backends.onnx.get_memory_info());
+        RedisModule_InfoAddFieldULongLong(ctx, "onnxruntime_memory_access_num",
+                                          RAI_backends.onnx.get_memory_access_num());
+    }
 
+    struct rusage self_ru, c_ru;
     // Return resource usage statistics for the calling process,
     // which is the sum of resources used by all threads in the
     // process
diff --git a/tests/flow/tests_onnx.py b/tests/flow/tests_onnx.py
index 849920560..b255f9ab2 100644
--- a/tests/flow/tests_onnx.py
+++ b/tests/flow/tests_onnx.py
@@ -4,12 +4,33 @@
 import redis
 from includes import *
 from RLTest import Env
+from functools import wraps
 
 '''
 python -m RLTest --test tests_onnx.py --module path/to/redisai.so
 '''
 
 
+def load_onnx_test_module(f):
+    @wraps(f)
+    def wrapper(env, *args, **kwargs):
+        goal_dir = os.path.join(os.path.dirname(__file__), "../module/OnnxAllocator.so")
+        TEST_MODULE_PATH = os.path.abspath(goal_dir)
+        con = env.getConnection()
+        modules = con.execute_command("MODULE", "LIST")
+        if b'RAI_onnxAllocator' in [module[1] for module in modules]:
+            return f(env, *args, **kwargs)
+        try:
+            ret = con.execute_command('MODULE', 'LOAD', TEST_MODULE_PATH)
+            env.assertEqual(ret, b'OK')
+        except Exception as e:
+            env.assertFalse(True)
+            env.debugPrint(str(e), force=True)
+            return
+        return f(env, *args, **kwargs)
+    return wrapper
+
+
 def test_onnx_modelrun_mnist(env):
     if not TEST_ONNX:
         env.debugPrint("skipping {} since TEST_ONNX=0".format(sys._getframe().f_code.co_name), force=True)
@@ -452,3 +473,41 @@ def test_parallelism():
                         for k in con.execute_command("INFO MODULES").decode().split("#")[3].split()[1:]}
     env.assertEqual(load_time_config["ai_inter_op_parallelism"], "2")
     env.assertEqual(load_time_config["ai_intra_op_parallelism"], "2")
+
+
+@load_onnx_test_module
+def test_onnx_use_custom_allocator(env):
+    if not TEST_ONNX:
+        env.debugPrint("skipping {} since TEST_ONNX=0".format(sys._getframe().f_code.co_name), force=True)
+        return
+
+    con = env.getConnection()
+    ret = con.execute_command("RAI_onnxAllocator.modelSet")
+    env.assertEquals(ret, b'OK')
+
+    test_data_path = os.path.join(os.path.dirname(__file__), 'test_data')
+    model_filename = os.path.join(test_data_path, 'mnist.onnx')
+    sample_filename = os.path.join(test_data_path, 'one.raw')
+
+    with open(model_filename, 'rb') as f:
+        model_pb = f.read()
+    with open(sample_filename, 'rb') as f:
+        sample_raw = f.read()
+
+    ret = con.execute_command('AI.MODELSET', 'm{1}', 'ONNX', DEVICE, 'BLOB', model_pb)
+    env.assertEqual(ret, b'OK')
+    con.execute_command('AI.TENSORSET', 'a{1}', 'FLOAT', 1, 1, 28, 28, 'BLOB', sample_raw)
+
+    # Expect 16 allocator's access from onnx during the run.
+    ai_memory_config = {k.split(":")[0]: k.split(":")[1]
+                        for k in con.execute_command("INFO MODULES").decode().split("#")[4].split()[1:]}
+    allocator_access_num_before = ai_memory_config["ai_onnxruntime_memory_access_num"]
+    con.execute_command('AI.MODELRUN', 'm{1}', 'INPUTS', 'a{1}', 'OUTPUTS', 'b{1}')
+    ai_memory_config = {k.split(":")[0]: k.split(":")[1]
+                        for k in con.execute_command("INFO MODULES").decode().split("#")[4].split()[1:]}
+    allocator_access_num_after = ai_memory_config["ai_onnxruntime_memory_access_num"]
+    env.assertEqual(int(allocator_access_num_after) - int(allocator_access_num_before), 16)
+
+    values = con.execute_command('AI.TENSORGET', 'b{1}', 'VALUES')
+    argmax = max(range(len(values)), key=lambda i: values[i])
+    env.assertEqual(argmax, 1)
diff --git a/tests/module/Makefile b/tests/module/Makefile
index 6673e0f9a..131e035d3 100644
--- a/tests/module/Makefile
+++ b/tests/module/Makefile
@@ -21,6 +21,8 @@ TEST_MODULES = \
     LLAPI.so \
     OnnxAllocator.so
 
+SAVE_PATH = @echo "const char* path = \"$(dir $(abspath $(lastword $(MAKEFILE_LIST))))\";" > path.h
+
 .PHONY: all
 
 all: $(TEST_MODULES)
@@ -29,6 +31,7 @@ all: $(TEST_MODULES)
 	$(MAKE) CFLAGS="-m32" LDFLAGS="-melf_i386"
 
 %.o: %.c
+	${SAVE_PATH}
 	$(CC) $(DEBUGFLAGS) -I../../src -I../../deps -DREDIS_MODULE_TARGET -DREDISMODULE_EXPERIMENTAL_API $(SHOBJ_CFLAGS) -fPIC -c $< -o $@
 
 LLAPI.so: LLAPI.o DAG_utils.o
diff --git a/tests/module/OnnxAllocator.c b/tests/module/OnnxAllocator.c
index cd5bb306c..7e1afe913 100644
--- a/tests/module/OnnxAllocator.c
+++ b/tests/module/OnnxAllocator.c
@@ -12,22 +12,25 @@ typedef unsigned int uint;
 #include "config.h"
 #include "err.h"
 #include <limits.h>
+#include "path.h"
 
-size_t _get_memory_usage(RedisModuleCtx *ctx) {
-    RedisModuleCallReply *reply;
-    reply = RedisModule_Call(ctx, "INFO", "c", "memory");
-    if (RedisModule_CallReplyType(reply) == REDISMODULE_REPLY_STRING) {
-        // Retrieve the used_memory field sub string and convert it to a number
-        size_t len;
-        const char *info_str = RedisModule_CallReplyStringPtr(reply, &len);
-        const char *used_memory_str = strchr(info_str, ':') + 1;
-        char *ptr;
-        size_t used_memory = strtoul(used_memory_str, &ptr, 10);
-        printf("used_memory: %zu\n", used_memory);
-        RedisModule_FreeCallReply(reply);
-        return used_memory;
-    }
-    return 0;
+typedef struct MemoryInfo {
+    size_t usage;
+    size_t access_counter;
+} MemoryInfo;
+
+MemoryInfo _get_memory_info() {
+
+    RedisModuleServerInfoData *server_info = RedisModule_GetServerInfo(NULL, "everything");
+    int onnx_info_error;
+    unsigned long long onnx_counter = RedisModule_ServerInfoGetFieldUnsigned(server_info,
+      "ai_onnxruntime_memory_access_num", &onnx_info_error);
+    RedisModule_Assert(onnx_info_error == 0);
+    unsigned long long onnx_mem = RedisModule_ServerInfoGetFieldUnsigned(server_info,
+      "ai_onnxruntime_memory", &onnx_info_error);
+    RedisModule_Assert(onnx_info_error == 0);
+    RedisModule_FreeServerInfo(NULL, server_info);
+    return (MemoryInfo){.usage = onnx_mem, .access_counter = onnx_counter};
 }
 
 int RAI_onnxAllocator_modelSet(RedisModuleCtx *ctx, RedisModuleString **argv, int argc) {
@@ -37,32 +40,49 @@ int RAI_onnxAllocator_modelSet(RedisModuleCtx *ctx, RedisModuleString **argv, in
         RedisModule_WrongArity(ctx);
         return REDISMODULE_OK;
     }
-    size_t memory_usage_before = _get_memory_usage(ctx);
-    RAI_ModelOpts opts = {0};
-    RAI_Error err = {0};
-    size_t model_len = 26454;
+    RedisAI_LoadDefaultBackend(ctx, RAI_BACKEND_ONNXRUNTIME);
+    RedisModule_Assert(_get_memory_info().usage == 0);
 
-    // Get onnx Model
+    // Load onnx model to local buffer.
+    size_t model_len = 130;
     FILE *fp;
-    char model[model_len];
+    char model_data[model_len];
     size_t index = 0;
-    char model_path[1000];
-    strcat(getcwd(model_path, sizeof(model_path)), "/../../RedisAI/tests/flow/test_data/linear_iris.onnx");
-    fp = fopen(model_path, "r");
+    char model_path[1024];
+    strcpy(model_path, path);
+    strcat(model_path, "/test_data/mul_1.onnx");
+    fp = fopen(model_path, "rb");
     char c;
-    while((c = getc(fp)) != EOF) {
-        model[index++] = c;
+    while((c = fgetc(fp)) != EOF) {
+        model_data[index++] = c;
     }
     fclose(fp);
-    model_len = index;
-    RedisAI_LoadDefaultBackend(ctx, RAI_BACKEND_ONNXRUNTIME);
-    RedisAI_ModelCreate(RAI_BACKEND_ONNXRUNTIME, "CPU", NULL, opts, 0, NULL, 0, NULL,
-      model, model_len, &err);
-    size_t memory_usage_after = _get_memory_usage(ctx);
-    printf("Redis allocated : %zu\n", memory_usage_after-memory_usage_before);
+    RedisModule_Assert(model_len == index);
+
+    // Create model and verify that onnx backend uses REDIS allocator.
+    RAI_ModelOpts opts = {0};
+    RAI_Error err = {0};
+    RAI_Model *model = RedisAI_ModelCreate(RAI_BACKEND_ONNXRUNTIME, "CPU", NULL, opts, 0, NULL, 0, NULL,
+      model_data, model_len, &err);
+    RedisModule_Assert(err.code == RAI_OK);
+
+    // Expect 3 allocation access by onnx backend for allocating model, input name
+    // and output name
+    MemoryInfo mem_info = _get_memory_info();
+    if (mem_info.access_counter != 3 || mem_info.usage == 0) {
+        return RedisModule_ReplyWithSimpleString(ctx, "model set allocation error");
+    }
+
+    RedisAI_ModelFree(model, &err);
+    RedisModule_Assert(err.code == RAI_OK);
+    mem_info = _get_memory_info();
+    if (mem_info.access_counter != 6 || mem_info.usage != 0) {
+        return RedisModule_ReplyWithSimpleString(ctx, "model free error");
+    }
     return RedisModule_ReplyWithSimpleString(ctx, "OK");
 }
 
+
 int RedisModule_OnLoad(RedisModuleCtx *ctx, RedisModuleString **argv, int argc) {
     REDISMODULE_NOT_USED(argv);
     REDISMODULE_NOT_USED(argc);
diff --git a/tests/module/onnx_redisai.h b/tests/module/onnx_redisai.h
index 2fbf65ae0..791a7800a 100644
--- a/tests/module/onnx_redisai.h
+++ b/tests/module/onnx_redisai.h
@@ -61,6 +61,9 @@ REDISAI_API RAI_Model *MODULE_API_FUNC(RedisAI_ModelCreate)(int backend, char *d
   const char **outputs,
   const char *modeldef, size_t modellen,
   RAI_Error *err);
+REDISAI_API void MODULE_API_FUNC(RedisAI_ModelFree)(RAI_Model *model, RAI_Error *err);
+REDISAI_API int MODULE_API_FUNC(RedisAI_ModelRun)(RAI_ModelRunCtx **mctx, long long n,
+  RAI_Error *err);
 
 REDISAI_API int MODULE_API_FUNC(RedisAI_GetLLAPIVersion)();
 
@@ -92,6 +95,8 @@ static int RedisAI_OnnxInitialize(RedisModuleCtx *ctx) {
 
     REDISAI_MODULE_INIT_FUNCTION(ctx, LoadDefaultBackend);
     REDISAI_MODULE_INIT_FUNCTION(ctx, ModelCreate);
+    REDISAI_MODULE_INIT_FUNCTION(ctx, ModelFree);
+    REDISAI_MODULE_INIT_FUNCTION(ctx, ModelRun);
 
     REDISAI_MODULE_INIT_FUNCTION(ctx, GetLLAPIVersion);
 
diff --git a/tests/module/test_data/mul_1.onnx b/tests/module/test_data/mul_1.onnx
new file mode 100644
index 000000000..c34899b6b
--- /dev/null
+++ b/tests/module/test_data/mul_1.onnx
@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:71f431c4e9321ec6fbeb158d02ed240459a7dcc98673fa79a4f439ce42efaf10
+size 130

From c00874eaf01115ee6654a223900cb8a17cfe5ba2 Mon Sep 17 00:00:00 2001
From: alonre24 <alonreshef24@gmail.com>
Date: Wed, 24 Feb 2021 22:51:15 +0200
Subject: [PATCH 09/32] Fix docker build

---
 opt/build/onnxruntime/pack.sh | 3 ++-
 1 file changed, 2 insertions(+), 1 deletion(-)

diff --git a/opt/build/onnxruntime/pack.sh b/opt/build/onnxruntime/pack.sh
index 7c689aec4..fdd1113b6 100755
--- a/opt/build/onnxruntime/pack.sh
+++ b/opt/build/onnxruntime/pack.sh
@@ -15,8 +15,9 @@ cd onnxruntime/
 git rev-parse HEAD > ../pack/GIT_COMMIT_ID
 cd ..
 cp onnxruntime/include/onnxruntime/core/session/onnxruntime_c_api.h pack/include/
+cp onnxruntime/include/onnxruntime/core/providers/cuda/cuda_provider_factory.h pack/include/
 cd pack/lib/
-ln -s libonnxruntime.so.${VER} libonnxruntimeo
+ln -s libonnxruntime.so.${VER} libonnxruntime.so
 cd ../..
 mv pack onnxruntime-linux-${PLATFORM}-${VER}
 tar czf onnxruntime-linux-${PLATFORM}-${VER}.tgz onnxruntime-linux-${PLATFORM}-${VER}/

From 67c40f81c1f68f954917b1407a00779e25e0df3e Mon Sep 17 00:00:00 2001
From: alonre24 <alonreshef24@gmail.com>
Date: Thu, 25 Feb 2021 10:46:36 +0200
Subject: [PATCH 10/32] Remove test model and use only flow test (the test
 module is redundant since onnx memory info is saved in redis' INFO report,
 that can be produced by calling the INFO MODULES command)

---
 src/CMakeLists.txt                          |   1 -
 src/backends/onnxruntime.c                  |  58 +++++------
 src/redisai.c                               |   7 +-
 src/tensor.h                                |   2 +-
 src/tensor_struct.h                         |   3 +-
 tests/flow/includes.py                      |   3 -
 tests/{module => flow}/test_data/mul_1.onnx |   0
 tests/flow/tests_onnx.py                    |  52 +++++-----
 tests/module/Makefile                       |  11 +-
 tests/module/OnnxAllocator.c                | 103 -------------------
 tests/module/onnx_redisai.h                 | 108 --------------------
 11 files changed, 62 insertions(+), 286 deletions(-)
 rename tests/{module => flow}/test_data/mul_1.onnx (100%)
 delete mode 100644 tests/module/OnnxAllocator.c
 delete mode 100644 tests/module/onnx_redisai.h

diff --git a/src/CMakeLists.txt b/src/CMakeLists.txt
index ac870f0df..22fceaa99 100644
--- a/src/CMakeLists.txt
+++ b/src/CMakeLists.txt
@@ -1,4 +1,3 @@
-set(CMAKE_BUILD_TYPE Debug)
 if (CMAKE_BUILD_TYPE STREQUAL Debug)
 	SET(DEBUG_SRC "${CMAKE_CURRENT_SOURCE_DIR}/../opt/readies/cetara/diag/gdb.c")
 endif()
diff --git a/src/backends/onnxruntime.c b/src/backends/onnxruntime.c
index d287b7699..fb9eefffb 100644
--- a/src/backends/onnxruntime.c
+++ b/src/backends/onnxruntime.c
@@ -17,6 +17,33 @@ OrtAllocator *global_allocator = NULL;
 unsigned long long OnnxMemory = 0;
 unsigned long long OnnxMemoryAccessCounter = 0;
 
+const OrtMemoryInfo *AllocatorInfo(const OrtAllocator *allocator) {
+    (void)allocator;
+    const OrtApi *ort = OrtGetApiBase()->GetApi(1);
+    OrtMemoryInfo *mem_info;
+    if (ort->CreateCpuMemoryInfo(OrtDeviceAllocator, OrtMemTypeDefault, &mem_info) != NULL) {
+        return NULL;
+    }
+    return mem_info;
+}
+
+void *AllocatorAlloc(OrtAllocator *ptr, size_t size) {
+    (void)ptr;
+    void *p = RedisModule_Alloc(size);
+    size_t allocated_size = RedisModule_MallocSize(p);
+    atomic_fetch_add(&OnnxMemory, allocated_size);
+    atomic_fetch_add(&OnnxMemoryAccessCounter, 1);
+    return p;
+}
+
+void AllocatorFree(OrtAllocator *ptr, void *p) {
+    (void)ptr;
+    size_t allocated_size = RedisModule_MallocSize(p);
+    atomic_fetch_add(&OnnxMemory, -1 * allocated_size);
+    atomic_fetch_add(&OnnxMemoryAccessCounter, 1);
+    return RedisModule_Free(p);
+}
+
 unsigned long long RAI_GetMemoryInfoORT() { return OnnxMemory; }
 
 unsigned long long RAI_GetMemoryAccessORT() { return OnnxMemoryAccessCounter; }
@@ -269,33 +296,6 @@ RAI_Tensor *RAI_TensorCreateFromOrtValue(OrtValue *v, size_t batch_offset, long
     return NULL;
 }
 
-const OrtMemoryInfo *myInfo(const OrtAllocator *allocator) {
-    (void)allocator;
-    const OrtApi *ort = OrtGetApiBase()->GetApi(1);
-    OrtMemoryInfo *mem_info;
-    if (ort->CreateCpuMemoryInfo(OrtDeviceAllocator, OrtMemTypeDefault, &mem_info) != NULL) {
-        return NULL;
-    }
-    return mem_info;
-}
-
-void *myAlloc(OrtAllocator *ptr, size_t size) {
-    (void)ptr;
-    void *p = RedisModule_Alloc(size);
-    size_t allocated_size = RedisModule_MallocSize(p);
-    atomic_fetch_add(&OnnxMemory, allocated_size);
-    atomic_fetch_add(&OnnxMemoryAccessCounter, 1);
-    return p;
-}
-
-void myFree(OrtAllocator *ptr, void *p) {
-    (void)ptr;
-    size_t allocated_size = RedisModule_MallocSize(p);
-    atomic_fetch_add(&OnnxMemory, -1 * allocated_size);
-    atomic_fetch_add(&OnnxMemoryAccessCounter, 1);
-    return RedisModule_Free(p);
-}
-
 RAI_Model *RAI_ModelCreateORT(RAI_Backend backend, const char *devicestr, RAI_ModelOpts opts,
                               const char *modeldef, size_t modellen, RAI_Error *error) {
 
@@ -308,8 +308,8 @@ RAI_Model *RAI_ModelCreateORT(RAI_Backend backend, const char *devicestr, RAI_Mo
 
     if (env == NULL) {
         ONNX_API(ort->CreateEnv(ORT_LOGGING_LEVEL_WARNING, "test", &env))
-        ONNX_API(ort->CreateCustomDeviceAllocator(ORT_API_VERSION, myAlloc, myFree, myInfo,
-                                                  &global_allocator))
+        ONNX_API(ort->CreateCustomDeviceAllocator(ORT_API_VERSION, AllocatorAlloc, AllocatorFree,
+                                                  AllocatorInfo, &global_allocator))
         ONNX_API(ort->RegisterCustomDeviceAllocator(env, global_allocator))
     }
 
diff --git a/src/redisai.c b/src/redisai.c
index 5a0669d2d..af1403e95 100644
--- a/src/redisai.c
+++ b/src/redisai.c
@@ -23,7 +23,6 @@
 #include <sys/resource.h>
 #include <sys/time.h>
 #include <unistd.h>
-#include "backends/onnxruntime.h"
 
 #include "rmutil/alloc.h"
 #include "rmutil/args.h"
@@ -1093,9 +1092,6 @@ static int RedisAI_RegisterApi(RedisModuleCtx *ctx) {
     REGISTER_API(DAGRunOpFree, ctx);
     REGISTER_API(DAGFree, ctx);
 
-    // For ORT test module
-    REGISTER_API(LoadDefaultBackend, ctx);
-
     return REDISMODULE_OK;
 }
 
@@ -1112,6 +1108,9 @@ void RAI_moduleInfoFunc(RedisModuleInfoCtx *ctx, int for_crash_report) {
                                           RAI_backends.onnx.get_memory_info());
         RedisModule_InfoAddFieldULongLong(ctx, "onnxruntime_memory_access_num",
                                           RAI_backends.onnx.get_memory_access_num());
+    } else {
+        RedisModule_InfoAddFieldULongLong(ctx, "onnxruntime_memory", 0);
+        RedisModule_InfoAddFieldULongLong(ctx, "onnxruntime_memory_access_num", 0);
     }
 
     struct rusage self_ru, c_ru;
diff --git a/src/tensor.h b/src/tensor.h
index 41b8a24e0..45062135c 100644
--- a/src/tensor.h
+++ b/src/tensor.h
@@ -11,7 +11,7 @@
 #define SRC_TENSOR_H_
 
 #include "config.h"
-//#include "dlpack/dlpack.h"
+#include "dlpack/dlpack.h"
 #include "err.h"
 #include "redismodule.h"
 #include "redisai.h"
diff --git a/src/tensor_struct.h b/src/tensor_struct.h
index d5df8288e..6203553b6 100644
--- a/src/tensor_struct.h
+++ b/src/tensor_struct.h
@@ -1,8 +1,7 @@
 #pragma once
 
 #include "config.h"
-#include "../deps/linux-x64-cpu/dlpack/include/dlpack/dlpack.h"
-//#include "dlpack/dlpack.h"
+#include "dlpack/dlpack.h"
 #include "limits.h"
 
 #define LEN_UNKOWN ULONG_MAX
diff --git a/tests/flow/includes.py b/tests/flow/includes.py
index fbe356990..55709700e 100755
--- a/tests/flow/includes.py
+++ b/tests/flow/includes.py
@@ -15,9 +15,6 @@
 import paella
 
 ROOT = os.environ.get("ROOT", None)
-if ROOT is None:
-    sys.stderr.write("ROOT was not defined in the environment.\n")
-    sys.exit(1)
 MAX_ITERATIONS = 2 if os.environ.get("MAX_ITERATIONS") == None else os.environ.get("MAX_ITERATIONS")
 TEST_TF = os.environ.get("TEST_TF") != "0" and os.environ.get("WITH_TF") != "0"
 TEST_TFLITE = os.environ.get("TEST_TFLITE") != "0" and os.environ.get("WITH_TFLITE") != "0"
diff --git a/tests/module/test_data/mul_1.onnx b/tests/flow/test_data/mul_1.onnx
similarity index 100%
rename from tests/module/test_data/mul_1.onnx
rename to tests/flow/test_data/mul_1.onnx
diff --git a/tests/flow/tests_onnx.py b/tests/flow/tests_onnx.py
index b255f9ab2..6e0607e60 100644
--- a/tests/flow/tests_onnx.py
+++ b/tests/flow/tests_onnx.py
@@ -4,33 +4,12 @@
 import redis
 from includes import *
 from RLTest import Env
-from functools import wraps
 
 '''
 python -m RLTest --test tests_onnx.py --module path/to/redisai.so
 '''
 
 
-def load_onnx_test_module(f):
-    @wraps(f)
-    def wrapper(env, *args, **kwargs):
-        goal_dir = os.path.join(os.path.dirname(__file__), "../module/OnnxAllocator.so")
-        TEST_MODULE_PATH = os.path.abspath(goal_dir)
-        con = env.getConnection()
-        modules = con.execute_command("MODULE", "LIST")
-        if b'RAI_onnxAllocator' in [module[1] for module in modules]:
-            return f(env, *args, **kwargs)
-        try:
-            ret = con.execute_command('MODULE', 'LOAD', TEST_MODULE_PATH)
-            env.assertEqual(ret, b'OK')
-        except Exception as e:
-            env.assertFalse(True)
-            env.debugPrint(str(e), force=True)
-            return
-        return f(env, *args, **kwargs)
-    return wrapper
-
-
 def test_onnx_modelrun_mnist(env):
     if not TEST_ONNX:
         env.debugPrint("skipping {} since TEST_ONNX=0".format(sys._getframe().f_code.co_name), force=True)
@@ -475,17 +454,40 @@ def test_parallelism():
     env.assertEqual(load_time_config["ai_intra_op_parallelism"], "2")
 
 
-@load_onnx_test_module
 def test_onnx_use_custom_allocator(env):
     if not TEST_ONNX:
         env.debugPrint("skipping {} since TEST_ONNX=0".format(sys._getframe().f_code.co_name), force=True)
         return
+    if DEVICE != 'CPU':
+        env.debugPrint("skipping {} since the custom allocator is for CPU".format(sys._getframe().f_code.co_name), force=True)
+        return
 
     con = env.getConnection()
-    ret = con.execute_command("RAI_onnxAllocator.modelSet")
-    env.assertEquals(ret, b'OK')
-
     test_data_path = os.path.join(os.path.dirname(__file__), 'test_data')
+    model_filename = os.path.join(test_data_path, 'mul_1.onnx')
+    with open(model_filename, 'rb') as f:
+        model_pb = f.read()
+    ai_memory_config = {k.split(":")[0]: k.split(":")[1]
+                            for k in con.execute_command("INFO MODULES").decode().split("#")[4].split()[1:]}
+    env.assertEqual(int(ai_memory_config["ai_onnxruntime_memory"]), 0)
+
+    # Expect using the allocator during model set for allocating the model, its input name and output name.
+    ret = con.execute_command('AI.MODELSET', 'm{1}', 'ONNX', DEVICE, 'BLOB', model_pb)
+    env.assertEqual(ret, b'OK')
+    ai_memory_config = {k.split(":")[0]: k.split(":")[1]
+                        for k in con.execute_command("INFO MODULES").decode().split("#")[4].split()[1:]}
+    env.assertTrue(int(ai_memory_config["ai_onnxruntime_memory"]) > 100)
+    env.assertEqual(int(ai_memory_config["ai_onnxruntime_memory_access_num"]), 3)
+
+    # Expect using the allocator free function when releasing the model.
+    con.execute_command('AI.MODELDEL', 'm{1}')
+    env.assertFalse(con.execute_command('EXISTS', 'm{1}'))
+    ai_memory_config = {k.split(":")[0]: k.split(":")[1]
+                        for k in con.execute_command("INFO MODULES").decode().split("#")[4].split()[1:]}
+    env.assertEqual(int(ai_memory_config["ai_onnxruntime_memory"]), 0)
+    env.assertEqual(int(ai_memory_config["ai_onnxruntime_memory_access_num"]), 6)
+
+    # test allocator in model run op
     model_filename = os.path.join(test_data_path, 'mnist.onnx')
     sample_filename = os.path.join(test_data_path, 'one.raw')
 
diff --git a/tests/module/Makefile b/tests/module/Makefile
index 131e035d3..ca282dd6b 100644
--- a/tests/module/Makefile
+++ b/tests/module/Makefile
@@ -17,11 +17,7 @@ else
 	SHOBJ_LDFLAGS ?= -bundle -undefined dynamic_lookup
 endif
 
-TEST_MODULES = \
-    LLAPI.so \
-    OnnxAllocator.so
-
-SAVE_PATH = @echo "const char* path = \"$(dir $(abspath $(lastword $(MAKEFILE_LIST))))\";" > path.h
+TEST_MODULES = LLAPI.so
 
 .PHONY: all
 
@@ -31,17 +27,12 @@ all: $(TEST_MODULES)
 	$(MAKE) CFLAGS="-m32" LDFLAGS="-melf_i386"
 
 %.o: %.c
-	${SAVE_PATH}
 	$(CC) $(DEBUGFLAGS) -I../../src -I../../deps -DREDIS_MODULE_TARGET -DREDISMODULE_EXPERIMENTAL_API $(SHOBJ_CFLAGS) -fPIC -c $< -o $@
 
 LLAPI.so: LLAPI.o DAG_utils.o
 	$(CC) -o $@ $^ $(SHOBJ_LDFLAGS) -lc -lm
 	chmod +x LLAPI.so
 
-OnnxAllocator.so: OnnxAllocator.o
-	$(CC) -o $@ $^ $(SHOBJ_LDFLAGS) -lc -lm
-	chmod +x OnnxAllocator.so
-
 .PHONY: clean
 
 clean:
diff --git a/tests/module/OnnxAllocator.c b/tests/module/OnnxAllocator.c
deleted file mode 100644
index 7e1afe913..000000000
--- a/tests/module/OnnxAllocator.c
+++ /dev/null
@@ -1,103 +0,0 @@
-#define REDISMODULE_MAIN
-
-typedef unsigned int uint;
-
-#include "err.h"
-#include <string.h>
-#include <stdio.h>
-#include <stdlib.h>
-#include <unistd.h>
-#include "onnx_redisai.h"
-#include "backends/onnxruntime.h"
-#include "config.h"
-#include "err.h"
-#include <limits.h>
-#include "path.h"
-
-typedef struct MemoryInfo {
-    size_t usage;
-    size_t access_counter;
-} MemoryInfo;
-
-MemoryInfo _get_memory_info() {
-
-    RedisModuleServerInfoData *server_info = RedisModule_GetServerInfo(NULL, "everything");
-    int onnx_info_error;
-    unsigned long long onnx_counter = RedisModule_ServerInfoGetFieldUnsigned(server_info,
-      "ai_onnxruntime_memory_access_num", &onnx_info_error);
-    RedisModule_Assert(onnx_info_error == 0);
-    unsigned long long onnx_mem = RedisModule_ServerInfoGetFieldUnsigned(server_info,
-      "ai_onnxruntime_memory", &onnx_info_error);
-    RedisModule_Assert(onnx_info_error == 0);
-    RedisModule_FreeServerInfo(NULL, server_info);
-    return (MemoryInfo){.usage = onnx_mem, .access_counter = onnx_counter};
-}
-
-int RAI_onnxAllocator_modelSet(RedisModuleCtx *ctx, RedisModuleString **argv, int argc) {
-    REDISMODULE_NOT_USED(argv);
-
-    if(argc > 1) {
-        RedisModule_WrongArity(ctx);
-        return REDISMODULE_OK;
-    }
-    RedisAI_LoadDefaultBackend(ctx, RAI_BACKEND_ONNXRUNTIME);
-    RedisModule_Assert(_get_memory_info().usage == 0);
-
-    // Load onnx model to local buffer.
-    size_t model_len = 130;
-    FILE *fp;
-    char model_data[model_len];
-    size_t index = 0;
-    char model_path[1024];
-    strcpy(model_path, path);
-    strcat(model_path, "/test_data/mul_1.onnx");
-    fp = fopen(model_path, "rb");
-    char c;
-    while((c = fgetc(fp)) != EOF) {
-        model_data[index++] = c;
-    }
-    fclose(fp);
-    RedisModule_Assert(model_len == index);
-
-    // Create model and verify that onnx backend uses REDIS allocator.
-    RAI_ModelOpts opts = {0};
-    RAI_Error err = {0};
-    RAI_Model *model = RedisAI_ModelCreate(RAI_BACKEND_ONNXRUNTIME, "CPU", NULL, opts, 0, NULL, 0, NULL,
-      model_data, model_len, &err);
-    RedisModule_Assert(err.code == RAI_OK);
-
-    // Expect 3 allocation access by onnx backend for allocating model, input name
-    // and output name
-    MemoryInfo mem_info = _get_memory_info();
-    if (mem_info.access_counter != 3 || mem_info.usage == 0) {
-        return RedisModule_ReplyWithSimpleString(ctx, "model set allocation error");
-    }
-
-    RedisAI_ModelFree(model, &err);
-    RedisModule_Assert(err.code == RAI_OK);
-    mem_info = _get_memory_info();
-    if (mem_info.access_counter != 6 || mem_info.usage != 0) {
-        return RedisModule_ReplyWithSimpleString(ctx, "model free error");
-    }
-    return RedisModule_ReplyWithSimpleString(ctx, "OK");
-}
-
-
-int RedisModule_OnLoad(RedisModuleCtx *ctx, RedisModuleString **argv, int argc) {
-    REDISMODULE_NOT_USED(argv);
-    REDISMODULE_NOT_USED(argc);
-
-    if(RedisModule_Init(ctx, "RAI_onnxAllocator", 1, REDISMODULE_APIVER_1) ==
-       REDISMODULE_ERR)
-        return REDISMODULE_ERR;
-
-    if(RedisAI_OnnxInitialize(ctx) != REDISMODULE_OK)
-        RedisModule_Log(ctx, "warning",
-          "could not initialize RedisAI api, running without AI support.");
-
-    if(RedisModule_CreateCommand(ctx, "RAI_onnxAllocator.modelSet",
-      RAI_onnxAllocator_modelSet, "", 0, 0, 0) == REDISMODULE_ERR)
-        return REDISMODULE_ERR;
-
-    return REDISMODULE_OK;
-}
\ No newline at end of file
diff --git a/tests/module/onnx_redisai.h b/tests/module/onnx_redisai.h
deleted file mode 100644
index 791a7800a..000000000
--- a/tests/module/onnx_redisai.h
+++ /dev/null
@@ -1,108 +0,0 @@
-#pragma once
-
-#include <stdbool.h>
-#include "redismodule.h"
-#include "model_struct.h"
-
-#define REDISAI_LLAPI_VERSION 1
-#define MODULE_API_FUNC(x)    (*x)
-
-#ifdef REDISAI_EXTERN
-#define REDISAI_API extern
-#endif
-
-#ifndef REDISAI_API
-#define REDISAI_API
-#endif
-
-#ifndef REDISAI_H_INCLUDE
-typedef struct RAI_Tensor RAI_Tensor;
-typedef struct RAI_Model RAI_Model;
-typedef struct RAI_Script RAI_Script;
-
-typedef struct RAI_Error RAI_Error;
-
-#endif
-
-#define REDISAI_BACKEND_TENSORFLOW  0
-#define REDISAI_BACKEND_TFLITE      1
-#define REDISAI_BACKEND_TORCH       2
-#define REDISAI_BACKEND_ONNXRUNTIME 3
-
-#define REDISAI_DEVICE_CPU 0
-#define REDISAI_DEVICE_GPU 1
-
-typedef enum RedisAI_ErrorCode {
-    RedisAI_ErrorCode_OK = 0,
-    RedisAI_ErrorCode_EMODELIMPORT,
-    RedisAI_ErrorCode_EMODELCONFIGURE,
-    RedisAI_ErrorCode_EMODELCREATE,
-    RedisAI_ErrorCode_EMODELRUN,
-    RedisAI_ErrorCode_EMODELSERIALIZE,
-    RedisAI_ErrorCode_EMODELFREE,
-    RedisAI_ErrorCode_ESCRIPTIMPORT,
-    RedisAI_ErrorCode_ESCRIPTCONFIGURE,
-    RedisAI_ErrorCode_ESCRIPTCREATE,
-    RedisAI_ErrorCode_ESCRIPTRUN,
-    RedisAI_ErrorCode_EUNSUPPORTEDBACKEND,
-    RedisAI_ErrorCode_EBACKENDNOTLOADED,
-    RedisAI_ErrorCode_ESCRIPTFREE,
-    RedisAI_ErrorCode_ETENSORSET,
-    RedisAI_ErrorCode_ETENSORGET,
-    RedisAI_ErrorCode_EDAGBUILDER,
-    RedisAI_ErrorCode_EDAGRUN,
-    RedisAI_ErrorCode_EFINISHCTX
-} RedisAI_ErrorCode;
-
-REDISAI_API int MODULE_API_FUNC(RedisAI_LoadDefaultBackend)(RedisModuleCtx *ctx, int backend);
-REDISAI_API RAI_Model *MODULE_API_FUNC(RedisAI_ModelCreate)(int backend, char *devicestr, char *tag,
-  RAI_ModelOpts opts, size_t ninputs,
-  const char **inputs, size_t noutputs,
-  const char **outputs,
-  const char *modeldef, size_t modellen,
-  RAI_Error *err);
-REDISAI_API void MODULE_API_FUNC(RedisAI_ModelFree)(RAI_Model *model, RAI_Error *err);
-REDISAI_API int MODULE_API_FUNC(RedisAI_ModelRun)(RAI_ModelRunCtx **mctx, long long n,
-  RAI_Error *err);
-
-REDISAI_API int MODULE_API_FUNC(RedisAI_GetLLAPIVersion)();
-
-#ifndef __cplusplus
-#define REDISAI_MODULE_INIT_FUNCTION(ctx, name)                                                    \
-    RedisAI_##name = RedisModule_GetSharedAPI(ctx, "RedisAI_" #name);                              \
-    if (!RedisAI_##name) {                                                                         \
-        RedisModule_Log(ctx, "warning", "could not initialize RedisAI_" #name "\r\n");             \
-        return REDISMODULE_ERR;                                                                    \
-    }
-#else
-#define REDISAI_MODULE_INIT_FUNCTION(ctx, name)                                                    \
-    RedisAI_##name = reinterpret_cast<decltype(RedisAI_##name)>(                                   \
-        RedisModule_GetSharedAPI((RedisModuleCtx *)(ctx), "RedisAI_" #name));                      \
-    if (!RedisAI_##name) {                                                                         \
-        RedisModule_Log(ctx, "warning", "could not initialize RedisAI_" #name "\r\n");             \
-        return REDISMODULE_ERR;                                                                    \
-    }
-#endif
-
-static int RedisAI_OnnxInitialize(RedisModuleCtx *ctx) {
-
-    if (!RedisModule_GetSharedAPI) {
-        RedisModule_Log(ctx, "warning",
-                        "redis version is not compatible with module shared api, "
-                        "use redis 5.0.4 or above.");
-        return REDISMODULE_ERR;
-    }
-
-    REDISAI_MODULE_INIT_FUNCTION(ctx, LoadDefaultBackend);
-    REDISAI_MODULE_INIT_FUNCTION(ctx, ModelCreate);
-    REDISAI_MODULE_INIT_FUNCTION(ctx, ModelFree);
-    REDISAI_MODULE_INIT_FUNCTION(ctx, ModelRun);
-
-    REDISAI_MODULE_INIT_FUNCTION(ctx, GetLLAPIVersion);
-
-    if (RedisAI_GetLLAPIVersion() < REDISAI_LLAPI_VERSION) {
-        return REDISMODULE_ERR;
-    }
-
-    return REDISMODULE_OK;
-}

From 79a613a61d493b2ae1912596d57fa1cf63d43dfd Mon Sep 17 00:00:00 2001
From: alonre24 <alonreshef24@gmail.com>
Date: Thu, 25 Feb 2021 13:38:44 +0200
Subject: [PATCH 11/32] Backward compatibility for GPU - when running on GPU,
 use the default allocator.

---
 opt/build/onnxruntime/Makefile |  8 ++++++++
 src/backends/onnxruntime.c     | 22 ++++++++++++++++++----
 2 files changed, 26 insertions(+), 4 deletions(-)

diff --git a/opt/build/onnxruntime/Makefile b/opt/build/onnxruntime/Makefile
index 5ab04f22e..c28472b1f 100755
--- a/opt/build/onnxruntime/Makefile
+++ b/opt/build/onnxruntime/Makefile
@@ -73,6 +73,13 @@ publish_x64:
 .PHONY: publish_x64
 endef
 
+define publish_x64-gpu # (1=arch, 2=tar-arch)
+publish_x64-gpu:
+	@aws s3 cp $(STEM)-$(2)-$(VERSION).tgz s3://$(S3_URL)/ --acl public-read
+
+.PHONY: publish_x64-gpu
+endef
+
 define publish_arm # (1=arch, 2=tar-arch)
 publish_$(1):
 	@aws s3 cp $(STEM)-$(2)-$(VERSION).tgz s3://$(S3_URL)/ --acl public-read
@@ -102,6 +109,7 @@ endif
 publish: $(PUBLISH_TARGETS)
 
 $(eval $(call publish_x64,x64,x64))
+$(eval $(call publish_x64-gpu,x64-gpu,x64-gpu))
 $(eval $(call publish_arm,arm64v8,arm64))
 $(eval $(call publish_arm,arm32v7,arm))
 
diff --git a/src/backends/onnxruntime.c b/src/backends/onnxruntime.c
index fb9eefffb..447cbbde5 100644
--- a/src/backends/onnxruntime.c
+++ b/src/backends/onnxruntime.c
@@ -12,6 +12,13 @@
     if ((status = (x)) != NULL)                                                                    \
         goto error;
 
+#ifdef RAI_ONNXRUNTIME_USE_CUDA
+#define GET_GLOBAL_ALLOCATOR
+ONNX_API(ort->GetAllocatorWithDefaultOptions(&global_allocator))
+#else
+#define GET_GLOBAL_ALLOCATOR
+#endif
+
 OrtEnv *env = NULL;
 OrtAllocator *global_allocator = NULL;
 unsigned long long OnnxMemory = 0;
@@ -171,6 +178,7 @@ OrtValue *RAI_OrtValueFromTensors(RAI_Tensor **ts, size_t count, RAI_Error *erro
     batched_shape[0] = batch_size;
 
     OrtValue *out;
+    GET_GLOBAL_ALLOCATOR
     if (count > 1) {
         ONNX_API(
             ort->CreateTensorAsOrtValue(global_allocator, batched_shape, t0->tensor.dl_tensor.ndim,
@@ -308,16 +316,21 @@ RAI_Model *RAI_ModelCreateORT(RAI_Backend backend, const char *devicestr, RAI_Mo
 
     if (env == NULL) {
         ONNX_API(ort->CreateEnv(ORT_LOGGING_LEVEL_WARNING, "test", &env))
+#ifndef RAI_ONNXRUNTIME_USE_CUDA
         ONNX_API(ort->CreateCustomDeviceAllocator(ORT_API_VERSION, AllocatorAlloc, AllocatorFree,
                                                   AllocatorInfo, &global_allocator))
         ONNX_API(ort->RegisterCustomDeviceAllocator(env, global_allocator))
+#endif
     }
 
     ONNX_API(ort->CreateSessionOptions(&session_options))
 
-    // These are required to ensure that onnx will use the registered REDIS allocator.
+#ifndef RAI_ONNXRUNTIME_USE_CUDA
+    // These are required to ensure that onnx will use the registered REDIS allocator (for CPU
+    // only).
     ONNX_API(ort->AddSessionConfigEntry(session_options, "session.use_env_allocators", "1"))
     ONNX_API(ort->DisableCpuMemArena(session_options))
+#endif
 
     // TODO: these options could be configured at the AI.CONFIG level
     ONNX_API(ort->SetSessionGraphOptimizationLevel(session_options, 1))
@@ -334,6 +347,7 @@ RAI_Model *RAI_ModelCreateORT(RAI_Backend backend, const char *devicestr, RAI_Mo
     size_t n_output_nodes;
     ONNX_API(ort->SessionGetOutputCount(session, &n_output_nodes))
 
+    GET_GLOBAL_ALLOCATOR
     inputs_ = array_new(char *, n_input_nodes);
     for (long long i = 0; i < n_input_nodes; i++) {
         char *input_name;
@@ -400,7 +414,7 @@ void RAI_ModelFreeORT(RAI_Model *model, RAI_Error *error) {
     const OrtApi *ort = OrtGetApiBase()->GetApi(1);
     OrtStatus *status = NULL;
 
-    RedisModule_Free(model->devicestr);
+    GET_GLOBAL_ALLOCATOR
     for (uint32_t i = 0; i < model->ninputs; i++) {
         ONNX_API(ort->AllocatorFree(global_allocator, model->inputs[i]))
     }
@@ -411,6 +425,7 @@ void RAI_ModelFreeORT(RAI_Model *model, RAI_Error *error) {
     }
     array_free(model->outputs);
 
+    RedisModule_Free(model->devicestr);
     RedisModule_Free(model->data);
     ort->ReleaseSession(model->session);
     model->model = NULL;
@@ -457,8 +472,7 @@ int RAI_ModelRunORT(RAI_ModelRunCtx **mctxs, RAI_Error *error) {
 
     size_t n_output_nodes;
     ONNX_API(ort->SessionGetOutputCount(session, &n_output_nodes))
-
-    {
+    GET_GLOBAL_ALLOCATOR {
         const char *input_names[n_input_nodes];
         const char *output_names[n_output_nodes];
 

From c83fccc920172f32a7d30c7cafa3e8043169fe90 Mon Sep 17 00:00:00 2001
From: alonre24 <alonreshef24@gmail.com>
Date: Thu, 25 Feb 2021 16:15:48 +0200
Subject: [PATCH 12/32] Addressing Rafi's comments.

---
 opt/build/onnxruntime/Dockerfile.arm  | 2 +-
 opt/build/onnxruntime/Dockerfile.arm7 | 2 +-
 opt/build/onnxruntime/pack.sh         | 6 +++---
 3 files changed, 5 insertions(+), 5 deletions(-)

diff --git a/opt/build/onnxruntime/Dockerfile.arm b/opt/build/onnxruntime/Dockerfile.arm
index b24c9bf54..979b3e852 100755
--- a/opt/build/onnxruntime/Dockerfile.arm
+++ b/opt/build/onnxruntime/Dockerfile.arm
@@ -48,6 +48,6 @@ RUN set -e ;\
 	./build.sh ${BUILDARGS} --build_shared_lib
 # RUN ./build.sh ${BUILDARGS} --enable_pybind --build_wheel
 
-RUN ./pack.sh ${ONNXRUNTIME_VER} ${ARCH}
+RUN ./pack.sh ${ONNXRUNTIME_VER}
 
 RUN [ "cross-build-end" ]
diff --git a/opt/build/onnxruntime/Dockerfile.arm7 b/opt/build/onnxruntime/Dockerfile.arm7
index 4de482563..4cfea0af0 100755
--- a/opt/build/onnxruntime/Dockerfile.arm7
+++ b/opt/build/onnxruntime/Dockerfile.arm7
@@ -42,6 +42,6 @@ RUN set -e ;\
 	./build.sh ${BUILDARGS} --build_shared_lib
 # RUN ./build.sh ${BUILDARGS} --enable_pybind --build_wheel
 
-RUN ./pack.sh ${ONNXRUNTIME_VER} ${ARCH}
+RUN ./pack.sh ${ONNXRUNTIME_VER}
 
 RUN [ "cross-build-end" ]
diff --git a/opt/build/onnxruntime/pack.sh b/opt/build/onnxruntime/pack.sh
index fdd1113b6..447174b70 100755
--- a/opt/build/onnxruntime/pack.sh
+++ b/opt/build/onnxruntime/pack.sh
@@ -2,7 +2,7 @@
 
 set -e
 VER="$1"
-PLATFORM="$2"
+PLATFORM=$($READIES/bin/platform -t)
 
 mkdir -p pack/include pack/lib
 cp onnxruntime/build/Linux/MinSizeRel/libonnxruntime.so.${VER} pack/lib/
@@ -19,5 +19,5 @@ cp onnxruntime/include/onnxruntime/core/providers/cuda/cuda_provider_factory.h p
 cd pack/lib/
 ln -s libonnxruntime.so.${VER} libonnxruntime.so
 cd ../..
-mv pack onnxruntime-linux-${PLATFORM}-${VER}
-tar czf onnxruntime-linux-${PLATFORM}-${VER}.tgz onnxruntime-linux-${PLATFORM}-${VER}/
+mv pack onnxruntime-${PLATFORM}-${VER}
+tar czf onnxruntime-${PLATFORM}-${VER}.tgz onnxruntime-${PLATFORM}-${VER}/

From 504f1eb38e3eb75e5e2e24be4be6aecc2eba4d4d Mon Sep 17 00:00:00 2001
From: alonre24 <alonreshef24@gmail.com>
Date: Sun, 28 Feb 2021 14:19:14 +0200
Subject: [PATCH 13/32] Add sanitizer test for GPU + Dvir's PR comments

---
 src/backends/onnxruntime.c | 91 ++++++++++++++++++++++----------------
 tests/flow/tests_onnx.py   | 28 ++++++++++++
 tests/module/Makefile      |  2 +-
 3 files changed, 81 insertions(+), 40 deletions(-)

diff --git a/src/backends/onnxruntime.c b/src/backends/onnxruntime.c
index 9f0b5155e..9313672d2 100644
--- a/src/backends/onnxruntime.c
+++ b/src/backends/onnxruntime.c
@@ -8,12 +8,17 @@
 
 #include "onnxruntime_c_api.h"
 
-#define ONNX_API(x)                                                                                \
+// Use as a wrapper for ORT api call. If ORT api hasn't returned null, it has failed.
+#define ONNX_VALIDATE_STATUS(x)                                                                    \
     if ((status = (x)) != NULL)                                                                    \
         goto error;
 
+// If we run on GPU, we do not use the custom allocator (redis allocator), but
+// the default allocator returned by ORT api. If we run on CPU, we do not need to
+// use this api, as the global allocator is already set to be the custom allocator.
 #ifdef RAI_ONNXRUNTIME_USE_CUDA
-#define GET_GLOBAL_ALLOCATOR ONNX_API(ort->GetAllocatorWithDefaultOptions(&global_allocator))
+#define GET_GLOBAL_ALLOCATOR                                                                       \
+    ONNX_VALIDATE_STATUS(ort->GetAllocatorWithDefaultOptions(&global_allocator))
 #else
 #define GET_GLOBAL_ALLOCATOR
 #endif
@@ -45,7 +50,7 @@ void *AllocatorAlloc(OrtAllocator *ptr, size_t size) {
 void AllocatorFree(OrtAllocator *ptr, void *p) {
     (void)ptr;
     size_t allocated_size = RedisModule_MallocSize(p);
-    atomic_fetch_add(&OnnxMemory, -1 * allocated_size);
+    atomic_fetch_sub(&OnnxMemory, allocated_size);
     atomic_fetch_add(&OnnxMemoryAccessCounter, 1);
     return RedisModule_Free(p);
 }
@@ -179,19 +184,19 @@ OrtValue *RAI_OrtValueFromTensors(RAI_Tensor **ts, size_t count, RAI_Error *erro
     OrtValue *out;
     GET_GLOBAL_ALLOCATOR
     if (count > 1) {
-        ONNX_API(
+        ONNX_VALIDATE_STATUS(
             ort->CreateTensorAsOrtValue(global_allocator, batched_shape, t0->tensor.dl_tensor.ndim,
                                         RAI_GetOrtDataTypeFromDL(t0->tensor.dl_tensor.dtype), &out))
 
         char *ort_data;
-        ONNX_API(ort->GetTensorMutableData(out, (void **)&ort_data))
+        ONNX_VALIDATE_STATUS(ort->GetTensorMutableData(out, (void **)&ort_data))
         size_t offset = 0;
         for (size_t i = 0; i < count; i++) {
             memcpy(ort_data + offset, RAI_TensorData(ts[i]), RAI_TensorByteSize(ts[i]));
             offset += RAI_TensorByteSize(ts[i]);
         }
     } else {
-        ONNX_API(ort->CreateTensorWithDataAsOrtValue(
+        ONNX_VALIDATE_STATUS(ort->CreateTensorWithDataAsOrtValue(
             global_allocator->Info(global_allocator), t0->tensor.dl_tensor.data,
             RAI_TensorByteSize(t0), t0->tensor.dl_tensor.shape, t0->tensor.dl_tensor.ndim,
             RAI_GetOrtDataTypeFromDL(t0->tensor.dl_tensor.dtype), &out))
@@ -213,7 +218,7 @@ RAI_Tensor *RAI_TensorCreateFromOrtValue(OrtValue *v, size_t batch_offset, long
     int64_t *strides = NULL;
 
     int is_tensor;
-    ONNX_API(ort->IsTensor(v, &is_tensor))
+    ONNX_VALIDATE_STATUS(ort->IsTensor(v, &is_tensor))
     if (!is_tensor) {
         // TODO: if not tensor, flatten the data structure (sequence or map) and store it in a
         // tensor. If return value is string, emit warning.
@@ -223,15 +228,15 @@ RAI_Tensor *RAI_TensorCreateFromOrtValue(OrtValue *v, size_t batch_offset, long
     ret = RAI_TensorNew();
     DLContext ctx = (DLContext){.device_type = kDLCPU, .device_id = 0};
     OrtTensorTypeAndShapeInfo *info;
-    ONNX_API(ort->GetTensorTypeAndShape(v, &info))
+    ONNX_VALIDATE_STATUS(ort->GetTensorTypeAndShape(v, &info))
 
     {
         size_t ndims;
-        ONNX_API(ort->GetDimensionsCount(info, &ndims))
+        ONNX_VALIDATE_STATUS(ort->GetDimensionsCount(info, &ndims))
         int64_t dims[ndims];
-        ONNX_API(ort->GetDimensions(info, dims, ndims))
+        ONNX_VALIDATE_STATUS(ort->GetDimensions(info, dims, ndims))
         enum ONNXTensorElementDataType ort_dtype;
-        ONNX_API(ort->GetTensorElementType(info, &ort_dtype))
+        ONNX_VALIDATE_STATUS(ort->GetTensorElementType(info, &ort_dtype))
         int64_t total_batch_size = dims[0];
         total_batch_size = total_batch_size > 0 ? total_batch_size : 1;
 
@@ -253,9 +258,9 @@ RAI_Tensor *RAI_TensorCreateFromOrtValue(OrtValue *v, size_t batch_offset, long
         DLDataType dtype = RAI_GetDLDataTypeFromORT(ort_dtype);
 #ifdef RAI_COPY_RUN_OUTPUT
         char *ort_data;
-        ONNX_API(ort->GetTensorMutableData(v, (void **)&ort_data))
+        ONNX_VALIDATE_STATUS(ort->GetTensorMutableData(v, (void **)&ort_data))
         size_t elem_count;
-        ONNX_API(ort->GetTensorShapeElementCount(info, &elem_count))
+        ONNX_VALIDATE_STATUS(ort->GetTensorShapeElementCount(info, &elem_count))
 
         const size_t len = dtype.bits * elem_count / 8;
         const size_t total_bytesize = len * sizeof(char);
@@ -314,50 +319,55 @@ RAI_Model *RAI_ModelCreateORT(RAI_Backend backend, const char *devicestr, RAI_Mo
     OrtStatus *status = NULL;
 
     if (env == NULL) {
-        ONNX_API(ort->CreateEnv(ORT_LOGGING_LEVEL_WARNING, "test", &env))
+        ONNX_VALIDATE_STATUS(ort->CreateEnv(ORT_LOGGING_LEVEL_WARNING, "test", &env))
 #ifndef RAI_ONNXRUNTIME_USE_CUDA
-        ONNX_API(ort->CreateCustomDeviceAllocator(ORT_API_VERSION, AllocatorAlloc, AllocatorFree,
-                                                  AllocatorInfo, &global_allocator))
-        ONNX_API(ort->RegisterCustomDeviceAllocator(env, global_allocator))
+        ONNX_VALIDATE_STATUS(ort->CreateCustomDeviceAllocator(
+            ORT_API_VERSION, AllocatorAlloc, AllocatorFree, AllocatorInfo, &global_allocator))
+        ONNX_VALIDATE_STATUS(ort->RegisterCustomDeviceAllocator(env, global_allocator))
 #endif
     }
 
-    ONNX_API(ort->CreateSessionOptions(&session_options))
+    ONNX_VALIDATE_STATUS(ort->CreateSessionOptions(&session_options))
 
 #ifndef RAI_ONNXRUNTIME_USE_CUDA
     // These are required to ensure that onnx will use the registered REDIS allocator (for CPU
     // only).
-    ONNX_API(ort->AddSessionConfigEntry(session_options, "session.use_env_allocators", "1"))
-    ONNX_API(ort->DisableCpuMemArena(session_options))
+    ONNX_VALIDATE_STATUS(
+        ort->AddSessionConfigEntry(session_options, "session.use_env_allocators", "1"))
+    ONNX_VALIDATE_STATUS(ort->DisableCpuMemArena(session_options))
 #endif
 
     // TODO: these options could be configured at the AI.CONFIG level
-    ONNX_API(ort->SetSessionGraphOptimizationLevel(session_options, 1))
-    ONNX_API(ort->SetIntraOpNumThreads(session_options, (int)opts.backends_intra_op_parallelism))
-    ONNX_API(ort->SetInterOpNumThreads(session_options, (int)opts.backends_inter_op_parallelism))
+    ONNX_VALIDATE_STATUS(ort->SetSessionGraphOptimizationLevel(session_options, 1))
+    ONNX_VALIDATE_STATUS(
+        ort->SetIntraOpNumThreads(session_options, (int)opts.backends_intra_op_parallelism))
+    ONNX_VALIDATE_STATUS(
+        ort->SetInterOpNumThreads(session_options, (int)opts.backends_inter_op_parallelism))
     if (!setDeviceId(devicestr, session_options, error)) {
+        ort->ReleaseSessionOptions(session_options);
         return NULL;
     }
 
-    ONNX_API(ort->CreateSessionFromArray(env, modeldef, modellen, session_options, &session))
+    ONNX_VALIDATE_STATUS(
+        ort->CreateSessionFromArray(env, modeldef, modellen, session_options, &session))
 
     size_t n_input_nodes;
-    ONNX_API(ort->SessionGetInputCount(session, &n_input_nodes))
+    ONNX_VALIDATE_STATUS(ort->SessionGetInputCount(session, &n_input_nodes))
     size_t n_output_nodes;
-    ONNX_API(ort->SessionGetOutputCount(session, &n_output_nodes))
+    ONNX_VALIDATE_STATUS(ort->SessionGetOutputCount(session, &n_output_nodes))
 
     GET_GLOBAL_ALLOCATOR
     inputs_ = array_new(char *, n_input_nodes);
     for (long long i = 0; i < n_input_nodes; i++) {
         char *input_name;
-        ONNX_API(ort->SessionGetInputName(session, i, global_allocator, &input_name))
+        ONNX_VALIDATE_STATUS(ort->SessionGetInputName(session, i, global_allocator, &input_name))
         inputs_ = array_append(inputs_, input_name);
     }
 
     outputs_ = array_new(char *, n_output_nodes);
     for (long long i = 0; i < n_output_nodes; i++) {
         char *output_name;
-        ONNX_API(ort->SessionGetOutputName(session, i, global_allocator, &output_name))
+        ONNX_VALIDATE_STATUS(ort->SessionGetOutputName(session, i, global_allocator, &output_name))
         outputs_ = array_append(outputs_, output_name);
     }
 
@@ -415,12 +425,12 @@ void RAI_ModelFreeORT(RAI_Model *model, RAI_Error *error) {
 
     GET_GLOBAL_ALLOCATOR
     for (uint32_t i = 0; i < model->ninputs; i++) {
-        ONNX_API(ort->AllocatorFree(global_allocator, model->inputs[i]))
+        ONNX_VALIDATE_STATUS(ort->AllocatorFree(global_allocator, model->inputs[i]))
     }
     array_free(model->inputs);
 
     for (uint32_t i = 0; i < model->noutputs; i++) {
-        ONNX_API(ort->AllocatorFree(global_allocator, model->outputs[i]))
+        ONNX_VALIDATE_STATUS(ort->AllocatorFree(global_allocator, model->outputs[i]))
     }
     array_free(model->outputs);
 
@@ -467,10 +477,10 @@ int RAI_ModelRunORT(RAI_ModelRunCtx **mctxs, RAI_Error *error) {
 
     OrtStatus *status = NULL;
     size_t n_input_nodes;
-    ONNX_API(ort->SessionGetInputCount(session, &n_input_nodes))
+    ONNX_VALIDATE_STATUS(ort->SessionGetInputCount(session, &n_input_nodes))
 
     size_t n_output_nodes;
-    ONNX_API(ort->SessionGetOutputCount(session, &n_output_nodes))
+    ONNX_VALIDATE_STATUS(ort->SessionGetOutputCount(session, &n_output_nodes))
     GET_GLOBAL_ALLOCATOR {
         const char *input_names[n_input_nodes];
         const char *output_names[n_output_nodes];
@@ -497,7 +507,8 @@ int RAI_ModelRunORT(RAI_ModelRunCtx **mctxs, RAI_Error *error) {
 
         for (size_t i = 0; i < n_input_nodes; i++) {
             char *input_name;
-            ONNX_API(ort->SessionGetInputName(session, i, global_allocator, &input_name))
+            ONNX_VALIDATE_STATUS(
+                ort->SessionGetInputName(session, i, global_allocator, &input_name))
             input_names[i] = input_name;
 
             RAI_Tensor *batched_input_tensors[nbatches];
@@ -514,23 +525,25 @@ int RAI_ModelRunORT(RAI_ModelRunCtx **mctxs, RAI_Error *error) {
 
         for (size_t i = 0; i < n_output_nodes; i++) {
             char *output_name;
-            ONNX_API(ort->SessionGetOutputName(session, i, global_allocator, &output_name))
+            ONNX_VALIDATE_STATUS(
+                ort->SessionGetOutputName(session, i, global_allocator, &output_name))
             output_names[i] = output_name;
             outputs[i] = NULL;
         }
 
         OrtRunOptions *run_options = NULL;
-        ONNX_API(ort->Run(session, run_options, input_names, (const OrtValue *const *)inputs,
-                          n_input_nodes, output_names, n_output_nodes, outputs))
+        ONNX_VALIDATE_STATUS(ort->Run(session, run_options, input_names,
+                                      (const OrtValue *const *)inputs, n_input_nodes, output_names,
+                                      n_output_nodes, outputs))
 
         for (size_t i = 0; i < n_output_nodes; i++) {
             if (nbatches > 1) {
                 OrtTensorTypeAndShapeInfo *info;
-                ONNX_API(ort->GetTensorTypeAndShape(outputs[i], &info))
+                ONNX_VALIDATE_STATUS(ort->GetTensorTypeAndShape(outputs[i], &info))
                 size_t ndims;
-                ONNX_API(ort->GetDimensionsCount(info, &ndims))
+                ONNX_VALIDATE_STATUS(ort->GetDimensionsCount(info, &ndims))
                 int64_t dims[ndims];
-                ONNX_API(ort->GetDimensions(info, dims, ndims))
+                ONNX_VALIDATE_STATUS(ort->GetDimensions(info, dims, ndims))
                 if (dims[0] != total_batch_size) {
                     RAI_SetError(error, RAI_EMODELRUN,
                                  "ERR Model did not generate the expected batch size");
diff --git a/tests/flow/tests_onnx.py b/tests/flow/tests_onnx.py
index 6e0607e60..26788e7ed 100644
--- a/tests/flow/tests_onnx.py
+++ b/tests/flow/tests_onnx.py
@@ -513,3 +513,31 @@ def test_onnx_use_custom_allocator(env):
     values = con.execute_command('AI.TENSORGET', 'b{1}', 'VALUES')
     argmax = max(range(len(values)), key=lambda i: values[i])
     env.assertEqual(argmax, 1)
+
+
+def test_onnx_use_custom_allocator_with_GPU(env):
+    if not TEST_ONNX:
+        env.debugPrint("skipping {} since TEST_ONNX=0".format(sys._getframe().f_code.co_name), force=True)
+        return
+    if DEVICE == 'CPU':
+        env.debugPrint("skipping {} since this test if for GPU only".format(sys._getframe().f_code.co_name), force=True)
+        return
+
+    con = env.getConnection()
+    test_data_path = os.path.join(os.path.dirname(__file__), 'test_data')
+    model_filename = os.path.join(test_data_path, 'mul_1.onnx')
+    with open(model_filename, 'rb') as f:
+        model_pb = f.read()
+    ai_memory_config = {k.split(":")[0]: k.split(":")[1]
+                        for k in con.execute_command("INFO MODULES").decode().split("#")[4].split()[1:]}
+    env.assertEqual(int(ai_memory_config["ai_onnxruntime_memory"]), 0)
+
+    # Create the same model, once for CPU and once for GPU, and verify that redis allocator was being used only for CPU
+    ret = con.execute_command('AI.MODELSET', 'm_gpu{1}', 'ONNX', DEVICE, 'BLOB', model_pb)
+    env.assertEqual(ret, b'OK')
+    ret = con.execute_command('AI.MODELSET', 'm_cpu{1}', 'ONNX', 'CPU', 'BLOB', model_pb)
+    env.assertEqual(ret, b'OK')
+    ai_memory_config = {k.split(":")[0]: k.split(":")[1]
+                        for k in con.execute_command("INFO MODULES").decode().split("#")[4].split()[1:]}
+    env.assertTrue(int(ai_memory_config["ai_onnxruntime_memory"]) > 100)
+    env.assertEqual(int(ai_memory_config["ai_onnxruntime_memory_access_num"]), 3)
diff --git a/tests/module/Makefile b/tests/module/Makefile
index ca282dd6b..454b7ace7 100644
--- a/tests/module/Makefile
+++ b/tests/module/Makefile
@@ -27,7 +27,7 @@ all: $(TEST_MODULES)
 	$(MAKE) CFLAGS="-m32" LDFLAGS="-melf_i386"
 
 %.o: %.c
-	$(CC) $(DEBUGFLAGS) -I../../src -I../../deps -DREDIS_MODULE_TARGET -DREDISMODULE_EXPERIMENTAL_API $(SHOBJ_CFLAGS) -fPIC -c $< -o $@
+	$(CC) $(DEBUGFLAGS) -I../../src -DREDIS_MODULE_TARGET -DREDISMODULE_EXPERIMENTAL_API $(SHOBJ_CFLAGS) -fPIC -c $< -o $@
 
 LLAPI.so: LLAPI.o DAG_utils.o
 	$(CC) -o $@ $^ $(SHOBJ_LDFLAGS) -lc -lm

From 2502d09fa245c9d1df3ab61f5a1bb23d9caa80ef Mon Sep 17 00:00:00 2001
From: alonre24 <alonreshef24@gmail.com>
Date: Sun, 28 Feb 2021 15:43:35 +0200
Subject: [PATCH 14/32] Use custom allocator also in GPU, for allocating input
 and output names.

---
 src/backends/onnxruntime.c | 26 +++++++-------------------
 tests/flow/tests_onnx.py   |  7 +++++--
 2 files changed, 12 insertions(+), 21 deletions(-)

diff --git a/src/backends/onnxruntime.c b/src/backends/onnxruntime.c
index 044932c1c..9df54aaaa 100644
--- a/src/backends/onnxruntime.c
+++ b/src/backends/onnxruntime.c
@@ -13,17 +13,11 @@
     if ((status = (x)) != NULL)                                                                    \
         goto error;
 
-// If we run on GPU, we do not use the custom allocator (redis allocator), but
-// the default allocator returned by ORT api. If we run on CPU, we do not need to
-// use this api, as the global allocator is already set to be the custom allocator.
-#ifdef RAI_ONNXRUNTIME_USE_CUDA
-#define GET_GLOBAL_ALLOCATOR                                                                       \
-    ONNX_VALIDATE_STATUS(ort->GetAllocatorWithDefaultOptions(&global_allocator))
-#else
-#define GET_GLOBAL_ALLOCATOR
-#endif
-
 OrtEnv *env = NULL;
+
+// For model that run on GPU, onnx will not use the custom allocator (redis allocator), but
+// the onnx allocator for GPU. But for the auxilery allocations of the input and output names,
+// we will use the custom global allocator for models that run on GPU as well
 OrtAllocator *global_allocator = NULL;
 unsigned long long OnnxMemory = 0;
 unsigned long long OnnxMemoryAccessCounter = 0;
@@ -182,7 +176,6 @@ OrtValue *RAI_OrtValueFromTensors(RAI_Tensor **ts, size_t count, RAI_Error *erro
     batched_shape[0] = batch_size;
 
     OrtValue *out;
-    GET_GLOBAL_ALLOCATOR
     if (count > 1) {
         ONNX_VALIDATE_STATUS(
             ort->CreateTensorAsOrtValue(global_allocator, batched_shape, t0->tensor.dl_tensor.ndim,
@@ -321,18 +314,16 @@ RAI_Model *RAI_ModelCreateORT(RAI_Backend backend, const char *devicestr, RAI_Mo
 
     if (env == NULL) {
         ONNX_VALIDATE_STATUS(ort->CreateEnv(ORT_LOGGING_LEVEL_WARNING, "test", &env))
-#ifndef RAI_ONNXRUNTIME_USE_CUDA
         ONNX_VALIDATE_STATUS(ort->CreateCustomDeviceAllocator(
             ORT_API_VERSION, AllocatorAlloc, AllocatorFree, AllocatorInfo, &global_allocator))
         ONNX_VALIDATE_STATUS(ort->RegisterCustomDeviceAllocator(env, global_allocator))
-#endif
     }
 
     ONNX_VALIDATE_STATUS(ort->CreateSessionOptions(&session_options))
 
 #ifndef RAI_ONNXRUNTIME_USE_CUDA
-    // These are required to ensure that onnx will use the registered REDIS allocator (for CPU
-    // only).
+    // These are required to ensure that onnx will use the registered REDIS allocator (for
+    // a model that defined to run on CPU).
     ONNX_VALIDATE_STATUS(
         ort->AddSessionConfigEntry(session_options, "session.use_env_allocators", "1"))
     ONNX_VALIDATE_STATUS(ort->DisableCpuMemArena(session_options))
@@ -357,7 +348,6 @@ RAI_Model *RAI_ModelCreateORT(RAI_Backend backend, const char *devicestr, RAI_Mo
     size_t n_output_nodes;
     ONNX_VALIDATE_STATUS(ort->SessionGetOutputCount(session, &n_output_nodes))
 
-    GET_GLOBAL_ALLOCATOR
     inputs_ = array_new(char *, n_input_nodes);
     for (long long i = 0; i < n_input_nodes; i++) {
         char *input_name;
@@ -424,7 +414,6 @@ void RAI_ModelFreeORT(RAI_Model *model, RAI_Error *error) {
     const OrtApi *ort = OrtGetApiBase()->GetApi(1);
     OrtStatus *status = NULL;
 
-    GET_GLOBAL_ALLOCATOR
     for (uint32_t i = 0; i < model->ninputs; i++) {
         ONNX_VALIDATE_STATUS(ort->AllocatorFree(global_allocator, model->inputs[i]))
     }
@@ -481,8 +470,7 @@ int RAI_ModelRunORT(RAI_ModelRunCtx **mctxs, RAI_Error *error) {
     ONNX_VALIDATE_STATUS(ort->SessionGetInputCount(session, &n_input_nodes))
 
     size_t n_output_nodes;
-    ONNX_VALIDATE_STATUS(ort->SessionGetOutputCount(session, &n_output_nodes))
-    GET_GLOBAL_ALLOCATOR {
+    ONNX_VALIDATE_STATUS(ort->SessionGetOutputCount(session, &n_output_nodes)) {
         const char *input_names[n_input_nodes];
         const char *output_names[n_output_nodes];
 
diff --git a/tests/flow/tests_onnx.py b/tests/flow/tests_onnx.py
index 26788e7ed..1946bf7fc 100644
--- a/tests/flow/tests_onnx.py
+++ b/tests/flow/tests_onnx.py
@@ -532,7 +532,9 @@ def test_onnx_use_custom_allocator_with_GPU(env):
                         for k in con.execute_command("INFO MODULES").decode().split("#")[4].split()[1:]}
     env.assertEqual(int(ai_memory_config["ai_onnxruntime_memory"]), 0)
 
-    # Create the same model, once for CPU and once for GPU, and verify that redis allocator was being used only for CPU
+    # Create the same model, once for CPU and once for GPU.
+    # Expect using the allocator during model set for allocating the model, its input name and output name in CPU,
+    # but for GPU, expcet using the allocator only for allocating input and output names.
     ret = con.execute_command('AI.MODELSET', 'm_gpu{1}', 'ONNX', DEVICE, 'BLOB', model_pb)
     env.assertEqual(ret, b'OK')
     ret = con.execute_command('AI.MODELSET', 'm_cpu{1}', 'ONNX', 'CPU', 'BLOB', model_pb)
@@ -540,4 +542,5 @@ def test_onnx_use_custom_allocator_with_GPU(env):
     ai_memory_config = {k.split(":")[0]: k.split(":")[1]
                         for k in con.execute_command("INFO MODULES").decode().split("#")[4].split()[1:]}
     env.assertTrue(int(ai_memory_config["ai_onnxruntime_memory"]) > 100)
-    env.assertEqual(int(ai_memory_config["ai_onnxruntime_memory_access_num"]), 3)
+    env.assertTrue(int(ai_memory_config["ai_onnxruntime_memory"]) < 400)
+    env.assertEqual(int(ai_memory_config["ai_onnxruntime_memory_access_num"]), 5)

From daa6b2e25256dc4c7fbaee71223e8bab55d9f81c Mon Sep 17 00:00:00 2001
From: DvirDukhan <dvir@redislabs.com>
Date: Mon, 1 Mar 2021 01:25:34 +0200
Subject: [PATCH 15/32] added gpu dockerfile

---
 opt/build/onnxruntime/Dockerfile.x64-gpu | 47 ++++++++++++++++++++++++
 opt/build/onnxruntime/Makefile           | 14 +++++++
 2 files changed, 61 insertions(+)
 create mode 100644 opt/build/onnxruntime/Dockerfile.x64-gpu

diff --git a/opt/build/onnxruntime/Dockerfile.x64-gpu b/opt/build/onnxruntime/Dockerfile.x64-gpu
new file mode 100644
index 000000000..4c0210325
--- /dev/null
+++ b/opt/build/onnxruntime/Dockerfile.x64-gpu
@@ -0,0 +1,47 @@
+ARG OS=ubuntu18.04
+ARG CUDA_VER=11.0-cudnn8
+
+FROM nvidia/cuda:11.0-cudnn8-devel-ubuntu18.04
+
+ARG ONNXRUNTIME_REPO=https://github.com/RedisAI/onnxruntime
+ARG ONNXRUNTIME_VER=1.6.0
+ARG ARCH=x64-gpu
+
+RUN apt-get -qq update
+RUN apt-get -qq install -y curl wget tar git
+RUN apt-get -qq install -y build-essential
+RUN apt-get -qq install -y libcurl4-openssl-dev libssl-dev libatlas-base-dev zlib1g-dev
+
+RUN apt-get -qq install -y python3 python3-pip python3-dev
+RUN pip3 install --upgrade pip setuptools wheel
+# RUN pip3 install numpy
+RUN apt-get -q install -y python3-numpy
+RUN apt-get install -y snapd
+RUN wget https://github.com/Kitware/CMake/releases/download/v3.15.2/cmake-3.15.2.tar.gz; \
+    tar -zxvf cmake-3.15.2.tar.gz; \
+    cd cmake-3.15.2; \
+    ./bootstrap; \
+    make -j4; \
+    make install
+
+
+ENV LANG=en_US.UTF-8
+RUN apt-get install -y locales && \
+    sed -i -e "s/# $LANG.*/$LANG UTF-8/" /etc/locale.gen && \
+    dpkg-reconfigure --frontend=noninteractive locales && \
+    update-locale LANG=$LANG
+
+WORKDIR /build
+
+ADD ./pack.sh /build/
+ARG BUILDTYPE=MinSizeRel
+ARG BUILDARGS="--config ${BUILDTYPE} --parallel"
+
+RUN git clone --single-branch --branch rel-${ONNXRUNTIME_VER} --recursive ${ONNXRUNTIME_REPO} onnxruntime
+
+
+RUN	cd onnxruntime ;\
+    ./build.sh ${BUILDARGS} --update --build --use_cuda --cudnn_home /usr/local/cuda --cuda_home /usr/local/cuda --build_shared_lib --parallel
+# RUN ./build.sh ${BUILDARGS} --enable_pybind --build_wheel
+
+RUN ./pack.sh ${ONNXRUNTIME_VER} ${ARCH}
diff --git a/opt/build/onnxruntime/Makefile b/opt/build/onnxruntime/Makefile
index c28472b1f..f37f10a78 100755
--- a/opt/build/onnxruntime/Makefile
+++ b/opt/build/onnxruntime/Makefile
@@ -48,6 +48,19 @@ build_x64:
 .PHONY: build_x64
 endef
 
+define build_x64-gpu # (1=arch, 2=tar-arch)
+IID_$(1)=$(1)_$(VERSION).iid
+CID_$(1)=$(1)_$(VERSION).cid
+
+build_x64-gpu:
+	@docker build --iidfile $$(IID_$(1)) -t redisfab/$(STEM)-x64:$(VERSION) -f Dockerfile.x64-gpu \
+		--build-arg OS=$(DOCKER_OS) $(ROOT)
+	@docker create --cidfile $$(CID_$(1)) `cat $$(IID_$(1))`
+	@docker cp `cat $$(CID_$(1))`:/build/$(STEM)-x64-$(VERSION).tgz .
+
+.PHONY: build_x64
+endef
+
 define build_arm # (1=arch, 2=tar-arch)
 IID_$(1)=$(1)_$(VERSION).iid
 CID_$(1)=$(1)_$(VERSION).cid
@@ -97,6 +110,7 @@ all: build publish
 build: $(BUILD_TARGETS)
 
 $(eval $(call build_x64,x64,x86_64))
+$(eval $(call build_x64-gpu,x64-gpu,x86_64-gpu))
 $(eval $(call build_arm,arm64v8,arm64))
 $(eval $(call build_arm,arm32v7,arm))
 

From fb7c7584d138447728a811eb5fda2cb0652b23bb Mon Sep 17 00:00:00 2001
From: alonre24 <alonreshef24@gmail.com>
Date: Tue, 2 Mar 2021 00:54:58 +0200
Subject: [PATCH 16/32] Fixes for building and publishing onnx for gpu.

---
 opt/build/onnxruntime/Makefile | 4 ++--
 opt/build/onnxruntime/pack.sh  | 6 +++---
 2 files changed, 5 insertions(+), 5 deletions(-)

diff --git a/opt/build/onnxruntime/Makefile b/opt/build/onnxruntime/Makefile
index f37f10a78..202a9b3ed 100755
--- a/opt/build/onnxruntime/Makefile
+++ b/opt/build/onnxruntime/Makefile
@@ -53,10 +53,10 @@ IID_$(1)=$(1)_$(VERSION).iid
 CID_$(1)=$(1)_$(VERSION).cid
 
 build_x64-gpu:
-	@docker build --iidfile $$(IID_$(1)) -t redisfab/$(STEM)-x64:$(VERSION) -f Dockerfile.x64-gpu \
+	@docker build --iidfile $$(IID_$(1)) -t redisfab/$(STEM)-$(1):$(VERSION) -f Dockerfile.x64-gpu \
 		--build-arg OS=$(DOCKER_OS) $(ROOT)
 	@docker create --cidfile $$(CID_$(1)) `cat $$(IID_$(1))`
-	@docker cp `cat $$(CID_$(1))`:/build/$(STEM)-x64-$(VERSION).tgz .
+	@docker cp `cat $$(CID_$(1))`:/build/$(STEM)-$(1)-$(VERSION).tgz .
 
 .PHONY: build_x64
 endef
diff --git a/opt/build/onnxruntime/pack.sh b/opt/build/onnxruntime/pack.sh
index 447174b70..fdd1113b6 100755
--- a/opt/build/onnxruntime/pack.sh
+++ b/opt/build/onnxruntime/pack.sh
@@ -2,7 +2,7 @@
 
 set -e
 VER="$1"
-PLATFORM=$($READIES/bin/platform -t)
+PLATFORM="$2"
 
 mkdir -p pack/include pack/lib
 cp onnxruntime/build/Linux/MinSizeRel/libonnxruntime.so.${VER} pack/lib/
@@ -19,5 +19,5 @@ cp onnxruntime/include/onnxruntime/core/providers/cuda/cuda_provider_factory.h p
 cd pack/lib/
 ln -s libonnxruntime.so.${VER} libonnxruntime.so
 cd ../..
-mv pack onnxruntime-${PLATFORM}-${VER}
-tar czf onnxruntime-${PLATFORM}-${VER}.tgz onnxruntime-${PLATFORM}-${VER}/
+mv pack onnxruntime-linux-${PLATFORM}-${VER}
+tar czf onnxruntime-linux-${PLATFORM}-${VER}.tgz onnxruntime-linux-${PLATFORM}-${VER}/

From 3d7be97b91b21ce667a113094b0f0854c169627c Mon Sep 17 00:00:00 2001
From: "Chayim I. Kirshen" <chayim@gnupower.net>
Date: Tue, 2 Mar 2021 08:37:15 +0200
Subject: [PATCH 17/32] speeding up the docker build by pararellizing the
 submodule fetch

Adding pack from the filesystem was moved to the bottom, so that it can be iterated on outside of the docker, without triggering a layer (and subsequent) layer rebuild. This speeds up the debug cycle
---
 opt/build/onnxruntime/Dockerfile.x64     | 24 +++++++--------
 opt/build/onnxruntime/Dockerfile.x64-gpu | 38 +++++++++++-------------
 2 files changed, 28 insertions(+), 34 deletions(-)

diff --git a/opt/build/onnxruntime/Dockerfile.x64 b/opt/build/onnxruntime/Dockerfile.x64
index 65cd9ed72..8f9d32682 100755
--- a/opt/build/onnxruntime/Dockerfile.x64
+++ b/opt/build/onnxruntime/Dockerfile.x64
@@ -10,14 +10,11 @@ ARG ONNXRUNTIME_VER=1.6.0
 ARG ARCH=x64
 
 RUN apt-get -qq update
-RUN apt-get -qq install -y curl wget tar git
-RUN apt-get -qq install -y build-essential cmake
-RUN apt-get -qq install -y libcurl4-openssl-dev libssl-dev libatlas-base-dev zlib1g-dev
-
-RUN apt-get -qq install -y python3 python3-pip python3-dev
+RUN apt-get -qq install -y curl wget tar git \
+            build-essential cmake \
+            libcurl4-openssl-dev libssl-dev libatlas-base-dev zlib1g-dev \
+            python3 python3-pip python3-dev python3-numpy
 RUN pip3 install --upgrade pip setuptools wheel
-# RUN pip3 install numpy
-RUN apt-get -q install -y python3-numpy
 
 ENV LANG=en_US.UTF-8
 RUN apt-get install -y locales && \
@@ -27,15 +24,16 @@ RUN apt-get install -y locales && \
 
 WORKDIR /build
 
-ADD ./pack.sh /build/
 ARG BUILDTYPE=MinSizeRel
 ARG BUILDARGS="--config ${BUILDTYPE} --parallel"
 
-RUN git clone --single-branch --branch rel-${ONNXRUNTIME_VER} --recursive ${ONNXRUNTIME_REPO} onnxruntime
-	
-RUN	cd onnxruntime ;\
-	./build.sh ${BUILDARGS} --update --build ;\
-	./build.sh ${BUILDARGS} --build_shared_lib
+RUN git clone --single-branch --branch rel-${ONNXRUNTIME_VER} ${ONNXRUNTIME_REPO} onnxruntime
+WORKDIR /build/onnxruntime
+RUN git fetch --recurse-submodules -j4
+RUN ./build.sh ${BUILDARGS} --update --build
+RUN ./build.sh ${BUILDARGS} --build_shared_lib
 # RUN ./build.sh ${BUILDARGS} --enable_pybind --build_wheel
 
+ADD ./pack.sh /build
+WORKDIR /build
 RUN ./pack.sh ${ONNXRUNTIME_VER} ${ARCH}
diff --git a/opt/build/onnxruntime/Dockerfile.x64-gpu b/opt/build/onnxruntime/Dockerfile.x64-gpu
index 4c0210325..e381c2a52 100644
--- a/opt/build/onnxruntime/Dockerfile.x64-gpu
+++ b/opt/build/onnxruntime/Dockerfile.x64-gpu
@@ -8,21 +8,18 @@ ARG ONNXRUNTIME_VER=1.6.0
 ARG ARCH=x64-gpu
 
 RUN apt-get -qq update
-RUN apt-get -qq install -y curl wget tar git
-RUN apt-get -qq install -y build-essential
-RUN apt-get -qq install -y libcurl4-openssl-dev libssl-dev libatlas-base-dev zlib1g-dev
+RUN apt-get -qq install -y curl \
+            wget tar git build-essential \
+            libcurl4-openssl-dev libssl-dev \
+            libatlas-base-dev zlib1g-dev python3 python3-pip \
+            python3-dev python3-numpy rsync
 
-RUN apt-get -qq install -y python3 python3-pip python3-dev
 RUN pip3 install --upgrade pip setuptools wheel
-# RUN pip3 install numpy
-RUN apt-get -q install -y python3-numpy
-RUN apt-get install -y snapd
-RUN wget https://github.com/Kitware/CMake/releases/download/v3.15.2/cmake-3.15.2.tar.gz; \
-    tar -zxvf cmake-3.15.2.tar.gz; \
-    cd cmake-3.15.2; \
-    ./bootstrap; \
-    make -j4; \
-    make install
+
+RUN wget -q https://github.com/Kitware/CMake/releases/download/v3.19.5/cmake-3.19.5-Linux-x86_64.tar.gz -O /tmp/cmake.tgz
+WORKDIR /tmp
+RUN tar -xpf cmake.tgz
+RUN rsync -aqH cmake*/* /usr
 
 
 ENV LANG=en_US.UTF-8
@@ -33,15 +30,14 @@ RUN apt-get install -y locales && \
 
 WORKDIR /build
 
-ADD ./pack.sh /build/
 ARG BUILDTYPE=MinSizeRel
 ARG BUILDARGS="--config ${BUILDTYPE} --parallel"
 
-RUN git clone --single-branch --branch rel-${ONNXRUNTIME_VER} --recursive ${ONNXRUNTIME_REPO} onnxruntime
-
+RUN git clone --single-branch --branch rel-${ONNXRUNTIME_VER} ${ONNXRUNTIME_REPO} onnxruntime
+WORKDIR /build/onnxruntime
+RUN git fetch --recurse-submodules -j4
+RUN ./build.sh ${BUILDARGS} --update --build --use_cuda --cudnn_home /usr/local/cuda --cuda_home /usr/local/cuda --build_shared_lib --parallel
 
-RUN	cd onnxruntime ;\
-    ./build.sh ${BUILDARGS} --update --build --use_cuda --cudnn_home /usr/local/cuda --cuda_home /usr/local/cuda --build_shared_lib --parallel
-# RUN ./build.sh ${BUILDARGS} --enable_pybind --build_wheel
-
-RUN ./pack.sh ${ONNXRUNTIME_VER} ${ARCH}
+ADD ./pack.sh /build/
+WORKDIR /build
+RUN /pack.sh ${ONNXRUNTIME_VER} ${ARCH}

From dd5bdc64d1381ee464b7fd4e4a8fa9e58a10f4d1 Mon Sep 17 00:00:00 2001
From: alonre24 <alonreshef24@gmail.com>
Date: Tue, 2 Mar 2021 16:06:38 +0200
Subject: [PATCH 18/32] for Debug COV test in CI

---
 tests/flow/tests_onnx.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/tests/flow/tests_onnx.py b/tests/flow/tests_onnx.py
index 1946bf7fc..b4ba8273d 100644
--- a/tests/flow/tests_onnx.py
+++ b/tests/flow/tests_onnx.py
@@ -16,7 +16,7 @@ def test_onnx_modelrun_mnist(env):
         return
 
     con = env.getConnection()
-
+    x=input()
     test_data_path = os.path.join(os.path.dirname(__file__), 'test_data')
     model_filename = os.path.join(test_data_path, 'mnist.onnx')
     wrong_model_filename = os.path.join(test_data_path, 'graph.pb')

From 4c49aa9b77b586f23ed6f06fe1494f20882e37b0 Mon Sep 17 00:00:00 2001
From: alonre24 <alonreshef24@gmail.com>
Date: Tue, 2 Mar 2021 18:47:16 +0200
Subject: [PATCH 19/32] Use custom allocator for models that created for CPU,
 even when CUDA is used.

---
 src/backends/onnxruntime.c | 15 +++++++--------
 tests/flow/tests_onnx.py   |  7 ++-----
 2 files changed, 9 insertions(+), 13 deletions(-)

diff --git a/src/backends/onnxruntime.c b/src/backends/onnxruntime.c
index 9df54aaaa..afa643245 100644
--- a/src/backends/onnxruntime.c
+++ b/src/backends/onnxruntime.c
@@ -320,14 +320,13 @@ RAI_Model *RAI_ModelCreateORT(RAI_Backend backend, const char *devicestr, RAI_Mo
     }
 
     ONNX_VALIDATE_STATUS(ort->CreateSessionOptions(&session_options))
-
-#ifndef RAI_ONNXRUNTIME_USE_CUDA
-    // These are required to ensure that onnx will use the registered REDIS allocator (for
-    // a model that defined to run on CPU).
-    ONNX_VALIDATE_STATUS(
-        ort->AddSessionConfigEntry(session_options, "session.use_env_allocators", "1"))
-    ONNX_VALIDATE_STATUS(ort->DisableCpuMemArena(session_options))
-#endif
+    if (strcasecmp(devicestr, "CPU") != 0) {
+        // These are required to ensure that onnx will use the registered REDIS allocator (for
+        // a model that defined to run on CPU).
+        ONNX_VALIDATE_STATUS(
+            ort->AddSessionConfigEntry(session_options, "session.use_env_allocators", "1"))
+        ONNX_VALIDATE_STATUS(ort->DisableCpuMemArena(session_options))
+    }
 
     // TODO: these options could be configured at the AI.CONFIG level
     ONNX_VALIDATE_STATUS(ort->SetSessionGraphOptimizationLevel(session_options, 1))
diff --git a/tests/flow/tests_onnx.py b/tests/flow/tests_onnx.py
index b4ba8273d..177e27723 100644
--- a/tests/flow/tests_onnx.py
+++ b/tests/flow/tests_onnx.py
@@ -458,9 +458,6 @@ def test_onnx_use_custom_allocator(env):
     if not TEST_ONNX:
         env.debugPrint("skipping {} since TEST_ONNX=0".format(sys._getframe().f_code.co_name), force=True)
         return
-    if DEVICE != 'CPU':
-        env.debugPrint("skipping {} since the custom allocator is for CPU".format(sys._getframe().f_code.co_name), force=True)
-        return
 
     con = env.getConnection()
     test_data_path = os.path.join(os.path.dirname(__file__), 'test_data')
@@ -472,7 +469,7 @@ def test_onnx_use_custom_allocator(env):
     env.assertEqual(int(ai_memory_config["ai_onnxruntime_memory"]), 0)
 
     # Expect using the allocator during model set for allocating the model, its input name and output name.
-    ret = con.execute_command('AI.MODELSET', 'm{1}', 'ONNX', DEVICE, 'BLOB', model_pb)
+    ret = con.execute_command('AI.MODELSET', 'm{1}', 'ONNX', 'CPU', 'BLOB', model_pb)
     env.assertEqual(ret, b'OK')
     ai_memory_config = {k.split(":")[0]: k.split(":")[1]
                         for k in con.execute_command("INFO MODULES").decode().split("#")[4].split()[1:]}
@@ -496,7 +493,7 @@ def test_onnx_use_custom_allocator(env):
     with open(sample_filename, 'rb') as f:
         sample_raw = f.read()
 
-    ret = con.execute_command('AI.MODELSET', 'm{1}', 'ONNX', DEVICE, 'BLOB', model_pb)
+    ret = con.execute_command('AI.MODELSET', 'm{1}', 'ONNX', 'CPU', 'BLOB', model_pb)
     env.assertEqual(ret, b'OK')
     con.execute_command('AI.TENSORSET', 'a{1}', 'FLOAT', 1, 1, 28, 28, 'BLOB', sample_raw)
 

From eac91a1936e029355043607a26b515bce0cf655b Mon Sep 17 00:00:00 2001
From: alonre24 <alonreshef24@gmail.com>
Date: Wed, 3 Mar 2021 16:22:30 +0200
Subject: [PATCH 20/32] Update build config of onnxruntime to RelWithDebInfo

---
 opt/build/onnxruntime/Dockerfile.x64 | 2 +-
 opt/build/onnxruntime/pack.sh        | 2 +-
 2 files changed, 2 insertions(+), 2 deletions(-)

diff --git a/opt/build/onnxruntime/Dockerfile.x64 b/opt/build/onnxruntime/Dockerfile.x64
index 65cd9ed72..eb80d670e 100755
--- a/opt/build/onnxruntime/Dockerfile.x64
+++ b/opt/build/onnxruntime/Dockerfile.x64
@@ -28,7 +28,7 @@ RUN apt-get install -y locales && \
 WORKDIR /build
 
 ADD ./pack.sh /build/
-ARG BUILDTYPE=MinSizeRel
+ARG BUILDTYPE=RelWithDebInfo
 ARG BUILDARGS="--config ${BUILDTYPE} --parallel"
 
 RUN git clone --single-branch --branch rel-${ONNXRUNTIME_VER} --recursive ${ONNXRUNTIME_REPO} onnxruntime
diff --git a/opt/build/onnxruntime/pack.sh b/opt/build/onnxruntime/pack.sh
index fdd1113b6..106aafdd7 100755
--- a/opt/build/onnxruntime/pack.sh
+++ b/opt/build/onnxruntime/pack.sh
@@ -5,7 +5,7 @@ VER="$1"
 PLATFORM="$2"
 
 mkdir -p pack/include pack/lib
-cp onnxruntime/build/Linux/MinSizeRel/libonnxruntime.so.${VER} pack/lib/
+cp onnxruntime/build/Linux/RelWithDebInfo/libonnxruntime.so.${VER} pack/lib/
 cp onnxruntime/docs/C_API.md pack/
 cp onnxruntime/LICENSE pack/
 cp onnxruntime/README.md pack/

From a0ec6c3edeac7431b00d99b9c34799f0348de9f9 Mon Sep 17 00:00:00 2001
From: alonre24 <alonreshef24@gmail.com>
Date: Thu, 4 Mar 2021 13:28:39 +0200
Subject: [PATCH 21/32] Debug coverage in CI

---
 .circleci/config.yml                 | 6 +++---
 opt/build/onnxruntime/Dockerfile.x64 | 2 +-
 opt/build/onnxruntime/pack.sh        | 2 +-
 3 files changed, 5 insertions(+), 5 deletions(-)

diff --git a/.circleci/config.yml b/.circleci/config.yml
index eadf0a1db..05fee460f 100644
--- a/.circleci/config.yml
+++ b/.circleci/config.yml
@@ -162,9 +162,9 @@ jobs:
       - run:
           name: Submodule checkout
           command: git submodule update --init --recursive
-      - restore_cache:
-          keys:
-          - build-dependencies-{{ checksum "get_deps.sh" }}
+      #- restore_cache:
+      #    keys:
+      #    - build-dependencies-{{ checksum "get_deps.sh" }}
           # If no exact match is found will get dependencies from source
       - setup-automation
       - run:
diff --git a/opt/build/onnxruntime/Dockerfile.x64 b/opt/build/onnxruntime/Dockerfile.x64
index eb80d670e..49a1233be 100755
--- a/opt/build/onnxruntime/Dockerfile.x64
+++ b/opt/build/onnxruntime/Dockerfile.x64
@@ -28,7 +28,7 @@ RUN apt-get install -y locales && \
 WORKDIR /build
 
 ADD ./pack.sh /build/
-ARG BUILDTYPE=RelWithDebInfo
+ARG BUILDTYPE=Debug
 ARG BUILDARGS="--config ${BUILDTYPE} --parallel"
 
 RUN git clone --single-branch --branch rel-${ONNXRUNTIME_VER} --recursive ${ONNXRUNTIME_REPO} onnxruntime
diff --git a/opt/build/onnxruntime/pack.sh b/opt/build/onnxruntime/pack.sh
index 106aafdd7..254135b80 100755
--- a/opt/build/onnxruntime/pack.sh
+++ b/opt/build/onnxruntime/pack.sh
@@ -5,7 +5,7 @@ VER="$1"
 PLATFORM="$2"
 
 mkdir -p pack/include pack/lib
-cp onnxruntime/build/Linux/RelWithDebInfo/libonnxruntime.so.${VER} pack/lib/
+cp onnxruntime/build/Linux/Debug/libonnxruntime.so.${VER} pack/lib/
 cp onnxruntime/docs/C_API.md pack/
 cp onnxruntime/LICENSE pack/
 cp onnxruntime/README.md pack/

From 1b6ba6ca75b9048899cb4c261654fab635d06ca7 Mon Sep 17 00:00:00 2001
From: alonre24 <alonreshef24@gmail.com>
Date: Thu, 4 Mar 2021 13:28:39 +0200
Subject: [PATCH 22/32] Debug coverage in CI

---
 .circleci/config.yml                 | 6 +++---
 opt/build/onnxruntime/Dockerfile.x64 | 2 +-
 opt/build/onnxruntime/pack.sh        | 2 +-
 src/backends/onnxruntime.c           | 2 +-
 4 files changed, 6 insertions(+), 6 deletions(-)

diff --git a/.circleci/config.yml b/.circleci/config.yml
index eadf0a1db..05fee460f 100644
--- a/.circleci/config.yml
+++ b/.circleci/config.yml
@@ -162,9 +162,9 @@ jobs:
       - run:
           name: Submodule checkout
           command: git submodule update --init --recursive
-      - restore_cache:
-          keys:
-          - build-dependencies-{{ checksum "get_deps.sh" }}
+      #- restore_cache:
+      #    keys:
+      #    - build-dependencies-{{ checksum "get_deps.sh" }}
           # If no exact match is found will get dependencies from source
       - setup-automation
       - run:
diff --git a/opt/build/onnxruntime/Dockerfile.x64 b/opt/build/onnxruntime/Dockerfile.x64
index eb80d670e..49a1233be 100755
--- a/opt/build/onnxruntime/Dockerfile.x64
+++ b/opt/build/onnxruntime/Dockerfile.x64
@@ -28,7 +28,7 @@ RUN apt-get install -y locales && \
 WORKDIR /build
 
 ADD ./pack.sh /build/
-ARG BUILDTYPE=RelWithDebInfo
+ARG BUILDTYPE=Debug
 ARG BUILDARGS="--config ${BUILDTYPE} --parallel"
 
 RUN git clone --single-branch --branch rel-${ONNXRUNTIME_VER} --recursive ${ONNXRUNTIME_REPO} onnxruntime
diff --git a/opt/build/onnxruntime/pack.sh b/opt/build/onnxruntime/pack.sh
index 106aafdd7..254135b80 100755
--- a/opt/build/onnxruntime/pack.sh
+++ b/opt/build/onnxruntime/pack.sh
@@ -5,7 +5,7 @@ VER="$1"
 PLATFORM="$2"
 
 mkdir -p pack/include pack/lib
-cp onnxruntime/build/Linux/RelWithDebInfo/libonnxruntime.so.${VER} pack/lib/
+cp onnxruntime/build/Linux/Debug/libonnxruntime.so.${VER} pack/lib/
 cp onnxruntime/docs/C_API.md pack/
 cp onnxruntime/LICENSE pack/
 cp onnxruntime/README.md pack/
diff --git a/src/backends/onnxruntime.c b/src/backends/onnxruntime.c
index cd3322a3d..052e806cb 100644
--- a/src/backends/onnxruntime.c
+++ b/src/backends/onnxruntime.c
@@ -329,7 +329,7 @@ RAI_Model *RAI_ModelCreateORT(RAI_Backend backend, const char *devicestr, RAI_Mo
     }
 
     // TODO: these options could be configured at the AI.CONFIG level
-    ONNX_VALIDATE_STATUS(ort->SetSessionGraphOptimizationLevel(session_options, 1))
+    //ONNX_VALIDATE_STATUS(ort->SetSessionGraphOptimizationLevel(session_options, 1))
     ONNX_VALIDATE_STATUS(
         ort->SetIntraOpNumThreads(session_options, (int)opts.backends_intra_op_parallelism))
     ONNX_VALIDATE_STATUS(

From f9e45ffed811150ab32ddec4c5663937bad7e129 Mon Sep 17 00:00:00 2001
From: alonre24 <alonreshef24@gmail.com>
Date: Fri, 5 Mar 2021 15:37:50 +0200
Subject: [PATCH 23/32] Disable graph optimization (to fix flaky tests in onnx)

---
 .circleci/config.yml                 | 6 +++---
 opt/build/onnxruntime/Dockerfile.x64 | 2 +-
 opt/build/onnxruntime/pack.sh        | 2 +-
 src/backends/onnxruntime.c           | 2 +-
 tests/flow/tests_onnx.py             | 2 +-
 5 files changed, 7 insertions(+), 7 deletions(-)

diff --git a/.circleci/config.yml b/.circleci/config.yml
index 05fee460f..eadf0a1db 100644
--- a/.circleci/config.yml
+++ b/.circleci/config.yml
@@ -162,9 +162,9 @@ jobs:
       - run:
           name: Submodule checkout
           command: git submodule update --init --recursive
-      #- restore_cache:
-      #    keys:
-      #    - build-dependencies-{{ checksum "get_deps.sh" }}
+      - restore_cache:
+          keys:
+          - build-dependencies-{{ checksum "get_deps.sh" }}
           # If no exact match is found will get dependencies from source
       - setup-automation
       - run:
diff --git a/opt/build/onnxruntime/Dockerfile.x64 b/opt/build/onnxruntime/Dockerfile.x64
index 3023fb4ba..b14bd80a3 100755
--- a/opt/build/onnxruntime/Dockerfile.x64
+++ b/opt/build/onnxruntime/Dockerfile.x64
@@ -24,7 +24,7 @@ RUN apt-get install -y locales && \
 
 WORKDIR /build
 
-ARG BUILDTYPE=RelWithDebInfo
+ARG BUILDTYPE=MinSizeRel
 
 ARG BUILDARGS="--config ${BUILDTYPE} --parallel"
 
diff --git a/opt/build/onnxruntime/pack.sh b/opt/build/onnxruntime/pack.sh
index 254135b80..fdd1113b6 100755
--- a/opt/build/onnxruntime/pack.sh
+++ b/opt/build/onnxruntime/pack.sh
@@ -5,7 +5,7 @@ VER="$1"
 PLATFORM="$2"
 
 mkdir -p pack/include pack/lib
-cp onnxruntime/build/Linux/Debug/libonnxruntime.so.${VER} pack/lib/
+cp onnxruntime/build/Linux/MinSizeRel/libonnxruntime.so.${VER} pack/lib/
 cp onnxruntime/docs/C_API.md pack/
 cp onnxruntime/LICENSE pack/
 cp onnxruntime/README.md pack/
diff --git a/src/backends/onnxruntime.c b/src/backends/onnxruntime.c
index 55dd0b9c5..faa8c08fa 100644
--- a/src/backends/onnxruntime.c
+++ b/src/backends/onnxruntime.c
@@ -329,7 +329,7 @@ RAI_Model *RAI_ModelCreateORT(RAI_Backend backend, const char *devicestr, RAI_Mo
     }
 
     // TODO: these options could be configured at the AI.CONFIG level
-    // ONNX_VALIDATE_STATUS(ort->SetSessionGraphOptimizationLevel(session_options, 1))
+    ONNX_VALIDATE_STATUS(ort->SetSessionGraphOptimizationLevel(session_options, ORT_DISABLE_ALL))
     ONNX_VALIDATE_STATUS(
         ort->SetIntraOpNumThreads(session_options, (int)opts.backends_intra_op_parallelism))
     ONNX_VALIDATE_STATUS(
diff --git a/tests/flow/tests_onnx.py b/tests/flow/tests_onnx.py
index 7efb632ab..202f0c44f 100644
--- a/tests/flow/tests_onnx.py
+++ b/tests/flow/tests_onnx.py
@@ -504,7 +504,7 @@ def test_onnx_use_custom_allocator(env):
     ai_memory_config = {k.split(":")[0]: k.split(":")[1]
                         for k in con.execute_command("INFO MODULES").decode().split("#")[4].split()[1:]}
     allocator_access_num_after = ai_memory_config["ai_onnxruntime_memory_access_num"]
-    env.assertEqual(int(allocator_access_num_after) - int(allocator_access_num_before), 16)
+    env.assertEqual(int(allocator_access_num_after) - int(allocator_access_num_before), 22)
 
     values = con.execute_command('AI.TENSORGET', 'b{1}', 'VALUES')
     argmax = max(range(len(values)), key=lambda i: values[i])

From 1628d381963b334ccfd3aa9dc5f3962decbf06d9 Mon Sep 17 00:00:00 2001
From: alonre24 <alonreshef24@gmail.com>
Date: Mon, 8 Mar 2021 13:30:47 +0200
Subject: [PATCH 24/32] Upgrade to onnx 1.7.1 (forked version with custom
 allocator)

---
 get_deps.sh                              | 2 +-
 opt/build/onnxruntime/Dockerfile.x64     | 2 +-
 opt/build/onnxruntime/Dockerfile.x64-gpu | 2 +-
 opt/build/onnxruntime/Makefile           | 2 +-
 opt/build/onnxruntime/pack.sh            | 2 +-
 src/backends/onnxruntime.c               | 2 +-
 tests/flow/tests_onnx.py                 | 2 +-
 7 files changed, 7 insertions(+), 7 deletions(-)

diff --git a/get_deps.sh b/get_deps.sh
index 0e5b0b417..a5d0f07e1 100755
--- a/get_deps.sh
+++ b/get_deps.sh
@@ -282,7 +282,7 @@ fi
 
 ################################################################################### ONNXRUNTIME
 
-ORT_VERSION="1.6.0"
+ORT_VERSION="1.7.1"
 
 if [[ $WITH_ORT != 0 ]]; then
 	[[ $FORCE == 1 ]] && rm -rf $ONNXRUNTIME
diff --git a/opt/build/onnxruntime/Dockerfile.x64 b/opt/build/onnxruntime/Dockerfile.x64
index b14bd80a3..439899269 100755
--- a/opt/build/onnxruntime/Dockerfile.x64
+++ b/opt/build/onnxruntime/Dockerfile.x64
@@ -6,7 +6,7 @@ ARG OS=debian:buster
 FROM ${OS}
 
 ARG ONNXRUNTIME_REPO=https://github.com/RedisAI/onnxruntime
-ARG ONNXRUNTIME_VER=1.6.0
+ARG ONNXRUNTIME_VER=1.7.1
 ARG ARCH=x64
 
 RUN apt-get -qq update
diff --git a/opt/build/onnxruntime/Dockerfile.x64-gpu b/opt/build/onnxruntime/Dockerfile.x64-gpu
index e381c2a52..2f3953f36 100644
--- a/opt/build/onnxruntime/Dockerfile.x64-gpu
+++ b/opt/build/onnxruntime/Dockerfile.x64-gpu
@@ -4,7 +4,7 @@ ARG CUDA_VER=11.0-cudnn8
 FROM nvidia/cuda:11.0-cudnn8-devel-ubuntu18.04
 
 ARG ONNXRUNTIME_REPO=https://github.com/RedisAI/onnxruntime
-ARG ONNXRUNTIME_VER=1.6.0
+ARG ONNXRUNTIME_VER=1.7.1
 ARG ARCH=x64-gpu
 
 RUN apt-get -qq update
diff --git a/opt/build/onnxruntime/Makefile b/opt/build/onnxruntime/Makefile
index 202a9b3ed..1faaabae9 100755
--- a/opt/build/onnxruntime/Makefile
+++ b/opt/build/onnxruntime/Makefile
@@ -1,7 +1,7 @@
 
 ROOT=.
 
-VERSION ?= 1.6.0
+VERSION ?= 1.7.1
 OSNICK ?= buster
 
 #----------------------------------------------------------------------------------------------
diff --git a/opt/build/onnxruntime/pack.sh b/opt/build/onnxruntime/pack.sh
index fdd1113b6..01b6bb5ed 100755
--- a/opt/build/onnxruntime/pack.sh
+++ b/opt/build/onnxruntime/pack.sh
@@ -6,7 +6,7 @@ PLATFORM="$2"
 
 mkdir -p pack/include pack/lib
 cp onnxruntime/build/Linux/MinSizeRel/libonnxruntime.so.${VER} pack/lib/
-cp onnxruntime/docs/C_API.md pack/
+cp onnxruntime/docs/C_API_Guidelines.md pack/
 cp onnxruntime/LICENSE pack/
 cp onnxruntime/README.md pack/
 cp onnxruntime/ThirdPartyNotices.txt pack/
diff --git a/src/backends/onnxruntime.c b/src/backends/onnxruntime.c
index faa8c08fa..be805fd45 100644
--- a/src/backends/onnxruntime.c
+++ b/src/backends/onnxruntime.c
@@ -329,7 +329,7 @@ RAI_Model *RAI_ModelCreateORT(RAI_Backend backend, const char *devicestr, RAI_Mo
     }
 
     // TODO: these options could be configured at the AI.CONFIG level
-    ONNX_VALIDATE_STATUS(ort->SetSessionGraphOptimizationLevel(session_options, ORT_DISABLE_ALL))
+    ONNX_VALIDATE_STATUS(ort->SetSessionGraphOptimizationLevel(session_options, ORT_ENABLE_BASIC))
     ONNX_VALIDATE_STATUS(
         ort->SetIntraOpNumThreads(session_options, (int)opts.backends_intra_op_parallelism))
     ONNX_VALIDATE_STATUS(
diff --git a/tests/flow/tests_onnx.py b/tests/flow/tests_onnx.py
index 202f0c44f..7efb632ab 100644
--- a/tests/flow/tests_onnx.py
+++ b/tests/flow/tests_onnx.py
@@ -504,7 +504,7 @@ def test_onnx_use_custom_allocator(env):
     ai_memory_config = {k.split(":")[0]: k.split(":")[1]
                         for k in con.execute_command("INFO MODULES").decode().split("#")[4].split()[1:]}
     allocator_access_num_after = ai_memory_config["ai_onnxruntime_memory_access_num"]
-    env.assertEqual(int(allocator_access_num_after) - int(allocator_access_num_before), 22)
+    env.assertEqual(int(allocator_access_num_after) - int(allocator_access_num_before), 16)
 
     values = con.execute_command('AI.TENSORGET', 'b{1}', 'VALUES')
     argmax = max(range(len(values)), key=lambda i: values[i])

From 84b7e215c0c118b186ac9888e2f4cece5554a733 Mon Sep 17 00:00:00 2001
From: alonre24 <alonreshef24@gmail.com>
Date: Mon, 8 Mar 2021 14:02:56 +0200
Subject: [PATCH 25/32] Debug CI

---
 .circleci/config.yml                 | 6 +++---
 opt/build/onnxruntime/Dockerfile.x64 | 2 +-
 opt/build/onnxruntime/pack.sh        | 2 +-
 3 files changed, 5 insertions(+), 5 deletions(-)

diff --git a/.circleci/config.yml b/.circleci/config.yml
index eadf0a1db..05fee460f 100644
--- a/.circleci/config.yml
+++ b/.circleci/config.yml
@@ -162,9 +162,9 @@ jobs:
       - run:
           name: Submodule checkout
           command: git submodule update --init --recursive
-      - restore_cache:
-          keys:
-          - build-dependencies-{{ checksum "get_deps.sh" }}
+      #- restore_cache:
+      #    keys:
+      #    - build-dependencies-{{ checksum "get_deps.sh" }}
           # If no exact match is found will get dependencies from source
       - setup-automation
       - run:
diff --git a/opt/build/onnxruntime/Dockerfile.x64 b/opt/build/onnxruntime/Dockerfile.x64
index 439899269..c0ad2898e 100755
--- a/opt/build/onnxruntime/Dockerfile.x64
+++ b/opt/build/onnxruntime/Dockerfile.x64
@@ -24,7 +24,7 @@ RUN apt-get install -y locales && \
 
 WORKDIR /build
 
-ARG BUILDTYPE=MinSizeRel
+ARG BUILDTYPE=Debug
 
 ARG BUILDARGS="--config ${BUILDTYPE} --parallel"
 
diff --git a/opt/build/onnxruntime/pack.sh b/opt/build/onnxruntime/pack.sh
index 01b6bb5ed..8325f94ac 100755
--- a/opt/build/onnxruntime/pack.sh
+++ b/opt/build/onnxruntime/pack.sh
@@ -5,7 +5,7 @@ VER="$1"
 PLATFORM="$2"
 
 mkdir -p pack/include pack/lib
-cp onnxruntime/build/Linux/MinSizeRel/libonnxruntime.so.${VER} pack/lib/
+cp onnxruntime/build/Linux/Debug/libonnxruntime.so.${VER} pack/lib/
 cp onnxruntime/docs/C_API_Guidelines.md pack/
 cp onnxruntime/LICENSE pack/
 cp onnxruntime/README.md pack/

From 049417e81b680392d5fb52c9df0ea0c352e4f59c Mon Sep 17 00:00:00 2001
From: alonre24 <alonreshef24@gmail.com>
Date: Mon, 8 Mar 2021 23:31:11 +0200
Subject: [PATCH 26/32] Align allocation for onnx to 32-byte

---
 .circleci/config.yml       |  6 +++---
 src/backends/onnxruntime.c | 17 ++++++++++++-----
 2 files changed, 15 insertions(+), 8 deletions(-)

diff --git a/.circleci/config.yml b/.circleci/config.yml
index 05fee460f..eadf0a1db 100644
--- a/.circleci/config.yml
+++ b/.circleci/config.yml
@@ -162,9 +162,9 @@ jobs:
       - run:
           name: Submodule checkout
           command: git submodule update --init --recursive
-      #- restore_cache:
-      #    keys:
-      #    - build-dependencies-{{ checksum "get_deps.sh" }}
+      - restore_cache:
+          keys:
+          - build-dependencies-{{ checksum "get_deps.sh" }}
           # If no exact match is found will get dependencies from source
       - setup-automation
       - run:
diff --git a/src/backends/onnxruntime.c b/src/backends/onnxruntime.c
index be805fd45..1ee514640 100644
--- a/src/backends/onnxruntime.c
+++ b/src/backends/onnxruntime.c
@@ -33,20 +33,27 @@ const OrtMemoryInfo *AllocatorInfo(const OrtAllocator *allocator) {
 }
 
 void *AllocatorAlloc(OrtAllocator *ptr, size_t size) {
+
     (void)ptr;
-    void *p = RedisModule_Alloc(size);
-    size_t allocated_size = RedisModule_MallocSize(p);
+    int offset = 31 + sizeof(void *);
+    void *p1 = (void *)RedisModule_Alloc(size + offset);
+    size_t allocated_size = RedisModule_MallocSize(p1);
     atomic_fetch_add(&OnnxMemory, allocated_size);
     atomic_fetch_add(&OnnxMemoryAccessCounter, 1);
-    return p;
+    void **p2 = (void **)(((uintptr_t)(p1) + offset) & (~31));
+    p2[-1] = p1;
+    return p2;
 }
 
 void AllocatorFree(OrtAllocator *ptr, void *p) {
     (void)ptr;
-    size_t allocated_size = RedisModule_MallocSize(p);
+    if (p == NULL)
+        return;
+    void *p1 = ((void **)p)[-1];
+    size_t allocated_size = RedisModule_MallocSize(p1);
     atomic_fetch_sub(&OnnxMemory, allocated_size);
     atomic_fetch_add(&OnnxMemoryAccessCounter, 1);
-    return RedisModule_Free(p);
+    return RedisModule_Free(p1);
 }
 
 unsigned long long RAI_GetMemoryInfoORT() { return OnnxMemory; }

From 8a6766228817c744d179853faa0a22fb5e5a0dd4 Mon Sep 17 00:00:00 2001
From: alonre24 <alonreshef24@gmail.com>
Date: Tue, 9 Mar 2021 00:02:45 +0200
Subject: [PATCH 27/32] Align allocation for onnx to 64-byte

---
 src/backends/onnxruntime.c | 4 ++--
 1 file changed, 2 insertions(+), 2 deletions(-)

diff --git a/src/backends/onnxruntime.c b/src/backends/onnxruntime.c
index 1ee514640..c6fc4f12f 100644
--- a/src/backends/onnxruntime.c
+++ b/src/backends/onnxruntime.c
@@ -35,12 +35,12 @@ const OrtMemoryInfo *AllocatorInfo(const OrtAllocator *allocator) {
 void *AllocatorAlloc(OrtAllocator *ptr, size_t size) {
 
     (void)ptr;
-    int offset = 31 + sizeof(void *);
+    int offset = 63 + sizeof(void *);
     void *p1 = (void *)RedisModule_Alloc(size + offset);
     size_t allocated_size = RedisModule_MallocSize(p1);
     atomic_fetch_add(&OnnxMemory, allocated_size);
     atomic_fetch_add(&OnnxMemoryAccessCounter, 1);
-    void **p2 = (void **)(((uintptr_t)(p1) + offset) & (~31));
+    void **p2 = (void **)(((uintptr_t)(p1) + offset) & (~63));
     p2[-1] = p1;
     return p2;
 }

From d4fba336b8049f28866fc580139ec8a2d372fcf1 Mon Sep 17 00:00:00 2001
From: alonre24 <alonreshef24@gmail.com>
Date: Tue, 9 Mar 2021 12:31:59 +0200
Subject: [PATCH 28/32] Fix test for GPU+CPU

---
 opt/build/onnxruntime/Dockerfile.x64     | 2 +-
 opt/build/onnxruntime/Dockerfile.x64-gpu | 2 +-
 opt/build/onnxruntime/pack.sh            | 2 +-
 src/backends/onnxruntime.c               | 4 +++-
 tests/flow/tests_onnx.py                 | 2 +-
 5 files changed, 7 insertions(+), 5 deletions(-)

diff --git a/opt/build/onnxruntime/Dockerfile.x64 b/opt/build/onnxruntime/Dockerfile.x64
index c0ad2898e..439899269 100755
--- a/opt/build/onnxruntime/Dockerfile.x64
+++ b/opt/build/onnxruntime/Dockerfile.x64
@@ -24,7 +24,7 @@ RUN apt-get install -y locales && \
 
 WORKDIR /build
 
-ARG BUILDTYPE=Debug
+ARG BUILDTYPE=MinSizeRel
 
 ARG BUILDARGS="--config ${BUILDTYPE} --parallel"
 
diff --git a/opt/build/onnxruntime/Dockerfile.x64-gpu b/opt/build/onnxruntime/Dockerfile.x64-gpu
index 2f3953f36..e691f7f4a 100644
--- a/opt/build/onnxruntime/Dockerfile.x64-gpu
+++ b/opt/build/onnxruntime/Dockerfile.x64-gpu
@@ -40,4 +40,4 @@ RUN ./build.sh ${BUILDARGS} --update --build --use_cuda --cudnn_home /usr/local/
 
 ADD ./pack.sh /build/
 WORKDIR /build
-RUN /pack.sh ${ONNXRUNTIME_VER} ${ARCH}
+RUN ./pack.sh ${ONNXRUNTIME_VER} ${ARCH}
diff --git a/opt/build/onnxruntime/pack.sh b/opt/build/onnxruntime/pack.sh
index 8325f94ac..01b6bb5ed 100755
--- a/opt/build/onnxruntime/pack.sh
+++ b/opt/build/onnxruntime/pack.sh
@@ -5,7 +5,7 @@ VER="$1"
 PLATFORM="$2"
 
 mkdir -p pack/include pack/lib
-cp onnxruntime/build/Linux/Debug/libonnxruntime.so.${VER} pack/lib/
+cp onnxruntime/build/Linux/MinSizeRel/libonnxruntime.so.${VER} pack/lib/
 cp onnxruntime/docs/C_API_Guidelines.md pack/
 cp onnxruntime/LICENSE pack/
 cp onnxruntime/README.md pack/
diff --git a/src/backends/onnxruntime.c b/src/backends/onnxruntime.c
index c6fc4f12f..33c1fedbe 100644
--- a/src/backends/onnxruntime.c
+++ b/src/backends/onnxruntime.c
@@ -32,6 +32,7 @@ const OrtMemoryInfo *AllocatorInfo(const OrtAllocator *allocator) {
     return mem_info;
 }
 
+// Allocate address with 64-byte alignment to cope with onnx optimizations.
 void *AllocatorAlloc(OrtAllocator *ptr, size_t size) {
 
     (void)ptr;
@@ -47,8 +48,9 @@ void *AllocatorAlloc(OrtAllocator *ptr, size_t size) {
 
 void AllocatorFree(OrtAllocator *ptr, void *p) {
     (void)ptr;
-    if (p == NULL)
+    if (p == NULL) {
         return;
+    }
     void *p1 = ((void **)p)[-1];
     size_t allocated_size = RedisModule_MallocSize(p1);
     atomic_fetch_sub(&OnnxMemory, allocated_size);
diff --git a/tests/flow/tests_onnx.py b/tests/flow/tests_onnx.py
index 7efb632ab..545284b84 100644
--- a/tests/flow/tests_onnx.py
+++ b/tests/flow/tests_onnx.py
@@ -538,5 +538,5 @@ def test_onnx_use_custom_allocator_with_GPU(env):
     ai_memory_config = {k.split(":")[0]: k.split(":")[1]
                         for k in con.execute_command("INFO MODULES").decode().split("#")[4].split()[1:]}
     env.assertTrue(int(ai_memory_config["ai_onnxruntime_memory"]) > 100)
-    env.assertTrue(int(ai_memory_config["ai_onnxruntime_memory"]) < 400)
+    env.assertTrue(int(ai_memory_config["ai_onnxruntime_memory"]) < 705)
     env.assertEqual(int(ai_memory_config["ai_onnxruntime_memory_access_num"]), 5)

From 688ac741be45b24b3df2a5fc6e35a3b0fd6d2747 Mon Sep 17 00:00:00 2001
From: alonre24 <alonreshef24@gmail.com>
Date: Wed, 10 Mar 2021 22:54:06 +0200
Subject: [PATCH 29/32] Elaborate documentation.

---
 src/backends/onnxruntime.c | 21 ++++++++++++++++++++-
 tests/flow/tests_onnx.py   | 30 +++++++++++++++++++++---------
 2 files changed, 41 insertions(+), 10 deletions(-)

diff --git a/src/backends/onnxruntime.c b/src/backends/onnxruntime.c
index 33c1fedbe..9c35af221 100644
--- a/src/backends/onnxruntime.c
+++ b/src/backends/onnxruntime.c
@@ -9,6 +9,7 @@
 #include "onnxruntime_c_api.h"
 
 // Use as a wrapper for ORT api call. If ORT api hasn't returned null, it has failed.
+// A label "error" must exist in every function that uses this macro.
 #define ONNX_VALIDATE_STATUS(x)                                                                    \
     if ((status = (x)) != NULL)                                                                    \
         goto error;
@@ -36,12 +37,19 @@ const OrtMemoryInfo *AllocatorInfo(const OrtAllocator *allocator) {
 void *AllocatorAlloc(OrtAllocator *ptr, size_t size) {
 
     (void)ptr;
+    // Allocate an additional 63 bytes to ensure that we can return an address which is
+    // 64-byte aligned, and an additional space in the size of a pointer to store
+    // the address that RedisModule_Alloc returns.
     int offset = 63 + sizeof(void *);
     void *p1 = (void *)RedisModule_Alloc(size + offset);
     size_t allocated_size = RedisModule_MallocSize(p1);
+    // Update the total number of bytes that onnx is using and the number of accesses
+    // that onnx made to the allocator.
     atomic_fetch_add(&OnnxMemory, allocated_size);
     atomic_fetch_add(&OnnxMemoryAccessCounter, 1);
-    void **p2 = (void **)(((uintptr_t)(p1) + offset) & (~63));
+    // This operation guarantees that p2 is the closest 64-aligned address to (p1+size_t).
+    void **p2 = (void **)(((size_t)(p1) + offset) & (~63));
+    // This stores the address p1 right before p2 (so we can retrieve it when we free).
     p2[-1] = p1;
     return p2;
 }
@@ -51,8 +59,12 @@ void AllocatorFree(OrtAllocator *ptr, void *p) {
     if (p == NULL) {
         return;
     }
+    // Retrieve the address that we originally received from RedisModule_Alloc
+    // (this is the address that we need to sent to RedisModule_Free).
     void *p1 = ((void **)p)[-1];
     size_t allocated_size = RedisModule_MallocSize(p1);
+    // Update the total number of bytes that onnx is using and the number of accesses
+    // that onnx made to the allocator.
     atomic_fetch_sub(&OnnxMemory, allocated_size);
     atomic_fetch_add(&OnnxMemoryAccessCounter, 1);
     return RedisModule_Free(p1);
@@ -321,6 +333,10 @@ RAI_Model *RAI_ModelCreateORT(RAI_Backend backend, const char *devicestr, RAI_Mo
     OrtSession *session = NULL;
     OrtStatus *status = NULL;
 
+    // In the first time we set a model for onnx, we create an environment and register
+    // an allocator to it that uses Redis allocator. This allocator is going to be used for
+    // allocating buffers when creating and running models that run on CPU, and for allocations of
+    // models inputs and outputs names (for both models that run on CPU and GPU)
     if (env == NULL) {
         ONNX_VALIDATE_STATUS(ort->CreateEnv(ORT_LOGGING_LEVEL_WARNING, "test", &env))
         ONNX_VALIDATE_STATUS(ort->CreateCustomDeviceAllocator(
@@ -343,6 +359,9 @@ RAI_Model *RAI_ModelCreateORT(RAI_Backend backend, const char *devicestr, RAI_Mo
         ort->SetIntraOpNumThreads(session_options, (int)opts.backends_intra_op_parallelism))
     ONNX_VALIDATE_STATUS(
         ort->SetInterOpNumThreads(session_options, (int)opts.backends_inter_op_parallelism))
+
+    // If the model is set for GPU, this will set CUDA provider for the session,
+    // so that onnx will use its own allocator for CUDA (not Redis allocator)
     if (!setDeviceId(devicestr, session_options, error)) {
         ort->ReleaseSessionOptions(session_options);
         return NULL;
diff --git a/tests/flow/tests_onnx.py b/tests/flow/tests_onnx.py
index 545284b84..0a36e5012 100644
--- a/tests/flow/tests_onnx.py
+++ b/tests/flow/tests_onnx.py
@@ -467,15 +467,21 @@ def test_onnx_use_custom_allocator(env):
                             for k in con.execute_command("INFO MODULES").decode().split("#")[4].split()[1:]}
     env.assertEqual(int(ai_memory_config["ai_onnxruntime_memory"]), 0)
 
-    # Expect using the allocator during model set for allocating the model, its input name and output name.
+    # Expect using the allocator during model set for allocating the model, its input name and output name:
+    # overall 3 allocations. The model raw size is 130B ,and the names are 2B each. In practice we allocate
+    # more than 134B as Redis allocator will use additional memory for its internal management and for the
+    # 64-Byte alignment. When the test runs with valgrind, redis will use malloc for the allocations
+    # (hence will not use additional memory).
     ret = con.execute_command('AI.MODELSET', 'm{1}', 'ONNX', 'CPU', 'BLOB', model_pb)
     env.assertEqual(ret, b'OK')
     ai_memory_config = {k.split(":")[0]: k.split(":")[1]
                         for k in con.execute_command("INFO MODULES").decode().split("#")[4].split()[1:]}
-    env.assertTrue(int(ai_memory_config["ai_onnxruntime_memory"]) > 100)
+
+    # Expect using at least 130+63+(size of an address) + 2*(2+63+(size of an address)) bytes.
+    env.assertTrue(int(ai_memory_config["ai_onnxruntime_memory"]) > 334)
     env.assertEqual(int(ai_memory_config["ai_onnxruntime_memory_access_num"]), 3)
 
-    # Expect using the allocator free function when releasing the model.
+    # Expect using the allocator free function when releasing the model and input and output names.
     con.execute_command('AI.MODELDEL', 'm{1}')
     env.assertFalse(con.execute_command('EXISTS', 'm{1}'))
     ai_memory_config = {k.split(":")[0]: k.split(":")[1]
@@ -483,7 +489,7 @@ def test_onnx_use_custom_allocator(env):
     env.assertEqual(int(ai_memory_config["ai_onnxruntime_memory"]), 0)
     env.assertEqual(int(ai_memory_config["ai_onnxruntime_memory_access_num"]), 6)
 
-    # test allocator in model run op
+    # test the use of Redis allocator in model run op.
     model_filename = os.path.join(test_data_path, 'mnist.onnx')
     sample_filename = os.path.join(test_data_path, 'one.raw')
 
@@ -496,7 +502,8 @@ def test_onnx_use_custom_allocator(env):
     env.assertEqual(ret, b'OK')
     con.execute_command('AI.TENSORSET', 'a{1}', 'FLOAT', 1, 1, 28, 28, 'BLOB', sample_raw)
 
-    # Expect 16 allocator's access from onnx during the run.
+    # Expect 16 allocator's access from onnx during the run (in addition to the allocations that were made while
+    # creating the model).
     ai_memory_config = {k.split(":")[0]: k.split(":")[1]
                         for k in con.execute_command("INFO MODULES").decode().split("#")[4].split()[1:]}
     allocator_access_num_before = ai_memory_config["ai_onnxruntime_memory_access_num"]
@@ -528,15 +535,20 @@ def test_onnx_use_custom_allocator_with_GPU(env):
                         for k in con.execute_command("INFO MODULES").decode().split("#")[4].split()[1:]}
     env.assertEqual(int(ai_memory_config["ai_onnxruntime_memory"]), 0)
 
-    # Create the same model, once for CPU and once for GPU.
-    # Expect using the allocator during model set for allocating the model, its input name and output name in CPU,
-    # but for GPU, expcet using the allocator only for allocating input and output names.
+    # Expect using the allocator during model set for allocating the model, its input name and output name:
+    # overall 3 allocations. The model raw size is 130B ,and the names are 2B each. In practice we allocate
+    # more than 134B as Redis allocator will use additional memory for its internal management and for the
+    # 64-Byte alignment. When the test runs with valgrind, redis will use malloc for the allocations.
     ret = con.execute_command('AI.MODELSET', 'm_gpu{1}', 'ONNX', DEVICE, 'BLOB', model_pb)
     env.assertEqual(ret, b'OK')
+
+    # but for GPU, expect using the allocator only for allocating input and output names (not the model itself).
     ret = con.execute_command('AI.MODELSET', 'm_cpu{1}', 'ONNX', 'CPU', 'BLOB', model_pb)
     env.assertEqual(ret, b'OK')
     ai_memory_config = {k.split(":")[0]: k.split(":")[1]
                         for k in con.execute_command("INFO MODULES").decode().split("#")[4].split()[1:]}
-    env.assertTrue(int(ai_memory_config["ai_onnxruntime_memory"]) > 100)
+
+    # Expect using at least 130+63+(size of an address) + 4*(2+63+(size of an address)) bytes.
+    env.assertTrue(int(ai_memory_config["ai_onnxruntime_memory"]) > 472)
     env.assertTrue(int(ai_memory_config["ai_onnxruntime_memory"]) < 705)
     env.assertEqual(int(ai_memory_config["ai_onnxruntime_memory_access_num"]), 5)

From 9bf31c2136e38b6fb69194de46c5fb822d306c3c Mon Sep 17 00:00:00 2001
From: alonre24 <alonreshef24@gmail.com>
Date: Thu, 11 Mar 2021 12:59:27 +0200
Subject: [PATCH 30/32] Extend GPU test

---
 src/backends/onnxruntime.c | 20 ++++++++++----------
 tests/flow/tests_onnx.py   | 16 ++++++++++++++++
 2 files changed, 26 insertions(+), 10 deletions(-)

diff --git a/src/backends/onnxruntime.c b/src/backends/onnxruntime.c
index 9c35af221..2c68d5d5d 100644
--- a/src/backends/onnxruntime.c
+++ b/src/backends/onnxruntime.c
@@ -41,33 +41,33 @@ void *AllocatorAlloc(OrtAllocator *ptr, size_t size) {
     // 64-byte aligned, and an additional space in the size of a pointer to store
     // the address that RedisModule_Alloc returns.
     int offset = 63 + sizeof(void *);
-    void *p1 = (void *)RedisModule_Alloc(size + offset);
-    size_t allocated_size = RedisModule_MallocSize(p1);
+    void *allocated_address = (void *)RedisModule_Alloc(size + offset);
+    size_t allocated_size = RedisModule_MallocSize(allocated_address);
     // Update the total number of bytes that onnx is using and the number of accesses
     // that onnx made to the allocator.
     atomic_fetch_add(&OnnxMemory, allocated_size);
     atomic_fetch_add(&OnnxMemoryAccessCounter, 1);
     // This operation guarantees that p2 is the closest 64-aligned address to (p1+size_t).
-    void **p2 = (void **)(((size_t)(p1) + offset) & (~63));
+    void **aligned_address = (void **)(((size_t)(allocated_address) + offset) & (~63));
     // This stores the address p1 right before p2 (so we can retrieve it when we free).
-    p2[-1] = p1;
-    return p2;
+    aligned_address[-1] = allocated_address;
+    return aligned_address;
 }
 
-void AllocatorFree(OrtAllocator *ptr, void *p) {
+void AllocatorFree(OrtAllocator *ptr, void *aligned_address) {
     (void)ptr;
-    if (p == NULL) {
+    if (aligned_address == NULL) {
         return;
     }
     // Retrieve the address that we originally received from RedisModule_Alloc
     // (this is the address that we need to sent to RedisModule_Free).
-    void *p1 = ((void **)p)[-1];
-    size_t allocated_size = RedisModule_MallocSize(p1);
+    void *allocated_address = ((void **)aligned_address)[-1];
+    size_t allocated_size = RedisModule_MallocSize(allocated_address);
     // Update the total number of bytes that onnx is using and the number of accesses
     // that onnx made to the allocator.
     atomic_fetch_sub(&OnnxMemory, allocated_size);
     atomic_fetch_add(&OnnxMemoryAccessCounter, 1);
-    return RedisModule_Free(p1);
+    return RedisModule_Free(allocated_address);
 }
 
 unsigned long long RAI_GetMemoryInfoORT() { return OnnxMemory; }
diff --git a/tests/flow/tests_onnx.py b/tests/flow/tests_onnx.py
index 0a36e5012..905db476c 100644
--- a/tests/flow/tests_onnx.py
+++ b/tests/flow/tests_onnx.py
@@ -552,3 +552,19 @@ def test_onnx_use_custom_allocator_with_GPU(env):
     env.assertTrue(int(ai_memory_config["ai_onnxruntime_memory"]) > 472)
     env.assertTrue(int(ai_memory_config["ai_onnxruntime_memory"]) < 705)
     env.assertEqual(int(ai_memory_config["ai_onnxruntime_memory_access_num"]), 5)
+
+    # Make sure that allocator is not used for running and freeing the GPU model.
+    sample_filename = os.path.join(test_data_path, 'one.raw')
+    with open(sample_filename, 'rb') as f:
+        sample_raw = f.read()
+    con.execute_command('AI.TENSORSET', 'a{1}', 'FLOAT', 1, 1, 28, 28, 'BLOB', sample_raw)
+    con.execute_command('AI.MODELRUN', 'm_gpu{1}', 'INPUTS', 'a{1}', 'OUTPUTS', 'b{1}')
+    values = con.execute_command('AI.TENSORGET', 'b{1}', 'VALUES')
+    argmax = max(range(len(values)), key=lambda i: values[i])
+    env.assertEqual(argmax, 1)
+    con.execute_command('AI.MODELDEL', 'm{1}')
+    env.assertFalse(con.execute_command('EXISTS', 'm_gpu{1}'))
+    ai_memory_config = {k.split(":")[0]: k.split(":")[1]
+                        for k in con.execute_command("INFO MODULES").decode().split("#")[4].split()[1:]}
+    env.assertTrue(int(ai_memory_config["ai_onnxruntime_memory"]) < 705)
+    env.assertEqual(int(ai_memory_config["ai_onnxruntime_memory_access_num"]), 5)

From dd19238671bdefe4b7d3dfc37dc1b7d9829064d7 Mon Sep 17 00:00:00 2001
From: alonre24 <alonreshef24@gmail.com>
Date: Thu, 11 Mar 2021 14:20:45 +0200
Subject: [PATCH 31/32] Fix GPU test

---
 tests/flow/tests_onnx.py | 10 +++-------
 1 file changed, 3 insertions(+), 7 deletions(-)

diff --git a/tests/flow/tests_onnx.py b/tests/flow/tests_onnx.py
index 905db476c..d8052c459 100644
--- a/tests/flow/tests_onnx.py
+++ b/tests/flow/tests_onnx.py
@@ -554,15 +554,11 @@ def test_onnx_use_custom_allocator_with_GPU(env):
     env.assertEqual(int(ai_memory_config["ai_onnxruntime_memory_access_num"]), 5)
 
     # Make sure that allocator is not used for running and freeing the GPU model.
-    sample_filename = os.path.join(test_data_path, 'one.raw')
-    with open(sample_filename, 'rb') as f:
-        sample_raw = f.read()
-    con.execute_command('AI.TENSORSET', 'a{1}', 'FLOAT', 1, 1, 28, 28, 'BLOB', sample_raw)
+    con.execute_command('AI.TENSORSET', 'a{1}', 'FLOAT', 3, 2, 'VALUES', 1.0, 2.0, 3.0, 4.0, 5.0, 6.0)
     con.execute_command('AI.MODELRUN', 'm_gpu{1}', 'INPUTS', 'a{1}', 'OUTPUTS', 'b{1}')
     values = con.execute_command('AI.TENSORGET', 'b{1}', 'VALUES')
-    argmax = max(range(len(values)), key=lambda i: values[i])
-    env.assertEqual(argmax, 1)
-    con.execute_command('AI.MODELDEL', 'm{1}')
+    env.assertEqual(values, [b'1', b'4', b'9', b'16', b'25', b'36'])
+    con.execute_command('AI.MODELDEL', 'm_gpu{1}')
     env.assertFalse(con.execute_command('EXISTS', 'm_gpu{1}'))
     ai_memory_config = {k.split(":")[0]: k.split(":")[1]
                         for k in con.execute_command("INFO MODULES").decode().split("#")[4].split()[1:]}

From f27405baf7b6e8f87557c2e10010946b7763095a Mon Sep 17 00:00:00 2001
From: alonre24 <alonreshef24@gmail.com>
Date: Thu, 11 Mar 2021 16:46:44 +0200
Subject: [PATCH 32/32] Remove onnx from valgrind suppressions, and fix leaks
 that were found. Extend tests to make a more precise check of the allocator
 usage.

---
 opt/redis_valgrind.sup     |  7 -------
 src/backends/onnxruntime.c | 19 +++++++++++++----
 tests/flow/tests_onnx.py   | 42 +++++++++++++++++++++++++++++---------
 3 files changed, 47 insertions(+), 21 deletions(-)

diff --git a/opt/redis_valgrind.sup b/opt/redis_valgrind.sup
index 5431acaf1..8b00a448c 100644
--- a/opt/redis_valgrind.sup
+++ b/opt/redis_valgrind.sup
@@ -12,13 +12,6 @@
    obj:*/libtensorflow_framework.so.*
 }
 
-{
-   ignore_unversioned_libs
-   Memcheck:Leak
-   ...
-   obj:*/libonnxruntime.so.*
-}
-
 {
    ignore_unversioned_libs
    Memcheck:Leak
diff --git a/src/backends/onnxruntime.c b/src/backends/onnxruntime.c
index 2c68d5d5d..21494d9b0 100644
--- a/src/backends/onnxruntime.c
+++ b/src/backends/onnxruntime.c
@@ -17,8 +17,9 @@
 OrtEnv *env = NULL;
 
 // For model that run on GPU, onnx will not use the custom allocator (redis allocator), but
-// the onnx allocator for GPU. But for the auxilery allocations of the input and output names,
-// we will use the custom global allocator for models that run on GPU as well
+// the onnx allocator for GPU. But for the auxiliary allocations of the input and output names,
+// we will use the custom global allocator for models that run on GPU as well.
+OrtMemoryInfo *mem_info = NULL;
 OrtAllocator *global_allocator = NULL;
 unsigned long long OnnxMemory = 0;
 unsigned long long OnnxMemoryAccessCounter = 0;
@@ -26,7 +27,9 @@ unsigned long long OnnxMemoryAccessCounter = 0;
 const OrtMemoryInfo *AllocatorInfo(const OrtAllocator *allocator) {
     (void)allocator;
     const OrtApi *ort = OrtGetApiBase()->GetApi(1);
-    OrtMemoryInfo *mem_info;
+    if (mem_info != NULL) {
+        return mem_info;
+    }
     if (ort->CreateCpuMemoryInfo(OrtDeviceAllocator, OrtMemTypeDefault, &mem_info) != NULL) {
         return NULL;
     }
@@ -369,6 +372,7 @@ RAI_Model *RAI_ModelCreateORT(RAI_Backend backend, const char *devicestr, RAI_Mo
 
     ONNX_VALIDATE_STATUS(
         ort->CreateSessionFromArray(env, modeldef, modellen, session_options, &session))
+    ort->ReleaseSessionOptions(session_options);
 
     size_t n_input_nodes;
     ONNX_VALIDATE_STATUS(ort->SessionGetInputCount(session, &n_input_nodes))
@@ -550,7 +554,14 @@ int RAI_ModelRunORT(RAI_ModelRunCtx **mctxs, RAI_Error *error) {
         OrtRunOptions *run_options = NULL;
         ONNX_VALIDATE_STATUS(ort->Run(session, run_options, input_names,
                                       (const OrtValue *const *)inputs, n_input_nodes, output_names,
-                                      n_output_nodes, outputs))
+                                      n_output_nodes, outputs));
+
+        for (uint32_t i = 0; i < ninputs; i++) {
+            status = ort->AllocatorFree(global_allocator, (void *)input_names[i]);
+        }
+        for (uint32_t i = 0; i < noutputs; i++) {
+            status = ort->AllocatorFree(global_allocator, (void *)output_names[i]);
+        }
 
         for (size_t i = 0; i < n_output_nodes; i++) {
             if (nbatches > 1) {
diff --git a/tests/flow/tests_onnx.py b/tests/flow/tests_onnx.py
index d8052c459..32b5b1fc7 100644
--- a/tests/flow/tests_onnx.py
+++ b/tests/flow/tests_onnx.py
@@ -478,16 +478,28 @@ def test_onnx_use_custom_allocator(env):
                         for k in con.execute_command("INFO MODULES").decode().split("#")[4].split()[1:]}
 
     # Expect using at least 130+63+(size of an address) + 2*(2+63+(size of an address)) bytes.
-    env.assertTrue(int(ai_memory_config["ai_onnxruntime_memory"]) > 334)
+    model_allocation_bytes_used = int(ai_memory_config["ai_onnxruntime_memory"])
+    env.assertTrue(model_allocation_bytes_used > 334)
     env.assertEqual(int(ai_memory_config["ai_onnxruntime_memory_access_num"]), 3)
+    con.execute_command('AI.TENSORSET', 'a_mul{1}', 'FLOAT', 3, 2, 'VALUES', 1.0, 2.0, 3.0, 4.0, 5.0, 6.0)
 
-    # Expect using the allocator free function when releasing the model and input and output names.
+    # Running the model should access the allocator 6 times: allocating+freeing input+output names,
+    # and allocating+freeing the output as OrtValue.
+    con.execute_command('AI.MODELRUN', 'm{1}', 'INPUTS', 'a_mul{1}', 'OUTPUTS', 'b{1}')
+    values = con.execute_command('AI.TENSORGET', 'b{1}', 'VALUES')
+    env.assertEqual(values, [b'1', b'4', b'9', b'16', b'25', b'36'])
+    ai_memory_config = {k.split(":")[0]: k.split(":")[1]
+                        for k in con.execute_command("INFO MODULES").decode().split("#")[4].split()[1:]}
+    env.assertEqual(int(ai_memory_config["ai_onnxruntime_memory_access_num"]), 9)
+    env.assertEqual(int(ai_memory_config["ai_onnxruntime_memory"]), model_allocation_bytes_used)
+
+    # Expect using the allocator free function 3 times: when releasing the model, input name and output name.
     con.execute_command('AI.MODELDEL', 'm{1}')
     env.assertFalse(con.execute_command('EXISTS', 'm{1}'))
     ai_memory_config = {k.split(":")[0]: k.split(":")[1]
                         for k in con.execute_command("INFO MODULES").decode().split("#")[4].split()[1:]}
     env.assertEqual(int(ai_memory_config["ai_onnxruntime_memory"]), 0)
-    env.assertEqual(int(ai_memory_config["ai_onnxruntime_memory_access_num"]), 6)
+    env.assertEqual(int(ai_memory_config["ai_onnxruntime_memory_access_num"]), 12)
 
     # test the use of Redis allocator in model run op.
     model_filename = os.path.join(test_data_path, 'mnist.onnx')
@@ -502,7 +514,7 @@ def test_onnx_use_custom_allocator(env):
     env.assertEqual(ret, b'OK')
     con.execute_command('AI.TENSORSET', 'a{1}', 'FLOAT', 1, 1, 28, 28, 'BLOB', sample_raw)
 
-    # Expect 16 allocator's access from onnx during the run (in addition to the allocations that were made while
+    # Expect 18 allocator's access from onnx during the run (in addition to the allocations that were made while
     # creating the model).
     ai_memory_config = {k.split(":")[0]: k.split(":")[1]
                         for k in con.execute_command("INFO MODULES").decode().split("#")[4].split()[1:]}
@@ -511,7 +523,7 @@ def test_onnx_use_custom_allocator(env):
     ai_memory_config = {k.split(":")[0]: k.split(":")[1]
                         for k in con.execute_command("INFO MODULES").decode().split("#")[4].split()[1:]}
     allocator_access_num_after = ai_memory_config["ai_onnxruntime_memory_access_num"]
-    env.assertEqual(int(allocator_access_num_after) - int(allocator_access_num_before), 16)
+    env.assertEqual(int(allocator_access_num_after) - int(allocator_access_num_before), 18)
 
     values = con.execute_command('AI.TENSORGET', 'b{1}', 'VALUES')
     argmax = max(range(len(values)), key=lambda i: values[i])
@@ -549,18 +561,28 @@ def test_onnx_use_custom_allocator_with_GPU(env):
                         for k in con.execute_command("INFO MODULES").decode().split("#")[4].split()[1:]}
 
     # Expect using at least 130+63+(size of an address) + 4*(2+63+(size of an address)) bytes.
-    env.assertTrue(int(ai_memory_config["ai_onnxruntime_memory"]) > 472)
-    env.assertTrue(int(ai_memory_config["ai_onnxruntime_memory"]) < 705)
+    model_allocation_bytes_used = int(ai_memory_config["ai_onnxruntime_memory"])
+    env.assertTrue(model_allocation_bytes_used > 472)
+    env.assertTrue(model_allocation_bytes_used < 705)
     env.assertEqual(int(ai_memory_config["ai_onnxruntime_memory_access_num"]), 5)
 
-    # Make sure that allocator is not used for running and freeing the GPU model.
+    # Make sure that allocator is not used for running and freeing the GPU model, except for
+    # the input and output names allocations (and deallocations).
     con.execute_command('AI.TENSORSET', 'a{1}', 'FLOAT', 3, 2, 'VALUES', 1.0, 2.0, 3.0, 4.0, 5.0, 6.0)
     con.execute_command('AI.MODELRUN', 'm_gpu{1}', 'INPUTS', 'a{1}', 'OUTPUTS', 'b{1}')
     values = con.execute_command('AI.TENSORGET', 'b{1}', 'VALUES')
     env.assertEqual(values, [b'1', b'4', b'9', b'16', b'25', b'36'])
+    # Expect that memory usage didn't change, and for another 4 accesses to the allocator (input and output names
+    # allocation and free)
+    ai_memory_config = {k.split(":")[0]: k.split(":")[1]
+                        for k in con.execute_command("INFO MODULES").decode().split("#")[4].split()[1:]}
+    env.assertEqual(int(ai_memory_config["ai_onnxruntime_memory"]), model_allocation_bytes_used)
+    env.assertEqual(int(ai_memory_config["ai_onnxruntime_memory_access_num"]), 9)
+
+    # Expect only 2 more accesses in delete - for deallocating input and output names
     con.execute_command('AI.MODELDEL', 'm_gpu{1}')
     env.assertFalse(con.execute_command('EXISTS', 'm_gpu{1}'))
     ai_memory_config = {k.split(":")[0]: k.split(":")[1]
                         for k in con.execute_command("INFO MODULES").decode().split("#")[4].split()[1:]}
-    env.assertTrue(int(ai_memory_config["ai_onnxruntime_memory"]) < 705)
-    env.assertEqual(int(ai_memory_config["ai_onnxruntime_memory_access_num"]), 5)
+    env.assertEqual(int(ai_memory_config["ai_onnxruntime_memory_access_num"]), 11)
+