From bfa17c8df078b423cbe6ea7e1cf6526c3b733df7 Mon Sep 17 00:00:00 2001 From: Zane Fink Date: Mon, 27 Jul 2026 15:13:45 -0700 Subject: [PATCH 1/2] feature: Add native Python signature API --- bindings/python/PythonBindings.cpp | 199 +++++++++++------- bindings/python/PythonBindings.h | 12 +- .../tests/test_gpu_cpp_launch_validation.py | 15 +- .../tests/test_gpu_cpp_pointer_validation.py | 4 +- bindings/python/tests/test_gpu_cpp_smoke.py | 9 +- .../python/tests/test_gpu_mlir_fatal_paths.py | 2 +- bindings/python/tests/test_gpu_mlir_smoke.py | 6 +- .../python/tests/test_gpu_mlir_validation.py | 10 +- bindings/python/tests/test_host_cpp_smoke.py | 11 +- .../python/tests/test_host_cpp_validation.py | 173 ++++++++++++++- .../python/tests/test_host_llvmir_smoke.py | 11 +- .../python/tests/test_proteus_init_exports.py | 6 + bindings/python/tests/test_wildcard_import.py | 2 + docs/user/llvmir-frontend.md | 5 +- docs/user/python-bindings.md | 96 +++++++++ mkdocs.yml | 1 + 16 files changed, 451 insertions(+), 111 deletions(-) create mode 100644 docs/user/python-bindings.md diff --git a/bindings/python/PythonBindings.cpp b/bindings/python/PythonBindings.cpp index d0c6e3750..bd43bb549 100644 --- a/bindings/python/PythonBindings.cpp +++ b/bindings/python/PythonBindings.cpp @@ -147,6 +147,9 @@ ArgStorage convertArg(Type Ty, py::object Value, bool DevicePointers) { // addresses. ArgStorage Storage; switch (Ty.Kind) { + case PyType::Void: + throw py::type_error( + "proteus.void is only valid as a signature return type"); case PyType::I8: Storage.set(Value.cast()); break; @@ -179,14 +182,70 @@ ArgStorage convertArg(Type Ty, py::object Value, bool DevicePointers) { class Module; class Function; -py::object typeToCType(py::handle ctypes, py::object Ty) { - if (Ty.is_none()) - return py::none(); +std::string typeRepr(Type Ty) { + switch (Ty.Kind) { + case PyType::Void: + return "proteus.void"; + case PyType::I8: + return "proteus.i8"; + case PyType::I32: + return "proteus.i32"; + case PyType::I64: + return "proteus.i64"; + case PyType::U32: + return "proteus.u32"; + case PyType::U64: + return "proteus.u64"; + case PyType::F32: + return "proteus.f32"; + case PyType::F64: + return "proteus.f64"; + case PyType::Ptr: + return "proteus.ptr"; + } + throw py::type_error("unsupported Proteus type descriptor"); +} + +std::string signatureRepr(const Signature &Sig) { + std::string Repr = typeRepr(Sig.RetType) + "("; + for (std::size_t I = 0; I < Sig.ArgTypes.size(); ++I) { + if (I != 0) + Repr += ", "; + Repr += typeRepr(Sig.ArgTypes[I]); + } + Repr += ")"; + return Repr; +} +void validateSignature(const Signature &Sig) { + for (Type ArgType : Sig.ArgTypes) + if (ArgType.Kind == PyType::Void) + throw py::type_error( + "proteus.void is only valid as a signature return type"); +} + +Signature makeSignature(Type RetType, py::args Args) { + std::vector ArgTypes; + ArgTypes.reserve(Args.size()); + for (py::handle Arg : Args) { + if (!py::isinstance(Arg)) + throw py::type_error( + "signature arguments must be Proteus type descriptors"); + Type ArgType = py::cast(Arg); + if (ArgType.Kind == PyType::Void) + throw py::type_error( + "proteus.void is only valid as a signature return type"); + ArgTypes.push_back(ArgType); + } + return Signature{RetType, std::move(ArgTypes)}; +} + +py::object typeToCType(py::handle ctypes, Type Ty) { // Build the ctypes signature from the lightweight descriptors exported by // this module. - Type T = Ty.cast(); - switch (T.Kind) { + switch (Ty.Kind) { + case PyType::Void: + return py::none(); case PyType::I8: return ctypes.attr("c_int8"); case PyType::I32: @@ -210,26 +269,26 @@ py::object typeToCType(py::handle ctypes, py::object Ty) { class Kernel { std::shared_ptr Mod; void *KernelFunc = nullptr; - std::vector ArgTypes; + Signature Sig; std::string Name; public: - Kernel(std::shared_ptr Mod, void *KernelFunc, - std::vector ArgTypes, std::string Name) - : Mod(std::move(Mod)), KernelFunc(KernelFunc), - ArgTypes(std::move(ArgTypes)), Name(std::move(Name)) {} + Kernel(std::shared_ptr Mod, void *KernelFunc, Signature Sig, + std::string Name) + : Mod(std::move(Mod)), KernelFunc(KernelFunc), Sig(std::move(Sig)), + Name(std::move(Name)) {} void launch(py::object Grid, py::object Block, py::sequence Args, uint64_t Shmem, py::object Stream) { - if (Args.size() != ArgTypes.size()) - throw py::type_error("kernel argument count does not match argtypes"); + if (Args.size() != Sig.ArgTypes.size()) + throw py::type_error("kernel argument count does not match signature"); // Keep the converted values alive for the duration of the kernel launch. std::vector Storage; - Storage.reserve(ArgTypes.size()); - for (std::size_t I = 0; I < ArgTypes.size(); ++I) + Storage.reserve(Sig.ArgTypes.size()); + for (std::size_t I = 0; I < Sig.ArgTypes.size(); ++I) Storage.push_back(convertArg( - ArgTypes[I], py::reinterpret_borrow(Args[I]), true)); + Sig.ArgTypes[I], py::reinterpret_borrow(Args[I]), true)); std::vector RawArgs; RawArgs.reserve(Storage.size()); @@ -252,15 +311,8 @@ class Kernel { } std::string repr() const { - std::string Repr = - "(); - } - Repr += "]>"; - return Repr; + return ""; } }; @@ -268,25 +320,23 @@ class Function { std::shared_ptr Mod; py::object Callable; std::string Name; - py::object RetType; - std::vector ArgTypes; + Signature Sig; public: Function(std::shared_ptr Mod, py::object Callable, - std::string Name, py::object RetType, std::vector ArgTypes) + std::string Name, Signature Sig) : Mod(std::move(Mod)), Callable(std::move(Callable)), - Name(std::move(Name)), RetType(std::move(RetType)), - ArgTypes(std::move(ArgTypes)) {} + Name(std::move(Name)), Sig(std::move(Sig)) {} py::object call(py::args Args) const { - if (Args.size() != ArgTypes.size()) - throw py::type_error("function argument count does not match argtypes"); + if (Args.size() != Sig.ArgTypes.size()) + throw py::type_error("function argument count does not match signature"); py::tuple CoercedArgs(Args.size()); py::module_ ctypes = py::module_::import("ctypes"); - for (std::size_t I = 0; I < ArgTypes.size(); ++I) { + for (std::size_t I = 0; I < Sig.ArgTypes.size(); ++I) { py::object Arg = py::reinterpret_borrow(Args[I]); - if (ArgTypes[I].Kind == PyType::Ptr) + if (Sig.ArgTypes[I].Kind == PyType::Ptr) // Force pointer-typed arguments through `c_void_p` so ctypes does not // reinterpret Python integers as narrower scalar values. CoercedArgs[I] = @@ -298,16 +348,8 @@ class Function { } std::string repr() const { - std::string Repr = "(); - } - Repr += "]>"; - return Repr; + return ""; } }; @@ -317,13 +359,15 @@ class Module { public: explicit Module(std::shared_ptr Impl) : Impl(std::move(Impl)) {} - Kernel getKernel(const std::string &Name, std::vector ArgTypes) { - return Kernel(Impl, Impl->getKernelAddress(Name), std::move(ArgTypes), - Name); + Kernel getKernel(const std::string &Name, Signature Sig) { + validateSignature(Sig); + if (Sig.RetType.Kind != PyType::Void) + throw py::type_error("kernel signatures must return proteus.void"); + return Kernel(Impl, Impl->getKernelAddress(Name), std::move(Sig), Name); } - Function getFunction(const std::string &Name, py::object RetType, - std::vector ArgTypes) { + Function getFunction(const std::string &Name, Signature Sig) { + validateSignature(Sig); if (!proteus::isHostTargetModel(Impl->getTargetModel())) throw py::value_error( "Target is a GPU model, cannot directly run functions, use launch()"); @@ -331,18 +375,17 @@ class Module { py::module_ ctypes = py::module_::import("ctypes"); // CFUNCTYPE expects the return type first, then the positional argument // types. - py::tuple CTypeArgs(ArgTypes.size() + 1); - CTypeArgs[0] = typeToCType(ctypes, RetType); - for (std::size_t I = 0; I < ArgTypes.size(); ++I) - CTypeArgs[I + 1] = typeToCType(ctypes, py::cast(ArgTypes[I])); + py::tuple CTypeArgs(Sig.ArgTypes.size() + 1); + CTypeArgs[0] = typeToCType(ctypes, Sig.RetType); + for (std::size_t I = 0; I < Sig.ArgTypes.size(); ++I) + CTypeArgs[I + 1] = typeToCType(ctypes, Sig.ArgTypes[I]); py::object FuncType = ctypes.attr("CFUNCTYPE")(*CTypeArgs); // Wrap the JIT symbol address in a Python callable with the requested // signature. py::object Callable = FuncType( py::int_(reinterpret_cast(Impl->getFunctionAddress(Name)))); - return Function(Impl, std::move(Callable), Name, std::move(RetType), - std::move(ArgTypes)); + return Function(Impl, std::move(Callable), Name, std::move(Sig)); } uintptr_t getFunctionAddress(const std::string &Name) { @@ -392,28 +435,24 @@ PYBIND11_MODULE(_proteus, M) { M.doc() = "Thin Python bindings for Proteus JIT frontends"; // Expose the builtin scalar/pointer descriptors as module-level singletons. - py::class_(M, "Type").def("__repr__", [](const Type &T) { - switch (T.Kind) { - case PyType::I8: - return "proteus.i8"; - case PyType::I32: - return "proteus.i32"; - case PyType::I64: - return "proteus.i64"; - case PyType::U32: - return "proteus.u32"; - case PyType::U64: - return "proteus.u64"; - case PyType::F32: - return "proteus.f32"; - case PyType::F64: - return "proteus.f64"; - case PyType::Ptr: - return "proteus.ptr"; - } - return "proteus.Type"; - }); - + py::class_(M, "Type") + .def("__call__", &makeSignature) + .def("__repr__", &typeRepr); + + py::class_(M, "Signature") + .def_property_readonly("restype", + [](const Signature &Sig) { return Sig.RetType; }) + .def_property_readonly("argtypes", + [](const Signature &Sig) { + py::tuple Result(Sig.ArgTypes.size()); + for (std::size_t I = 0; I < Sig.ArgTypes.size(); + ++I) + Result[I] = py::cast(Sig.ArgTypes[I]); + return Result; + }) + .def("__repr__", &signatureRepr); + + M.attr("void") = Type{PyType::Void}; M.attr("i8") = Type{PyType::I8}; M.attr("i32") = Type{PyType::I32}; M.attr("i64") = Type{PyType::I64}; @@ -425,10 +464,10 @@ PYBIND11_MODULE(_proteus, M) { // Modules own compiled code; kernels capture both the code and arg schema. py::class_(M, "Module") - .def("get_kernel", &Module::getKernel, py::arg("name"), - py::arg("argtypes")) - .def("get_function", &Module::getFunction, py::arg("name"), - py::arg("restype"), py::arg("argtypes")) + .def("get_kernel", &Module::getKernel, py::arg("name"), py::kw_only(), + py::arg("signature")) + .def("get_function", &Module::getFunction, py::arg("name"), py::kw_only(), + py::arg("signature")) .def("get_function_address", &Module::getFunctionAddress, py::arg("name")); diff --git a/bindings/python/PythonBindings.h b/bindings/python/PythonBindings.h index 65be062dc..8a6bb5031 100644 --- a/bindings/python/PythonBindings.h +++ b/bindings/python/PythonBindings.h @@ -11,14 +11,20 @@ namespace proteus_python { -// Scalar and pointer element kinds exposed through the Python bindings. -enum class PyType { I8, I32, I64, U32, U64, F32, F64, Ptr }; +// Return, scalar, and pointer kinds exposed through the Python bindings. +enum class PyType { Void, I8, I32, I64, U32, U64, F32, F64, Ptr }; -// Lightweight type descriptor used when marshalling Python-side values. +// Lightweight type descriptor used in signatures and argument marshalling. struct Type { PyType Kind; }; +// Exact native function signature assembled from the builtin descriptors. +struct Signature { + Type RetType; + std::vector ArgTypes; +}; + // Common execution interface implemented by the C++ and MLIR-backed modules. class ModuleBase { public: diff --git a/bindings/python/tests/test_gpu_cpp_launch_validation.py b/bindings/python/tests/test_gpu_cpp_launch_validation.py index 119a7ef1a..f3ddfca7d 100644 --- a/bindings/python/tests/test_gpu_cpp_launch_validation.py +++ b/bindings/python/tests/test_gpu_cpp_launch_validation.py @@ -13,19 +13,28 @@ def main(): device_value = runtime.malloc_i32() mod = proteus.compile(source, frontend="cpp", target=target) - kernel = mod.get_kernel("write_int", [proteus.ptr, proteus.i32]) + kernel = mod.get_kernel( + "write_int", signature=proteus.void(proteus.ptr, proteus.i32) + ) + expect_raises( + TypeError, + lambda: mod.get_kernel( + "write_int", signature=proteus.i32(proteus.ptr, proteus.i32) + ), + "kernel signatures must return proteus.void", + ) expect_raises( ValueError, lambda: mod.get_function( - "write_int", restype=None, argtypes=[proteus.ptr, proteus.i32] + "write_int", signature=proteus.void(proteus.ptr, proteus.i32) ), "Target is a GPU model, cannot directly run functions, use launch()", ) expect_raises( TypeError, lambda: kernel.launch(grid=1, block=1, args=[device_value]), - "kernel argument count does not match argtypes", + "kernel argument count does not match signature", ) expect_raises( TypeError, diff --git a/bindings/python/tests/test_gpu_cpp_pointer_validation.py b/bindings/python/tests/test_gpu_cpp_pointer_validation.py index a9c8162ff..7b79260a1 100644 --- a/bindings/python/tests/test_gpu_cpp_pointer_validation.py +++ b/bindings/python/tests/test_gpu_cpp_pointer_validation.py @@ -20,7 +20,9 @@ def main(): runtime.copy_h2d_i32(device_value, 0) mod = proteus.compile(source, frontend="cpp", target=target) - kernel = mod.get_kernel("write_int", [proteus.ptr, proteus.i32]) + kernel = mod.get_kernel( + "write_int", signature=proteus.void(proteus.ptr, proteus.i32) + ) kernel.launch(grid=(1, 1, 1), block=1, args=[DataPtr(device_value), 8]) runtime.sync() diff --git a/bindings/python/tests/test_gpu_cpp_smoke.py b/bindings/python/tests/test_gpu_cpp_smoke.py index a43933179..79d058793 100644 --- a/bindings/python/tests/test_gpu_cpp_smoke.py +++ b/bindings/python/tests/test_gpu_cpp_smoke.py @@ -14,8 +14,13 @@ def main(): runtime.copy_h2d_i32(device_value, 0) mod = proteus.compile(source, frontend="cpp", target=target) - kernel = mod.get_kernel("write_int", [proteus.ptr, proteus.i32]) - assert repr(kernel) == "" + kernel = mod.get_kernel( + "write_int", signature=proteus.void(proteus.ptr, proteus.i32) + ) + assert ( + repr(kernel) + == "" + ) kernel.launch(grid=1, block=(1, 1, 1), args=[device_value, 7]) runtime.sync() assert runtime.copy_d2h_i32(device_value) == 7 diff --git a/bindings/python/tests/test_gpu_mlir_fatal_paths.py b/bindings/python/tests/test_gpu_mlir_fatal_paths.py index a26abe935..a6f32c4b1 100644 --- a/bindings/python/tests/test_gpu_mlir_fatal_paths.py +++ b/bindings/python/tests/test_gpu_mlir_fatal_paths.py @@ -13,7 +13,7 @@ def main(): source = {source!r} target = {target!r} mod = proteus.compile(source, frontend="mlir", target=target, verify=True) -mod.get_kernel("missing_kernel", [proteus.ptr]) +mod.get_kernel("missing_kernel", signature=proteus.void(proteus.ptr)) """ ) assert missing_kernel.returncode != 0 diff --git a/bindings/python/tests/test_gpu_mlir_smoke.py b/bindings/python/tests/test_gpu_mlir_smoke.py index e478ef420..325f253a3 100644 --- a/bindings/python/tests/test_gpu_mlir_smoke.py +++ b/bindings/python/tests/test_gpu_mlir_smoke.py @@ -19,7 +19,7 @@ def main(): mod = proteus.compile(source, frontend="mlir", target=target, verify=True) assert mod.get_function_address("write42") != 0 - kernel = mod.get_kernel("write42", [proteus.ptr]) + kernel = mod.get_kernel("write42", signature=proteus.void(proteus.ptr)) kernel.launch(grid=1, block=1, args=[device_value]) runtime.sync() assert runtime.copy_d2h_i32(device_value) == 42 @@ -28,7 +28,9 @@ def main(): path = pathlib.Path(tmpdir) / "kernel.mlir" path.write_text(source) file_mod = proteus.compile(path, frontend="mlir", target=target, verify=True) - file_kernel = file_mod.get_kernel("write42", [proteus.ptr]) + file_kernel = file_mod.get_kernel( + "write42", signature=proteus.void(proteus.ptr) + ) runtime.copy_h2d_i32(device_value, 0) file_kernel.launch(grid=[1], block=(1, 1, 1), args=[int(device_value)]) runtime.sync() diff --git a/bindings/python/tests/test_gpu_mlir_validation.py b/bindings/python/tests/test_gpu_mlir_validation.py index de3deeb14..5a3345bb8 100644 --- a/bindings/python/tests/test_gpu_mlir_validation.py +++ b/bindings/python/tests/test_gpu_mlir_validation.py @@ -30,12 +30,18 @@ def main(): runtime = create_runtime() device_value = runtime.malloc_i32() mod = proteus.compile(source, frontend="mlir", target=target, verify=True) - kernel = mod.get_kernel("write42", [proteus.ptr]) + kernel = mod.get_kernel("write42", signature=proteus.void(proteus.ptr)) + + expect_raises( + TypeError, + lambda: mod.get_kernel("write42", signature=proteus.i32(proteus.ptr)), + "kernel signatures must return proteus.void", + ) expect_raises( TypeError, lambda: kernel.launch(grid=1, block=1, args=[]), - "kernel argument count does not match argtypes", + "kernel argument count does not match signature", ) expect_raises( TypeError, diff --git a/bindings/python/tests/test_host_cpp_smoke.py b/bindings/python/tests/test_host_cpp_smoke.py index 87c461a9f..a845a9bed 100644 --- a/bindings/python/tests/test_host_cpp_smoke.py +++ b/bindings/python/tests/test_host_cpp_smoke.py @@ -9,8 +9,11 @@ def main(): mod = proteus.compile(source, frontend="cpp", target="host") assert mod.get_function_address("forty_two") != 0 - plus1 = mod.get_function("plus1", restype=proteus.i32, argtypes=[proteus.i32]) - assert repr(plus1) == "" + plus1 = mod.get_function("plus1", signature=proteus.i32(proteus.i32)) + assert ( + repr(plus1) + == "" + ) assert plus1(41) == 42 assert proteus.compile(source, frontend="cpp").get_function_address("forty_two") != 0 @@ -19,7 +22,9 @@ def main(): path.write_text(source) mod = proteus.compile(path, frontend="cpp", target="host") assert mod.get_function_address("forty_two") != 0 - assert mod.get_function("plus1", restype=proteus.i32, argtypes=[proteus.i32])(41) == 42 + assert mod.get_function( + "plus1", signature=proteus.i32(proteus.i32) + )(41) == 42 print("python_host_cpp_smoke: ok") diff --git a/bindings/python/tests/test_host_cpp_validation.py b/bindings/python/tests/test_host_cpp_validation.py index 49347bdeb..12190da0c 100644 --- a/bindings/python/tests/test_host_cpp_validation.py +++ b/bindings/python/tests/test_host_cpp_validation.py @@ -6,6 +6,7 @@ def main(): + assert repr(proteus.void) == "proteus.void" assert repr(proteus.i8) == "proteus.i8" assert repr(proteus.i32) == "proteus.i32" assert repr(proteus.i64) == "proteus.i64" @@ -14,6 +15,72 @@ def main(): assert repr(proteus.f32) == "proteus.f32" assert repr(proteus.f64) == "proteus.f64" assert repr(proteus.ptr) == "proteus.ptr" + + arg_descriptors = ( + proteus.i8, + proteus.i32, + proteus.i64, + proteus.u32, + proteus.u64, + proteus.f32, + proteus.f64, + proteus.ptr, + ) + for descriptor in (proteus.void, *arg_descriptors): + assert isinstance(descriptor, proteus.Type) + assert isinstance(descriptor(), proteus.Signature) + + signature = proteus.f64(proteus.f64, proteus.f64) + assert repr(signature.restype) == "proteus.f64" + assert tuple(map(repr, signature.argtypes)) == ( + "proteus.f64", + "proteus.f64", + ) + assert isinstance(signature.argtypes, tuple) + assert repr(signature) == "proteus.f64(proteus.f64, proteus.f64)" + + zero_arg_signature = proteus.i32() + assert repr(zero_arg_signature.restype) == "proteus.i32" + assert zero_arg_signature.argtypes == () + assert repr(zero_arg_signature) == "proteus.i32()" + + all_arg_signature = proteus.void(*arg_descriptors) + assert repr(all_arg_signature.restype) == "proteus.void" + assert tuple(map(repr, all_arg_signature.argtypes)) == tuple( + map(repr, arg_descriptors) + ) + + expect_raises(TypeError, lambda: proteus.Signature()) + expect_raises( + AttributeError, + lambda: setattr(signature, "restype", proteus.i32), + ) + expect_raises( + AttributeError, + lambda: setattr(signature, "argtypes", (proteus.i32,)), + ) + + invalid_argtypes = ( + int, + float, + ctypes.c_int32, + ctypes.c_int32(), + object(), + 1, + None, + ) + for invalid in invalid_argtypes: + expect_raises( + TypeError, + lambda invalid=invalid: proteus.i32(invalid), + "signature arguments must be Proteus type descriptors", + ) + expect_raises( + TypeError, + lambda: proteus.i32(proteus.void), + "proteus.void is only valid as a signature return type", + ) + assert isinstance(proteus.has_cuda, bool) assert isinstance(proteus.has_hip, bool) assert isinstance(proteus.has_mlir, bool) @@ -45,26 +112,112 @@ def main(): ), "MLIR frontend does not support extra_args", ) - plus1 = proteus.compile( + mod = proteus.compile( r''' -extern "C" int plus1(int x) { return x + 1; } -extern "C" int load0(const int *xs) { return xs[0]; } -extern "C" void store0(int *xs, int value) { xs[0] = value; } +#include + +extern "C" std::int32_t forty_two() { return 42; } +extern "C" std::int32_t plus1(std::int32_t x) { return x + 1; } +extern "C" std::int8_t echo_i8(std::int8_t x) { return x; } +extern "C" std::int32_t echo_i32(std::int32_t x) { return x; } +extern "C" std::int64_t echo_i64(std::int64_t x) { return x; } +extern "C" std::uint32_t echo_u32(std::uint32_t x) { return x; } +extern "C" std::uint64_t echo_u64(std::uint64_t x) { return x; } +extern "C" float echo_f32(float x) { return x; } +extern "C" double add_f64(double x, double y) { return x + y; } +extern "C" void *echo_ptr(void *x) { return x; } +extern "C" std::int32_t load0(const std::int32_t *xs) { return xs[0]; } +extern "C" void store0(std::int32_t *xs, std::int32_t value) { xs[0] = value; } ''', frontend="cpp", target="host", ) - plus1_fn = plus1.get_function("plus1", restype=proteus.i32, argtypes=[proteus.i32]) - load0 = plus1.get_function("load0", restype=proteus.i32, argtypes=[proteus.ptr]) - store0 = plus1.get_function( - "store0", restype=None, argtypes=[proteus.ptr, proteus.i32] + + expect_raises( + TypeError, + lambda: mod.get_function("plus1", proteus.i32(proteus.i32)), + ) + expect_raises( + TypeError, + lambda: mod.get_function( + "plus1", restype=proteus.i32, argtypes=[proteus.i32] + ), + ) + expect_raises( + TypeError, + lambda: mod.get_function("plus1", signature=int), + ) + expect_raises( + TypeError, + lambda: mod.get_kernel( + "missing_kernel", signature=proteus.i32(proteus.i32) + ), + "kernel signatures must return proteus.void", + ) + expect_raises( + TypeError, + lambda: mod.get_kernel("plus1", proteus.void(proteus.i32)), + ) + expect_raises( + TypeError, + lambda: mod.get_kernel("plus1", argtypes=[proteus.i32]), + ) + + forty_two = mod.get_function("forty_two", signature=proteus.i32()) + plus1_fn = mod.get_function( + "plus1", signature=proteus.i32(proteus.i32) ) - assert repr(plus1_fn) == "" + echo_i8 = mod.get_function("echo_i8", signature=proteus.i8(proteus.i8)) + echo_i32 = mod.get_function( + "echo_i32", signature=proteus.i32(proteus.i32) + ) + echo_i64 = mod.get_function( + "echo_i64", signature=proteus.i64(proteus.i64) + ) + echo_u32 = mod.get_function( + "echo_u32", signature=proteus.u32(proteus.u32) + ) + echo_u64 = mod.get_function( + "echo_u64", signature=proteus.u64(proteus.u64) + ) + echo_f32 = mod.get_function( + "echo_f32", signature=proteus.f32(proteus.f32) + ) + add_f64 = mod.get_function( + "add_f64", signature=proteus.f64(proteus.f64, proteus.f64) + ) + echo_ptr = mod.get_function( + "echo_ptr", signature=proteus.ptr(proteus.ptr) + ) + load0 = mod.get_function("load0", signature=proteus.i32(proteus.ptr)) + store0 = mod.get_function( + "store0", signature=proteus.void(proteus.ptr, proteus.i32) + ) + assert ( + repr(plus1_fn) + == "" + ) + assert forty_two() == 42 assert plus1_fn(41) == 42 + expect_raises( + TypeError, + lambda: plus1_fn(), + "function argument count does not match signature", + ) + + assert echo_i8(-(2**7)) == -(2**7) + assert echo_i32(-(2**31)) == -(2**31) + assert echo_i64(-(2**63)) == -(2**63) + assert echo_u32(2**32 - 1) == 2**32 - 1 + assert echo_u64(2**64 - 1) == 2**64 - 1 + assert echo_f32(1.5) == 1.5 + assert add_f64(1.25, 2.5) == 3.75 values = (ctypes.c_int * 2)(11, 22) + assert echo_ptr(values) == ctypes.addressof(values) + assert echo_ptr(None) is None assert load0(values) == 11 - store0(values, 33) + assert store0(values, 33) is None assert values[0] == 33 scalar = ctypes.c_int(44) diff --git a/bindings/python/tests/test_host_llvmir_smoke.py b/bindings/python/tests/test_host_llvmir_smoke.py index 6f0a4b33b..6ebd25bd6 100644 --- a/bindings/python/tests/test_host_llvmir_smoke.py +++ b/bindings/python/tests/test_host_llvmir_smoke.py @@ -20,8 +20,11 @@ def main(): mod = proteus.compile(source, frontend="llvmir", target="host") assert mod.get_function_address("forty_two") != 0 - plus1 = mod.get_function("plus1", restype=proteus.i32, argtypes=[proteus.i32]) - assert repr(plus1) == "" + plus1 = mod.get_function("plus1", signature=proteus.i32(proteus.i32)) + assert ( + repr(plus1) + == "" + ) assert plus1(41) == 42 assert proteus.compile(source, frontend="llvmir").get_function_address("forty_two") != 0 @@ -30,7 +33,9 @@ def main(): path.write_text(source) mod = proteus.compile(path, frontend="llvmir", target="host") assert mod.get_function_address("forty_two") != 0 - assert mod.get_function("plus1", restype=proteus.i32, argtypes=[proteus.i32])(41) == 42 + assert mod.get_function( + "plus1", signature=proteus.i32(proteus.i32) + )(41) == 42 print("python_host_llvmir_smoke: ok") diff --git a/bindings/python/tests/test_proteus_init_exports.py b/bindings/python/tests/test_proteus_init_exports.py index a78081779..289846ddc 100644 --- a/bindings/python/tests/test_proteus_init_exports.py +++ b/bindings/python/tests/test_proteus_init_exports.py @@ -20,6 +20,8 @@ def main(): native_module.__file__ = "/tmp/fake_proteus_backend.so" native_module.compile = lambda *args, **kwargs: None native_module.i32 = object() + native_module.void = object() + native_module.Signature = type("Signature", (), {}) native_module.has_cuda = False native_module.has_hip = False @@ -44,6 +46,10 @@ def main(): assert namespace["active_backend_variant"] == "host_llvm22" assert "compile" in namespace assert "i32" in namespace + assert "void" in namespace + assert "Signature" in namespace + assert exported.void is native_module.void + assert exported.Signature is native_module.Signature assert exported.__doc__ == "fake native module" assert exported.__file__ == "/tmp/fake_proteus_backend.so" diff --git a/bindings/python/tests/test_wildcard_import.py b/bindings/python/tests/test_wildcard_import.py index 573305d55..e570fa69c 100644 --- a/bindings/python/tests/test_wildcard_import.py +++ b/bindings/python/tests/test_wildcard_import.py @@ -5,6 +5,8 @@ def main(): assert namespace["active_backend_variant"], namespace assert callable(namespace["compile"]), namespace.keys() assert "i32" in namespace, namespace.keys() + assert "void" in namespace, namespace.keys() + assert "Signature" in namespace, namespace.keys() print("test_wildcard_import: ok") diff --git a/docs/user/llvmir-frontend.md b/docs/user/llvmir-frontend.md index c57b785aa..287726b6c 100644 --- a/docs/user/llvmir-frontend.md +++ b/docs/user/llvmir-frontend.md @@ -126,9 +126,12 @@ entry: """ mod = proteus.compile(source, frontend="llvmir", target="host") -plus1 = mod.get_function("plus1", restype=proteus.i32, argtypes=[proteus.i32]) +plus1 = mod.get_function("plus1", signature=proteus.i32(proteus.i32)) assert plus1(41) == 42 ``` In the current Python API, LLVM IR input is provided as text through a string or a path to a `.ll` file. +See [Python Bindings](python-bindings.md) for native signature descriptors, +GPU kernel lookup, +and ABI matching requirements. diff --git a/docs/user/python-bindings.md b/docs/user/python-bindings.md new file mode 100644 index 000000000..702dbdc59 --- /dev/null +++ b/docs/user/python-bindings.md @@ -0,0 +1,96 @@ +# Python Bindings + +The Python API compiles C++, LLVM IR, or MLIR source with `proteus.compile()`. +After compilation, +retrieve a native entry point by supplying its exact ABI signature. + +## Host Functions + +Build a signature by calling the return-type descriptor with the argument-type +descriptors: + +```python +import proteus + +source = r''' +extern "C" int plus1(int x) { + return x + 1; +} +''' + +mod = proteus.compile(source, frontend="cpp", target="host") +plus1 = mod.get_function( + "plus1", + signature=proteus.i32(proteus.i32), +) + +assert plus1(41) == 42 +``` + +The `signature` parameter is keyword-only. +A function with no arguments uses an empty descriptor call, +such as `proteus.i32()`. +Host functions may return any supported scalar or pointer descriptor, +or `proteus.void`. + +## GPU Kernels + +CUDA and HIP kernel signatures must return `proteus.void`: + +```python +kernel = mod.get_kernel( + "write_int", + signature=proteus.void(proteus.ptr, proteus.i32), +) + +kernel.launch( + grid=1, + block=256, + args=[device_output, 42], +) +``` + +`proteus.void` is valid only as a return descriptor. +It cannot appear in the argument list. + +## Supported ABI Descriptors + +Signatures preserve the exact width and signedness of every descriptor: + +| Descriptor | Native ABI type | +| --- | --- | +| `proteus.i8` | signed 8-bit integer | +| `proteus.i32` | signed 32-bit integer | +| `proteus.i64` | signed 64-bit integer | +| `proteus.u32` | unsigned 32-bit integer | +| `proteus.u64` | unsigned 64-bit integer | +| `proteus.f32` | 32-bit floating point | +| `proteus.f64` | 64-bit floating point | +| `proteus.ptr` | opaque pointer | +| `proteus.void` | return only; no value | + +A constructed `proteus.Signature` exposes read-only `restype` and `argtypes` +properties. +`argtypes` is always a tuple. +For example, +`repr(proteus.f64(proteus.f64, proteus.f64))` is +`proteus.f64(proteus.f64, proteus.f64)`. + +Proteus does not infer native types from Python builtins, +function annotations, +decorators, +or runtime argument values. +Python types such as `int` and `float`, +and `ctypes` types, +are not valid signature descriptors. + +!!! warning + + Proteus cannot verify that a supplied signature matches the compiled + native symbol. + A mismatch in return type, + argument order, + width, + signedness, + or pointer position can cross the native ABI incorrectly and cause + corrupted results or a process failure. diff --git a/mkdocs.yml b/mkdocs.yml index 9c754449c..36630dd1f 100644 --- a/mkdocs.yml +++ b/mkdocs.yml @@ -9,6 +9,7 @@ nav: - User Interface: - Overview: user/interface.md - Code Annotations: user/annotations.md + - Python Bindings: user/python-bindings.md - C++ Frontend API: user/cpp-frontend.md - LLVM IR Frontend API: user/llvmir-frontend.md - MLIR Frontend API: user/mlir-frontend.md From 118cb01f10a0a05d2f258b68de384178510fa51d Mon Sep 17 00:00:00 2001 From: Zane Fink Date: Mon, 27 Jul 2026 15:37:02 -0700 Subject: [PATCH 2/2] fix: Make Python type descriptors comparable --- bindings/python/PythonBindings.cpp | 12 ++++++++++- .../python/tests/test_host_cpp_validation.py | 20 +++++++++---------- 2 files changed, 21 insertions(+), 11 deletions(-) diff --git a/bindings/python/PythonBindings.cpp b/bindings/python/PythonBindings.cpp index bd43bb549..4577bcc5b 100644 --- a/bindings/python/PythonBindings.cpp +++ b/bindings/python/PythonBindings.cpp @@ -437,7 +437,17 @@ PYBIND11_MODULE(_proteus, M) { // Expose the builtin scalar/pointer descriptors as module-level singletons. py::class_(M, "Type") .def("__call__", &makeSignature) - .def("__repr__", &typeRepr); + .def("__repr__", &typeRepr) + .def( + "__eq__", + [](const Type &LHS, const Type &RHS) { return LHS.Kind == RHS.Kind; }, + py::is_operator()) + .def( + "__ne__", + [](const Type &LHS, const Type &RHS) { return LHS.Kind != RHS.Kind; }, + py::is_operator()) + .def("__hash__", + [](const Type &T) { return static_cast(T.Kind); }); py::class_(M, "Signature") .def_property_readonly("restype", diff --git a/bindings/python/tests/test_host_cpp_validation.py b/bindings/python/tests/test_host_cpp_validation.py index 12190da0c..f82947718 100644 --- a/bindings/python/tests/test_host_cpp_validation.py +++ b/bindings/python/tests/test_host_cpp_validation.py @@ -29,26 +29,26 @@ def main(): for descriptor in (proteus.void, *arg_descriptors): assert isinstance(descriptor, proteus.Type) assert isinstance(descriptor(), proteus.Signature) + assert descriptor == descriptor + assert proteus.i32 != proteus.i64 + assert proteus.i32 != object() signature = proteus.f64(proteus.f64, proteus.f64) - assert repr(signature.restype) == "proteus.f64" - assert tuple(map(repr, signature.argtypes)) == ( - "proteus.f64", - "proteus.f64", - ) + assert signature.restype == proteus.f64 + assert signature.argtypes == (proteus.f64, proteus.f64) + assert hash(signature.restype) == hash(proteus.f64) + assert len({signature.restype, proteus.f64}) == 1 assert isinstance(signature.argtypes, tuple) assert repr(signature) == "proteus.f64(proteus.f64, proteus.f64)" zero_arg_signature = proteus.i32() - assert repr(zero_arg_signature.restype) == "proteus.i32" + assert zero_arg_signature.restype == proteus.i32 assert zero_arg_signature.argtypes == () assert repr(zero_arg_signature) == "proteus.i32()" all_arg_signature = proteus.void(*arg_descriptors) - assert repr(all_arg_signature.restype) == "proteus.void" - assert tuple(map(repr, all_arg_signature.argtypes)) == tuple( - map(repr, arg_descriptors) - ) + assert all_arg_signature.restype == proteus.void + assert all_arg_signature.argtypes == arg_descriptors expect_raises(TypeError, lambda: proteus.Signature()) expect_raises(