Skip to content

Register the locality-domain shim under the _C module's own name #31

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

Merged
zyongye merged 1 commit into from
Oct 12, 2026
Merged

Register the locality-domain shim under the _C module's own name #31

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
5 changes: 3 additions & 2 deletions deep_gemm/_C.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -632,7 +632,7 @@ def slice_input_buffers(buffer):


# Memory homed in the GPU's locality domains, re-exported by `deep_gemm.locality_domain`
locality_domain = types.ModuleType('deep_gemm._C.locality_domain', "Memory homed in the GPU's locality domains")
locality_domain = types.ModuleType(f'{__name__}.locality_domain', "Memory homed in the GPU's locality domains")
locality_domain.get_num_locality_domains = _torch_ops.locality_domain_get_num_locality_domains
locality_domain.get_granularity = _torch_ops.locality_domain_get_granularity
locality_domain.is_localization_available = _torch_ops.locality_domain_is_localization_available
Expand All @@ -642,7 +642,8 @@ def slice_input_buffers(buffer):
locality_domain.get_sm_locality_domains = _torch_ops.locality_domain_get_sm_locality_domains
locality_domain.get_balanced_sm_locality_domains = _torch_ops.locality_domain_get_balanced_sm_locality_domains
locality_domain.release_mlopart = _torch_ops.locality_domain_release_mlopart
# NOTES: registered so that `from ._C.locality_domain import ...` resolves as with the former pybind submodule
# NOTES: registered so that `from ._C.locality_domain import ...` resolves as with the former pybind submodule,
# named after this module so the import also resolves when the package is vendored under another name
sys.modules[locality_domain.__name__] = locality_domain


Expand Down