Extend GPU test

Author: alonre24

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

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)