diff --git a/python/operator_runtime/_runtime/__init__.py b/python/operator_runtime/_internal/__init__.py similarity index 95% rename from python/operator_runtime/_runtime/__init__.py rename to python/operator_runtime/_internal/__init__.py index 908842d..a6311d1 100644 --- a/python/operator_runtime/_runtime/__init__.py +++ b/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, diff --git a/python/operator_runtime/_runtime/ctypes_bindings.py b/python/operator_runtime/_internal/bindings.py similarity index 99% rename from python/operator_runtime/_runtime/ctypes_bindings.py rename to python/operator_runtime/_internal/bindings.py index 1297190..a4aeaf9 100644 --- a/python/operator_runtime/_runtime/ctypes_bindings.py +++ b/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) - diff --git a/python/operator_runtime/_runtime/loader.py b/python/operator_runtime/_internal/loader.py similarity index 94% rename from python/operator_runtime/_runtime/loader.py rename to python/operator_runtime/_internal/loader.py index c894794..2db7349 100644 --- a/python/operator_runtime/_runtime/loader.py +++ b/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}") - diff --git a/python/operator_runtime/_runtime/prepared.py b/python/operator_runtime/_internal/prepared.py similarity index 95% rename from python/operator_runtime/_runtime/prepared.py rename to python/operator_runtime/_internal/prepared.py index 05e7119..a8d3c40 100644 --- a/python/operator_runtime/_runtime/prepared.py +++ b/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 - diff --git a/python/operator_runtime/_runtime/tensor_view.py b/python/operator_runtime/_internal/tensor_view.py similarity index 99% rename from python/operator_runtime/_runtime/tensor_view.py rename to python/operator_runtime/_internal/tensor_view.py index d2a6eb4..30b5358 100644 --- a/python/operator_runtime/_runtime/tensor_view.py +++ b/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) - diff --git a/python/operator_runtime/ops/_common.py b/python/operator_runtime/ops/_common.py new file mode 100644 index 0000000..b6a1231 --- /dev/null +++ b/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) diff --git a/python/operator_runtime/ops/copy.py b/python/operator_runtime/ops/copy.py index db25695..49a3677 100644 --- a/python/operator_runtime/ops/copy.py +++ b/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) - diff --git a/python/operator_runtime/ops/reduce_sum.py b/python/operator_runtime/ops/reduce_sum.py index 7082f8a..597216c 100644 --- a/python/operator_runtime/ops/reduce_sum.py +++ b/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) - diff --git a/python/operator_runtime/ops/softmax.py b/python/operator_runtime/ops/softmax.py index fa1d0ed..4382b61 100644 --- a/python/operator_runtime/ops/softmax.py +++ b/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) - diff --git a/python/operator_runtime/ops/vector_add.py b/python/operator_runtime/ops/vector_add.py index 997be87..706ee4b 100644 --- a/python/operator_runtime/ops/vector_add.py +++ b/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) - diff --git a/python/operator_runtime/testing/__init__.py b/python/operator_runtime_testing/__init__.py similarity index 100% rename from python/operator_runtime/testing/__init__.py rename to python/operator_runtime_testing/__init__.py diff --git a/python/operator_runtime/testing/assertions.py b/python/operator_runtime_testing/assertions.py similarity index 100% rename from python/operator_runtime/testing/assertions.py rename to python/operator_runtime_testing/assertions.py diff --git a/python/operator_runtime/testing/benchmark.py b/python/operator_runtime_testing/benchmark.py similarity index 100% rename from python/operator_runtime/testing/benchmark.py rename to python/operator_runtime_testing/benchmark.py diff --git a/python/operator_runtime/testing/profiler.py b/python/operator_runtime_testing/profiler.py similarity index 100% rename from python/operator_runtime/testing/profiler.py rename to python/operator_runtime_testing/profiler.py diff --git a/tests/bench/copy.py b/tests/bench/copy.py index c161a49..dd966aa 100644 --- a/tests/bench/copy.py +++ b/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 diff --git a/tests/bench/reduce_sum.py b/tests/bench/reduce_sum.py index d9f7b22..7716c6e 100644 --- a/tests/bench/reduce_sum.py +++ b/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 diff --git a/tests/bench/softmax.py b/tests/bench/softmax.py index c8015bf..044e80e 100644 --- a/tests/bench/softmax.py +++ b/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 diff --git a/tests/bench/vector_add.py b/tests/bench/vector_add.py index 199cc6c..6ed1a5f 100644 --- a/tests/bench/vector_add.py +++ b/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 diff --git a/tests/ops/test_copy.py b/tests/ops/test_copy.py index e1d7ba9..4b8b893 100644 --- a/tests/ops/test_copy.py +++ b/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 diff --git a/tests/ops/test_reduce_sum.py b/tests/ops/test_reduce_sum.py index 64ede1c..be59c05 100644 --- a/tests/ops/test_reduce_sum.py +++ b/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 diff --git a/tests/ops/test_softmax.py b/tests/ops/test_softmax.py index 8b67e31..9c94811 100644 --- a/tests/ops/test_softmax.py +++ b/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 diff --git a/tests/ops/test_vector_add.py b/tests/ops/test_vector_add.py index 67e483a..0cdf21d 100644 --- a/tests/ops/test_vector_add.py +++ b/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 diff --git a/tests/test_tilelang_copy_templates.py b/tests/test_tilelang_copy_templates.py index ed02bac..dbc5ddc 100644 --- a/tests/test_tilelang_copy_templates.py +++ b/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,