[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
Sean Silva 2020-11-23 16:30:49 -08:00
parent f13994fdf7
commit 0b7c443256
2 changed files with 2 additions and 0 deletions

View File

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

View File

@ -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);