Skip to content

Api

Modules:

Name Description
array
buffer
buffer_cuda
dtypes
kernels
transpiler
utils

Classes:

Name Description
Array
Buffer
BufferCuda
PythonModule

Functions:

Name Description
compile_str

Transpile a python function into a C source code string.
ref

Equal to doing &x in C. Returning [x] is just for type hinting purposes.

Attributes:

Name Type Description
element_wise_module

element_wise_module = element_wise_module.compile() module-attribute

Array

Methods:

Name Description
__getitem__
__init__
__repr__
arange
binary_op
from_iterable
reduction_op
reshape
squeeze
to_python
transpose
unary_op

Attributes:

Name Type Description
T 'Array'
__add__
__div__
__mul__
__sub__
__truediv__
acos
asin
atan
atan2
cos
cosh
data
device Device
exp
exp2
is_contiguous
log
log10
log2
mT 'Array'
max
min
ndim
offset
prod
relu
shape
sin
sinh
size int
sqrt
square
strides
sum
tan
tanh
Source code in src/simplendarray/array.py 
 61
 62
 63
 64
 65
 66
 67
 68
 69
 70
 71
 72
 73
 74
 75
 76
 77
 78
 79
 80
 81
 82
 83
 84
 85
 86
 87
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
class Array:
    @classmethod
    def from_iterable(cls, data: NestedIterable | Scalar, dtype: str | type[DType], device: Device = "cpu"):
        flat, shape = flatten_and_get_shape(data)
        _typecode = typecode(get_dtype(dtype))
        buffer = buf_cls[device].from_iterable(flat, _typecode)
        strides = contiguous_strides(shape)
        offset = 0
        return cls(buffer, shape, strides, offset)

    @classmethod
    def arange(cls, numel: int, dtype: str | type[DType], device: Device = "cpu") -> "Array":
        if numel < 0:
            raise ValueError("numel must be >= 0")
        _typecode = typecode(get_dtype(dtype))
        buffer = buf_cls[device].empty(numel, _typecode)
        dispatch_arange(buffer, 0, 1, numel)
        return cls(buffer, (numel,), (1,), 0)

    def __init__(self, buffer: BufferType, shape: tuple[int, ...], strides: tuple[int, ...], offset: int):
        self.data = buffer
        self.shape = shape
        self.strides = strides
        self.offset = offset

    @property
    def ndim(self):
        return len(self.shape)

    @property
    def device(self) -> Device:
        return self.data.device

    @property
    def size(self) -> int:
        return product(self.shape)

    @property
    def is_contiguous(self):
        if self.ndim == 0 or 0 in self.shape:
            return True
        real_shape = tuple(x for x in self.shape if x > 1)
        real_stride = tuple(y for x, y in zip(self.shape, self.strides) if x > 1)
        return real_stride == contiguous_strides(real_shape)

    def __repr__(self) -> str:
        return f"Array({self.to_python()}, shape={self.shape}, strides={self.strides}, offset={self.offset})"

    def to_python(self) -> NestedIterable:
        if self.ndim == 0:
            return self.data.data[self.offset]
        nested = []
        for i in range(self.shape[0]):
            # Does [self[0, :, :, ...], self[1, :, :, ...], ..., self[shape[0] - 1, :, :, ...]]
            # Then recursively calls to_python on each of these children, until base case reached
            indexed = self[i, *(slice(None) for _ in range(self.ndim - 1))].squeeze(0)
            nested.append(indexed.to_python())
        return nested

    def squeeze(self, dims: int | Iterable[int]) -> "Array":
        if isinstance(dims, int):
            dims = [dims]
        if self.ndim == 0 or any(self.shape[d] != 1 for d in dims):
            raise ValueError("Can only squeeze a non scalar array with length 1 dims, but shape is", self.shape)
        dims_set = set(dims)
        new_shape = tuple(s for i, s in enumerate(self.shape) if i not in dims_set)
        new_strides = tuple(stride for i, stride in enumerate(self.strides) if i not in dims_set)
        return Array(self.data, new_shape, new_strides, self.offset)

    def __getitem__(self, items: tuple[int | slice, ...] | int | slice):
        if not isinstance(items, tuple):
            items = (items,)
        if len(items) != self.ndim:
            raise ValueError("Must index the same number of dimensions as the array")
        new_shape = []
        new_strides = []
        new_offset = self.offset
        for shape, stride, item in zip(self.shape, self.strides, items):
            if isinstance(item, int):
                item = slice(item, item + 1).indices(shape)[0]
                new_shape.append(1)
                new_strides.append(0)
                new_offset += stride * item
            elif isinstance(item, slice):
                start, stop, step = item.indices(shape)
                if step > 0:
                    # Num elements in [start, stop) = stop - start
                    new_shape.append(max(0, ceildiv(stop - start, step)))
                else:
                    new_shape.append(max(0, ceildiv(start - stop, -step)))
                new_strides.append(stride * step)
                new_offset += stride * start
            else:
                raise TypeError(f"Index must be int or slice, got {type(item).__name__}")
        return Array(self.data, tuple(new_shape), tuple(new_strides), new_offset)

    def transpose(self, dims: Iterable[int]) -> "Array":
        new_shape = tuple(self.shape[i] for i in dims)
        new_strides = tuple(self.strides[i] for i in dims)
        return Array(self.data, new_shape, new_strides, self.offset)

    @property
    def T(self) -> "Array":
        return self.transpose(range(self.ndim - 1, -1, -1))

    @property
    def mT(self) -> "Array":
        if self.ndim < 2:
            raise ValueError("matrix transpose with ndim < 2 is undefined")
        dims = tuple(range(self.ndim - 2)) + (-1, -2)
        return self.transpose(dims)

    def reshape(self, new_shape: int | tuple[int, ...]) -> "Array":
        numel = product(self.shape)
        if isinstance(new_shape, int):
            new_shape = (new_shape,)
        num_n1s = sum(1 for s in new_shape if s == -1)
        for s in new_shape:
            if s < -1:
                raise ValueError("Dimension must be non-negative, or with a single -1")
        if num_n1s > 1:
            raise ValueError("only one dimension can be -1")
        if num_n1s == 1:
            known = product(x for x in new_shape if x != -1)
            if numel % known != 0:
                raise ValueError(
                    f"Cannot reshape array of size {numel} into shape "
                    f"{tuple(numel // known if s == -1 else s for s in new_shape)}"
                )
            replacement = numel // known
            new_shape = tuple(replacement if s == -1 else s for s in new_shape)
        new_strides = reshape_strides(self.shape, new_shape, self.strides)
        if new_strides is None:
            out_ndim = len(new_shape)
            new_strides = contiguous_strides(new_shape)
            new_buffer = buf_cls[self.device].empty(numel, self.data.typecode)
            inp_shape_buffer = buf_cls[self.device].from_iterable(self.shape, "l")
            inp_strides_buffer = buf_cls[self.device].from_iterable(self.strides, "l")
            out_shape_buffer = buf_cls[self.device].from_iterable(new_shape, "l")
            out_strides_buffer = buf_cls[self.device].from_iterable(new_strides, "l")
            inp_work_buffer = buf_cls[self.device].empty(self.ndim, "l")
            out_work_buffer = buf_cls[self.device].empty(out_ndim, "l")

            dispatch_reshape_copy(
                self,
                new_shape,
                new_buffer,
                inp_shape_buffer,
                inp_strides_buffer,
                inp_work_buffer,
                out_shape_buffer,
                out_strides_buffer,
                out_work_buffer,
            )
            return Array(new_buffer, new_shape, new_strides, 0)
        return Array(self.data, new_shape, new_strides, self.offset)

    @staticmethod
    def unary_op(op: str):
        def fn(self: Array):
            buf = buf_cls[self.device].empty(product(self.shape), self.data.typecode)
            out = Array(buf, self.shape, contiguous_strides(self.shape), 0)
            dispatch_element_wise_unary(self, out, op)
            return out

        return fn

    relu = unary_op("relu")
    exp = unary_op("exp")
    exp2 = unary_op("exp2")
    log = unary_op("log")
    log2 = unary_op("log2")
    log10 = unary_op("log10")
    relu = unary_op("relu")
    square = unary_op("square")
    sqrt = unary_op("sqrt")
    sin = unary_op("sin")
    cos = unary_op("cos")
    tan = unary_op("tan")
    asin = unary_op("asin")
    acos = unary_op("acos")
    atan = unary_op("atan")
    sinh = unary_op("sinh")
    cosh = unary_op("cosh")
    tanh = unary_op("tanh")

    @staticmethod
    def binary_op(op: str):
        def fn(self: Array, other: Array) -> Array:
            buf = buf_cls[self.device].empty(product(self.shape), self.data.typecode)
            out = Array(buf, self.shape, contiguous_strides(self.shape), 0)
            dispatch_element_wise_binary(self, other, out, op)
            return out

        return fn

    __add__ = binary_op("add")
    __sub__ = binary_op("sub")
    __mul__ = binary_op("mul")
    __div__ = binary_op("div")
    __truediv__ = binary_op("div")
    atan2 = binary_op("atan2")

    @staticmethod
    def reduction_op(op: str):
        def fn(self: Array, dims: tuple[int]):
            reduction_size = product(x for i, x in enumerate(self.shape) if i in dims)
            buf = buf_cls[self.device].empty(self.size // reduction_size, self.data.typecode)
            out = Array(buf, (self.size // reduction_size,), (1,), 0)
            dispatch_reduction(self, out, op, dims)
            return out

        return fn

    sum = reduction_op("add")
    min = reduction_op("min")
    max = reduction_op("max")
    prod = reduction_op("mul")

T: 'Array' property

__add__ = binary_op('add') class-attribute instance-attribute

__div__ = binary_op('div') class-attribute instance-attribute

__mul__ = binary_op('mul') class-attribute instance-attribute

__sub__ = binary_op('sub') class-attribute instance-attribute

__truediv__ = binary_op('div') class-attribute instance-attribute

acos = unary_op('acos') class-attribute instance-attribute

asin = unary_op('asin') class-attribute instance-attribute

atan = unary_op('atan') class-attribute instance-attribute

atan2 = binary_op('atan2') class-attribute instance-attribute

cos = unary_op('cos') class-attribute instance-attribute

cosh = unary_op('cosh') class-attribute instance-attribute

data = buffer instance-attribute

device: Device property

exp = unary_op('exp') class-attribute instance-attribute

exp2 = unary_op('exp2') class-attribute instance-attribute

is_contiguous property

log = unary_op('log') class-attribute instance-attribute

log10 = unary_op('log10') class-attribute instance-attribute

log2 = unary_op('log2') class-attribute instance-attribute

mT: 'Array' property

max = reduction_op('max') class-attribute instance-attribute

min = reduction_op('min') class-attribute instance-attribute

ndim property

offset = offset instance-attribute

prod = reduction_op('mul') class-attribute instance-attribute

relu = unary_op('relu') class-attribute instance-attribute

shape = shape instance-attribute

sin = unary_op('sin') class-attribute instance-attribute

sinh = unary_op('sinh') class-attribute instance-attribute

size: int property

sqrt = unary_op('sqrt') class-attribute instance-attribute

square = unary_op('square') class-attribute instance-attribute

strides = strides instance-attribute

sum = reduction_op('add') class-attribute instance-attribute

tan = unary_op('tan') class-attribute instance-attribute

tanh = unary_op('tanh') class-attribute instance-attribute

__getitem__(items: tuple[int | slice, ...] | int | slice)

Source code in src/simplendarray/array.py 
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
def __getitem__(self, items: tuple[int | slice, ...] | int | slice):
    if not isinstance(items, tuple):
        items = (items,)
    if len(items) != self.ndim:
        raise ValueError("Must index the same number of dimensions as the array")
    new_shape = []
    new_strides = []
    new_offset = self.offset
    for shape, stride, item in zip(self.shape, self.strides, items):
        if isinstance(item, int):
            item = slice(item, item + 1).indices(shape)[0]
            new_shape.append(1)
            new_strides.append(0)
            new_offset += stride * item
        elif isinstance(item, slice):
            start, stop, step = item.indices(shape)
            if step > 0:
                # Num elements in [start, stop) = stop - start
                new_shape.append(max(0, ceildiv(stop - start, step)))
            else:
                new_shape.append(max(0, ceildiv(start - stop, -step)))
            new_strides.append(stride * step)
            new_offset += stride * start
        else:
            raise TypeError(f"Index must be int or slice, got {type(item).__name__}")
    return Array(self.data, tuple(new_shape), tuple(new_strides), new_offset)

__init__(buffer: BufferType, shape: tuple[int, ...], strides: tuple[int, ...], offset: int)

Source code in src/simplendarray/array.py 
80
81
82
83
84
def __init__(self, buffer: BufferType, shape: tuple[int, ...], strides: tuple[int, ...], offset: int):
    self.data = buffer
    self.shape = shape
    self.strides = strides
    self.offset = offset

__repr__() -> str

Source code in src/simplendarray/array.py 
106
107
def __repr__(self) -> str:
    return f"Array({self.to_python()}, shape={self.shape}, strides={self.strides}, offset={self.offset})"

arange(numel: int, dtype: str | type[DType], device: Device = 'cpu') -> 'Array' classmethod

Source code in src/simplendarray/array.py 
71
72
73
74
75
76
77
78
@classmethod
def arange(cls, numel: int, dtype: str | type[DType], device: Device = "cpu") -> "Array":
    if numel < 0:
        raise ValueError("numel must be >= 0")
    _typecode = typecode(get_dtype(dtype))
    buffer = buf_cls[device].empty(numel, _typecode)
    dispatch_arange(buffer, 0, 1, numel)
    return cls(buffer, (numel,), (1,), 0)

binary_op(op: str) staticmethod

Source code in src/simplendarray/array.py 
247
248
249
250
251
252
253
254
255
@staticmethod
def binary_op(op: str):
    def fn(self: Array, other: Array) -> Array:
        buf = buf_cls[self.device].empty(product(self.shape), self.data.typecode)
        out = Array(buf, self.shape, contiguous_strides(self.shape), 0)
        dispatch_element_wise_binary(self, other, out, op)
        return out

    return fn

from_iterable(data: NestedIterable | Scalar, dtype: str | type[DType], device: Device = 'cpu') classmethod

Source code in src/simplendarray/array.py 
62
63
64
65
66
67
68
69
@classmethod
def from_iterable(cls, data: NestedIterable | Scalar, dtype: str | type[DType], device: Device = "cpu"):
    flat, shape = flatten_and_get_shape(data)
    _typecode = typecode(get_dtype(dtype))
    buffer = buf_cls[device].from_iterable(flat, _typecode)
    strides = contiguous_strides(shape)
    offset = 0
    return cls(buffer, shape, strides, offset)

reduction_op(op: str) staticmethod

Source code in src/simplendarray/array.py 
264
265
266
267
268
269
270
271
272
273
@staticmethod
def reduction_op(op: str):
    def fn(self: Array, dims: tuple[int]):
        reduction_size = product(x for i, x in enumerate(self.shape) if i in dims)
        buf = buf_cls[self.device].empty(self.size // reduction_size, self.data.typecode)
        out = Array(buf, (self.size // reduction_size,), (1,), 0)
        dispatch_reduction(self, out, op, dims)
        return out

    return fn

reshape(new_shape: int | tuple[int, ...]) -> 'Array'

Source code in src/simplendarray/array.py 
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
def reshape(self, new_shape: int | tuple[int, ...]) -> "Array":
    numel = product(self.shape)
    if isinstance(new_shape, int):
        new_shape = (new_shape,)
    num_n1s = sum(1 for s in new_shape if s == -1)
    for s in new_shape:
        if s < -1:
            raise ValueError("Dimension must be non-negative, or with a single -1")
    if num_n1s > 1:
        raise ValueError("only one dimension can be -1")
    if num_n1s == 1:
        known = product(x for x in new_shape if x != -1)
        if numel % known != 0:
            raise ValueError(
                f"Cannot reshape array of size {numel} into shape "
                f"{tuple(numel // known if s == -1 else s for s in new_shape)}"
            )
        replacement = numel // known
        new_shape = tuple(replacement if s == -1 else s for s in new_shape)
    new_strides = reshape_strides(self.shape, new_shape, self.strides)
    if new_strides is None:
        out_ndim = len(new_shape)
        new_strides = contiguous_strides(new_shape)
        new_buffer = buf_cls[self.device].empty(numel, self.data.typecode)
        inp_shape_buffer = buf_cls[self.device].from_iterable(self.shape, "l")
        inp_strides_buffer = buf_cls[self.device].from_iterable(self.strides, "l")
        out_shape_buffer = buf_cls[self.device].from_iterable(new_shape, "l")
        out_strides_buffer = buf_cls[self.device].from_iterable(new_strides, "l")
        inp_work_buffer = buf_cls[self.device].empty(self.ndim, "l")
        out_work_buffer = buf_cls[self.device].empty(out_ndim, "l")

        dispatch_reshape_copy(
            self,
            new_shape,
            new_buffer,
            inp_shape_buffer,
            inp_strides_buffer,
            inp_work_buffer,
            out_shape_buffer,
            out_strides_buffer,
            out_work_buffer,
        )
        return Array(new_buffer, new_shape, new_strides, 0)
    return Array(self.data, new_shape, new_strides, self.offset)

squeeze(dims: int | Iterable[int]) -> 'Array'

Source code in src/simplendarray/array.py 
120
121
122
123
124
125
126
127
128
def squeeze(self, dims: int | Iterable[int]) -> "Array":
    if isinstance(dims, int):
        dims = [dims]
    if self.ndim == 0 or any(self.shape[d] != 1 for d in dims):
        raise ValueError("Can only squeeze a non scalar array with length 1 dims, but shape is", self.shape)
    dims_set = set(dims)
    new_shape = tuple(s for i, s in enumerate(self.shape) if i not in dims_set)
    new_strides = tuple(stride for i, stride in enumerate(self.strides) if i not in dims_set)
    return Array(self.data, new_shape, new_strides, self.offset)

to_python() -> NestedIterable

Source code in src/simplendarray/array.py 
109
110
111
112
113
114
115
116
117
118
def to_python(self) -> NestedIterable:
    if self.ndim == 0:
        return self.data.data[self.offset]
    nested = []
    for i in range(self.shape[0]):
        # Does [self[0, :, :, ...], self[1, :, :, ...], ..., self[shape[0] - 1, :, :, ...]]
        # Then recursively calls to_python on each of these children, until base case reached
        indexed = self[i, *(slice(None) for _ in range(self.ndim - 1))].squeeze(0)
        nested.append(indexed.to_python())
    return nested

transpose(dims: Iterable[int]) -> 'Array'

Source code in src/simplendarray/array.py 
157
158
159
160
def transpose(self, dims: Iterable[int]) -> "Array":
    new_shape = tuple(self.shape[i] for i in dims)
    new_strides = tuple(self.strides[i] for i in dims)
    return Array(self.data, new_shape, new_strides, self.offset)

unary_op(op: str) staticmethod

Source code in src/simplendarray/array.py 
218
219
220
221
222
223
224
225
226
@staticmethod
def unary_op(op: str):
    def fn(self: Array):
        buf = buf_cls[self.device].empty(product(self.shape), self.data.typecode)
        out = Array(buf, self.shape, contiguous_strides(self.shape), 0)
        dispatch_element_wise_unary(self, out, op)
        return out

    return fn

Buffer

Methods:

Name Description
__init__
__repr__
empty
from_iterable

Attributes:

Name Type Description
data
device Device
typecode
Source code in src/simplendarray/buffer.py 
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
class Buffer:
    def __init__(self, data: array.array):
        self.data = data
        self.address, self.num_bytes = data.buffer_info()
        self.typecode = data.typecode
        self.num_bytes *= data.itemsize
        self.device: Device = "cpu"

    @classmethod
    def empty(cls, size: int, dtype: str) -> Buffer:
        # Not empty, full of zeros.
        return cls(array.array(dtype, [0]) * size)

    @classmethod
    def from_iterable(cls, data: Iterable[int | float | bool], dtype: str) -> Buffer:
        data = array.array(dtype, data)
        return cls(data)

    def __repr__(self) -> str:
        return repr(self.data)

data = data instance-attribute

device: Device = 'cpu' instance-attribute

typecode = data.typecode instance-attribute

__init__(data: array.array)

Source code in src/simplendarray/buffer.py 
11
12
13
14
15
16
def __init__(self, data: array.array):
    self.data = data
    self.address, self.num_bytes = data.buffer_info()
    self.typecode = data.typecode
    self.num_bytes *= data.itemsize
    self.device: Device = "cpu"

__repr__() -> str

Source code in src/simplendarray/buffer.py 
28
29
def __repr__(self) -> str:
    return repr(self.data)

empty(size: int, dtype: str) -> Buffer classmethod

Source code in src/simplendarray/buffer.py 
18
19
20
21
@classmethod
def empty(cls, size: int, dtype: str) -> Buffer:
    # Not empty, full of zeros.
    return cls(array.array(dtype, [0]) * size)

from_iterable(data: Iterable[int | float | bool], dtype: str) -> Buffer classmethod

Source code in src/simplendarray/buffer.py 
23
24
25
26
@classmethod
def from_iterable(cls, data: Iterable[int | float | bool], dtype: str) -> Buffer:
    data = array.array(dtype, data)
    return cls(data)

BufferCuda

Methods:

Name Description
__del__
__init__
__repr__
copy_from_host

Copy data from a CPU buffer or array.array to this GPU buffer.
copy_to_host

Copy data from this GPU buffer to a new CPU array.array.
empty
from_iterable

Attributes:

Name Type Description
address
data array

Return a CPU copy of the GPU data for compatibility.
device Device
dtype
num_bytes
size
typecode
Source code in src/simplendarray/buffer_cuda.py 
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
class BufferCuda:
    def __init__(self, size: int, dtype: str):
        self.dtype = dtype
        self.typecode = dtype
        self.size = size
        self.num_bytes = size * array.array(dtype).itemsize
        self.device: Device = "gpu"

        # Allocate device memory
        ptr = buffer_cuda_module.cuda_malloc(self.num_bytes)
        self.address = ptr
        self._owns_memory = True

    @classmethod
    def empty(cls, size: int, dtype: str) -> BufferCuda:
        return cls(size, dtype)

    @classmethod
    def from_iterable(cls, data: array.array | Iterable[int | float | bool], dtype: str) -> BufferCuda:
        if isinstance(data, array.array):
            cpu_buf = data
        else:
            cpu_buf = array.array(dtype, data)
        gpu_buf = cls(len(cpu_buf), dtype)
        gpu_buf.copy_from_host(cpu_buf)
        return gpu_buf

    def copy_from_host(self, cpu_buffer: array.array | HasDataAndAddress) -> None:
        """Copy data from a CPU buffer or array.array to this GPU buffer."""
        if isinstance(cpu_buffer, array.array):
            src_data = cpu_buffer
            src_addr = cpu_buffer.buffer_info()[0]
        else:
            src_data = cpu_buffer.data
            src_addr = cpu_buffer.address
        src_bytes = len(src_data) * src_data.itemsize
        if src_bytes != self.num_bytes:
            raise ValueError("Buffer size mismatch")
        buffer_cuda_module.cuda_memcpy_h2d(self.address, src_addr, self.num_bytes)

    def copy_to_host(self) -> array.array:
        """Copy data from this GPU buffer to a new CPU array.array."""
        cpu_data = array.array(self.typecode, [0]) * self.size
        buffer_cuda_module.cuda_memcpy_d2h(cpu_data.buffer_info()[0], self.address, self.num_bytes)
        return cpu_data

    def __repr__(self) -> str:
        cpu_data = self.copy_to_host()
        return f"BufferCuda({repr(cpu_data)})"

    def __del__(self):  # pragma: no cover
        owns = getattr(self, "_owns_memory", False)
        if owns and hasattr(self, "address") and self.address:
            buffer_cuda_module.cuda_free(self.address)
            self._owns_memory = False

    @property
    def data(self) -> array.array:
        """Return a CPU copy of the GPU data for compatibility."""
        return self.copy_to_host()  # pragma: no cover

_owns_memory = True instance-attribute

address = ptr instance-attribute

data: array.array property

Return a CPU copy of the GPU data for compatibility.

device: Device = 'gpu' instance-attribute

dtype = dtype instance-attribute

num_bytes = size * array.array(dtype).itemsize instance-attribute

size = size instance-attribute

typecode = dtype instance-attribute

__del__()

Source code in src/simplendarray/buffer_cuda.py 
72
73
74
75
76
def __del__(self):  # pragma: no cover
    owns = getattr(self, "_owns_memory", False)
    if owns and hasattr(self, "address") and self.address:
        buffer_cuda_module.cuda_free(self.address)
        self._owns_memory = False

__init__(size: int, dtype: str)

Source code in src/simplendarray/buffer_cuda.py 
23
24
25
26
27
28
29
30
31
32
33
def __init__(self, size: int, dtype: str):
    self.dtype = dtype
    self.typecode = dtype
    self.size = size
    self.num_bytes = size * array.array(dtype).itemsize
    self.device: Device = "gpu"

    # Allocate device memory
    ptr = buffer_cuda_module.cuda_malloc(self.num_bytes)
    self.address = ptr
    self._owns_memory = True

__repr__() -> str

Source code in src/simplendarray/buffer_cuda.py 
68
69
70
def __repr__(self) -> str:
    cpu_data = self.copy_to_host()
    return f"BufferCuda({repr(cpu_data)})"

copy_from_host(cpu_buffer: array.array | HasDataAndAddress) -> None

Copy data from a CPU buffer or array.array to this GPU buffer.

Source code in src/simplendarray/buffer_cuda.py 
49
50
51
52
53
54
55
56
57
58
59
60
def copy_from_host(self, cpu_buffer: array.array | HasDataAndAddress) -> None:
    """Copy data from a CPU buffer or array.array to this GPU buffer."""
    if isinstance(cpu_buffer, array.array):
        src_data = cpu_buffer
        src_addr = cpu_buffer.buffer_info()[0]
    else:
        src_data = cpu_buffer.data
        src_addr = cpu_buffer.address
    src_bytes = len(src_data) * src_data.itemsize
    if src_bytes != self.num_bytes:
        raise ValueError("Buffer size mismatch")
    buffer_cuda_module.cuda_memcpy_h2d(self.address, src_addr, self.num_bytes)

copy_to_host() -> array.array

Copy data from this GPU buffer to a new CPU array.array.

Source code in src/simplendarray/buffer_cuda.py 
62
63
64
65
66
def copy_to_host(self) -> array.array:
    """Copy data from this GPU buffer to a new CPU array.array."""
    cpu_data = array.array(self.typecode, [0]) * self.size
    buffer_cuda_module.cuda_memcpy_d2h(cpu_data.buffer_info()[0], self.address, self.num_bytes)
    return cpu_data

empty(size: int, dtype: str) -> BufferCuda classmethod

Source code in src/simplendarray/buffer_cuda.py 
35
36
37
@classmethod
def empty(cls, size: int, dtype: str) -> BufferCuda:
    return cls(size, dtype)

from_iterable(data: array.array | Iterable[int | float | bool], dtype: str) -> BufferCuda classmethod

Source code in src/simplendarray/buffer_cuda.py 
39
40
41
42
43
44
45
46
47
@classmethod
def from_iterable(cls, data: array.array | Iterable[int | float | bool], dtype: str) -> BufferCuda:
    if isinstance(data, array.array):
        cpu_buf = data
    else:
        cpu_buf = array.array(dtype, data)
    gpu_buf = cls(len(cpu_buf), dtype)
    gpu_buf.copy_from_host(cpu_buf)
    return gpu_buf

PythonModule

Methods:

Name Description
__getattr__
__init__
compile
compile_fn
Source code in src/simplendarray/transpiler/runtime.py 
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
287
288
289
290
291
292
293
294
295
296
297
298
299
300
301
302
303
304
305
306
307
308
309
310
311
312
313
314
315
316
317
318
319
320
321
322
323
324
325
326
327
328
329
330
331
332
333
334
335
336
337
338
339
340
341
342
343
344
345
346
347
348
349
350
351
352
353
354
355
356
357
358
359
360
361
362
363
364
365
366
367
368
369
370
371
372
373
374
375
376
377
378
379
380
381
382
383
384
385
386
387
388
389
390
391
392
393
394
395
396
397
398
399
400
401
402
403
404
405
406
407
408
409
410
411
412
413
414
415
416
417
418
419
420
421
422
423
424
425
426
427
428
429
430
431
432
433
434
435
436
437
438
439
440
441
442
443
444
445
446
447
448
449
450
451
452
453
454
455
456
457
458
459
460
461
462
463
464
465
466
467
468
469
470
471
472
473
474
475
476
477
478
479
480
481
482
483
484
485
486
487
488
489
490
491
492
493
494
495
496
497
498
499
500
501
502
503
504
505
506
507
508
509
510
511
512
513
514
515
516
517
518
519
520
521
522
523
524
525
526
527
528
529
530
531
532
533
534
535
536
537
538
539
540
541
542
543
544
545
546
547
548
549
550
551
552
553
554
class PythonModule:
    def __init__(self, includes: list[str] | None = None, stub_path: str | None = None, stub_var: str | None = None):
        self._funcs: list[CFunction] = []
        self._includes = includes or []
        self._module = None
        self._compiled = False
        self._stub_path = stub_path
        self._stub_var = stub_var

    def _register(self, func, c_attrs: list[str], pybind: bool, group: str | None = None):
        src = dedent(inspect.getsource(func))
        tree = ast.parse(src)
        fn = tree.body[0]
        if not isinstance(fn, ast.FunctionDef):
            raise ValueError("only function definitions can be compiled")
        name = fn.name

        if fn.returns is None:
            fn.returns = ast.Name("void")

        self._register_ast(name, fn, c_attrs, pybind, group=group)

    def _register_ast(
        self,
        name: str,
        fn: ast.FunctionDef,
        c_attrs: list[str],
        pybind: bool,
        group: str | None = None,
        dispatch_key=None,
    ):
        c_source = _stmt(fn)

        params = []
        for a in fn.args.args:
            annotation = cast(ast.AST, a.annotation)
            is_bool = isinstance(annotation, ast.Name) and annotation.id == "bool"
            is_string = isinstance(annotation, ast.Name) and annotation.id == "str"
            ct = _c_type(annotation)
            params.append({"name": a.arg, "c_type": ct, "is_bool": is_bool, "is_string": is_string})

        ret = cast(ast.AST, fn.returns)
        ret_is_bool = isinstance(ret, ast.Name) and ret.id == "bool"
        ret_c_type = _c_type(ret)

        entry = CFunction(
            name=name,
            c_source=c_source,
            params=params,
            ret_c_type=ret_c_type,
            ret_is_bool=ret_is_bool,
            c_attrs=c_attrs or [],
            pybind=pybind,
            group=group,
            dispatch_key=dispatch_key,
        )
        self._funcs.append(entry)

    _cache_dir_override: Path | None = None

    @staticmethod
    def _cache_dir() -> Path:
        if PythonModule._cache_dir_override is not None:
            return PythonModule._cache_dir_override
        env = os.environ.get("SNDA_CACHE_HOME")
        if env:
            base = Path(env)
        else:
            base = Path.home() / ".cache" / "simplendarray"
        base.mkdir(parents=True, exist_ok=True)
        return base

    @contextmanager
    def _cache_lock(self, cache_dir: Path):
        lock_path = cache_dir / ".lock"
        lock_path.parent.mkdir(parents=True, exist_ok=True)
        fd = os.open(lock_path, os.O_WRONLY | os.O_CREAT | os.O_TRUNC)
        try:
            fcntl.flock(fd, fcntl.LOCK_EX)
            yield
        finally:
            os.close(fd)

    def _cache_key(self, ext_src: str, compiler: str, cflags: list[str] | None, ldflags: list[str] | None) -> str:
        meta = json.dumps({"compiler": compiler, "cflags": cflags, "ldflags": ldflags}, sort_keys=True)
        raw = ext_src + meta
        return hashlib.sha256(raw.encode()).hexdigest()

    def _load_from_so(self, module_name: str, so_path: Path) -> bool:
        spec = importlib.util.spec_from_file_location(module_name, so_path)
        if spec is None or spec.loader is None:
            return False
        self._module = importlib.util.module_from_spec(spec)
        spec.loader.exec_module(self._module)
        self._compiled = True
        return True

    def compile(self, compiler="gcc", cflags=None, ldflags=None) -> Self:
        if not self._funcs:
            raise ValueError("no functions registered")
        if self._compiled:
            raise RuntimeError("module already compiled")

        self._write_source_stub()

        if compiler == "nvcc" and shutil.which("nvcc") is None:
            return self

        module_name = _NDARRAY_MODULE_NAME
        ext_src = self._generate_extension(module_name)
        ext_suffix = sysconfig.get_config_var("EXT_SUFFIX") or ".so"
        src_ext = ".cu" if compiler == "nvcc" else ".c"
        cache_dir = self._cache_dir()
        key = self._cache_key(ext_src, compiler, cflags, ldflags)
        build_dir = cache_dir / key
        build_dir.mkdir(parents=True, exist_ok=True)
        so_path = build_dir / f"{key}{ext_suffix}"

        if so_path.exists() and self._load_from_so(module_name, so_path):
            self._build_dispatch_dicts()
            return self

        with self._cache_lock(build_dir):
            src_path = build_dir / f"{key}{src_ext}"
            src_path.write_text(ext_src)

            meta = {"compiler": compiler, "cflags": cflags, "ldflags": ldflags}
            meta_path = build_dir / f"{key}.json"
            meta_path.write_text(json.dumps(meta, indent=2))

            py_include = sysconfig.get_config_var("INCLUDEPY")
            if not py_include:
                py_include = sysconfig.get_path("include")

            cmd = [compiler, "-shared"]
            if compiler == "nvcc":  # pragma: no cover
                cmd.extend(["-Xcompiler", "-fPIC"])  # pragma: no cover
            else:
                cmd.append("-fPIC")
            cmd.append(f"-I{py_include}")
            if cflags:
                cmd.extend(cflags)
            cmd.extend(["-o", str(so_path), str(src_path)])
            if ldflags:
                cmd.extend(ldflags)
            if sys.platform == "darwin":
                cmd.append("-undefined")  # pragma: no cover
                cmd.append("dynamic_lookup")  # pragma: no cover

            result = subprocess.run(cmd, stdout=subprocess.PIPE, stderr=subprocess.PIPE)  # noqa: S603
            if result.returncode != 0:
                msg = result.stderr.decode()
                raise RuntimeError(f"compilation failed:\n{msg}")

            if not self._load_from_so(module_name, so_path):
                raise RuntimeError(f"failed to create module spec for {so_path}")
            self._build_dispatch_dicts()

        return self

    def _generate_forward_decl(self, func: CFunction) -> str:
        ret = func.ret_c_type
        name = func.name
        param_strs = [p["c_type"] for p in func.params]
        sig = f"{ret} {name}({', '.join(param_strs)})"
        attrs = func.c_attrs
        if attrs:
            sig = f"{' '.join(attrs)} {sig}"
        return f"{sig};"

    def _generate_extension(self, module_name):
        c_functions = []
        wrapper_functions = []
        method_defs = []
        forward_decls = []

        for func in self._funcs:
            lines = func.c_source.split("\n", 1)
            lines[0] = f"{' '.join(func.c_attrs)} {lines[0]}"
            c_source = "\n".join(lines)
            c_functions.append(c_source)
            if func.pybind:
                wrapper_functions.append(self._generate_wrapper(func))
                method_defs.append(
                    f'    {{"{func.name}", {func.name}_wrapper, METH_VARARGS, "Transpiled function {func.name}"}}'
                )
            forward_decls.append(self._generate_forward_decl(func))

        if not method_defs:
            raise ValueError("Useless module, no python bindings present", module_name)

        c_funcs_str = "\n\n".join(c_functions)
        wrapper_str = "\n\n".join(wrapper_functions)
        method_defs_str = ",\n".join(method_defs)
        forward_decls_str = "\n".join(forward_decls)
        includes_str = "\n".join(self._includes)

        return _EXTENSION_TEMPLATE.format(
            includes=includes_str,
            forward_decls=forward_decls_str,
            c_funcs=c_funcs_str,
            wrapper=wrapper_str,
            method_defs=method_defs_str,
            module_name=module_name,
        )

    def _generate_wrapper(self, func: CFunction):
        name = func.name
        params = func.params
        ret_c_type = func.ret_c_type
        ret_is_bool = func.ret_is_bool

        fmt_chars = []
        for p in params:
            if p["is_bool"]:
                fmt_chars.append("p")
            elif p.get("is_string") and p["c_type"] in ("char*", "unsigned char*"):
                fmt_chars.append("s")
            elif p["c_type"] in ("char*", "unsigned char*"):
                fmt_chars.append("K")
            else:
                f = _C_PARSE_FMT.get(p["c_type"])
                if f is None:
                    raise ValueError(f"unsupported param type for Python wrapper: {p['c_type']}")
                fmt_chars.append(f)

        fmt = "".join(fmt_chars)

        var_lines = []
        parse_args = []
        for p in params:
            if p["is_bool"]:
                vt = "int"
            elif p.get("is_string") and p["c_type"] in ("char*", "unsigned char*"):
                vt = "char*"
            elif p["c_type"] in ("char*", "unsigned char*"):
                vt = "unsigned long long"
            else:
                vt = _C_PARSE_TYPE.get(p["c_type"], p["c_type"])
            var_lines.append(f"{T}{vt} {p['name']};")
            parse_args.append(f"&{p['name']}")

        call_args = []
        for p in params:
            if p.get("is_string") and p["c_type"] in ("char*", "unsigned char*"):
                call_args.append(p["name"])
            elif p["c_type"] in _POINTER_TYPES:
                call_args.append(f"({p['c_type']}){p['name']}")
            else:
                call_args.append(p["name"])
        arg_str = ", ".join(call_args)

        if ret_c_type == "void":
            call_line = f"{T}{name}({arg_str});"
            ret_line = f"{T}Py_RETURN_NONE;"
        else:
            decl_ret_type = "int" if ret_is_bool else ret_c_type
            call_line = f"{T}{decl_ret_type} result = {name}({arg_str});"
            if ret_is_bool:
                ret_line = f"{T}return PyBool_FromLong(result);"
            else:
                builder = _C_RETURN_BUILD.get(ret_c_type)
                if builder is None:
                    raise ValueError(f"unsupported return type for Python wrapper: {ret_c_type}")
                ret_line = f"{T}return {builder}(result);"

        var_block = "\n".join(var_lines) if var_lines else f"{T}(void)self;"

        if parse_args:
            parse_call = f'if (!PyArg_ParseTuple(args, "{fmt}", {", ".join(parse_args)})) {{'
        else:
            parse_call = 'if (!PyArg_ParseTuple(args, "")) {'

        return _WRAPPER_TEMPLATE.format(
            name=name,
            var_block=var_block,
            parse_call=parse_call,
            call_line=call_line,
            ret_line=ret_line,
        )

    def _write_source_stub(self):
        if self._stub_path is not None and self._stub_var is not None:
            stem = Path(self._stub_path).stem
            parent = Path(self._stub_path).parent
            stub_file = parent / f"_{stem}_stubs.py"
            class_name = f"_{self._stub_var.title().replace('_', '')}Class"

            has_dispatch = any(
                func.pybind and func.group is not None and func.dispatch_key is not None for func in self._funcs
            )

            lines = [
                "# Auto-generated stub for compiled functions.",
                "from __future__ import annotations",
                "from typing import TYPE_CHECKING",
            ]
            if has_dispatch:
                lines.append("from typing import Callable, ClassVar")
            lines.append("")
            lines.append("from simplendarray.transpiler.runtime import PythonModule")
            lines.append("")
            lines.append(f"class {class_name}(PythonModule):")
            lines.append(f"{T}if TYPE_CHECKING:")

            seen_groups: set[str] = set()
            for func in self._funcs:
                if (
                    func.pybind
                    and func.group is not None
                    and func.dispatch_key is not None
                    and func.group not in seen_groups
                ):
                    seen_groups.add(func.group)
                    lines.append(
                        f"{T * 2}DISPATCH_DICT_{func.group}: ClassVar[dict[tuple, Callable[..., None]]]"  # noqa: E501
                    )
            if seen_groups:
                lines.append("")

            for func in self._funcs:
                if func.pybind:
                    params = ", ".join(f"{p['name']}: {_C_TYPE_TO_PY.get(p['c_type'], 'int')}" for p in func.params)
                    return_type = _C_TYPE_TO_PY.get(func.ret_c_type, "None")
                    lines.append(f"{T * 2}def {func.name}(self, {params}) -> {return_type}: ...")
            lines.append(f"{T * 2}pass")
            lines.append("")
            content = "\n".join(lines)
            fd, tmp = tempfile.mkstemp(dir=parent, suffix=".py")
            with os.fdopen(fd, "w") as f:
                f.write(content)
            os.replace(tmp, stub_file)

    def _build_dispatch_dicts(self):
        groups: dict[str, dict[str, str]] = {}
        for func in self._funcs:
            if func.pybind and func.group is not None and func.dispatch_key is not None:
                groups.setdefault(func.group, {})[func.dispatch_key] = func.name
        for group_name, mapping in groups.items():
            dict_name = f"DISPATCH_DICT_{group_name}"
            dispatch_dict = {key: getattr(self, mangled_name) for key, mangled_name in mapping.items()}
            setattr(self, dict_name, dispatch_dict)

    def __getattr__(self, name):
        if self._compiled and self._module is not None and hasattr(self._module, name):
            return getattr(self._module, name)
        raise AttributeError(f"Module has no attribute '{name}'")

    def compile_fn(
        self, types: Iterable[SpecItem] | None = None, c_attrs: list[str] | None = None, pybind: bool = False
    ):
        if c_attrs is None:
            c_attrs = []

        class IndexableFunction[**P, R]:  # pragma: no cover
            # Purely for type hinting
            def __init__(self, fn: Callable[P, R]):
                self.fn = fn

            def __call__(self, *args: P.args, **kwargs: P.kwargs):
                return self.fn(*args, **kwargs)

            def __getitem__(self, _idx):
                return self

        def decorator[R, **P](func: Callable[P, R]) -> IndexableFunction[P, R]:
            if types:
                src = dedent(inspect.getsource(func))
                tree = ast.parse(src)
                fn = tree.body[0]
                if not isinstance(fn, ast.FunctionDef):
                    raise ValueError("only function definitions can be compiled")
                for spec in types:
                    fn_name = spec.fn_name
                    fn_copy = copy.deepcopy(fn)
                    substituter = _TypeSubstituter(spec)
                    specialized = substituter.visit(fn_copy)
                    ast.fix_missing_locations(specialized)
                    group = func.__name__
                    dispatch_key = tuple(spec.mapping.items())
                    self._register_ast(fn_name, specialized, c_attrs, pybind, group=group, dispatch_key=dispatch_key)
            else:
                group = func.__name__
                self._register(func, c_attrs, pybind, group=group)
            return IndexableFunction(func)

        return decorator

_cache_dir_override: Path | None = None class-attribute instance-attribute

_compiled = False instance-attribute

_funcs: list[CFunction] = [] instance-attribute

_includes = includes or [] instance-attribute

_module = None instance-attribute

_stub_path = stub_path instance-attribute

_stub_var = stub_var instance-attribute

__getattr__(name)

Source code in src/simplendarray/transpiler/runtime.py 
511
512
513
514
def __getattr__(self, name):
    if self._compiled and self._module is not None and hasattr(self._module, name):
        return getattr(self._module, name)
    raise AttributeError(f"Module has no attribute '{name}'")

__init__(includes: list[str] | None = None, stub_path: str | None = None, stub_var: str | None = None)

Source code in src/simplendarray/transpiler/runtime.py 
169
170
171
172
173
174
175
def __init__(self, includes: list[str] | None = None, stub_path: str | None = None, stub_var: str | None = None):
    self._funcs: list[CFunction] = []
    self._includes = includes or []
    self._module = None
    self._compiled = False
    self._stub_path = stub_path
    self._stub_var = stub_var

_build_dispatch_dicts()

Source code in src/simplendarray/transpiler/runtime.py 
501
502
503
504
505
506
507
508
509
def _build_dispatch_dicts(self):
    groups: dict[str, dict[str, str]] = {}
    for func in self._funcs:
        if func.pybind and func.group is not None and func.dispatch_key is not None:
            groups.setdefault(func.group, {})[func.dispatch_key] = func.name
    for group_name, mapping in groups.items():
        dict_name = f"DISPATCH_DICT_{group_name}"
        dispatch_dict = {key: getattr(self, mangled_name) for key, mangled_name in mapping.items()}
        setattr(self, dict_name, dispatch_dict)

_cache_dir() -> Path staticmethod

Source code in src/simplendarray/transpiler/runtime.py 
228
229
230
231
232
233
234
235
236
237
238
@staticmethod
def _cache_dir() -> Path:
    if PythonModule._cache_dir_override is not None:
        return PythonModule._cache_dir_override
    env = os.environ.get("SNDA_CACHE_HOME")
    if env:
        base = Path(env)
    else:
        base = Path.home() / ".cache" / "simplendarray"
    base.mkdir(parents=True, exist_ok=True)
    return base

_cache_key(ext_src: str, compiler: str, cflags: list[str] | None, ldflags: list[str] | None) -> str

Source code in src/simplendarray/transpiler/runtime.py 
251
252
253
254
def _cache_key(self, ext_src: str, compiler: str, cflags: list[str] | None, ldflags: list[str] | None) -> str:
    meta = json.dumps({"compiler": compiler, "cflags": cflags, "ldflags": ldflags}, sort_keys=True)
    raw = ext_src + meta
    return hashlib.sha256(raw.encode()).hexdigest()

_cache_lock(cache_dir: Path)

Source code in src/simplendarray/transpiler/runtime.py 
240
241
242
243
244
245
246
247
248
249
@contextmanager
def _cache_lock(self, cache_dir: Path):
    lock_path = cache_dir / ".lock"
    lock_path.parent.mkdir(parents=True, exist_ok=True)
    fd = os.open(lock_path, os.O_WRONLY | os.O_CREAT | os.O_TRUNC)
    try:
        fcntl.flock(fd, fcntl.LOCK_EX)
        yield
    finally:
        os.close(fd)

_generate_extension(module_name)

Source code in src/simplendarray/transpiler/runtime.py 
338
339
340
341
342
343
344
345
346
347
348
349
350
351
352
353
354
355
356
357
358
359
360
361
362
363
364
365
366
367
368
369
370
371
372
def _generate_extension(self, module_name):
    c_functions = []
    wrapper_functions = []
    method_defs = []
    forward_decls = []

    for func in self._funcs:
        lines = func.c_source.split("\n", 1)
        lines[0] = f"{' '.join(func.c_attrs)} {lines[0]}"
        c_source = "\n".join(lines)
        c_functions.append(c_source)
        if func.pybind:
            wrapper_functions.append(self._generate_wrapper(func))
            method_defs.append(
                f'    {{"{func.name}", {func.name}_wrapper, METH_VARARGS, "Transpiled function {func.name}"}}'
            )
        forward_decls.append(self._generate_forward_decl(func))

    if not method_defs:
        raise ValueError("Useless module, no python bindings present", module_name)

    c_funcs_str = "\n\n".join(c_functions)
    wrapper_str = "\n\n".join(wrapper_functions)
    method_defs_str = ",\n".join(method_defs)
    forward_decls_str = "\n".join(forward_decls)
    includes_str = "\n".join(self._includes)

    return _EXTENSION_TEMPLATE.format(
        includes=includes_str,
        forward_decls=forward_decls_str,
        c_funcs=c_funcs_str,
        wrapper=wrapper_str,
        method_defs=method_defs_str,
        module_name=module_name,
    )

_generate_forward_decl(func: CFunction) -> str

Source code in src/simplendarray/transpiler/runtime.py 
328
329
330
331
332
333
334
335
336
def _generate_forward_decl(self, func: CFunction) -> str:
    ret = func.ret_c_type
    name = func.name
    param_strs = [p["c_type"] for p in func.params]
    sig = f"{ret} {name}({', '.join(param_strs)})"
    attrs = func.c_attrs
    if attrs:
        sig = f"{' '.join(attrs)} {sig}"
    return f"{sig};"

_generate_wrapper(func: CFunction)

Source code in src/simplendarray/transpiler/runtime.py 
374
375
376
377
378
379
380
381
382
383
384
385
386
387
388
389
390
391
392
393
394
395
396
397
398
399
400
401
402
403
404
405
406
407
408
409
410
411
412
413
414
415
416
417
418
419
420
421
422
423
424
425
426
427
428
429
430
431
432
433
434
435
436
437
438
439
440
441
442
443
444
445
446
447
def _generate_wrapper(self, func: CFunction):
    name = func.name
    params = func.params
    ret_c_type = func.ret_c_type
    ret_is_bool = func.ret_is_bool

    fmt_chars = []
    for p in params:
        if p["is_bool"]:
            fmt_chars.append("p")
        elif p.get("is_string") and p["c_type"] in ("char*", "unsigned char*"):
            fmt_chars.append("s")
        elif p["c_type"] in ("char*", "unsigned char*"):
            fmt_chars.append("K")
        else:
            f = _C_PARSE_FMT.get(p["c_type"])
            if f is None:
                raise ValueError(f"unsupported param type for Python wrapper: {p['c_type']}")
            fmt_chars.append(f)

    fmt = "".join(fmt_chars)

    var_lines = []
    parse_args = []
    for p in params:
        if p["is_bool"]:
            vt = "int"
        elif p.get("is_string") and p["c_type"] in ("char*", "unsigned char*"):
            vt = "char*"
        elif p["c_type"] in ("char*", "unsigned char*"):
            vt = "unsigned long long"
        else:
            vt = _C_PARSE_TYPE.get(p["c_type"], p["c_type"])
        var_lines.append(f"{T}{vt} {p['name']};")
        parse_args.append(f"&{p['name']}")

    call_args = []
    for p in params:
        if p.get("is_string") and p["c_type"] in ("char*", "unsigned char*"):
            call_args.append(p["name"])
        elif p["c_type"] in _POINTER_TYPES:
            call_args.append(f"({p['c_type']}){p['name']}")
        else:
            call_args.append(p["name"])
    arg_str = ", ".join(call_args)

    if ret_c_type == "void":
        call_line = f"{T}{name}({arg_str});"
        ret_line = f"{T}Py_RETURN_NONE;"
    else:
        decl_ret_type = "int" if ret_is_bool else ret_c_type
        call_line = f"{T}{decl_ret_type} result = {name}({arg_str});"
        if ret_is_bool:
            ret_line = f"{T}return PyBool_FromLong(result);"
        else:
            builder = _C_RETURN_BUILD.get(ret_c_type)
            if builder is None:
                raise ValueError(f"unsupported return type for Python wrapper: {ret_c_type}")
            ret_line = f"{T}return {builder}(result);"

    var_block = "\n".join(var_lines) if var_lines else f"{T}(void)self;"

    if parse_args:
        parse_call = f'if (!PyArg_ParseTuple(args, "{fmt}", {", ".join(parse_args)})) {{'
    else:
        parse_call = 'if (!PyArg_ParseTuple(args, "")) {'

    return _WRAPPER_TEMPLATE.format(
        name=name,
        var_block=var_block,
        parse_call=parse_call,
        call_line=call_line,
        ret_line=ret_line,
    )

_load_from_so(module_name: str, so_path: Path) -> bool

Source code in src/simplendarray/transpiler/runtime.py 
256
257
258
259
260
261
262
263
def _load_from_so(self, module_name: str, so_path: Path) -> bool:
    spec = importlib.util.spec_from_file_location(module_name, so_path)
    if spec is None or spec.loader is None:
        return False
    self._module = importlib.util.module_from_spec(spec)
    spec.loader.exec_module(self._module)
    self._compiled = True
    return True

_register(func, c_attrs: list[str], pybind: bool, group: str | None = None)

Source code in src/simplendarray/transpiler/runtime.py 
177
178
179
180
181
182
183
184
185
186
187
188
def _register(self, func, c_attrs: list[str], pybind: bool, group: str | None = None):
    src = dedent(inspect.getsource(func))
    tree = ast.parse(src)
    fn = tree.body[0]
    if not isinstance(fn, ast.FunctionDef):
        raise ValueError("only function definitions can be compiled")
    name = fn.name

    if fn.returns is None:
        fn.returns = ast.Name("void")

    self._register_ast(name, fn, c_attrs, pybind, group=group)

_register_ast(name: str, fn: ast.FunctionDef, c_attrs: list[str], pybind: bool, group: str | None = None, dispatch_key=None)

Source code in src/simplendarray/transpiler/runtime.py 
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
def _register_ast(
    self,
    name: str,
    fn: ast.FunctionDef,
    c_attrs: list[str],
    pybind: bool,
    group: str | None = None,
    dispatch_key=None,
):
    c_source = _stmt(fn)

    params = []
    for a in fn.args.args:
        annotation = cast(ast.AST, a.annotation)
        is_bool = isinstance(annotation, ast.Name) and annotation.id == "bool"
        is_string = isinstance(annotation, ast.Name) and annotation.id == "str"
        ct = _c_type(annotation)
        params.append({"name": a.arg, "c_type": ct, "is_bool": is_bool, "is_string": is_string})

    ret = cast(ast.AST, fn.returns)
    ret_is_bool = isinstance(ret, ast.Name) and ret.id == "bool"
    ret_c_type = _c_type(ret)

    entry = CFunction(
        name=name,
        c_source=c_source,
        params=params,
        ret_c_type=ret_c_type,
        ret_is_bool=ret_is_bool,
        c_attrs=c_attrs or [],
        pybind=pybind,
        group=group,
        dispatch_key=dispatch_key,
    )
    self._funcs.append(entry)

_write_source_stub()

Source code in src/simplendarray/transpiler/runtime.py 
449
450
451
452
453
454
455
456
457
458
459
460
461
462
463
464
465
466
467
468
469
470
471
472
473
474
475
476
477
478
479
480
481
482
483
484
485
486
487
488
489
490
491
492
493
494
495
496
497
498
499
def _write_source_stub(self):
    if self._stub_path is not None and self._stub_var is not None:
        stem = Path(self._stub_path).stem
        parent = Path(self._stub_path).parent
        stub_file = parent / f"_{stem}_stubs.py"
        class_name = f"_{self._stub_var.title().replace('_', '')}Class"

        has_dispatch = any(
            func.pybind and func.group is not None and func.dispatch_key is not None for func in self._funcs
        )

        lines = [
            "# Auto-generated stub for compiled functions.",
            "from __future__ import annotations",
            "from typing import TYPE_CHECKING",
        ]
        if has_dispatch:
            lines.append("from typing import Callable, ClassVar")
        lines.append("")
        lines.append("from simplendarray.transpiler.runtime import PythonModule")
        lines.append("")
        lines.append(f"class {class_name}(PythonModule):")
        lines.append(f"{T}if TYPE_CHECKING:")

        seen_groups: set[str] = set()
        for func in self._funcs:
            if (
                func.pybind
                and func.group is not None
                and func.dispatch_key is not None
                and func.group not in seen_groups
            ):
                seen_groups.add(func.group)
                lines.append(
                    f"{T * 2}DISPATCH_DICT_{func.group}: ClassVar[dict[tuple, Callable[..., None]]]"  # noqa: E501
                )
        if seen_groups:
            lines.append("")

        for func in self._funcs:
            if func.pybind:
                params = ", ".join(f"{p['name']}: {_C_TYPE_TO_PY.get(p['c_type'], 'int')}" for p in func.params)
                return_type = _C_TYPE_TO_PY.get(func.ret_c_type, "None")
                lines.append(f"{T * 2}def {func.name}(self, {params}) -> {return_type}: ...")
        lines.append(f"{T * 2}pass")
        lines.append("")
        content = "\n".join(lines)
        fd, tmp = tempfile.mkstemp(dir=parent, suffix=".py")
        with os.fdopen(fd, "w") as f:
            f.write(content)
        os.replace(tmp, stub_file)

compile(compiler='gcc', cflags=None, ldflags=None) -> Self

Source code in src/simplendarray/transpiler/runtime.py 
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
287
288
289
290
291
292
293
294
295
296
297
298
299
300
301
302
303
304
305
306
307
308
309
310
311
312
313
314
315
316
317
318
319
320
321
322
323
324
325
326
def compile(self, compiler="gcc", cflags=None, ldflags=None) -> Self:
    if not self._funcs:
        raise ValueError("no functions registered")
    if self._compiled:
        raise RuntimeError("module already compiled")

    self._write_source_stub()

    if compiler == "nvcc" and shutil.which("nvcc") is None:
        return self

    module_name = _NDARRAY_MODULE_NAME
    ext_src = self._generate_extension(module_name)
    ext_suffix = sysconfig.get_config_var("EXT_SUFFIX") or ".so"
    src_ext = ".cu" if compiler == "nvcc" else ".c"
    cache_dir = self._cache_dir()
    key = self._cache_key(ext_src, compiler, cflags, ldflags)
    build_dir = cache_dir / key
    build_dir.mkdir(parents=True, exist_ok=True)
    so_path = build_dir / f"{key}{ext_suffix}"

    if so_path.exists() and self._load_from_so(module_name, so_path):
        self._build_dispatch_dicts()
        return self

    with self._cache_lock(build_dir):
        src_path = build_dir / f"{key}{src_ext}"
        src_path.write_text(ext_src)

        meta = {"compiler": compiler, "cflags": cflags, "ldflags": ldflags}
        meta_path = build_dir / f"{key}.json"
        meta_path.write_text(json.dumps(meta, indent=2))

        py_include = sysconfig.get_config_var("INCLUDEPY")
        if not py_include:
            py_include = sysconfig.get_path("include")

        cmd = [compiler, "-shared"]
        if compiler == "nvcc":  # pragma: no cover
            cmd.extend(["-Xcompiler", "-fPIC"])  # pragma: no cover
        else:
            cmd.append("-fPIC")
        cmd.append(f"-I{py_include}")
        if cflags:
            cmd.extend(cflags)
        cmd.extend(["-o", str(so_path), str(src_path)])
        if ldflags:
            cmd.extend(ldflags)
        if sys.platform == "darwin":
            cmd.append("-undefined")  # pragma: no cover
            cmd.append("dynamic_lookup")  # pragma: no cover

        result = subprocess.run(cmd, stdout=subprocess.PIPE, stderr=subprocess.PIPE)  # noqa: S603
        if result.returncode != 0:
            msg = result.stderr.decode()
            raise RuntimeError(f"compilation failed:\n{msg}")

        if not self._load_from_so(module_name, so_path):
            raise RuntimeError(f"failed to create module spec for {so_path}")
        self._build_dispatch_dicts()

    return self

compile_fn(types: Iterable[SpecItem] | None = None, c_attrs: list[str] | None = None, pybind: bool = False)

Source code in src/simplendarray/transpiler/runtime.py 
516
517
518
519
520
521
522
523
524
525
526
527
528
529
530
531
532
533
534
535
536
537
538
539
540
541
542
543
544
545
546
547
548
549
550
551
552
553
554
def compile_fn(
    self, types: Iterable[SpecItem] | None = None, c_attrs: list[str] | None = None, pybind: bool = False
):
    if c_attrs is None:
        c_attrs = []

    class IndexableFunction[**P, R]:  # pragma: no cover
        # Purely for type hinting
        def __init__(self, fn: Callable[P, R]):
            self.fn = fn

        def __call__(self, *args: P.args, **kwargs: P.kwargs):
            return self.fn(*args, **kwargs)

        def __getitem__(self, _idx):
            return self

    def decorator[R, **P](func: Callable[P, R]) -> IndexableFunction[P, R]:
        if types:
            src = dedent(inspect.getsource(func))
            tree = ast.parse(src)
            fn = tree.body[0]
            if not isinstance(fn, ast.FunctionDef):
                raise ValueError("only function definitions can be compiled")
            for spec in types:
                fn_name = spec.fn_name
                fn_copy = copy.deepcopy(fn)
                substituter = _TypeSubstituter(spec)
                specialized = substituter.visit(fn_copy)
                ast.fix_missing_locations(specialized)
                group = func.__name__
                dispatch_key = tuple(spec.mapping.items())
                self._register_ast(fn_name, specialized, c_attrs, pybind, group=group, dispatch_key=dispatch_key)
        else:
            group = func.__name__
            self._register(func, c_attrs, pybind, group=group)
        return IndexableFunction(func)

    return decorator

compile_str(func) -> str

Transpile a python function into a C source code string.

Source code in src/simplendarray/transpiler/transpiler.py 
337
338
339
340
341
342
343
344
345
def compile_str(func) -> str:
    """Transpile a python function into a C source code string."""
    src = dedent(inspect.getsource(func))
    tree = ast.parse(src)
    fn = tree.body[0]
    if not isinstance(fn, ast.FunctionDef):
        raise ValueError("only function definitions can be compiled")
    c_str = _stmt(fn)
    return c_str

ref(x: T) -> list[T]

Equal to doing &x in C. Returning [x] is just for type hinting purposes.

Source code in src/simplendarray/transpiler/transpiler.py 
 9
10
11
def ref[T](x: T) -> list[T]:
    """Equal to doing &x in C. Returning [x] is just for type hinting purposes."""
    return [x]