purlin 0.1.0__tar.gz

purlin-0.1.0/LICENSE

@@ -0,0 +1,28 @@
+BSD 3-Clause License
+
+Copyright (c) 2026, Osayamen Jonathan Aimuyo
+
+Redistribution and use in source and binary forms, with or without
+modification, are permitted provided that the following conditions are met:
+
+1. Redistributions of source code must retain the above copyright notice, this
+   list of conditions and the following disclaimer.
+
+2. Redistributions in binary form must reproduce the above copyright notice,
+   this list of conditions and the following disclaimer in the documentation
+   and/or other materials provided with the distribution.
+
+3. Neither the name of the copyright holder nor the names of its
+   contributors may be used to endorse or promote products derived from
+   this software without specific prior written permission.
+
+THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS"
+AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
+IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE
+DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE
+FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL
+DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR
+SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER
+CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY,
+OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
+OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE

purlin-0.1.0/PKG-INFO

@@ -0,0 +1,17 @@
+Metadata-Version: 2.4
+Name: purlin
+Version: 0.1.0
+Summary: Adaptable, template-based substrate for high-performance inter-GPU communication
+Author-email: Osayamen Jonathan Aimuyo <osayamen@stanford.edu>
+License-Expression: BSD-3-Clause
+Project-URL: Homepage, https://github.com/osayamenja/purlin
+Project-URL: Repository, https://github.com/osayamenja/purlin
+Requires-Python: >=3.10
+Description-Content-Type: text/markdown
+License-File: LICENSE
+Requires-Dist: torch>=2.8.0
+Dynamic: license-file
+
+# Purlin
+Library that provides composable primitives for constructing high-performance and _adaptable_ device-initiated 
+communication operations.

purlin-0.1.0/README.md

@@ -0,0 +1,3 @@
+# Purlin
+Library that provides composable primitives for constructing high-performance and _adaptable_ device-initiated 
+communication operations.

purlin-0.1.0/purlin/CMakeLists.txt

@@ -0,0 +1,102 @@
+cmake_minimum_required(VERSION 3.27)
+project(purlin_v010 LANGUAGES CXX CUDA)
+
+set(CMAKE_CXX_STANDARD 20)
+set(CMAKE_CXX_STANDARD_REQUIRED ON)
+set(CMAKE_CXX_EXTENSIONS OFF)
+
+set(CMAKE_CUDA_STANDARD 20)
+set(CMAKE_CUDA_STANDARD_REQUIRED ON)
+set(CMAKE_CUDA_EXTENSIONS OFF)
+
+# flags
+add_library(purlin_warnings INTERFACE)
+
+target_compile_options(purlin_warnings INTERFACE
+        $<$<COMPILE_LANGUAGE:CXX>:
+        -Wall -Wextra
+        -fno-strict-aliasing
+        -Wno-unknown-pragmas
+        -Wnull-dereference
+        -Wnarrowing
+        -Wno-switch
+        -Wduplicated-branches
+        -Wformat=2
+        -Wno-unused-but-set-parameter
+        -Wno-sign-compare
+        >
+
+        $<$<COMPILE_LANGUAGE:CUDA>:
+        -Xcompiler=-Wall,-Wextra
+        -Xcompiler=-fno-strict-aliasing
+        -Xcompiler=-Wno-unknown-pragmas,-Wnull-dereference,-Wnarrowing
+        -Xcompiler=-Wno-switch,-Wduplicated-branches,-Wformat=2
+        -Xcompiler=-Wno-unused-but-set-parameter
+        -Xcudafe --display_error_number
+        -Xcompiler=-Wno-sign-compare
+        >
+)
+
+find_package(CUDAToolkit REQUIRED)
+
+if (NOT DEFINED GENERATED_SRC)
+  message(FATAL_ERROR "Need -DGENERATED_SRC")
+endif()
+
+if (NOT DEFINED TARGET_MODULE_NAME)
+  message(FATAL_ERROR "Need -DTARGET_MODULE_NAME")
+endif()
+
+if (NOT DEFINED ARCH)
+  message(FATAL_ERROR "Need -DARCH")
+endif()
+
+if(${ARCH} LESS 70)
+  message(FATAL_ERROR "Unsupported ARCH")
+endif()
+
+math(EXPR GPU_ARCH "${ARCH} * 10" OUTPUT_FORMAT DECIMAL)
+set(ARCH_TAG "${ARCH}")
+if(ARCH GREATER_EQUAL 90)
+    string(APPEND ARCH_TAG "a") # accelerated arch-specific instruction set
+endif()
+
+# ---- CPM ----
+include(${CMAKE_CURRENT_LIST_DIR}/CPM.cmake)
+
+#pybind11
+set(PYBIND11_FINDPYTHON ON)
+CPMAddPackage(
+  NAME pybind11
+  GITHUB_REPOSITORY pybind/pybind11
+  VERSION 3.0.4
+)
+if(NOT COMMAND pybind11_add_module)
+    message(FATAL_ERROR "pybind11 not found")
+endif()
+pybind11_add_module(${TARGET_MODULE_NAME} MODULE ${GENERATED_SRC})
+
+# -------------------------------------------------------
+# dependencies
+# -------------------------------------------------------
+CPMAddPackage(
+    NAME purlin
+    GITHUB_REPOSITORY osayamenja/purlin
+    GIT_TAG main
+)
+if (DEFINED PURLIN_SOURCE_DIR AND EXISTS "${PURLIN_SOURCE_DIR}/include")
+    target_include_directories(${TARGET_MODULE_NAME} PRIVATE "${PURLIN_SOURCE_DIR}/include")
+else()
+    target_link_libraries(${TARGET_MODULE_NAME} PRIVATE purlin::purlin)
+endif()
+
+target_compile_options(${TARGET_MODULE_NAME} PRIVATE
+            $<$<COMPILE_LANGUAGE:CXX>:-O3>
+            $<$<COMPILE_LANGUAGE:CUDA>:SHELL:-Xfatbin -compress-all>
+            $<$<COMPILE_LANGUAGE:CUDA>:--expt-relaxed-constexpr; -Xptxas -v; -t0>
+            $<$<COMPILE_LANGUAGE:CUDA>:SHELL:-gencode=arch=compute_${ARCH_TAG},code=sm_${ARCH_TAG}>
+)
+target_compile_definitions(${TARGET_MODULE_NAME} PRIVATE
+            ARCH=${GPU_ARCH}
+            PURLIN_NVTX=1
+)

