Every way to write and call a kernel

Every way to write and call a kernel#

A kernel is a plain function over one sample. This page is a tour of the forms: what you can return, how the outputs get their names, what a call looks like, what happens on the first build, and what hawk refuses. Each line says what it does. It runs on the CPU; a GPU only changes where the arrays live.

Colab or Kaggle: run this first cell (it installs the packages and downloads the helper files this notebook uses). Elsewhere it only checks the imports.

Hide code cell source

 
# Colab setup, generated by python -m raptor_docs.nb colab_setup (do not edit).
import os

try:
    import google.colab  # noqa: F401
    HOSTED = True  # a Colab runtime
except ImportError:
    HOSTED = "KAGGLE_KERNEL_RUN_TYPE" in os.environ  # a Kaggle kernel

if HOSTED:
    import importlib.metadata, pathlib, re, shutil, subprocess, sys, urllib.request

    gpu = bool(shutil.which("nvidia-smi")) and subprocess.run(["nvidia-smi", "-L"], capture_output=True).returncode == 0
    extra = ""
    if gpu:
        # the CUDA major the runtime already has: a CuPy that is installed, else torch's, else the driver's
        major = None
        for m in ("13", "12"):
            try:
                importlib.metadata.version(f"cupy-cuda{m}x")
                major = m
                break
            except importlib.metadata.PackageNotFoundError:
                pass
        if major is None:
            try:
                import torch
                major = "13" if str(torch.version.cuda).startswith("13") else None
            except ImportError:
                pass
        if major is None:
            smi = re.search(r"CUDA Version: (\d+)", subprocess.run(["nvidia-smi"], capture_output=True, text=True).stdout)
            major = "13" if smi and int(smi.group(1)) >= 13 else "12"
        extra = f"[cuda{major}]"
        print(f"GPU found: installing the {extra} extras to match this runtime.")
    pre = ["--pre"] if os.environ.get("RAPTOR_PRE") else []  # optional: test a release candidate
    ref = os.environ.get("RAPTOR_REF", "main")  # optional: fetch the helper files from this git ref
    if pre or "RAPTOR_REF" in os.environ:
        print(f"Pre-release test: installing with --pre, helper files from {ref}.")
    # our packages with --upgrade, so a runtime that kept an older release gets the current one
    subprocess.check_call([sys.executable, "-m", "pip", "install", "-q", "--upgrade", *pre, f"raptor-hawk{extra}"])
    if not gpu:
        print("Running on the CPU (no GPU in this runtime). For a GPU: " + ("Settings → Accelerator → GPU T4 x2." if "KAGGLE_KERNEL_RUN_TYPE" in os.environ else "Runtime → Change runtime type → T4 GPU."))

try:
    import hawk  # noqa: F401
except ImportError as err:
    raise ImportError(
        "hawk not found. Install with `pip install raptor-hawk`. For a GPU, add [cuda12] (or [cuda13]) to raptor-hawk."
    ) from err

import logging
import os
import sys
import tempfile

os.environ["HAWK_CACHE_DIR"] = tempfile.mkdtemp()      # keep this page's build in a folder of its own
log = logging.getLogger("hawk")                        # show the one line a first build logs
log.setLevel(logging.INFO)
handler = logging.StreamHandler(sys.stdout)
handler.setFormatter(type("Short", (logging.Formatter,), {
    "format": lambda self, r: super(type(self), self).format(r).replace(os.environ["HAWK_CACHE_DIR"], "<cache folder>")})())
log.addHandler(handler)                                # (this page prints the folder as <cache folder>)

import numpy as np
import hawk
from hawk import Mutable, Param, Scalar

x = np.arange(4.0)

Call it like a function#

k(x, a, b) runs the kernel over the whole array. The first call builds it (the log line below says where; delete that folder to rebuild); later calls, and later runs of your program, reuse the cache.

@hawk.kernel
def scale(x: Scalar, a: Param, b: Param) -> Scalar:    # `-> Scalar` is optional; when written, it is checked
    return a * x + b                                   # the returned value is the output

print(scale(x, 2.0, 1.0))          # positional, in the order the kernel declares
print(scale(x=x, a=2.0, b=1.0))    # or by name, in any order

building scale (host) in <cache folder> — delete the directory to rebuild
building scale (cuda) in <cache folder> — delete the directory to rebuild
[1. 3. 5. 7.]
[1. 3. 5. 7.]

What a kernel can return#

One value, a tuple, a list, a dict, or any nesting of them. The call gives back the same shape, with an array in place of each value.

@hawk.kernel
def one(x: Scalar):
    return x + 1.0                                     # one value: you get one array

@hawk.kernel
def pair(x: Scalar):
    return x + 1.0, x * 3.0                            # a tuple: you get a tuple

@hawk.kernel
def many(x: Scalar):
    return [x + 1.0, x * 3.0]                          # a list: you get a list

@hawk.kernel
def named(x: Scalar):
    return {"up": x + 1.0, "triple": x * 3.0}          # a dict: you get a dict, with your keys

@hawk.kernel
def nested(x: Scalar):
    up = x + 1.0
    return {"up": up, "rest": [x * 3.0, (x * 4.0, x * 5.0)]}   # any nesting of the above

for k in (one, pair, many, named, nested):
    print(k.name, "->", k(x))

