Skip to content

Device Functions

@tack.func

Functions decorated with @tack.func are inlined into kernels at the AST level. They let you write reusable device-side code:

@tack.func
def lerp(a, b, t):
    return a + t * (b - a)

@tack.kernel
def interpolate(src, dst, alpha, n):
    for i in range(n):
        dst[i] = lerp(src[i], 1.0, alpha)

When the kernel is compiled, lerp is inlined — no function call overhead.

Calls use the defining kernel or device function's Python bindings. You can import a device function under an alias or call it through an imported module. Two modules may define the same function name without interfering. Nested calls use the callee's namespace, and closure-bound device functions are supported. Template methods keep their defining bindings too.

Bindings are resolved when the IR specialization is first built. Cached kernels keep the inlined body; recreate a kernel to capture a changed callable binding. Runtime function parameters, ordinary Python callables, and recursive device calls are unsupported. This does not enable capturing numerical Python variables: pass those values as kernel arguments.

Multiple Return Values

@tack.func supports returning tuples:

@tack.func
def polar_to_cart(r, theta):
    x = r * cos(theta)
    y = r * sin(theta)
    return x, y

@tack.kernel
def convert(r_array, theta_array, x_out, y_out, n):
    for i in range(n):
        x, y = polar_to_cart(r_array[i], theta_array[i])
        x_out[i] = x
        y_out[i] = y

The values may be vectors and matrices as well as scalars, in any mix, and the targets may be array elements:

@tack.func
def closest_hit(origin, direction):
    ...
    return distance, normal, color              # a scalar and two vectors

@tack.func
def polar(F):
    ...
    return R, S                                 # two matrices

@tack.kernel
def shade(origins, directions, depth, normals, n):
    for i in range(n):
        depth[i], normals[i], color = closest_hit(origins[i], directions[i])

A result made of several values is unpacked into one target each. A device function may also return another call's result as it is (return closest_hit(o, d)), and a name may hold it until it is unpacked (hit = closest_hit(o, d), then depth, normal, color = hit); using such a name as one value is rejected.

Nested Calls

@tack.func functions can call other @tack.func functions:

@tack.func
def square(x):
    return x * x

@tack.func
def length(x, y, z):
    return sqrt(square(x) + square(y) + square(z))

Passing Arrays

Arrays can be passed to @tack.func — they work just like in kernels:

@tack.func
def safe_load(data, i, default_val):
    result = default_val
    if 0 <= i < data.shape[0]:
        result = data[i]
    return result

Texture Sampling in Functions

tex.sample() works inside @tack.func. Tack automatically propagates texture metadata through the inlining:

@tack.func
def sample_bilinear(tex, u, v, w):
    return tex.sample(u, v, w)

@tack.kernel
def render(tex, output, n):
    for i in range(n):
        output[i] = sample_bilinear(tex, 0.5, 0.5, float(i) / float(n))

Local Arrays in Functions

Local arrays (tack.local_array) can be passed to @tack.func. The function accesses the caller's array directly — no copy:

@tack.func
def fill_weights(w, ct, pc0, pc1, pc2):
    for v in range(ct.num_points):
        w[v] = ct.weight(v, pc0, pc1, pc2)

@tack.kernel
def interp(ct, data, out, n):
    for i in range(n):
        w = tack.local_array(tack.f32, ct.num_points)
        fill_weights(w, ct, 0.5, 0.5, 0.5)  # fills w in place
        # ... use w[v] ...

This also works with @tack.device_class methods — see Advanced Features for the cell set pattern.

Reuse a numerical operation across element types

A device function is specialized on the values or arrays it receives. The same linear interpolation expression works for scalars and vectors:

@tack.func
def mix_values(a, b, t):
    return a * (1.0 - t) + b * t

The dye tutorial builds a bilinear sampler from four array loads and applies it to RGB vectors. The MPM tutorial uses matrix-valued arithmetic, and N-body shares one interaction function between the plain and tiled kernels.

Keep all return paths compatible in shape and type. Early returns in a sequential loop of a device function remain unsupported; use a result local and break, then return after the loop. Recursion is unsupported. Local arrays passed to a function are scratch storage shared with that caller iteration, whereas vector and matrix arguments behave as values.