Tiling N-body forces¶
Each body accumulates the gravitational contribution of every other body. The plain implementation is already parallel over bodies. A GPU version can reuse positions within a workgroup by loading them into shared memory in tiles.
Origin: adapted from NVIDIA Warp's tile N-body example, Copyright © NVIDIA CORPORATION & AFFILIATES, Apache-2.0. The algorithm follows Lars Nyland, Mark Harris and Jan Prins, Fast N-Body Simulation with CUDA, GPU Gems 3 chapter 31. This teaching version calculates one acceleration field, uses seeded normal positions, and implements both ordinary and shared-memory kernels rather than Warp's tile API. It does not integrate trajectories.
Begin with the portable kernel¶
The pair interaction is r / (dot(r,r) + epsilon²)^(3/2) for unit masses.
Softening keeps the denominator nonzero. The self contribution is zero because
its displacement is zero.
@tack.kernel
def acceleration_plain(pos, acceleration, n):
for i in range(n):
a = tack.vec3(0.0, 0.0, 0.0)
for j in range(n):
a += interaction(pos[i], pos[j])
acceleration[i] = a
Iteration i owns acceleration[i]. The inner loop is sequential. This costs
O(n²) interactions and works on every backend.
Cooperate on each tile¶
Each 256-lane workgroup processes a tile of 256 positions. Every lane loads one position; every lane then reads all 256 positions from shared storage:
TILE = tack.constant(256)
@tack.kernel
def acceleration_tiled(pos, acceleration, n):
for i in range(n):
xyz = tack.shared(tack.f32, TILE * 3) # folded to 768 scalar components
lane = tack.thread_id()
a = tack.vec3(0.0, 0.0, 0.0)
for tile in range(n // TILE):
q = pos[tile * TILE + lane]
xyz[3 * lane], xyz[3 * lane + 1], xyz[3 * lane + 2] = q.x, q.y, q.z
tack.barrier() # finish filling before reading
for j in range(TILE):
a += interaction(pos[i], [xyz[3 * j], xyz[3 * j + 1], xyz[3 * j + 2]])
tack.barrier() # finish reading before reusing
acceleration[i] = a
The first barrier ensures the tile is filled before any lane reads it. The second ensures every lane has finished reading before the next tile overwrites it. Both are outside lane-dependent conditions. Keep every lane participating; an early exit around either barrier violates the collective contract.
The shared array packs three coordinates per position. TILE * 3 is a constant
expression, so Tack resolves it to 768 scalars before generating GPU code. It
is workgroup storage, unlike a local_array, which would give each lane a
private array and no reuse. TILE = 256 matches Tack's default 256-lane
workgroup; changing the constant alone does not change that launch configuration.
Choose a capability, not a backend name¶
if get_backend().supports_workgroups:
if n % TILE:
raise ValueError("the tiled example needs a multiple of 256 bodies")
acceleration_tiled(pos, acceleration, n)
else:
acceleration_plain(pos, acceleration, n)
CPU has no workgroup execution model and runs the plain kernel. The GPU path requires a body count divisible by 256. This example rejects partial tiles; it does not silently drop remaining bodies or introduce padding. A general padded implementation must keep all lanes at both barriers and give padded bodies zero physical contribution.
Validate before timing¶
uv run python docs/examples/nbody.py --arch cpu --check
uv run python docs/examples/nbody.py --arch metal --check
uv run --with vtk python docs/examples/nbody.py --output nbody.png
The check compares the acceleration with a vectorized NumPy float64 reference.
On a GPU it also compares the tiled result with the plain kernel. A CPU check
alone does not execute shared memory or barriers. The VTK view shows an XY projection
of 3-D positions and a subset of projected acceleration arrows.
For performance work, compile both kernels first, reuse their arrays, time warm
calls at the same n, and record the selected device. GPU launches are
queued, so call tack.sync() before reading the clock (see
timing). Shared memory reduces
repeated position loads within a workgroup, but it also adds barriers and uses
resources. Measure the net effect; this tutorial promises no universal speedup.
The arithmetic remains O(n²).
Complete program¶
Download nbody.py. The adapted code is Apache-2.0 licensed.
# Copyright (c) 2026 Kitware, Inc.
# SPDX-License-Identifier: Apache-2.0
"""All-pairs and tiled N-body forces, adapted from NVIDIA Warp's tile example.
Copyright (c) 2022 NVIDIA CORPORATION & AFFILIATES. Apache-2.0;
see LICENSE-APACHE-2.0.txt.
Source: https://github.com/NVIDIA/warp/blob/main/warp/examples/tile/example_tile_nbody.py
Algorithm: Lars Nyland, Mark Harris, Jan Prins, GPU Gems 3, chapter 31.
"""
import argparse
import numpy as np
import tack
from tack.runtime.dispatch import get_backend
@tack.func
def interaction(p, q):
r = q - p
inv = 1.0 / sqrt(r.dot(r) + 0.01)
return r * inv * inv * inv
@tack.kernel
def acceleration_plain(pos, acceleration, n):
for i in range(n):
a = tack.vec3(0.0, 0.0, 0.0)
for j in range(n):
a += interaction(pos[i], pos[j])
acceleration[i] = a
TILE = tack.constant(256)
@tack.kernel
def acceleration_tiled(pos, acceleration, n):
for i in range(n):
xyz = tack.shared(tack.f32, TILE * 3) # folded to 768 scalar components
lane = tack.thread_id()
a = tack.vec3(0.0, 0.0, 0.0)
for tile in range(n // TILE):
q = pos[tile * TILE + lane]
xyz[3 * lane], xyz[3 * lane + 1], xyz[3 * lane + 2] = q.x, q.y, q.z
tack.barrier() # finish filling before reading
for j in range(TILE):
a += interaction(pos[i], [xyz[3 * j], xyz[3 * j + 1], xyz[3 * j + 2]])
tack.barrier() # finish reading before reusing
acceleration[i] = a
def run(arch="cpu", check=False):
tack.set_default_device(arch)
n = 256 if check else 512
xyz = np.random.default_rng(42).normal(size=(n, 3)).astype(np.float32)
pos = tack.empty((n,), tack.vec3(tack.f32))
acceleration = tack.empty((n,), tack.vec3(tack.f32))
tack.copyto(pos, xyz)
if get_backend().supports_workgroups:
if n % TILE:
raise ValueError("the tiled example needs a multiple of 256 bodies")
acceleration_tiled(pos, acceleration, n)
else:
acceleration_plain(pos, acceleration, n)
result = acceleration.numpy()
if check:
r = xyz.astype(np.float64)[None, :, :] - xyz.astype(np.float64)[:, None, :]
distance_sq = (r * r).sum(axis=-1) + 0.01
reference = (r / distance_sq[..., None] ** 1.5).sum(axis=1)
np.testing.assert_allclose(result, reference, atol=8e-4, rtol=3e-5)
if get_backend().supports_workgroups:
baseline = tack.empty((n,), tack.vec3(tack.f32))
acceleration_plain(pos, baseline, n)
np.testing.assert_allclose(result, baseline.numpy(), atol=8e-4, rtol=3e-5)
return xyz, result
def save_figure(result, output):
import vtkmodules.vtkRenderingFreeType
import vtkmodules.vtkRenderingOpenGL2 # noqa: F401
from vtkmodules.vtkCommonCore import VTK_TEXT_CENTERED, VTK_TEXT_TOP, vtkFloatArray, vtkPoints
from vtkmodules.vtkCommonDataModel import vtkPolyData
from vtkmodules.vtkFiltersCore import VTK_SCALE_BY_VECTOR, VTK_USE_VECTOR, vtkGlyph3D
from vtkmodules.vtkFiltersGeneral import vtkVertexGlyphFilter
from vtkmodules.vtkFiltersSources import vtkArrowSource
from vtkmodules.vtkIOImage import vtkPNGWriter
from vtkmodules.vtkRenderingCore import (
VTK_NORMALIZED_VIEWPORT,
vtkActor,
vtkPolyDataMapper,
vtkRenderer,
vtkRenderWindow,
vtkTextActor,
vtkWindowToImageFilter,
)
pos, acceleration = result
# Project both positions and forces into the x-y plane.
xyz = pos.copy()
xyz[:, 2] = 0
# vtkFloatArray keeps the float32 positions: vtkPoints(data=numpy) would store doubles.
bodies = vtkPolyData(points=vtkPoints(data=vtkFloatArray(xyz)))
mapper = vtkPolyDataMapper()
bodies >> vtkVertexGlyphFilter() >> mapper
actor = vtkActor(mapper=mapper)
actor.property.color = (0.15, 0.39, 0.92)
actor.property.point_size = 4
actor.property.render_points_as_spheres = True
samples = vtkPolyData(points=vtkPoints(data=vtkFloatArray(np.ascontiguousarray(xyz[::8]))))
scaled = acceleration / (np.linalg.norm(acceleration, axis=1).max() + 1e-12)
scaled[:, 2] = 0
# Attribute data has no settable vectors property: SetVectors makes them active.
samples.point_data.SetVectors(vtkFloatArray(np.ascontiguousarray(scaled[::8])))
glyphs = vtkGlyph3D(vector_mode=VTK_USE_VECTOR, scale_mode=VTK_SCALE_BY_VECTOR,
scale_factor=0.8, orient=True)
arrow = vtkArrowSource(shaft_radius=0.04, tip_radius=0.15)
# The glyph source is input port 1, which has no property or >> form.
glyphs.SetSourceConnection(arrow.output_port)
arrow_mapper = vtkPolyDataMapper(scalar_visibility=False)
samples >> glyphs >> arrow_mapper
arrows = vtkActor(mapper=arrow_mapper)
arrows.property.color = (0.92, 0.35, 0.05)
arrows.property.lighting = False
renderer = vtkRenderer(background=(1, 1, 1))
renderer.AddActor(actor)
renderer.AddActor(arrows)
camera = renderer.active_camera
camera.position = (0, 0, 1)
camera.focal_point = (0, 0, 0)
camera.parallel_projection = True
camera.parallel_scale = float(np.abs(xyz).max()) * 1.15
title = vtkTextActor(input="Bodies and gravitational acceleration (x-y)", position=(0.5, 0.98))
title.position_coordinate.coordinate_system = VTK_NORMALIZED_VIEWPORT
text = title.text_property
text.justification = VTK_TEXT_CENTERED
text.vertical_justification = VTK_TEXT_TOP
text.font_size = 20
text.color = (0, 0, 0)
renderer.AddViewProp(title)
window = vtkRenderWindow(off_screen_rendering=True, multi_samples=0, size=(640, 640))
window.AddRenderer(renderer)
window.Render()
capture = vtkWindowToImageFilter(input=window, read_front_buffer=False)
writer = vtkPNGWriter(file_name=str(output))
capture >> writer
writer.Write() # calling the writer, or updating it, writes raw pixels instead of a PNG
window.Finalize()
if __name__ == "__main__":
parser = argparse.ArgumentParser(description=__doc__)
parser.add_argument("--arch", default="cpu", choices=["cpu", "metal", "cuda", "hip", "level_zero"])
parser.add_argument("--check", action="store_true")
parser.add_argument("--output", help="Save a figure (requires VTK)")
args = parser.parse_args()
result = run(args.arch, args.check)
if args.output:
save_figure(result, args.output)
print("N-body: check passed" if args.check else "N-body: force calculation complete")