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.