purlin-0.1.0/purlin/CPM.cmake

@@ -0,0 +1,24 @@
+# SPDX-License-Identifier: MIT
+#
+# SPDX-FileCopyrightText: Copyright (c) 2019-2023 Lars Melchior and contributors
+
+set(CPM_DOWNLOAD_VERSION 0.42.1)
+set(CPM_HASH_SUM "f3a6dcc6a04ce9e7f51a127307fa4f699fb2bade357a8eb4c5b45df76e1dc6a5")
+
+if(CPM_SOURCE_CACHE)
+  set(CPM_DOWNLOAD_LOCATION "${CPM_SOURCE_CACHE}/cpm/CPM_${CPM_DOWNLOAD_VERSION}.cmake")
+elseif(DEFINED ENV{CPM_SOURCE_CACHE})
+  set(CPM_DOWNLOAD_LOCATION "$ENV{CPM_SOURCE_CACHE}/cpm/CPM_${CPM_DOWNLOAD_VERSION}.cmake")
+else()
+  set(CPM_DOWNLOAD_LOCATION "${CMAKE_BINARY_DIR}/cmake/CPM_${CPM_DOWNLOAD_VERSION}.cmake")
+endif()
+
+# Expand relative path. This is important if the provided path contains a tilde (~)
+get_filename_component(CPM_DOWNLOAD_LOCATION ${CPM_DOWNLOAD_LOCATION} ABSOLUTE)
+
+file(DOWNLOAD
+     https://github.com/cpm-cmake/CPM.cmake/releases/download/v${CPM_DOWNLOAD_VERSION}/CPM.cmake
+     ${CPM_DOWNLOAD_LOCATION} EXPECTED_HASH SHA256=${CPM_HASH_SUM}
+)
+
+include(${CPM_DOWNLOAD_LOCATION})

purlin-0.1.0/purlin/__init__.py

