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

[Runtime] Allow query of available device memory through DeviceAPI #16994

Merged

Conversation

Lunderberg
Copy link
Contributor

Prior to this commit, the total device memory could be queried through the DeviceAPI interface, but the currently available device memory could not. This functionality may be useful for debugging, or for validating available memory prior to model execution.

This commit implements the property Device.available_global_memory, which queries the DeviceAttrKind::kAvailableGlobalMemory. Support for this query, like all device attribute queries, may vary across different backends, and will return None for backends that do not support this query. This commit only currently implements support for kAvailableGlobalMemory for TVM's Cuda backend.

Prior to this commit, the total device memory could be queried through
the `DeviceAPI` interface, but the currently available device memory
could not.  This functionality may be useful for debugging, or for
validating available memory prior to model execution.

This commit implements the property `Device.available_global_memory`,
which queries the `DeviceAttrKind::kAvailableGlobalMemory`.  Support
for this query, like all device attribute queries, may vary across
different backends, and will return `None` for backends that do not
support this query.  This commit only currently implements support for
`kAvailableGlobalMemory` for TVM's Cuda backend.
Return None if the device does not support this feature.
"""
return self._GetDeviceAttr(self.device_type, self.device_id, 14)

@property
def available_global_memory(self):
"""Return size of the total global memory.
Copy link
Member

Choose a reason for hiding this comment

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

Not total

Copy link
Contributor Author

Choose a reason for hiding this comment

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

Thank you, and fixed!

@Lunderberg Lunderberg merged commit 3cd6673 into apache:main May 19, 2024
18 checks passed
@Lunderberg Lunderberg deleted the runtime_query_available_device_memory branch May 19, 2024 13:29
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment
Labels
None yet
Projects
None yet
Development

Successfully merging this pull request may close these issues.

None yet

2 participants