refactor: split runtime internals and testing package

Co-authored-by: wawahejun <hejunlbbc@gmail.com>

python/operator_runtime/_internal/__init__.py

@@ -1,6 +1,6 @@
 from .loader import load_library
 from .tensor_view import TensorView, tensor_view, dtype_to_oprt, current_stream_ptr, OPRT_MAX_DIMS
-from .ctypes_bindings import (
+from .bindings import (
     CFunctions,
     Descriptor,
     OperatorRuntimeError,

python/operator_runtime/_internal/bindings.py

@@ -123,4 +123,3 @@ def bind_reduce_like(name: str) -> CFunctions:
     destroy.argtypes = [Descriptor]
     destroy.restype = Status
     return CFunctions(create, workspace, execute, destroy)
-

python/operator_runtime/_internal/loader.py

@@ -12,7 +12,7 @@ def _candidate_library_paths() -> list[Path]:
     if build_dir:
         root = Path(build_dir)
         candidates.extend([root / "libcamp_ops.so", root / "ops" / "libcamp_ops.so"])
-    repo_root = Path(__file__).resolve().parents[2]
+    repo_root = Path(__file__).resolve().parents[3]
     candidates.extend(
         [
             repo_root / "build" / "libcamp_ops.so",
@@ -29,4 +29,3 @@ def load_library() -> ctypes.CDLL:
             return ctypes.CDLL(str(path))
     searched = ", ".join(str(path) for path in _candidate_library_paths())
     raise FileNotFoundError(f"libcamp_ops.so not found; searched: {searched}")
-

python/operator_runtime/_internal/prepared.py

@@ -6,7 +6,7 @@ from typing import Any
 
 import torch
 
-from .ctypes_bindings import CFunctions, Descriptor, check_status
+from .bindings import CFunctions, Descriptor, check_status
 
 
 @dataclass
@@ -46,4 +46,3 @@ class PreparedOp:
             self.destroy()
         except Exception:
             pass
-

python/operator_runtime/_internal/tensor_view.py

@@ -52,4 +52,3 @@ def current_stream_ptr(tensor: torch.Tensor | None = None) -> ctypes.c_void_p:
     device = tensor.device if tensor is not None and tensor.is_cuda else None
     stream = torch.cuda.current_stream(device=device)
     return ctypes.c_void_p(stream.cuda_stream)
-

python/operator_runtime/ops/_common.py

@@ -0,0 +1,24 @@
+from __future__ import annotations
+
+import ctypes
+
+import torch
+
+from operator_runtime._internal import CFunctions, Descriptor, PreparedOp, check_status
+
+
+def build_prepared_op(
+    funcs: CFunctions,
+    create_args: tuple[object, ...],
+    tensors: tuple[torch.Tensor, ...],
+    stream_tensor: torch.Tensor,
+) -> PreparedOp:
+    desc = Descriptor()
+    check_status(funcs.create(ctypes.byref(desc), *create_args))
+    workspace_size = ctypes.c_size_t()
+    check_status(funcs.workspace(desc, ctypes.byref(workspace_size)))
+    workspace = None
+    if workspace_size.value:
+        workspace = torch.empty(workspace_size.value, dtype=torch.uint8, device=stream_tensor.device)
+    runner_args = tuple(ctypes.c_void_p(tensor.data_ptr()) for tensor in tensors)
+    return PreparedOp(funcs, desc, workspace, runner_args, stream_tensor)

python/operator_runtime/ops/copy.py

@@ -5,7 +5,8 @@ import ctypes
 import torch
 
 from operator_runtime.backend import Backend, normalize_backend
-from operator_runtime._runtime import Descriptor, bind_unary, check_status, PreparedOp, tensor_view
+from operator_runtime._internal import PreparedOp, bind_unary, tensor_view
+from operator_runtime.ops._common import build_prepared_op
 
 
 def _check(out: torch.Tensor, src: torch.Tensor) -> None:
@@ -30,15 +31,10 @@ def prepare_copy(out: torch.Tensor, src: torch.Tensor, backend: str | Backend =
         raise NotImplementedError(f"backend {backend.value} is not runnable")
 
     funcs = bind_unary("copy")
-    desc = Descriptor()
     out_view = tensor_view(out)
     src_view = tensor_view(src)
-    check_status(funcs.create(ctypes.byref(desc), ctypes.byref(out_view), ctypes.byref(src_view)))
-    workspace_size = ctypes.c_size_t()
-    check_status(funcs.workspace(desc, ctypes.byref(workspace_size)))
-    workspace = torch.empty(workspace_size.value, dtype=torch.uint8, device=out.device) if workspace_size.value else None
-    args = (ctypes.c_void_p(out.data_ptr()), ctypes.c_void_p(src.data_ptr()))
-    return PreparedOp(funcs, desc, workspace, args, out)
+    create_args = (ctypes.byref(out_view), ctypes.byref(src_view))
+    return build_prepared_op(funcs, create_args, (out, src), out)
 
 
 def copy_(out: torch.Tensor, src: torch.Tensor, backend: str | Backend = Backend.NVIDIA) -> torch.Tensor:
@@ -50,4 +46,3 @@ def copy_(out: torch.Tensor, src: torch.Tensor, backend: str | Backend = Backend
 def copy(src: torch.Tensor, backend: str | Backend = Backend.NVIDIA) -> torch.Tensor:
     out = torch.empty_like(src)
     return copy_(out, src, backend)
-

python/operator_runtime/ops/reduce_sum.py

@@ -5,7 +5,8 @@ import ctypes
 import torch
 
 from operator_runtime.backend import Backend, normalize_backend
-from operator_runtime._runtime import Descriptor, bind_reduce_like, check_status, PreparedOp, tensor_view
+from operator_runtime._internal import PreparedOp, bind_reduce_like, tensor_view
+from operator_runtime.ops._common import build_prepared_op
 
 
 def _check(out: torch.Tensor, src: torch.Tensor, dim: int) -> None:
@@ -37,15 +38,10 @@ def prepare_reduce_sum(
         raise NotImplementedError(f"backend {backend.value} is not runnable")
 
     funcs = bind_reduce_like("reduce_sum")
-    desc = Descriptor()
     out_view = tensor_view(out)
     src_view = tensor_view(src)
-    check_status(funcs.create(ctypes.byref(desc), ctypes.byref(out_view), ctypes.byref(src_view), ctypes.c_int64(dim)))
-    workspace_size = ctypes.c_size_t()
-    check_status(funcs.workspace(desc, ctypes.byref(workspace_size)))
-    workspace = torch.empty(workspace_size.value, dtype=torch.uint8, device=out.device) if workspace_size.value else None
-    args = (ctypes.c_void_p(out.data_ptr()), ctypes.c_void_p(src.data_ptr()))
-    return PreparedOp(funcs, desc, workspace, args, out)
+    create_args = (ctypes.byref(out_view), ctypes.byref(src_view), ctypes.c_int64(dim))
+    return build_prepared_op(funcs, create_args, (out, src), out)
 
 
 def reduce_sum_(out: torch.Tensor, src: torch.Tensor, dim: int = 1, backend: str | Backend = Backend.NVIDIA) -> torch.Tensor:
@@ -59,4 +55,3 @@ def reduce_sum(src: torch.Tensor, dim: int = 1, backend: str | Backend = Backend
         raise ValueError("reduce_sum v1 supports 2D row-wise reduction over dim=1")
     out = torch.empty((src.shape[0],), dtype=src.dtype, device=src.device)
     return reduce_sum_(out, src, dim, backend)
-

python/operator_runtime/ops/softmax.py

@@ -5,7 +5,8 @@ import ctypes
 import torch
 
 from operator_runtime.backend import Backend, normalize_backend
-from operator_runtime._runtime import Descriptor, bind_reduce_like, check_status, PreparedOp, tensor_view
+from operator_runtime._internal import PreparedOp, bind_reduce_like, tensor_view
+from operator_runtime.ops._common import build_prepared_op
 
 
 def _check(out: torch.Tensor, src: torch.Tensor, dim: int) -> None:
@@ -37,15 +38,10 @@ def prepare_softmax(
         raise NotImplementedError(f"backend {backend.value} is not runnable")
 
     funcs = bind_reduce_like("softmax")
-    desc = Descriptor()
     out_view = tensor_view(out)
     src_view = tensor_view(src)
-    check_status(funcs.create(ctypes.byref(desc), ctypes.byref(out_view), ctypes.byref(src_view), ctypes.c_int64(dim)))
-    workspace_size = ctypes.c_size_t()
-    check_status(funcs.workspace(desc, ctypes.byref(workspace_size)))
-    workspace = torch.empty(workspace_size.value, dtype=torch.uint8, device=out.device) if workspace_size.value else None
-    args = (ctypes.c_void_p(out.data_ptr()), ctypes.c_void_p(src.data_ptr()))
-    return PreparedOp(funcs, desc, workspace, args, out)
+    create_args = (ctypes.byref(out_view), ctypes.byref(src_view), ctypes.c_int64(dim))
+    return build_prepared_op(funcs, create_args, (out, src), out)
 
 
 def softmax_(out: torch.Tensor, src: torch.Tensor, dim: int = 1, backend: str | Backend = Backend.NVIDIA) -> torch.Tensor:
@@ -57,4 +53,3 @@ def softmax_(out: torch.Tensor, src: torch.Tensor, dim: int = 1, backend: str |
 def softmax(src: torch.Tensor, dim: int = 1, backend: str | Backend = Backend.NVIDIA) -> torch.Tensor:
     out = torch.empty_like(src)
     return softmax_(out, src, dim, backend)
-

python/operator_runtime/ops/vector_add.py

@@ -5,7 +5,8 @@ import ctypes
 import torch
 
 from operator_runtime.backend import Backend, normalize_backend
-from operator_runtime._runtime import Descriptor, bind_binary, check_status, PreparedOp, tensor_view
+from operator_runtime._internal import PreparedOp, bind_binary, tensor_view
+from operator_runtime.ops._common import build_prepared_op
 
 
 def _check(out: torch.Tensor, a: torch.Tensor, b: torch.Tensor) -> None:
@@ -35,16 +36,11 @@ def prepare_vector_add(
         raise NotImplementedError(f"backend {backend.value} is not runnable")
 
     funcs = bind_binary("vector_add")
-    desc = Descriptor()
     out_view = tensor_view(out)
     a_view = tensor_view(a)
     b_view = tensor_view(b)
-    check_status(funcs.create(ctypes.byref(desc), ctypes.byref(out_view), ctypes.byref(a_view), ctypes.byref(b_view)))
-    workspace_size = ctypes.c_size_t()
-    check_status(funcs.workspace(desc, ctypes.byref(workspace_size)))
-    workspace = torch.empty(workspace_size.value, dtype=torch.uint8, device=out.device) if workspace_size.value else None
-    args = (ctypes.c_void_p(out.data_ptr()), ctypes.c_void_p(a.data_ptr()), ctypes.c_void_p(b.data_ptr()))
-    return PreparedOp(funcs, desc, workspace, args, out)
+    create_args = (ctypes.byref(out_view), ctypes.byref(a_view), ctypes.byref(b_view))
+    return build_prepared_op(funcs, create_args, (out, a, b), out)
 
 
 def vector_add_(out: torch.Tensor, a: torch.Tensor, b: torch.Tensor, backend: str | Backend = Backend.NVIDIA) -> torch.Tensor:
@@ -56,4 +52,3 @@ def vector_add_(out: torch.Tensor, a: torch.Tensor, b: torch.Tensor, backend: st
 def vector_add(a: torch.Tensor, b: torch.Tensor, backend: str | Backend = Backend.NVIDIA) -> torch.Tensor:
     out = torch.empty_like(a)
     return vector_add_(out, a, b, backend)
-

tests/bench/copy.py

@@ -13,7 +13,7 @@ if str(ROOT) not in sys.path:
 import torch
 
 from operator_runtime import copy
-from operator_runtime.testing import cuda_time_ms, PerformanceResult
+from operator_runtime_testing import cuda_time_ms, PerformanceResult
 from tests.cases import copy as copy_cases
 
 

tests/bench/reduce_sum.py

@@ -13,7 +13,7 @@ if str(ROOT) not in sys.path:
 import torch
 
 from operator_runtime import reduce_sum
-from operator_runtime.testing import cuda_time_ms, PerformanceResult
+from operator_runtime_testing import cuda_time_ms, PerformanceResult
 from tests.cases import reduce_sum as reduce_sum_cases
 
 

tests/bench/softmax.py

@@ -13,7 +13,7 @@ if str(ROOT) not in sys.path:
 import torch
 
 from operator_runtime import softmax
-from operator_runtime.testing import cuda_time_ms, PerformanceResult
+from operator_runtime_testing import cuda_time_ms, PerformanceResult
 from tests.cases import softmax as softmax_cases
 
 

tests/bench/vector_add.py

@@ -13,7 +13,7 @@ if str(ROOT) not in sys.path:
 import torch
 
 from operator_runtime import vector_add
-from operator_runtime.testing import cuda_time_ms, PerformanceResult
+from operator_runtime_testing import cuda_time_ms, PerformanceResult
 from tests.cases import vector_add as vector_add_cases
 
 

tests/ops/test_copy.py

@@ -4,7 +4,7 @@ import pytest
 import torch
 
 from operator_runtime import copy, copy_
-from operator_runtime.testing import assert_close, require_cuda
+from operator_runtime_testing import assert_close, require_cuda
 from tests.cases import copy as copy_cases
 
 

tests/ops/test_reduce_sum.py

@@ -4,7 +4,7 @@ import pytest
 import torch
 
 from operator_runtime import reduce_sum, reduce_sum_
-from operator_runtime.testing import assert_close, require_cuda
+from operator_runtime_testing import assert_close, require_cuda
 from tests.cases import reduce_sum as reduce_sum_cases
 
 

tests/ops/test_softmax.py

@@ -4,7 +4,7 @@ import pytest
 import torch
 
 from operator_runtime import softmax, softmax_
-from operator_runtime.testing import assert_close, require_cuda
+from operator_runtime_testing import assert_close, require_cuda
 from tests.cases import softmax as softmax_cases
 
 

tests/ops/test_vector_add.py

@@ -5,7 +5,7 @@ import torch
 
 from operator_runtime import vector_add, vector_add_
 from operator_runtime.ops.vector_add import prepare_vector_add
-from operator_runtime.testing import assert_close, require_cuda
+from operator_runtime_testing import assert_close, require_cuda
 from tests.cases import vector_add as vector_add_cases
 
 

tests/test_tilelang_copy_templates.py

@@ -5,7 +5,7 @@ import importlib.util
 import pytest
 import torch
 
-from operator_runtime.testing import require_cuda
+from operator_runtime_testing import require_cuda
 
 pytestmark = pytest.mark.skipif(
     importlib.util.find_spec("tilelang") is None,