@@ -0,0 +1,107 @@
+from enum import IntEnum
+import torch
+
+class ContextHandle:
+    __slots__ = ("mod", "ctx", "buf", "hdl", "sig_buf", "sig_hdl")
+
+    def __init__(self, mod, ctx, sym_buf, sym_hdl, sig_buf, sig_hdl):
+        self.mod = mod
+        self.ctx = ctx
+        self.buf = sym_buf
+        self.hdl = sym_hdl
+        self.sig_buf = sig_buf
+        self.sig_hdl = sig_hdl
+
+
+STAGING_BUFFER_SIZE = 256 * 1024 * 1024
+PACKET_BUFFER_SIZE = 2 * 512 * 1024
+
+
+class DataType(IntEnum):
+    BF16 = 0
+    FP16 = 1
+    FP32 = 2
+    FP64 = 3
+
+
+def normalize_arch(arch: int):
+    if arch >= 100:
+        return 1000
+    if arch >= 90:
+        return 900
+    if arch >= 80:
+        return 800
+    return 700
+
+
+def buffer_type(t: torch.dtype):
+    if t == torch.float16:
+        return DataType.FP16
+    if t == torch.bfloat16:
+        return DataType.BF16
+    if t == torch.float32:
+        return DataType.FP32
+    if t == torch.float64:
+        return DataType.FP64
+    assert False, "invalid type"
+
+
+def initialize(group, device: torch.device, arch: int, stream_ptr: int):
+    import torch.distributed._symmetric_memory as sym_mem
+    from . import jit
+    from .bindings import purlin_bindings
+    rank = torch.distributed.get_rank(group)
+    world = torch.distributed.get_world_size(group)
+    mod_prefix = "purlin"
+    n_arch = normalize_arch(arch)
+    mod_name = "purlin_{}".format(n_arch)
+    src = purlin_bindings.substitute(mod_name=mod_name)
+    mod = jit.get_compiled(arch, src, mod_prefix, mod_name)
+    # staging buffers
+    staging_size = 2 * (STAGING_BUFFER_SIZE + (world * PACKET_BUFFER_SIZE))
+    t = sym_mem.empty(staging_size, dtype=torch.uint8, device=device)
+    t.zero_()
+    # signal pads
+    t1 = sym_mem.empty(2 * world, dtype=torch.uint64, device=device)
+    t1.zero_()
+    hdl = sym_mem.rendezvous(t, group)
+    hdl1 = sym_mem.rendezvous(t1, group)
+    ctx = mod.initialize(rank, world, hdl.buffer_ptrs, hdl1.buffer_ptrs, stream_ptr)
+    return ContextHandle(mod, ctx, t, hdl, t1, hdl1)
+
+
+def finalize(handle: ContextHandle, stream_ptr: int):
+    if handle.ctx is None:
+        return
+    handle.mod.finalize(handle.ctx, stream_ptr)
+    handle.ctx = None
+    handle.hdl = None
+    handle.buf = None
+    handle.sig_buf = None
+    handle.sig_hdl = None
+
+
+def all_gather(in_tensor: torch.Tensor, out_tensor: torch.Tensor, handle: ContextHandle, stream_ptr: int):
+    assert in_tensor.is_contiguous()
+    assert out_tensor.is_contiguous()
+    handle.mod.all_gather(in_tensor.data_ptr(), out_tensor.data_ptr(), in_tensor.nbytes, handle.ctx, stream_ptr)
+
+
+def all_reduce(in_tensor: torch.Tensor, out_tensor: torch.Tensor, handle: ContextHandle, stream_ptr: int):
+    assert in_tensor.is_contiguous()
+    assert out_tensor.is_contiguous()
+    bt = buffer_type(in_tensor.dtype)
+    handle.mod.all_reduce(in_tensor.data_ptr(), out_tensor.data_ptr(), in_tensor.nbytes, bt, handle.ctx, stream_ptr)
+
+
+def all_to_all(in_tensor: torch.Tensor, out_tensor: torch.Tensor, handle: ContextHandle, stream_ptr: int):
+    assert in_tensor.is_contiguous()
+    assert out_tensor.is_contiguous()
+    handle.mod.all_to_all(in_tensor.data_ptr(), out_tensor.data_ptr(), in_tensor.nbytes, handle.ctx, stream_ptr)
+
+
+def reduce_scatter(in_tensor: torch.Tensor, out_tensor: torch.Tensor, handle: ContextHandle, stream_ptr: int):
+    assert in_tensor.is_contiguous()
+    assert out_tensor.is_contiguous()
+    bt = buffer_type(in_tensor.dtype)
+    handle.mod.reduce_scatter(in_tensor.data_ptr(), out_tensor.data_ptr(), in_tensor.nbytes, bt, handle.ctx, stream_ptr)

