When hawk can’t spell it#
Three ways to reach past hawk’s own arithmetic, in the order to try them.
1. A custom primitive (the common case)#
If the function is a composition of ops hawk already has — log,
exp, norm, dot, the arithmetic operators — and you know its
derivative in closed form, hawk.ext.primitive
is the seam — the extension point where your own code plugs into hawk’s:
register the forward, supply vjp=/jvp=, done. The forward is inlined
(no function call, no separate compiled translation unit), and the
supplied rule replaces
the table’s for that subgraph. This covers almost everything a kernel
author reaches for: a smoothed activation, a regularised norm, a custom
loss term.
2. hawk.raw_device (concept only — not a runnable cell here)#
hawk.raw_device(text, access=..., reads=..., writes=..., returns=...)
splices literal C++ text into the generated kernel body, with the access
class (sample_local, cross_sample_write, …) declared explicitly — never
inferred, because spliced text is opaque to the walk (the pass that
inspects a kernel’s IR to infer things like this automatically). It
exists for
arithmetic hawk’s vocabulary cannot express at all: a hand-tuned numerical
trick, a call into a library hawk has no binding for.
Two honest limits, today:
It is not differentiable.
hawk.diff.vjp/jvprefuse a kernel that calls araw_deviceblock, by name — there is no rule for opaque text, and there will not be one (the point of the seam is that the text is outside hawk’s own IR).There is no end-to-end compile-and-run test for it yet. Its own test suite (
tests/test_trace_raw_device.py) stops at the trace and IR level — it checks that an unannotated block is refused and that a declared access class folds into the walk’s inferred class, but no test compiles and runs a kernel that calls one. This page shows no runnable example because none is tested yet.
Reach for raw_device only when a primitive genuinely cannot express what
you need, and expect to lose autodiff on that kernel.
3. Contribute a native op (the two-repo path)#
If the function is not a composition of existing ops at all — it needs a
new aether::math call, with its own host and device implementation — a
primitive cannot help: there is nothing to compose it from. This is a
contributor change, not an authoring one: a new name across aether
and hawk, six files, two repositories. See
Add a native function end to end for the exact
steps.
I need a function hawk's kernels don't have
|
Is it a COMPOSITION of ops hawk already has
(log, exp, norm, dot, the operators...)?
/ \
yes no -- it needs its
| own aether::math call
Do I know its derivative |
in closed form? contribute a native op
/ \ (3. above --
yes no two repos, six files)
| |
hawk.ext.primitive hawk.raw_device
(1. above -- (2. above -- concept only,
the common case) not differentiable, no
runnable cell here)
Next#
Add a native function end to end — the
contributor path, six files across aether and hawk. Back to
the capstone.