From a749eef07fd9e7d205956e0dc546501aecd50c0f Mon Sep 17 00:00:00 2001 From: Andre Masella Date: Fri, 21 Jul 2023 15:45:59 -0400 Subject: [PATCH] Add stack allocation handles Create handles for simple stack allocation and allocation of an array. --- dispyatcher/allocation.py | 93 +++++++++++++++++++++++++++++++++++++++ dispyatcher/cpython.py | 7 ++- docs/index.rst | 3 ++ tests/test_cpython.py | 14 ++++-- 4 files changed, 113 insertions(+), 4 deletions(-) create mode 100644 dispyatcher/allocation.py diff --git a/dispyatcher/allocation.py b/dispyatcher/allocation.py new file mode 100644 index 0000000..9185e47 --- /dev/null +++ b/dispyatcher/allocation.py @@ -0,0 +1,93 @@ +import itertools +import llvmlite.ir +from llvmlite.ir import Value as IRValue +from typing import Generic, Tuple, Sequence, Union + +from dispyatcher import F, Handle, Type, ReturnManagement, ArgumentManagement, TemporaryValue, Pointer +from dispyatcher.general import UncheckedArray + + +class ToStackPointer(Handle[F], Generic[F]): + """ + Creates a stack allocation for a value and provides a pointer to it. + + The allocation can take ownership of the value and the value will be freed when the pointer is dropped. + """ + __type: Type + __transfer: ReturnManagement + + def __init__(self, ty: Type, transfer: ReturnManagement = ReturnManagement.TRANSFER): + """ + Create a new stack allocation handle for a particular type. + + :param ty: the type the allocation contains; the result will be a pointer to this type + :param transfer: whether to move the value into the allocation (``TRANSFER``) or copy it and rely on the + original value for memory management (``BORROW``). + """ + super().__init__() + self.__type = ty + self.__transfer = transfer + + def __str__(self) -> str: + return f"ToStackPointer[{self.__transfer.name}] → {self.__type}" + + def generate_handle_ir(self, flow: F, args: Sequence[IRValue]) -> Union[TemporaryValue, IRValue]: + ty = self.__type.machine_type() + (arg, ) = args + alloc = flow.builder.alloca(ty) + flow.builder.store(arg, alloc) + return alloc + + def handle_arguments(self) -> Sequence[Tuple[Type, ArgumentManagement]]: + return (self.__type, + ArgumentManagement.TRANSFER_CAPTURE_PARENTS + if self.__transfer == ReturnManagement.TRANSFER + else ArgumentManagement.BORROW_CAPTURE_PARENTS), + + def handle_return(self) -> Tuple[Type, ReturnManagement]: + return Pointer(self.__type), self.__transfer + + +class CollectIntoArray(Handle[F], Generic[F]): + """ + Stores a arguments into a stack-allocated fixed-length array. + + This is meant for producing arrays like ``argv`` from individual values + """ + __type: Type + __count: int + __null_terminated: bool + + def __init__(self, ty: Type, count: int, null_terminated: bool = False): + """ + Creates a new stack-allocated array collector handle. + + :param ty: the type of the elements in the array + :param count: the number of items to collect + :param null_terminated: if true, the array will be allocated with an extra slot at the end that will have a null + or zero value; otherwise, the array will only be filled with the arguments. + """ + super().__init__() + assert count > 0, "Array length must be positive number" + self.__type = ty + self.__count = count + self.__null_terminated = null_terminated + + def __str__(self) -> str: + return f"CollectIntoArray({self.__type} * {self.__count} {'+null' if self.__null_terminated else ''})" + + def generate_handle_ir(self, flow: F, args: Sequence[IRValue]) -> Union[TemporaryValue, IRValue]: + i32 = llvmlite.ir.IntType(32) + array = flow.builder.alloca(llvmlite.ir.ArrayType(self.__type.machine_type(), + self.__count + self.__null_terminated)) + for idx, arg in enumerate(args): + flow.builder.store(arg, flow.builder.gep(array, [i32(0), i32(idx)])) + if self.__null_terminated: + flow.builder.store(self.__type.machine_type()(None), flow.builder.gep(array, [i32(0), i32(self.__count)])) + return flow.builder.bitcast(array, self.__type.machine_type().as_pointer()) + + def handle_arguments(self) -> Sequence[Tuple[Type, ArgumentManagement]]: + return list(itertools.repeat((self.__type, ArgumentManagement.BORROW_CAPTURE), self.__count)) + + def handle_return(self) -> Tuple[Type, ReturnManagement]: + return UncheckedArray(self.__type), ReturnManagement.BORROW diff --git a/dispyatcher/cpython.py b/dispyatcher/cpython.py index 8841bd2..80327a6 100644 --- a/dispyatcher/cpython.py +++ b/dispyatcher/cpython.py @@ -9,12 +9,13 @@ from dispyatcher import ArgumentManagement, BaseTransferUnaryHandle, ControlFlow, DerefPointer, F, FlowState, Handle, \ Identity, Type, ReturnManagement, TemporaryValue from dispyatcher.accessors import GetElementPointer -from dispyatcher.general import BaseIndirectFunction, CurrentProcessFunction, MachineType +from dispyatcher.general import BaseIndirectFunction, CurrentProcessFunction, MachineType, UncheckedArray from dispyatcher.permute import implode_args from dispyatcher.repacking import Repacker, RepackingDispatcher, RepackingState INT_RESULT_TYPE = llvmlite.ir.types.IntType(ctypes.sizeof(ctypes.c_int) * 8) SIZE_T_TYPE = llvmlite.ir.types.IntType(ctypes.sizeof(ctypes.c_size_t) * 8) +CHAR_ARRAY = UncheckedArray(MachineType(llvmlite.ir.IntType(8))) class PythonControlFlow(dispyatcher.ControlFlow): @@ -804,3 +805,7 @@ def handle_return(self) -> Tuple[Type, ReturnManagement]: ReturnManagement.TRANSFER, "PyDict_Size", (PY_DICT_TYPE, ArgumentManagement.BORROW_TRANSIENT)) +PY_UNICODE_FROM_STRING = CurrentProcessFunction(PyObjectType(str), + ReturnManagement.TRANSFER, + "PyUnicode_FromString", + (CHAR_ARRAY, ArgumentManagement.BORROW_TRANSIENT)) diff --git a/docs/index.rst b/docs/index.rst index 8f93df3..c4bb35e 100644 --- a/docs/index.rst +++ b/docs/index.rst @@ -48,6 +48,9 @@ Once a handle is constructed, creating a ``Callsite`` will compile it and turn i .. automodule:: dispyatcher.accessors :members: +.. automodule:: dispyatcher.allocation + :members: + .. automodule:: dispyatcher.cpython :members: diff --git a/tests/test_cpython.py b/tests/test_cpython.py index 0aef6a5..b0e75b8 100644 --- a/tests/test_cpython.py +++ b/tests/test_cpython.py @@ -6,14 +6,15 @@ import dispyatcher import dispyatcher.general from dispyatcher import CallSite, Identity, IgnoreArguments -from dispyatcher.general import SimpleConstant +from dispyatcher.allocation import CollectIntoArray +from dispyatcher.general import MachineType, SimpleConstant from dispyatcher.cpython import TupleUnpackingDispatcher, PY_OBJECT_TYPE, PythonControlFlowType class RepackTests(unittest.TestCase): def test_tuple_unpack_int(self): - i32 = dispyatcher.general.MachineType(llvmlite.ir.IntType(32)) + i32 = MachineType(llvmlite.ir.IntType(32)) tud = TupleUnpackingDispatcher(IgnoreArguments(SimpleConstant(i32, -1), 0, PY_OBJECT_TYPE), 0) tud.append(Identity(i32), None) callsite = CallSite(tud, PythonControlFlowType()) @@ -24,7 +25,7 @@ def test_tuple_unpack_int(self): self.assertEqual(callsite(ctypes.py_object((7.5,))), -1) def test_tuple_unpack_double(self): - dbl = dispyatcher.general.MachineType(llvmlite.ir.DoubleType()) + dbl = MachineType(llvmlite.ir.DoubleType()) tud = TupleUnpackingDispatcher(IgnoreArguments(SimpleConstant(dbl, -1), 0, PY_OBJECT_TYPE), 0) tud.append(Identity(dbl), None) callsite = CallSite(tud, PythonControlFlowType()) @@ -61,3 +62,10 @@ def test_value(self): self_handle = dispyatcher.cpython.Value(self, transfer=dispyatcher.ReturnManagement.TRANSFER) callsite = CallSite(self_handle, PythonControlFlowType()) self.assertEqual(self, callsite()) + + def test_array_collection(self): + # This is really an allocation test, but we use Python infrastructure to make it pleasant + i8 = MachineType(llvmlite.ir.IntType(8)) + handle = CollectIntoArray(i8, 2, null_terminated=True) + dispyatcher.cpython.PY_UNICODE_FROM_STRING + callsite = CallSite(handle, PythonControlFlowType()) + self.assertEqual(callsite(ord('h'), ord('i')), "hi")