building one (host) in <cache folder> — delete the directory to rebuild
building one (cuda) in <cache folder> — delete the directory to rebuild
one -> [1. 2. 3. 4.]
building pair (cuda) in <cache folder> — delete the directory to rebuild
building pair (host) in <cache folder> — delete the directory to rebuild
pair -> (array([1., 2., 3., 4.]), array([0., 3., 6., 9.]))
building many (cuda) in <cache folder> — delete the directory to rebuild
building many (host) in <cache folder> — delete the directory to rebuild
many -> [array([1., 2., 3., 4.]), array([0., 3., 6., 9.])]
building named (cuda) in <cache folder> — delete the directory to rebuild
building named (host) in <cache folder> — delete the directory to rebuild
named -> {'up': array([1., 2., 3., 4.]), 'triple': array([0., 3., 6., 9.])}
building nested (cuda) in <cache folder> — delete the directory to rebuild
building nested (host) in <cache folder> — delete the directory to rebuild
nested -> {'up': array([1., 2., 3., 4.]), 'rest': [array([0., 3., 6., 9.]), (array([ 0.,  4.,  8., 12.]), array([ 0.,  5., 10., 15.]))]}

How the outputs are named#

Each returned value is an output slot of the kernel, named after the variable it was returned from (up above is up). A value that is an expression has no name, so it is out when it is alone, and out<i>, numbered by its place in the return, when there are several. The names matter for two things: you can hand an array of your own in by that name, and a gradient is keyed by them. A name may be returned once, and may not repeat a parameter.

def outputs(k):
    return [name for role, name in k.arg_spec if role == "mutable"]

print(outputs(one))        # an expression: `out`
print(outputs(pair))       # several expressions: `out0`, `out1`
print(outputs(nested))     # `up` is named, the others are numbered by place

mine = np.zeros(4)
print(one(x, out=mine))    # your own array for the slot `out`: written in place and returned
print(mine)

['out']
['out0', 'out1']
['out1', 'out2', 'out3', 'up']
[1. 2. 3. 4.]
[1. 2. 3. 4.]

A Mutable is updated in place#

A Mutable parameter is an array the kernel reads and writes. It needs no return; you pass your array and it holds the answer. Mixed with a return, you get the returned values back and the Mutable updated.

@hawk.kernel
def bump(x: Scalar, acc: Mutable[Scalar]):
    acc = acc + x                                      # update in place: no return, the call gives back None

@hawk.kernel
def bump_and_double(x: Scalar, acc: Mutable[Scalar]):
    acc = acc + x                                      # in place, as above
    twice = 2.0 * x
    return twice                                       # and one returned output, named `twice`

acc = np.zeros(4)
print(bump(x, acc), acc)                       # None, and acc has moved
print(bump_and_double(x, acc), acc)            # the returned array, and acc moved again

building bump (host) in <cache folder> — delete the directory to rebuild
building bump (cuda) in <cache folder> — delete the directory to rebuild
[0. 1. 2. 3.] [0. 1. 2. 3.]
building bump_and_double (host) in <cache folder> — delete the directory to rebuild
building bump_and_double (cuda) in <cache folder> — delete the directory to rebuild
[0. 2. 4. 6.] [0. 2. 4. 6.]

The -> annotation#

Optional. When you write it, hawk checks it against what the body returns, at decoration, so a mismatch fails when you define the kernel.

@hawk.kernel
def checked(x: Scalar) -> tuple[Scalar, Scalar]:      # says: two values
    return x + 1.0, x * 3.0

print(checked(x))

building checked (host) in <cache folder> — delete the directory to rebuild
building checked (cuda) in <cache folder> — delete the directory to rebuild
(array([1., 2., 3., 4.]), array([0., 3., 6., 9.]))

What it refuses#

Each refusal names the rule and the way out.

def refused(label, define):
    try:
        define()
    except hawk.HawkError as e:
        print(f"{label}: {e}\n")

def annotation_mismatch():
    @hawk.kernel
    def bad(x: Scalar) -> tuple[Scalar, Scalar]:       # promises two values ...
        return x + 1.0                                 # ... returns one

def class_instance():
    class P:
        def __repr__(self): return "P()"
    @hawk.kernel
    def bad(x: Scalar):
        return P()                                     # only values, tuples, lists and dicts

def name_reuse():
    @hawk.kernel
    def bad(x: Scalar):
        return x                                       # returning a parameter under its own name

def buffer_parameter():
    @hawk.kernel
    def bad(x: Scalar, buffer: Scalar):                # `buffer` is the graph builder's keyword
        return x + 1.0

for label, f in (("annotation mismatch", annotation_mismatch), ("class instance", class_instance),
                 ("a parameter returned as itself", name_reuse), ("a parameter named buffer", buffer_parameter)):
    refused(label, f)

annotation mismatch: kernel 'bad': `-> tuple[Scalar, Scalar]` declares a tuple, but the return is Scalar

class instance: kernel 'bad': the body returned P P(), not a traced value; a class instance is not supported — return a value, a tuple/list, or a dict of them

a parameter returned as itself: kernel 'bad': the return names an output 'x' after a parameter or another output; use another name (x_next)

a parameter named buffer: kernel 'bad': `buffer` is a reserved parameter name — `buffer=` is the keyword a kernel call takes inside an eagle graph builder for the array an output is written to; rename the parameter

Where to read more#