purlin-0.1.0/purlin/bindings.py

@@ -0,0 +1,161 @@
+from string import Template
+purlin_bindings = Template(r"""
+//
+// Created by Osayamen on 6/2/26.
+//
+#include <cstdint>
+#include <cstdio>
+#include <cuda_runtime.h>
+
+#include <pybind11/pybind11.h>
+#include <pybind11/stl.h>
+#include <vector>
+
+#include <purlin/host.cuh>
+
+namespace py = pybind11;
+
+static std::uintptr_t purlin_initialize(const int& rank,
+  const int& world,
+  const std::vector<std::uintptr_t>& staging_table,
+  const std::vector<std::uintptr_t>& signal_table,
+  const std::uintptr_t& stream_ptr) {
+  auto stream = reinterpret_cast<cudaStream_t>(stream_ptr);
+  // allocate pointer tables
+  void* stagingTR = nullptr;
+  void* stagingLR = nullptr;
+  void* signals = nullptr;
+  void* gatherSignals = nullptr;
+
+  CHECK_CUDA(cudaMallocAsync(&stagingTR, sizeof(cuda::std::byte*) * world, stream));
+  static_assert(sizeof(uintptr_t) == sizeof(cuda::std::byte*));
+  CHECK_CUDA(cudaMemcpyAsync(stagingTR, staging_table.data(), sizeof(uintptr_t) * world,
+    cudaMemcpyHostToDevice, stream));
+  CHECK_CUDA(cudaMallocAsync(&stagingLR, sizeof(cuda::std::byte*) * world, stream));
+  CHECK_CUDA(cudaMallocAsync(&signals, sizeof(uint64_t*) * world, stream));
+  CHECK_CUDA(cudaMemcpyAsync(signals, signal_table.data(), sizeof(uintptr_t) * world,
+    cudaMemcpyHostToDevice, stream))
+  CHECK_CUDA(cudaMallocAsync(&gatherSignals, sizeof(uint64_t*) * world, stream));
+
+  std::vector<uintptr_t> stagingStash(world);
+  constexpr auto offsetTR = 2 * purlin::STAGING_BUFFER_SIZE_;
+  for (int i = 0; i < world; i++) {
+    stagingStash[i] = reinterpret_cast<uintptr_t>(reinterpret_cast<cuda::std::byte*>(staging_table[i]) + offsetTR);
+  }
+  CHECK_CUDA(cudaMemcpyAsync(stagingLR, stagingStash.data(), sizeof(uintptr_t) * world,
+    cudaMemcpyHostToDevice, stream));
+  std::vector<uintptr_t> signalStash(world);
+  const auto offsetSig = world;
+  for (int i = 0; i < world; i++) {
+    signalStash[i] = reinterpret_cast<uintptr_t>(reinterpret_cast<uint64_t*>(signal_table[i]) + offsetSig);
+  }
+  CHECK_CUDA(cudaMemcpyAsync(gatherSignals, signalStash.data(), sizeof(uintptr_t) * world,
+    cudaMemcpyHostToDevice, stream));
+
+  const auto ctx = purlin::initialize(rank, world,
+    static_cast<cuda::std::byte**>(stagingLR),
+    static_cast<cuda::std::byte**>(stagingTR),
+    static_cast<uint64_t**>(signals),
+    static_cast<uint64_t**>(gatherSignals), stream);
+  CHECK_CUDA(cudaStreamSynchronize(stream));
+  auto* pyCtx = new purlin::Context(ctx);
+  return reinterpret_cast<uintptr_t>(pyCtx);
+}
+
+static void purlin_finalize(const uintptr_t& raw_ctx, const uintptr_t& stream_ptr) {
+  auto stream = reinterpret_cast<cudaStream_t>(stream_ptr);
+  const auto* ctx = reinterpret_cast<purlin::Context*>(raw_ctx);
+  if (!ctx) return;
+  purlin::finalize(*ctx, stream);
+  CHECK_CUDA(cudaFreeAsync(ctx->staging, stream));
+  CHECK_CUDA(cudaFreeAsync(ctx->stagingLR, stream));
+  CHECK_CUDA(cudaFreeAsync(ctx->signals, stream));
+  CHECK_CUDA(cudaFreeAsync(ctx->gatherSignals, stream));
+  CHECK_CUDA(cudaStreamSynchronize(stream));
+  delete ctx;
+}
+
+static void all_gather(const uintptr_t& src, const uintptr_t& dst, const size_t& bytes,
+  const uintptr_t& raw_ctx, const uintptr_t& stream_ptr) {
+  purlin::allGather(
+    reinterpret_cast<cuda::std::byte*>(src),
+    reinterpret_cast<cuda::std::byte*>(dst),
+    bytes, *reinterpret_cast<purlin::Context*>(raw_ctx), reinterpret_cast<cudaStream_t>(stream_ptr));
+}
+
+static void all_reduce(const uintptr_t& src, const uintptr_t& dst, const size_t& bytes,
+  const int& buffer_type, const uintptr_t& raw_ctx, const uintptr_t& stream_ptr) {
+  switch (buffer_type) {
+    case purlin::TensorType::fp16: {
+      purlin::allReduce<__half>(reinterpret_cast<cuda::std::byte*>(src),
+        reinterpret_cast<cuda::std::byte*>(dst),
+        bytes, *reinterpret_cast<purlin::Context*>(raw_ctx), reinterpret_cast<cudaStream_t>(stream_ptr));
+    }
+      break;
+    case purlin::TensorType::bf16: {
+      purlin::allReduce<__nv_bfloat16>(reinterpret_cast<cuda::std::byte*>(src),
+        reinterpret_cast<cuda::std::byte*>(dst),
+        bytes, *reinterpret_cast<purlin::Context*>(raw_ctx), reinterpret_cast<cudaStream_t>(stream_ptr));
+    }
+      break;
+    case purlin::TensorType::fp32: {
+      purlin::allReduce<float>(reinterpret_cast<cuda::std::byte*>(src),
+      reinterpret_cast<cuda::std::byte*>(dst),
+      bytes, *reinterpret_cast<purlin::Context*>(raw_ctx), reinterpret_cast<cudaStream_t>(stream_ptr));
+    }
+      break;
+    default: {
+      // fp64
+      purlin::allReduce<double>(reinterpret_cast<cuda::std::byte*>(src),
+      reinterpret_cast<cuda::std::byte*>(dst),
+      bytes, *reinterpret_cast<purlin::Context*>(raw_ctx), reinterpret_cast<cudaStream_t>(stream_ptr));
+    }
+  }
+}
+
+static void all_to_all(const uintptr_t& src, const uintptr_t& dst, const size_t& bytes,
+  const uintptr_t& raw_ctx, const uintptr_t& stream_ptr) {
+  purlin::all2all(
+    reinterpret_cast<cuda::std::byte*>(src),
+    reinterpret_cast<cuda::std::byte*>(dst),
+    bytes, *reinterpret_cast<purlin::Context*>(raw_ctx), reinterpret_cast<cudaStream_t>(stream_ptr));
+}
+
+static void reduce_scatter(const uintptr_t& src, const uintptr_t& dst, const size_t& bytes,
+  const int& buffer_type, const uintptr_t& raw_ctx, const uintptr_t& stream_ptr) {
+  const auto ctx = *reinterpret_cast<purlin::Context*>(raw_ctx);
+  auto stream = reinterpret_cast<cudaStream_t>(stream_ptr);
+  const auto localBytes = bytes / ctx.world_l;
+  switch (buffer_type) {
+    case purlin::TensorType::fp16: {
+      purlin::reduceScatter<__half>(reinterpret_cast<cuda::std::byte*>(src),
+        reinterpret_cast<cuda::std::byte*>(dst), localBytes, ctx, stream);
+    }
+      break;
+    case purlin::TensorType::bf16: {
+      purlin::reduceScatter<__nv_bfloat16>(reinterpret_cast<cuda::std::byte*>(src),
+        reinterpret_cast<cuda::std::byte*>(dst), localBytes, ctx, stream);
+    }
+      break;
+    case purlin::TensorType::fp32: {
+      purlin::reduceScatter<float>(reinterpret_cast<cuda::std::byte*>(src),
+      reinterpret_cast<cuda::std::byte*>(dst), localBytes, ctx, stream);
+    }
+      break;
+    default: {
+      // fp64
+      purlin::reduceScatter<double>(reinterpret_cast<cuda::std::byte*>(src),
+      reinterpret_cast<cuda::std::byte*>(dst), localBytes, ctx, stream);
+    }
+  }
+}
+
+PYBIND11_MODULE($mod_name, m) {
+  m.def("initialize", &purlin_initialize);
+  m.def("finalize", &purlin_finalize);
+  m.def("all_gather", &all_gather);
+  m.def("all_reduce", &all_reduce);
+  m.def("all_to_all", &all_to_all);
+  m.def("reduce_scatter", &reduce_scatter);
+}
+""")

