mirror of https://github.com/llvm/torch-mlir
[RefBackend] Properly initialize refbackrt::Tensor refcount.
Although `refCount` is initialized as `std::atomic<int> refCount{0};` in the definition of Tensor, our tail-allocating malloc would ignore it, resulting in bogus values that led to leaks. Caught with LeakSanitizer, but I added an assertion that the refcount is non-negative to begin with, which should catch this bug in the future fairly consistently (assuming the garbage refcount is negative half the time).pull/130/head
parent
f13994fdf7
commit
0b7c443256
|
@ -34,6 +34,7 @@ public:
|
||||||
// Creates a Ref and increments the refcount by 1.
|
// Creates a Ref and increments the refcount by 1.
|
||||||
// rawPtr must be allocated with std::malloc.
|
// rawPtr must be allocated with std::malloc.
|
||||||
Ref(T *rawPtr) {
|
Ref(T *rawPtr) {
|
||||||
|
assert(rawPtr->refCount >= 0 && "expected non-negative refcount to start!");
|
||||||
ptr = rawPtr;
|
ptr = rawPtr;
|
||||||
ptr->refCount += 1;
|
ptr->refCount += 1;
|
||||||
}
|
}
|
||||||
|
|
|
@ -46,6 +46,7 @@ Tensor *Tensor::createRaw(ArrayRef<std::int32_t> extents, ElementType type,
|
||||||
auto *tensor = static_cast<Tensor *>(
|
auto *tensor = static_cast<Tensor *>(
|
||||||
std::malloc(sizeof(Tensor) + extents.size() * sizeof(std::int32_t)));
|
std::malloc(sizeof(Tensor) + extents.size() * sizeof(std::int32_t)));
|
||||||
|
|
||||||
|
tensor->refCount.store(0);
|
||||||
tensor->elementType = type;
|
tensor->elementType = type;
|
||||||
tensor->rank = extents.size();
|
tensor->rank = extents.size();
|
||||||
auto byteSize = getElementTypeByteSize(type) * totalElements(extents);
|
auto byteSize = getElementTypeByteSize(type) * totalElements(extents);
|
||||||
|
|
Loading…
Reference in New Issue