IMPORTANT: To view this page as Markdown, append `.md` to the URL (e.g. /max/get-started.md). For the complete documentation index, see llms.txt.
Skip to main content
For the complete documentation index, see llms.txt. Markdown versions of all pages are available by appending .md to any URL (e.g. /max/get-started.md).

Mojo function

register

def register(name: StringSlice[StaticConstantOrigin], type: StringSlice[StaticConstantOrigin] = StringSlice(""), api: StringSlice[StaticConstantOrigin] = StringSlice(""), arch: StringSlice[StaticConstantOrigin] = StringSlice(""), model: StringSlice[StaticConstantOrigin] = StringSlice(""))

Registers a struct as a kernel implementation for a MAX Graph op.

At compile time, the Graph Compiler selects the most specific registered kernel that matches the runtime device. The four device fields narrow the target from coarse to fine:

  • type β€” broad device category: "cpu", "gpu", or a custom label.
  • api β€” compute backend: "cuda", "hip", "metal", etc.
  • arch β€” microarchitecture: "sm_100a", "sm_90a", "gfx942", etc.
  • model β€” exact hardware model: "NVIDIA B200", "AMD Instinct MI355X", etc.

An empty string for any field acts as a wildcard β€” it matches any device value for that field. Fields are ordered from least to most specific: type < api < arch < model. Among all matching candidates, the most specific registration wins: a registration that sets arch is more specific than one that sets only type and api, because arch sits higher in the hierarchy.

If a user-defined kernel and a built-in kernel are equally specific for a device, the user-defined kernel takes precedence. Registering the same op with identical device fields more than once β€” whether within the same library or across multiple user libraries β€” is an error reported by the Graph Compiler.

Example β€” device field combinations and their reach:

# Matches all devices (backward-compatible default).
@compiler.register("mo.matmul")

# Matches all CPU devices.
@compiler.register("mo.matmul", type="cpu")

# Matches all GPU devices regardless of vendor or architecture.
@compiler.register("mo.matmul", type="gpu")

# Matches all CUDA GPUs (any architecture).
@compiler.register("mo.matmul", type="gpu", api="cuda")

# Matches only the NVIDIA SM100A architecture.
@compiler.register("mo.matmul", type="gpu", api="cuda", arch="sm_100a")

# Matches only the NVIDIA B200.
@compiler.register("mo.matmul", type="gpu", api="cuda", arch="sm_100a", model="NVIDIA B200")

Example β€” selection with multiple registrations in scope:

@compiler.register("mo.matmul")                                          # wildcard
@compiler.register("mo.matmul", type="cpu")                              # least specific
@compiler.register("mo.matmul", type="gpu")                              # least specific
@compiler.register("mo.matmul", type="gpu", api="cuda", arch="sm_100a")  # more specific
  • SM100A CUDA GPU β†’ all four match; arch registration wins (most specific).
  • SM90A CUDA GPU β†’ wildcard and type="gpu" match; type="gpu" wins.
  • CPU β†’ wildcard and type="cpu" match; type="cpu" wins.
  • Some NPU β†’ only the wildcard matches; wildcard wins.

Args:

  • ​name (StringSlice[StaticConstantOrigin]): The MAX Graph op name to register (e.g. "mo.matmul").
  • ​type (StringSlice[StaticConstantOrigin]): Broad device category β€” "cpu", "gpu", or a custom accelerator label such as "npu-xxx". Corresponds to the label used in DeviceRef (e.g. DeviceRef.GPU()). Empty matches all.
  • ​api (StringSlice[StaticConstantOrigin]): Programming API or compute backend for the device β€” for example "cuda" (NVIDIA CUDA), "hip" (AMD ROCm/HIP), or "metal" (Apple Metal). Empty matches all.
  • ​arch (StringSlice[StaticConstantOrigin]): Microarchitecture identifier β€” for example "sm_90a" (Hopper), "sm_100a" (Blackwell), or "gfx942" (AMD CDNA3). Empty matches all.
  • ​model (StringSlice[StaticConstantOrigin]): Specific device model β€” for example "NVIDIA B200" or "AMD Instinct MI355X". Empty matches all.