purlin-0.1.0/purlin/jit.py

@@ -0,0 +1,160 @@
+from __future__ import annotations
+
+import threading
+from pathlib import Path
+
+def _verify_dirs() -> None:
+    from pathlib import Path
+    root = Path(__file__).resolve().parent
+
+    if not (root / "CMakeLists.txt").exists():
+        raise RuntimeError("JIT CMakeLists.txt not found at package root")
+
+def _load_ext(mod_name: str, so_path: Path):
+    import importlib.util
+    spec = importlib.util.spec_from_file_location(mod_name, so_path)
+    if spec is None or spec.loader is None:
+        raise ImportError(f"Could not load {mod_name} from {so_path}")
+    mod = importlib.util.module_from_spec(spec)
+    spec.loader.exec_module(mod)
+    return mod
+
+def get_compiled(arch: int, src: str, mod_prefix: str, mod_name: str):
+    import hashlib
+    import os
+    import shutil
+    import socket
+    import subprocess
+    import sys
+    import time
+
+    _verify_dirs()
+
+    cache = Path(os.environ.get("PURLIN_CACHE_DIR", str(Path.home() / ".cache" / "purlin_jit")))
+    cache.mkdir(parents=True, exist_ok=True)
+
+    key = hashlib.sha256(f"{mod_name}|py{sys.version_info[:2]}|{src}".encode()).hexdigest()[:16]
+
+    build_root = cache / f"{key}"
+    build_root.mkdir(parents=True, exist_ok=True)
+
+    so_path = build_root / f"{mod_name}.so"
+    lock_path = build_root / ".build.lock"
+
+    # Fast path
+    if so_path.exists():
+        return _load_ext(mod_name, so_path)
+
+    # Process-unique tag for temp dirs
+    host = socket.gethostname()
+    pid = os.getpid()
+    tid = threading.get_ident()
+    uniq = f"{host}_tid{tid}_pid{pid}"
+
+    gen_dir = build_root / f"gen_{uniq}"
+    bdir = build_root / f"build_{uniq}"
+    gen_dir.mkdir(exist_ok=True)
+    bdir.mkdir(exist_ok=True)
+
+    generated = gen_dir / f"{mod_prefix}_bindings.cu"
+    generated.write_text(src)
+
+    cmake_source_dir = Path(__file__).resolve().parent
+
+    def _try_acquire_lock() -> bool:
+        try:
+            fd = os.open(lock_path, os.O_CREAT | os.O_EXCL | os.O_WRONLY)
+            with os.fdopen(fd, "w") as f:
+                f.write(f"host={host}\npid={pid}\ntid={tid}\ntime={time.time()}\n")
+            return True
+        except FileExistsError:
+            return False
+
+    def _lock_is_stale() -> bool:
+        try:
+            fields = dict(
+                line.split("=", 1)
+                for line in lock_path.read_text().splitlines()
+                if "=" in line
+            )
+        except OSError:
+            return False
+
+        lock_host = fields.get("host")
+        lock_pid = fields.get("pid")
+        if not lock_host or not lock_pid:
+            return True
+        if lock_host != host:
+            return False
+        try:
+            os.kill(int(lock_pid), 0)
+        except ProcessLookupError:
+            return True
+        except (PermissionError, ValueError):
+            return False
+        return False
+
+    def _release_lock() -> None:
+        try:
+            lock_path.unlink()
+        except FileNotFoundError:
+            pass
+
+    def _wait_for_artifact(timeout_s: float = 1800.0, poll_s: float = 0.1):
+        start = time.time()
+        while True:
+            if so_path.exists():
+                return _load_ext(mod_name, so_path)
+
+            if time.time() - start > timeout_s:
+                raise TimeoutError(
+                    f"Timed out waiting for JIT artifact {so_path} while another process was building it."
+                )
+
+            time.sleep(poll_s)
+
+    # package_dir = Path(__file__).resolve().parent
+    # repo_root = package_dir.parent
+    # csrc = repo_root / "csrc"
+    # Try to become the builder
+    have_lock = _try_acquire_lock()
+
+    if not have_lock and _lock_is_stale():
+        _release_lock()
+        have_lock = _try_acquire_lock()
+
+    if not have_lock:
+        # Another process is building. Wait for the final .so to appear.
+        return _wait_for_artifact()
+
+    try:
+        # Double-check after lock acquisition in case another process finished just before us
+        if so_path.exists():
+            return _load_ext(mod_name, so_path)
+
+        subprocess.run([
+            "cmake", "-S", str(cmake_source_dir), "-B", str(bdir), "-G", "Ninja",
+            # f"-DPURLIN_SOURCE_DIR={csrc}",
+            f"-DGENERATED_SRC={generated}",
+            f"-DTARGET_MODULE_NAME={mod_name}",
+            f"-DCMAKE_CUDA_ARCHITECTURES={arch}",
+            f"-DCPM_SOURCE_CACHE={Path.home() / '.cache' / 'cpm'}",
+            "-DCMAKE_BUILD_TYPE=Release",
+            f"-DARCH={arch}"
+        ], check=True)
+
+        subprocess.run([
+            "cmake", "--build", str(bdir), "--parallel"
+        ], check=True)
+
+        built = next(bdir.glob(mod_name + "*.so"))
+
+        # Copy into a temp path in build_root, then atomically replace final path
+        tmp_so = build_root / f".{mod_name}.{uniq}.tmp.so"
+        shutil.copy2(built, tmp_so)
+        tmp_so.replace(so_path)
+
+    finally:
+        _release_lock()
+
+    return _load_ext(mod_name, so_path)

