Support MLOPart imports from vendored package names

csrc/apis/config.hpp

@@ -3,12 +3,14 @@
 #include <torch/library.h>
 
 #include "../jit_kernels/heuristics/runtime.hpp"
+#include "../runtime/python_package.hpp"
 #include "../runtime/runtime.hpp"
 
 TORCH_LIBRARY_FRAGMENT(deep_gemm, m) {
     using namespace deep_gemm;
-    m.def("init(str library_root_path) -> ()", [&](const std::string& library_root_path) {
+    m.def("init(str library_root_path, str python_package_name='deep_gemm') -> ()", [&](const std::string& library_root_path, const std::string& python_package_name) {
         init_jit(library_root_path);
+        python_package::set_package_name(python_package_name);
     });
     m.def("shutdown() -> ()", []() {
         runtime = deep_jit::LazyInit<Runtime>(nullptr);

csrc/runtime/mlopart.hpp

@@ -11,6 +11,7 @@
 
 #include "../utils/exception.hpp"
 #include "jit.hpp"
+#include "python_package.hpp"
 
 namespace deep_gemm::mlopart {
 
@@ -69,7 +70,7 @@ static PyObjectRef checked(PyObject* ptr) {
     return PyObjectRef(ptr);
 }
 
-// `deep_gemm.utils.mlopart.<name>(*args)`, stealing the references of `args`
+// `<python package>.utils.mlopart.<name>(*args)`, stealing the references of `args`
 // NOTES: the caller must hold the GIL
 static PyObjectRef call(const char* name, const std::initializer_list<PyObject*>& args) {
     PyObjectRef tuple(PyTuple_New(static_cast<Py_ssize_t>(args.size())));
@@ -85,7 +86,8 @@ static PyObjectRef call(const char* name, const std::initializer_list<PyObject*>
     check_python_error();
     DG_HOST_ASSERT(tuple.get() != nullptr);
 
-    const auto module = checked(PyImport_ImportModule("deep_gemm.utils.mlopart"));
+    const auto module_name = python_package::module_name("utils.mlopart");
+    const auto module = checked(PyImport_ImportModule(module_name.c_str()));
     const auto func = checked(PyObject_GetAttrString(module.get(), name));
     return checked(PyObject_Call(func.get(), tuple.get(), nullptr));
 }

csrc/runtime/python_package.hpp

@@ -0,0 +1,17 @@
+#pragma once
+
+#include <string>
+
+namespace deep_gemm::python_package {
+
+inline std::string package_name = "deep_gemm";
+
+inline void set_package_name(const std::string& name) {
+    package_name = name;
+}
+
+inline std::string module_name(const char* relative_name) {
+    return package_name + "." + relative_name;
+}
+
+}  // namespace deep_gemm::python_package

deep_gemm/_C.py

@@ -36,7 +36,7 @@ def _bind_guarded_ops(*names):
     globals().update({name: getattr(_torch_ops, name) for name in names})
 
 
-init = _torch_ops.init
+_cpp_init = _torch_ops.init
 shutdown = _torch_ops.shutdown
 set_num_sms = _torch_ops.set_num_sms
 get_num_sms = _torch_ops.get_num_sms
@@ -53,6 +53,10 @@ def _bind_guarded_ops(*names):
 cublaslt_gemm_tt = _torch_ops.cublaslt_gemm_tt
 
 
+def init(library_root_path):
+    _cpp_init(library_root_path, __package__)
+
+
 def set_block_size_multiple_of(value):
     if isinstance(value, int):
         return _torch_ops.set_block_size_multiple_of([value])
@@ -624,7 +628,7 @@ def slice_input_buffers(buffer):
 
 # GEMM epilogue classes for the `epilogue` argument, as `deep_gemm.epilogue.X`
 # NOTES: `alpha` and the FP8 `(d, sfd)` output pair are the shorthand for `Alpha` and `FP8Quantization`
-epilogue = types.ModuleType('deep_gemm.epilogue', 'GEMM epilogue classes')
+epilogue = types.ModuleType(f'{__package__}.epilogue', 'GEMM epilogue classes')
 epilogue.Identity = lambda: _torch_ops.epilogue_identity()
 epilogue.Alpha = lambda alpha: _torch_ops.epilogue_alpha(alpha)
 epilogue.FP8Quantization = lambda sfd: _torch_ops.epilogue_fp8_quantization(sfd)

deep_gemm/utils/mlopart.py

@@ -19,6 +19,10 @@
 PR_SET_PDEATHSIG = 1  # `<linux/prctl.h>`
 
 
+class MLOPartUnavailable(RuntimeError):
+    pass
+
+
 class MLOPart:
     """
     An MLOPart client that allocates localized memory.
@@ -96,7 +100,8 @@ def serve(sock: socket.socket) -> None:
     check(cuda.cuInit(0))
     num_devices = ctypes.c_int()
     check(cuda.cuDeviceGetCount(ctypes.byref(num_devices)))
-    assert num_devices.value == NUM_LOCALITY_DOMAINS, f'{num_devices.value} MLOPart devices'
+    if num_devices.value != NUM_LOCALITY_DOMAINS:
+        raise MLOPartUnavailable(f'{num_devices.value} MLOPart devices')
     sock.sendall(READY)
 
     while request := sock.recv(REQUEST.size, socket.MSG_WAITALL):
@@ -131,4 +136,7 @@ def main(sock_fd: int, device: str) -> None:
 
 
 if __name__ == '__main__':
-    main(int(sys.argv[1]), sys.argv[2])
+    try:
+        main(int(sys.argv[1]), sys.argv[2])
+    except MLOPartUnavailable:
+        pass

tests/test_locality_domain.py

@@ -135,9 +135,10 @@ def test_without_mlopart() -> None:
     # MPS predating MLOPart ignores `-mlopart`: emulate it by stripping the flag before the real control CLI
     # NOTES: through the torch dispatcher, the Python `AssertionError` from MLOPart surfaces as a `RuntimeError`
     real = shutil.which('nvidia-cuda-mps-control')
-    script = ('import warnings, torch, deep_gemm\n'
+    script = ('import warnings, torch\n'
               'with warnings.catch_warnings(record=True) as caught:\n'
               '    warnings.simplefilter("always")\n'
+              '    import deep_gemm\n'
               '    assert not deep_gemm.is_localization_available()\n'
               'assert len(caught) == 1 and "MLOPart is unavailable" in str(caught[0].message)\n'
               'assert deep_gemm.locality_domain.get_sm_locality_domains().numel() == torch.cuda.get_device_properties(0).multi_processor_count\n'
@@ -153,7 +154,8 @@ def test_without_mlopart() -> None:
         fake.chmod(0o755)
         result = subprocess.run([sys.executable, '-c', script], env={**os.environ, 'PATH': f'{bin_dir}:{os.environ["PATH"]}'}, stderr=subprocess.PIPE, text=True)
     assert result.returncode == 0, result.stderr
-    assert '1 MLOPart devices' in result.stderr, result.stderr
+    assert 'Traceback' not in result.stderr, result.stderr
+    assert 'MLOPartUnavailable' not in result.stderr, result.stderr
     print(' > detected, `localize` refused, SM table available')
     print()