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.

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

Add allocated CPU and GPU memory reporting #81

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

Conversation

jhalakpatel
Copy link
Collaborator

@jhalakpatel jhalakpatel commented Aug 9, 2024

Add RuntimeClient python bindings to report allocated CPU and GPU memory usage.

@jhalakpatel jhalakpatel marked this pull request as ready for review August 9, 2024 21:52
@jhalakpatel jhalakpatel added enhancement New feature or request mlir-tensorrt Pull request for the mlir-tensorrt project labels Aug 9, 2024
@christopherbate
Copy link
Collaborator

It looks fine, but without knowledge of the use case it's hard to review. This is only giving you the amount of memory tracked by the RuntimeClient at any given point in time. I hope you wouldn't be calling this often (since you're adding it up on the fly). If you need specific memory stats we could build something a bit more robust

@@ -773,6 +773,9 @@ class AllocTracker {
/// Return true if the tracker's map contains `ptr`.
bool contains(uintptr_t ptr) const;

/// Report total CPU and GPU memory allocated by runtime client.
std::pair<int64_t, int64_t> reportAllocatedMemory() const;
Copy link
Collaborator

Choose a reason for hiding this comment

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

There actually could are more types than just these two, so I'd prefer if we separate it into a struct or array. Array could be indexed by all the potential values of PointerType.

@@ -429,6 +429,24 @@ PointerInfo AllocTracker::lookupOrDefault(uintptr_t ptr) const {
return map.at(ptr);
}

std::pair<int64_t, int64_t> AllocTracker::reportAllocatedMemory() const {
int64_t totalCpuMemory = 0;
Copy link
Collaborator

Choose a reason for hiding this comment

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

We should use uint64_t here

MTRT_Status s = mtrtReportAllocatedMemory(self, &totalCpuMemory, &totalGpuMemory);
THROW_IF_MTRT_ERROR(s);
py::object namedtuple = py::module::import("collections").attr("namedtuple");
py::object MemoryUsage = namedtuple("MemoryUsage", "cpu_memory gpu_memory");
Copy link
Collaborator

Choose a reason for hiding this comment

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

You'll need to update the stubs so users can see this type information in the IDE.

MLIR_CAPI_EXPORTED MTRT_RuntimeClient
mtrtMemRefGetClient(MTRT_MemRefValue memref);

/// Retrieve the runtime client allocated cpu and gpu memory.
Copy link
Collaborator

Choose a reason for hiding this comment

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

This isn't quite accurate. You're reporting the CPU/GPU memory that is being tracked by the RuntimeClient. It can track buffers that are externally allocated.

/// Retrieve the runtime client allocated cpu and gpu memory.
MTRT_Status mtrtReportAllocatedMemory(MTRT_RuntimeClient client,
int64_t *totalCpuMemory,
int64_t *totalGpuMemory);
Copy link
Collaborator

Choose a reason for hiding this comment

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

Let's use uint64_t or size_t here.


for (const auto &entry : map) {
const PointerInfo &info = entry.second;
if (info.isExternallyManaged())
Copy link
Collaborator Author

Choose a reason for hiding this comment

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

@christopherbate Is this sufficient for tracking only internally managed/allocated pointers?

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment
Labels
enhancement New feature or request mlir-tensorrt Pull request for the mlir-tensorrt project
Projects
None yet
Development

Successfully merging this pull request may close these issues.

2 participants