Example: one Kind, two force models

Example: one Kind, two force models#

One task: declare a kind once and reuse it for more than one kernel — the point of factoring a vocabulary out of an authored kernel at all. Two different accelerations share the same declared shape (a Vector[3] velocity in, a returned Vector[3] acceleration out) and therefore the same Kind.

Time: ~3 min · Runs on: CPU · You need: Writing a vocabulary I

Hide code cell source

 
import sys, pathlib
sys.path.insert(0, str(pathlib.Path.cwd().parent / "_shared"))
from nb_helpers import plot_style
plot_style()

from hawk import Param, Vector
from hawk.ext import KernelKind, Output
from hawk.math import norm


class Accel(KernelKind, slug="accel"):
    output = Output.returned(Vector[3], slot="acc")


@Accel
def linear_drag(velocity: Vector[3], k: Param):
    return (-k) * velocity


linear_drag

<Kernel linear_drag slots=3>

A second model on the same Kind:

@Accel
def quadratic_drag(velocity: Vector[3], k: Param):
    return (-k * norm(velocity)) * velocity


quadratic_drag

<Kernel quadratic_drag slots=3>

Build both kernels together, then run each:

import pathlib, tempfile

import hawk

work_dir = pathlib.Path(tempfile.mkdtemp())
hawk.build([linear_drag, quadratic_drag], work_dir, targets=("host",));

import numpy as np

velocity = np.array([[2.0, 4.0], [0.0, 0.0], [0.0, 0.0]])

acc_linear = np.zeros((3, 2))
hawk.run(hawk.load(work_dir, "linear_drag"), velocity=velocity, k=0.5, acc=acc_linear)

acc_quad = np.zeros((3, 2))
hawk.run(hawk.load(work_dir, "quadratic_drag"), velocity=velocity, k=0.5, acc=acc_quad)

print("linear drag:   ", acc_linear[0])
print("quadratic drag:", acc_quad[0])

linear drag:    [-1. -2.]
quadratic drag: [-2. -8.]

Both kernels were traced (turned into IR), emitted (that IR rendered as C++/CUDA source) and compiled through the exact same Kind — the same output set, sink policy (how a kernel’s results commit — here, one returned Vector[3] slot) and guard (which samples a launch runs; Accel takes the default, which gates on a bound terminated slot IF one is declared — Accel declares none, so every sample runs unconditionally) — so a new force model in this family costs one function body, not a redeclaration of its shape.