Add allocated CPU and GPU memory reporting #81
Conversation
jhalakpatel
added
enhancement
|
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; | |||
There was a problem hiding this comment.
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; | |||
There was a problem hiding this comment.
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"); |
There was a problem hiding this comment.
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. |
There was a problem hiding this comment.
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); |
There was a problem hiding this comment.
Let's use uint64_t or size_t here.
|
|
||
| for (const auto &entry : map) { | ||
| const PointerInfo &info = entry.second; | ||
| if (info.isExternallyManaged()) |
There was a problem hiding this comment.
@christopherbate Is this sufficient for tracking only internally managed/allocated pointers?
Add RuntimeClient python bindings to report allocated CPU and GPU memory usage.