Skip to content

Commit 049417e

Browse files
committed
Align allocation for onnx to 32-byte
1 parent 84b7e21 commit 049417e

File tree

2 files changed

+15
-8
lines changed

2 files changed

+15
-8
lines changed

.circleci/config.yml

Lines changed: 3 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -162,9 +162,9 @@ jobs:
162162
- run:
163163
name: Submodule checkout
164164
command: git submodule update --init --recursive
165-
#- restore_cache:
166-
# keys:
167-
# - build-dependencies-{{ checksum "get_deps.sh" }}
165+
- restore_cache:
166+
keys:
167+
- build-dependencies-{{ checksum "get_deps.sh" }}
168168
# If no exact match is found will get dependencies from source
169169
- setup-automation
170170
- run:

src/backends/onnxruntime.c

Lines changed: 12 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -33,20 +33,27 @@ const OrtMemoryInfo *AllocatorInfo(const OrtAllocator *allocator) {
3333
}
3434

3535
void *AllocatorAlloc(OrtAllocator *ptr, size_t size) {
36+
3637
(void)ptr;
37-
void *p = RedisModule_Alloc(size);
38-
size_t allocated_size = RedisModule_MallocSize(p);
38+
int offset = 31 + sizeof(void *);
39+
void *p1 = (void *)RedisModule_Alloc(size + offset);
40+
size_t allocated_size = RedisModule_MallocSize(p1);
3941
atomic_fetch_add(&OnnxMemory, allocated_size);
4042
atomic_fetch_add(&OnnxMemoryAccessCounter, 1);
41-
return p;
43+
void **p2 = (void **)(((uintptr_t)(p1) + offset) & (~31));
44+
p2[-1] = p1;
45+
return p2;
4246
}
4347

4448
void AllocatorFree(OrtAllocator *ptr, void *p) {
4549
(void)ptr;
46-
size_t allocated_size = RedisModule_MallocSize(p);
50+
if (p == NULL)
51+
return;
52+
void *p1 = ((void **)p)[-1];
53+
size_t allocated_size = RedisModule_MallocSize(p1);
4754
atomic_fetch_sub(&OnnxMemory, allocated_size);
4855
atomic_fetch_add(&OnnxMemoryAccessCounter, 1);
49-
return RedisModule_Free(p);
56+
return RedisModule_Free(p1);
5057
}
5158

5259
unsigned long long RAI_GetMemoryInfoORT() { return OnnxMemory; }

0 commit comments

Comments
 (0)