purlin-0.1.0/purlin.egg-info/PKG-INFO

@@ -0,0 +1,17 @@
+Metadata-Version: 2.4
+Name: purlin
+Version: 0.1.0
+Summary: Adaptable, template-based substrate for high-performance inter-GPU communication
+Author-email: Osayamen Jonathan Aimuyo <osayamen@stanford.edu>
+License-Expression: BSD-3-Clause
+Project-URL: Homepage, https://github.com/osayamenja/purlin
+Project-URL: Repository, https://github.com/osayamenja/purlin
+Requires-Python: >=3.10
+Description-Content-Type: text/markdown
+License-File: LICENSE
+Requires-Dist: torch>=2.8.0
+Dynamic: license-file
+
+# Purlin
+Library that provides composable primitives for constructing high-performance and _adaptable_ device-initiated 
+communication operations.

purlin-0.1.0/purlin.egg-info/SOURCES.txt

@@ -0,0 +1,13 @@
+LICENSE
+README.md
+pyproject.toml
+purlin/CMakeLists.txt
+purlin/CPM.cmake
+purlin/__init__.py
+purlin/bindings.py
+purlin/jit.py
+purlin.egg-info/PKG-INFO
+purlin.egg-info/SOURCES.txt
+purlin.egg-info/dependency_links.txt
+purlin.egg-info/requires.txt
+purlin.egg-info/top_level.txt

