Skip to content

src: fix memory leaks using TACO runtime API from reference cycle #520

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
wants to merge 1 commit into
base: master
Choose a base branch
from
Open

src: fix memory leaks using TACO runtime API from reference cycle #520

Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
23 changes: 17 additions & 6 deletions src/tensor.cpp
View file
Original file line number Diff line number Diff line change
@@ -476,10 +476,10 @@ static inline map<TensorVar, TensorBase> getTensors(const IndexExpr& expr);
struct AccessTensorNode : public AccessNode {
AccessTensorNode(TensorBase tensor, const std::vector<IndexVar>& indices)
: AccessNode(tensor.getTensorVar(), indices, {}, false),
tensor(tensor) {}
tensorPtr(tensor.content) {}

AccessTensorNode(TensorBase tensor, const std::vector<std::shared_ptr<IndexVarInterface>>& indices)
: AccessNode(tensor.getTensorVar()), tensor(tensor) {
: AccessNode(tensor.getTensorVar()), tensorPtr(tensor.content) {
// Create the vector of IndexVar to assign to this->indexVars.
std::vector<IndexVar> ivars(indices.size());
for (size_t i = 0; i < indices.size(); i++) {
@@ -517,10 +517,21 @@ struct AccessTensorNode : public AccessNode {
this->indexVars = std::move(ivars);
}

TensorBase tensor;
// We hold a weak_ptr to the accessed TensorBase to avoid creating a reference
// cycle between the accessed TensorBase and this AccessTensorNode, since the
// TensorBase will store the AccessTensorNode (as part of an IndexExpr) as a
// field on itself. Not using a weak pointer results in leaking TensorBases.
std::weak_ptr<TensorBase::Content> tensorPtr;
TensorBase getTensor() const {
TensorBase tensor;
tensor.content = tensorPtr.lock();
return tensor;
}

virtual void setAssignment(const Assignment& assignment) {
tensor.syncDependentTensors();
auto tensor = this->getTensor();

tensor.syncDependentTensors();
Assignment assign = makeReductionNotation(assignment);

tensor.setNeedsPack(false);
@@ -751,7 +762,7 @@ static inline map<TensorVar, TensorBase> getTensors(const IndexExpr& expr) {
taco_iassert(isa<AccessTensorNode>(node)) << "Unknown subexpression";

if (!util::contains(arguments, node->tensorVar)) {
arguments.insert({node->tensorVar, to<AccessTensorNode>(node)->tensor});
arguments.insert({node->tensorVar, to<AccessTensorNode>(node)->getTensor()});
}

// Also add any tensors backing index sets of tensor accesses.
@@ -763,7 +774,7 @@ static inline map<TensorVar, TensorBase> getTensors(const IndexExpr& expr) {
}

// TODO (rohany): This seems like dead code.
TensorBase tensor = to<AccessTensorNode>(node)->tensor;
TensorBase tensor = to<AccessTensorNode>(node)->getTensor();
if (!util::contains(inserted, tensor)) {
inserted.insert(tensor);
operands.push_back(tensor);