Add multi-device system support for AOT compilation (#133)

* Use `pathlib.Path.suffix` to get file extension in `ninetoothed.aot`

* Add C++ class `ThreadSafeUnorderedMap` content string

* Change generated kernel source files from C to C++

* Add multi-device system support for AOT

* Add a test case for AOT's multi-device system support

Author: voltjia

src/ninetoothed/aot.py

@@ -3,6 +3,7 @@
 import re
 import subprocess
 import tempfile
+import textwrap
 import uuid
 
 import ninetoothed.dtype
@@ -33,7 +34,7 @@ def aot(
     output_contents = _aot(func, caller, kernel_name, num_warps, num_stages)
 
     for output_name, output_content in output_contents.items():
-        output_path = output_dir / f"{kernel_name}{output_name[-2:]}"
+        output_path = output_dir / f"{kernel_name}{pathlib.Path(output_name).suffix}"
 
         with open(output_path, "w") as f:
             f.write(output_content)
@@ -121,24 +122,48 @@ def _find_tensor_by_source_name(tensors, name):
     pattern = rf"\({', '.join(rf'(.*) {param}' for param in param_strings)}\)"
     c_param_type_strings = re.search(pattern, c_header_file).groups()
 
+    kernel_name_with_hash = f"{kernel_name}_{signature_hash}"
+
     unparser = _Unparser(c_param_type_strings)
 
     launch_func_unparsed = unparser.unparse(launch_func)
+    launch_func_unparsed_lines = launch_func_unparsed.splitlines()
+    launch_func_unparsed_lines.insert(1, f"{_INDENTATION}cuCtxGetId(NULL, &ctx_id);\n")
+    launch_func_unparsed_lines.insert(1, f"{_INDENTATION}unsigned long long ctx_id;")
+    launch_func_unparsed = "\n".join(launch_func_unparsed_lines)
     launch_func_unparsed = launch_func_unparsed.replace(
-        func.__name__, f"{kernel_name}_{signature_hash}"
+        func.__name__, f"kernels[ctx_id].{kernel_name_with_hash}"
     )
 
-    c_source_file = f"{c_source_file}\n{launch_func_unparsed}\n"
     c_source_file = c_source_file.replace("<stdint.h>", f'"{_HEADER_PATH}"')
     output_contents[c_source_file_name] = c_source_file
 
     c_header_file = f'{c_header_file}\n#ifdef __cplusplus\nextern "C" {unparser.header};\n#else\n{unparser.header};\n#endif\n'
     c_header_file = c_header_file.replace("<stdint.h>", f'"{_HEADER_PATH}"')
     output_contents[c_header_file_name] = c_header_file
 
+    kernel_start = c_source_file.find("//")
+    kernel_end = len(c_source_file)
+    cpp_source_file = (
+        c_source_file[:kernel_start]
+        + f"namespace {kernel_name_with_hash} {{\n"
+        + "struct Kernel {\n"
+        + textwrap.indent(c_source_file[kernel_start:kernel_end], _INDENTATION)
+        + "};\n"
+        + textwrap.indent(c_source_file[kernel_end:], _INDENTATION)
+        + "}\n"
+        + f"\nstatic ThreadSafeUnorderedMap<unsigned long long, {kernel_name_with_hash}::Kernel> kernels;\n"
+        + f'\nextern "C" {launch_func_unparsed}\n'
+    )
+    cpp_source_file_name = f"{kernel_name}.{signature_hash}.cpp"
+    output_contents[cpp_source_file_name] = cpp_source_file
+    output_contents.pop(c_source_file_name)
+
     return output_contents
 
 
+_INDENTATION = "    "
+
 _MACRO_MAPPING = {True: ("NINETOOTHED_TRUE", 1), False: ("NINETOOTHED_FALSE", 0)}
 
 _DTYPE_MAPPING = {
@@ -163,6 +188,41 @@ def _find_tensor_by_source_name(tensors, name):
 
 _DATA_TYPE_BODY_CONTENT = ",\n    ".join(_DTYPE_MAPPING.values())
 
+_THREAD_SAFE_UNORDERED_MAP_CONTENT = """#include <mutex>
+#include <shared_mutex>
+#include <unordered_map>
+
+template<typename Key, typename Value>
+class ThreadSafeUnorderedMap {
+public:
+    Value& operator[](const Key& key) {
+        {
+            std::shared_lock lock{mutex};
+
+            auto iter = map.find(key);
+
+            if (iter != map.end()) {
+                return iter->second;
+            }
+        }
+
+        std::unique_lock lock{mutex};
+
+        auto iter = map.find(key);
+
+        if (iter == map.end()) {
+            iter = map.emplace(key, Value{}).first;
+        }
+
+        return iter->second;
+    }
+
+private:
+    std::unordered_map<Key, Value> map;
+
+    mutable std::shared_mutex mutex;
+};"""
+
 _HEADER_CONTENT = f"""#ifndef NINETOOTHED_H
 #define NINETOOTHED_H
 
@@ -184,6 +244,10 @@ def _find_tensor_by_source_name(tensors, name):
 
 typedef int NineToothedResult;
 
+#ifdef __cplusplus
+{_THREAD_SAFE_UNORDERED_MAP_CONTENT}
+#endif
+
 #endif // NINETOOTHED_H
 """
 

src/ninetoothed/build.py

@@ -80,15 +80,15 @@ def build(premake, configs, *, caller=None, kernel_name=None, output_dir=None):
         f"{type} {param}" for param, type in zip(param_names, param_types)
     )
 
-    source_file_name = f"{kernel_name}.c"
+    source_file_name = f"{kernel_name}.cpp"
     header_file_name = f"{kernel_name}.h"
 
     func_sig = f"NineToothedResult launch_{kernel_name}({param_decls})"
 
     joined_launches = "\n".join(launches)
 
     op_decl = f'#ifdef __cplusplus\nextern "C" {func_sig};\n#else\n{func_sig};\n#endif'
-    op_def = f"""{func_sig} {{
+    op_def = f"""extern "C" {func_sig} {{
 {joined_launches}
     return 1;
 }}"""

tests/test_aot.py

@@ -24,7 +24,8 @@
     "dtype, ninetoothed_dtype", ((torch.bfloat16, ninetoothed.bfloat16),)
 )
 @pytest.mark.parametrize("size", (45327,))
-def test_add(size, dtype, device, ninetoothed_dtype):
+@pytest.mark.parametrize("test_multi_device", (False, True))
+def test_add(test_multi_device, size, dtype, device, ninetoothed_dtype):
     def _arrangement(input, other, output):
         def _arrange(tensor):
             return tensor.tile((256,))
@@ -52,15 +53,25 @@ def _application(input, other, output):
 
     shape = (size,)
 
-    input = torch.randn(shape, dtype=dtype, device=device)
-    other = torch.randn(shape, dtype=dtype, device=device)
-    output = torch.empty_like(input)
+    if test_multi_device:
+        if torch.cuda.device_count() < 2:
+            pytest.skip("multi-device testing requires at least 2 devices")
 
-    _run_launch_func(launch_func, input, other, output)
+        devices = (f"{device}:0", f"{device}:1")
+    else:
+        devices = (device,)
 
-    expected = torch.add(input, other)
+    for device in devices:
+        with torch.cuda.Stream(device=device):
+            input = torch.randn(shape, dtype=dtype, device=device)
+            other = torch.randn(shape, dtype=dtype, device=device)
+            output = torch.empty_like(input)
 
-    assert torch.allclose(output, expected)
+            _run_launch_func(launch_func, input, other, output)
+
+            expected = torch.add(input, other)
+
+            assert torch.allclose(output, expected)
 
 
 @pytest.mark.parametrize("device", get_available_devices())
@@ -346,7 +357,7 @@ def _compile_library(kernel_name, output_dir):
         "-lcuda",
         "-o",
         output_dir / f"{kernel_name}.so",
-    ] + list(output_dir.glob(f"{kernel_name}*.c"))
+    ] + list(output_dir.glob(f"{kernel_name}*.cpp"))
 
     subprocess.run(command, check=True)