purlin-0.1.0/purlin.egg-info/dependency_links.txt

@@ -0,0 +1 @@
+

purlin-0.1.0/purlin.egg-info/requires.txt

@@ -0,0 +1 @@
+torch>=2.8.0

purlin-0.1.0/purlin.egg-info/top_level.txt

@@ -0,0 +1 @@
+purlin

purlin-0.1.0/pyproject.toml

@@ -0,0 +1,31 @@
+[build-system]
+requires = ["setuptools>=77"]
+build-backend = "setuptools.build_meta"
+
+[project]
+name = "purlin"
+version = "0.1.0"
+readme = "README.md"
+license = "BSD-3-Clause"
+license-files = ["LICENSE"]
+description = "Adaptable, template-based substrate for high-performance inter-GPU communication"
+requires-python = ">=3.10"
+authors = [
+    { name = "Osayamen Jonathan Aimuyo", email = "osayamen@stanford.edu" },
+]
+dependencies = [
+      "torch>=2.8.0",
+]
+
+[project.urls]
+Homepage = "https://github.com/osayamenja/purlin"
+Repository = "https://github.com/osayamenja/purlin"
+
+[tool.setuptools]
+packages = ["purlin"]
+
+[tool.setuptools.package-data]
+purlin = [
+    "CMakeLists.txt",
+    "CPM.cmake"
+]

purlin-0.1.0/setup.cfg

@@ -0,0 +1,4 @@
+[egg_info]
+tag_build = 
+tag_date = 0
+