Skip to content

Commit e141da8

Browse files
authored
[mlir][ExecutionEngine] fix default free function in OwningMemRef. (#153133)
`basePtr` should be freed instead of `data` because it is the one which is storing the output of `malloc`. In `allocAligned()`, the `data` is malloced and then assigned to `basePtr`.
1 parent 1f2fb8e commit e141da8

File tree

2 files changed

+8
-2
lines changed

2 files changed

+8
-2
lines changed

mlir/include/mlir/ExecutionEngine/MemRefUtils.h

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -151,7 +151,7 @@ class OwningMemRef {
151151
AllocFunType allocFun = &::malloc,
152152
std::function<void(StridedMemRefType<T, Rank>)> freeFun =
153153
[](StridedMemRefType<T, Rank> descriptor) {
154-
::free(descriptor.data);
154+
::free(descriptor.basePtr);
155155
})
156156
: freeFunc(freeFun) {
157157
if (shapeAlloc.empty())

mlir/unittests/ExecutionEngine/Invoke.cpp

Lines changed: 7 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -205,7 +205,13 @@ TEST(NativeMemRefJit, SKIP_WITHOUT_JIT(BasicMemref)) {
205205
};
206206
int64_t shape[] = {k, m};
207207
int64_t shapeAlloc[] = {k + 1, m + 1};
208-
OwningMemRef<float, 2> a(shape, shapeAlloc, init);
208+
// Use a large alignment to stress the case where the memref data/basePtr are
209+
// disjoint.
210+
int alignment = 8192;
211+
OwningMemRef<float, 2> a(shape, shapeAlloc, init, alignment);
212+
ASSERT_EQ(
213+
(void *)(((uintptr_t)a->basePtr + alignment - 1) & ~(alignment - 1)),
214+
a->data);
209215
ASSERT_EQ(a->sizes[0], k);
210216
ASSERT_EQ(a->sizes[1], m);
211217
ASSERT_EQ(a->strides[0], m + 1);

0 commit comments

Comments
 (0)