From Python to a Tack program¶
A Tack application contains two kinds of code. Host code is ordinary Python:
it allocates arrays, prepares data, chooses which kernels to run, repeats timesteps,
and displays the results. Device code is the body of a @tack.kernel or
@tack.func: Tack reads that source and compiles its supported numerical operations
for the selected backend. On the CPU, device code still means compiled code.
Understanding that boundary makes it much easier to turn a NumPy calculation into a program that runs on a GPU.
Start with one operation per element¶
Suppose you want out = alpha * x + y. NumPy expresses the whole array operation
at once. A Tack kernel expresses what happens to one element and provides the
iteration space:
import numpy as np
import tack
tack.set_default_device("cpu")
x = tack.array(np.arange(8, dtype=np.float32))
y = tack.array(np.ones(8, dtype=np.float32))
out = tack.empty((8,))
@tack.kernel
def saxpy(x, y, out, alpha):
for i in range(x.shape[0]):
out[i] = alpha * x[i] + y[i]
saxpy(x, y, out, 2.0)
np.testing.assert_array_equal(out.numpy(), np.arange(8, dtype=np.float32) * 2 + 1)
The allocation, call and assertion run in Python. The arithmetic in saxpy runs
as compiled code. The loop has eight independent iterations. Tack may schedule
them in any order; the program is correct because iteration i owns out[i].
A call describes a whole launch¶
A kernel has one parallel for directly in its body. Nested loops run sequentially
within an iteration. For a row sum, assign one iteration to a row and let that
iteration walk its columns:
@tack.kernel
def row_sums(matrix, sums, rows, columns):
for row in range(rows):
total = 0.0
for column in range(columns):
total += matrix[row, column]
sums[row] = total
Here total belongs to one iteration. A local variable does not communicate
between threads. Sequential statement order inside an iteration is preserved;
it does not establish an order between rows.
For a serial recurrence, use for _ in range(1) with an ordinary loop inside it.
For multidimensional work, use tack.ndindex. See Kernels for
stepped ranges, interior ranges and variable-length inner loops.
Separate stages into separate kernels¶
Consider particle motion: first compute every force from the old positions, then update every position. Writing positions while other iterations still read them mixes old and new state. Put the two stages in separate kernels:
@tack.kernel
def compute_force(pos, force, n):
for i in range(n):
force[i] = -pos[i] # independent restoring force
@tack.kernel
def integrate(pos, vel, force, dt, n):
for i in range(n):
vel[i] += dt * force[i]
pos[i] += dt * vel[i]
for step in range(100): # Python controls the simulation
compute_force(pos, force, n)
integrate(pos, vel, force, 0.01, n)
Every launch sees the results of the launches before it, and every host
access to an array -- numpy(), tack.copyto, indexing, reductions -- sees
the results of every launch before it. So the second kernel sees the
completed force array, and so does a read after the loop. On the CPU a
call completes before Python continues; on CUDA, HIP, Metal and Level Zero
it is queued and returns at once, and the host access waits. The program
above behaves the same either way, and there is no stream or event to
manage. tack.sync() waits for everything launched so far on every device, for timing
kernels or for memory another library reads without Tack.
Waiting from asyncio¶
A host access that waits blocks its thread. In an application with an event loop -- a trame or web server, a Jupyter kernel -- that thread is the loop, and nothing else runs while the device works. Two coroutines wait in a worker thread instead:
async def frame():
advance(state, dt) # queued; returns at once
await tack.sync_async() # the loop runs while it works
pixels = await image.numpy_async() # wait, then copy, off the loop
return pixels
tack.sync_async() waits for the launches made on every device
before the call, not for ones another task queues while it waits. numpy_async returns every
write queued before the call; one queued later by another task may or may
not be in it, so order the two with await if it matters. A launch's error
is raised at the await. On CUDA, HIP, Metal and Level Zero the wait runs
in a worker thread; on the CPU, whose launches complete before returning,
both return at once.
A tack.barrier() synchronizes only a GPU workgroup. It cannot turn two stages
over a whole grid into one ordered launch. Use completed kernel calls for that.
The memory chapter explains stencil buffers,
scatter operations, and the cases where atomics are appropriate.
Choose what changes between calls¶
There are three common ways to give device code a numerical value:
| Kind of value | How to supply it | Example |
|---|---|---|
| A coefficient changing each call | Scalar argument | advance(pos, vel, dt, n) |
| A fixed model parameter | tack.constant |
GRAVITY = tack.constant(9.8) |
| An array of state | Array argument or device-class attribute | Particle positions, a grid, a counter |
An ordinary Python global such as gravity = 9.8 is not captured in device code.
Declare a constant or pass it. Changing a runtime argument changes the next call
without changing the kernel source. Constants are fixed when the specialization
is first lowered; rebinding a Python name later does not update cached code.
@tack.device_class gathers arrays and parameters into an object. Its instance
scalars are runtime arguments; its class constants and method structure help
define the compiled specialization. The Templates chapter
shows how to use this to organize a solver.
Shapes, values and compilation¶
The first call includes compilation. Later compatible calls reuse the compiled variant. Changing a scalar value usually reuses it. Changing the scalar's dtype, array dtype, vector width, template structure or a baked-in dimension may create another variant.
This distinction matters for a grid. A 1-D length used only as the launch bound can vary without recompilation; multidimensional indexing uses row strides that depend on the shape. Two grids with different shapes may need different code even when the Python kernel is identical. Prefer fixed working arrays during a simulation; allocate once and update them each frame.
See Kernel Caching for examples and Debugging and timing for how to measure warm calls.
The kernel language uses Python syntax¶
Many familiar constructs work: arithmetic, if, sequential for and while,
vectors, matrices and device-function calls. Arbitrary Python objects, NumPy
calls and file operations do not run inside a kernel. For example, generate an
initial state with NumPy on the host, upload it, then work on the array:
rng = np.random.default_rng(42)
initial = rng.normal(size=1024).astype(np.float32) # Python / NumPy
data = tack.array(initial)
Use tack.random for random draws within a kernel. Use @tack.func for reusable
compiled helpers. Such helpers can be called from kernels and other device
functions, rather than as ordinary host functions.
Write complete examples in a .py file: Tack needs the decorated function's
source to compile it. Dynamically generated functions and interactive contexts
that cannot supply their source are poor starting points for learning the API.
Before tackling a larger program, run Heat diffusion. It combines allocation, a multidimensional kernel, a Python timestep loop and a NumPy reference in one small script.