Skip to content
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

[WIP] Change ExecutionContext.hpp/MemoryManager in runtime to use smart pointers #1284

Draft
wants to merge 1 commit into
base: main
Choose a base branch
from
Draft

[WIP] Change ExecutionContext.hpp/MemoryManager in runtime to use smart pointers #1284

Show file tree
Hide file tree
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
41 changes: 31 additions & 10 deletions runtime/lib/capi/ExecutionContext.hpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -40,38 +40,59 @@ extern "C" void __catalyst_inactive_callback(int64_t identifier, int64_t argc, i

class MemoryManager final {
private:
std::unordered_set<void *> _impl;
std::unordered_set<std::shared_ptr<void>> _impl;
std::mutex mu; // To guard the memory manager

public:
explicit MemoryManager() { _impl.reserve(1024); };

~MemoryManager()
std::shared_ptr<void> create(size_t size)
{
// Lock the mutex to protect _impl free
std::lock_guard<std::mutex> lock(mu);
for (auto allocation : _impl) {
free(allocation);
}
std::shared_ptr<void> p(std::malloc(size), [](void *ptr) { free(ptr); });
this->insert(p);
return p;
}

void insert(void *ptr)
std::shared_ptr<void> create_aligned(size_t alignment, size_t size)
{
std::shared_ptr<void> p(std::aligned_alloc(alignment, size), [](void *ptr) { free(ptr); });
this->insert(p);
return p;
}

void insert(std::shared_ptr<void> ptr)
{
// Lock the mutex to protect _impl update
std::lock_guard<std::mutex> lock(mu);
_impl.insert(ptr);
}

void erase(void *ptr)
{
// Lock the mutex to protect _impl update
std::lock_guard<std::mutex> lock(mu);
_impl.erase(ptr);

std::shared_ptr<void> target;
for (std::shared_ptr<void> sharedP : _impl) {
if (sharedP.get() == ptr) {
target = sharedP;
}
}
_impl.erase(target);
Copy link
Contributor

@erick-xanadu erick-xanadu Nov 6, 2024 •
edited
Loading

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Doesn't the implementation of erase will call the destructor of element equivalent to target? And since target is a std::shared_ptr it will decrement its reference count? This goes against what we want erase to do, which is to erase the pointer of the set which needs to be deleted. I.e., erase will make sure that the pointer is not free'd.

Copy link
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Ah, so this class is kind of like the anti-smart_ptr, where we hold a "dead" pointer, waiting for it to potentially be recycled, and manually free the pointer once the holder itself goes out of scope?

I wasn't familiar with the context here and was just trying to fix clang tidy warnings, but now that I know I think this is one instance where we have to manually silence clang tidy. I reckon this can be closed then?

Copy link
Contributor

@erick-xanadu erick-xanadu Nov 13, 2024 •
edited
Loading

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Ah, so this class is kind of like the anti-smart_ptr, where we hold a "dead" pointer, waiting for it to potentially be recycled,

No. It is like a unique_ptr. But with a way to not free the memory. If we knew they were dead, we would delete them. Essentially:

  1. Elements in the set may or may not be alive.
  2. When mlir_memref_free is called, the pointer is removed from this set and freed.
  3. When the set is destroyed, everything in the set is freed. (Imagine it like every element in the set is reference counted by a single counter, and here we decrement that counter).
  4. You can remove elements from the set without free-ing. This is useful to transfer the ownership of memory (and technically not what an arena allocator is like).

I call this an "arena-like" allocator, but it is not exactly an arena allocator.

and manually free the pointer once the holder itself goes out of scope?

Yes.

Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Just as a note, this design you are proposing may work if you return a shared_pointer but erase it from the set. It is then up to the caller to keep the pointer alive.

}

bool contains(void *ptr)
{
// Lock the mutex to protect _impl update
std::lock_guard<std::mutex> lock(mu);
return _impl.contains(ptr);

bool result = false;
for (std::shared_ptr<void> sharedP : _impl) {
if (sharedP.get() == ptr) {
result = true;
}
}
return result;
}
};

Expand Down
16 changes: 5 additions & 11 deletions runtime/lib/capi/RuntimeCAPI.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -156,16 +156,14 @@ void __catalyst__host__rt__unrecoverable_error()

void *_mlir_memref_to_llvm_alloc(size_t size)
{
void *ptr = malloc(size);
CTX->getMemoryManager()->insert(ptr);
return ptr;
std::shared_ptr<void> ptr = CTX->getMemoryManager()->create(size);
return ptr.get();
}

void *_mlir_memref_to_llvm_aligned_alloc(size_t alignment, size_t size)
{
void *ptr = aligned_alloc(alignment, size);
CTX->getMemoryManager()->insert(ptr);
return ptr;
std::shared_ptr<void> ptr = CTX->getMemoryManager()->create_aligned(alignment, size);
return ptr.get();
}

bool _mlir_memory_transfer(void *ptr)
Expand All @@ -177,11 +175,7 @@ bool _mlir_memory_transfer(void *ptr)
return true;
}

void _mlir_memref_to_llvm_free(void *ptr)
{
CTX->getMemoryManager()->erase(ptr);
free(ptr);
}
void _mlir_memref_to_llvm_free(void *ptr) { CTX->getMemoryManager()->erase(ptr); }

void __catalyst__rt__print_string(char *string)
{
Expand Down
2 changes: 0 additions & 2 deletions runtime/tests/Test_LightningCoreQIS.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -491,7 +491,6 @@ TEST_CASE("Test memory transfer in rt", "[CoreQIS]")
bool is_in_rt = _mlir_memory_transfer(a);
CHECK(is_in_rt);
__catalyst__rt__finalize();
free(a);
}

TEST_CASE("Test memory transfer not in rt", "[CoreQIS]")
Expand All @@ -501,7 +500,6 @@ TEST_CASE("Test memory transfer not in rt", "[CoreQIS]")
bool is_in_rt = _mlir_memory_transfer(a);
CHECK(!is_in_rt);
__catalyst__rt__finalize();
free(a);
}

TEST_CASE("Test __catalyst__qis__Measure", "[CoreQIS]")
Expand Down
Loading