Adding a Primitive

This guide explains how to implement a new primitive. It will primarily focus on how to do this. See the internals article for more information on how primitives work.

In general, there are two main reasons to add a new primitive:

The operation is not expressible using existing primitives.
The operation (or, for example, its derivative) would benefit from a dedicated implementation.

In order to be able to add a new primitive, it needs to be expressible in the {stablehlo} package. There are various scenarios:

The operation is expressible in the StableHLO language and:
1. All required operations are already implemented in the {stablehlo} R package.
  
  $\rightarrow$ This is the simplest case and the one we assume in this guide.
2. One or more required operations are not already implemented in the {stablehlo} R package.
  
  $\rightarrow$ You first need to implement the missing operations in the {stablehlo} R package. See this issue for which operations are missing and the StableHLO specification for which are available.
The operation is not available/cannot be efficiently expressed in StableHLO, because:
1. It requires shape dynamism (output shape depends on data, such as subsetting with a boolean mask computed from the inputs):
  
  $\rightarrow$ This is currently not possible, but we hope we can add support for this in the future.
2. It cannot be expressed or can only be expressed inefficiently:
  
  $\rightarrow$ You can implement a StableHLO custom call, see the custom print operation in pjrt. The {pjrt} package has a dedicated “Adding Custom Calls” tutorial. This mechanism is used, for example, to implement some of the linear algebra primitives such as SVD: on CPU these call into LAPACK, and on GPU they use NVIDIA’s cuSOLVER library. In the future it will also be possible to call into user-provided custom CUDA kernels, but there is no infrastructure for this yet – currently we only call into predefined kernels shipped with the CUDA toolkit.

Adding a Primitive: Practical Example

Let’s add a new primitive step by step. We’ll implement prim_repeat_along – a primitive that repeats an array multiple times along a specified axis.

For example, repeating c(1, 2, 3) twice along axis 1 gives c(1, 2, 3, 1, 2, 3).

This primitive has a dynamic input (an array) and two static parameters (how many times to repeat and which axis).

Step 1: Define the Primitive

Primitives are created with new_primitive(): it builds the AnvlPrimitiveDef metadata object that holds the rules, wraps the body with jit(), and attaches the metadata. The returned callable becomes the primitive and is bound to a prim_<name> R symbol. The name passed to new_primitive() is that same <name>, not the StableHLO op the primitive lowers to, so an error or a printed graph naming a primitive names a function the reader can look up.

library(anvl)
prim_repeat_along <- new_primitive(
  "repeat_along",
  function(x, times, axis) {
    infer_fn <- function(x, times, axis) {
      if (!checkmate::test_integerish(axis, lower = 1, upper = naxes(x), len = 1L)) {
        cli::cli_abort("{.arg axis} must be between 1 and {naxes(x)}, but is {.val {axis}}.")
      }
      if (!checkmate::test_integerish(times, lower = 1, len = 1L)) {
        cli::cli_abort("{.arg times} must be a positive integer, but is {.val {times}}.")
      }
      new_shape <- shape(x)
      new_shape[axis] <- new_shape[axis] * times
      list(AbstractArray(
        dtype = dtype(x),
        shape = Shape(new_shape)
      ))
    }

    graph_desc_add(
      self,                       # lexically bound to the AnvlPrimitiveDef
      list(x = x),                # Dynamic inputs (arrays)
      params = list(              # Static parameters
        times = times,
        axis = axis
      ),
      infer_fn = infer_fn
    )[[1L]]  # Extract single output from list
  },
  static = c("times", "axis")
)

The primitive is now callable directly as prim_repeat_along(x, times, axis).

Key points:

Pass the lexically-bound self (the [AnvlPrimitiveDef]) as the first argument to graph_desc_add(). new_primitive() installs self into the enclosing environment of the body, so this is always available inside a primitive.
The inference function receives abstract arrays (types, not values) for dynamic inputs, and actual values for static parameters. It must verify that the arguments to the function are valid, and should throw clear error messages — {anvl} programs are otherwise hard to debug.
graph_desc_add() returns a list of outputs; use [[1L]] for single-output primitives.

Every argument that is not a dynamic array must be listed in static. new_primitive calls jit() internally, so the primitive runs on whichever backend is active when it is called.

Here we define the inference function inline to keep the example self-contained. Inside {anvl}, inference rules live together in R/rules-inference.R, named infer_<primitive>(). Its formals are the primitive’s own – every operand and every static parameter, including ones the rule has no use for – because graph_desc_add() calls it as do.call(infer_fn, c(avals_in, params)). It returns a plain list() of abstract arrays, named when the primitive has named outputs.

If your primitive lowers to a StableHLO operation, implement the constraints from the StableHLO specification directly, in {anvl}’s vocabulary: arrays rather than tensors, axes rather than dimensions, and axis numbers 1-based. Many rules need nothing of their own – an elementwise operation reuses infer_generic_biv(), infer_float_uni() and the like.

Handling of R inputs

One question that the primitive needs to answer is how it materializes R inputs: The default rule is to materialize an R input at its default data type, which on pjrt is f32 for doubles and i32 for integers, and bool for logicals everywhere (see default_dtypes()).

For functions taking a single dynamic input, as does prim_repeat_along from above, this is almost always the right thing to do. The inputs section of the printed graph shows both halves of it: the input’s data type is the default the value materialized at, and the <- double records that the caller supplies an R double rather than an array that already had one.

trace_fn(prim_repeat_along, list(nv_aval("double", c(2, 3)), 2, 1))
#> <AnvlGraph> (%x1: f32[2,3] <- double) {
#>   %1: f32[4,3] = repeat_along [times = 2, axis = 1] (%x1)
#>   return %1
#> }

Where the default goes wrong is when a primitive has several operands that must agree. Materializing each at its own default would make prim_add(1, nv_scalar(2, "f64")) an error – the literal would become the default float and could meet an f64 – where it should give f64. In such cases, you should call apply_promotion() in the primitive and pass the rules you want to apply. This needs to be done before calling graph_desc_add(). These should not change the data type of any AnvlArray inputs, which is what the promotion_rdata_common() rule is for. See the Data Types and Promotion Rules article for more information on these rules.

Special case: primitives with 0 dynamic inputs

Array constructors (e.g. prim_fill, prim_iota) have no dynamic array inputs – every argument is static. Two things change in this case:

All formals must be listed in static. Otherwise {anvl} would try to interpret a scalar argument like value or shape as an AnvlArray input.
The device cannot be inferred from inputs, because there are no array inputs. Instead, add a device formal to the function (it goes in static like every other one) and pass it on to graph_desc_add(), which declares it into the trace. The device the caller asked for – a device of the active backend, or a device name – then counts for device inference exactly like the device of an array input.

For example, the real definition of prim_fill looks like this:

prim_fill <- new_primitive(
  "fill",
  function(value, shape, dtype, device = NULL) {
    # ... graph_desc_add(self, ..., device = device) ...
  },
  static = 1:4
)

See graph_desc_add() for the detailed semantics.

Shortcut helpers for common shapes

prim_repeat_along has its own parameters and needs a custom body. Many primitives are simpler – elementwise unary or binary ops, reductions, and comparisons all share a standard shape. R/primitives.R exposes factories that generate the body for you:

make_unary_op(infer_fn) – elementwise unary (e.g. prim_abs, prim_negate).
make_binary_op(infer_fn) – elementwise binary (e.g. prim_add, prim_mul).
make_reduce_op(infer_fn) – reductions with axes / drop parameters (e.g. prim_sum).
make_compare_op(direction) – comparison ops with a fixed direction string (e.g. prim_eq, prim_lt).

Used together with the matching inference rule, the primitive definition collapses to a one-liner:

prim_add <- new_primitive("add", make_binary_op(infer_generic_biv))
prim_negate <- new_primitive("negate", make_unary_op(infer_numeric_uni))
prim_sum <- new_primitive("sum", make_reduce_op(), static = 2:3)

Reach for the manual graph_desc_add() form when the primitive has extra static parameters, a custom inference function, or multiple outputs.

Step 2: Add the StableHLO Rule

The StableHLO rule defines how to lower this primitive to actual operations, which is used in the stablehlo() lowering pass. We implement repeat_along using concatenation:

prim_repeat_along[["stablehlo"]] <- function(x, times, axis) {
  xs <- rep(list(x), times)
  list(rlang::exec(stablehlo::hlo_concatenate, !!!xs, dimension = axis - 1L))
}

The rule receives:

Dynamic inputs as stablehlo::FuncValues.
Static parameters as their R values.

It must return a list of stablehlo::FuncValues, even if there is only one output.

Important: StableHLO uses 0-based indexing, while {anvl} uses R’s 1-based indexing. Always convert axis indices by subtracting 1. Nothing translates StableHLO’s error messages back into {anvl}’s terms: an input that only StableHLO rejects surfaces as an MLIR error in its own vocabulary, with 0-based axes. This is why the inference function has to check every constraint the StableHLO operation imposes, so that the lowering rule only ever sees valid inputs.

Step 3: Add the Reverse Rule

If the operation should support automatic differentiation, attach a reverse rule built with rule_reverse(). The idea here is the following, where we assume the input x has shape (s_1, ..., s_n), which means that the output (and therefore its gradient) has shape (s_1, ..., s_{axis-1}, s_axis * times, s_{axis+1}, ..., s_n).

Reshape the gradient to (s_1, ..., s_{axis-1}, s_axis, times, s_{axis+1}, ..., s_n).
Sum over the times axis and drop the times axis.

prim_repeat_along[["reverse"]] <- rule_reverse(function(inputs, outputs, grads, params, required) {
  if (!required[[1L]]) {
    return(list(NULL))
  }

  grad <- grads[[1L]]
  x <- inputs[[1L]]
  axis <- params$axis
  times <- params$times

  old_shape <- shape(x)
  grad_shape <- shape(grad)

  new_shape <- grad_shape
  new_shape[axis] <- old_shape[axis]
  new_shape <- append(new_shape, times, after = axis)

  grad_reshaped <- prim_reshape(grad, new_shape)
  grad_summed <- prim_sum(grad_reshaped, axes = axis + 1L, drop = TRUE)
  list(grad_summed)
})

The wrapped backward receives:

inputs: Input GraphValues from the forward pass
outputs: Output GraphValues from the forward pass
grads: Gradients flowing back from downstream (one per output)
params: Named list of the statement’s static parameters (here: params$axis, params$times)
required: Logical vector indicating which input gradients are needed

It returns a list with one gradient per input (or NULL if not required).

Note that the reverse rule will only be called if at least one input gradient is required, so this can be assumed.

Optional: alternative-forward rule

For most primitives the backward-only form above is enough. If the gradient computation can be made significantly more efficient by running an alternative forward pass that exposes intermediate values for reuse, use the alternative-forward form: rule_reverse(forward = ...). The forward hook receives (inputs, params, required) – required says which inputs need a gradient –, emits whatever forward primitives it likes via the normal prim_* callables, and returns a list with the primal outputs and a backward closure. Intermediates flow from forward to backward via R’s lexical scoping:

prim_<name>[["reverse"]] <- rule_reverse(forward = function(inputs, params, required) {
  x <- inputs[[1L]]
  y <- prim_some_op(x)        # forward emit 1
  z <- prim_other_op(x, y)    # forward emit 2; both `x` and `y` captured
  list(
    outputs  = list(z),       # one box per output of the original statement
    backward = function(inputs, outputs, grads, params, required) {
      # `x` is also available via `inputs`; what lexical capture buys us is
      # access to intermediates like `y` that the framework doesn't pass in.
      ...
    }
  )
})

The backward closure shares the same signature as the backward-only form. Intermediates that aren’t passed in (like y above) flow in via lexical capture. The forward must return one output box per output of the original statement, with matching shape/dtype.

Optional: a quickr rule

prim_<name>[["stablehlo"]] and prim_<name>[["reverse"]] are the two rules you almost always want. A primitive can optionally also carry a quickr rule, which lowers it to plain R code for the quickr backend (see R/rules-quickr.R). The quickr rule is only required if you want the primitive to run under local_backend("quickr"); if you skip it, the primitive will still work on the pjrt backend.

Step 4: Add an `nv_` API Function

In {anvl}, we also offer convenience wrappers around the primitives. An example is prim_add vs nv_add, where the latter calls into the former after optionally broadcasting (scalar) inputs:

nv_add(1L, nv_array(2:3))
#> AnvlArray
#>  3
#>  4
#> [ CPUi32{2} ]
prim_add(1L, nv_array(2:3))
#> Error in `prim_add()`:
#> ! `lhs` and `rhs` must have the same array type.
#> ✖ Got i32[] and i32[2].

In our case, no such convenience is needed and the functionality is not too low-level (for it to be generally useful), so we can just reassign the prim_* function to an nv_* function:

nv_repeat_along <- prim_repeat_along

Note that in the nv_* wrapper function, the input may still be a bare R value – one that has no data type yet (see ?RData). shape() and naxes() answer for it as they do for an array. For the data type there is peek_dtype(), which reports the type the value would materialize at; dtype() errors on an R value, because there is nothing to report until it is used.

Using Your Primitive

You can now use the primitive, both eagerly and inside a larger JIT-compiled function:

x <- nv_array(c(1, 2, 3), shape = c(3, 1))
# Eager call -- works because prim_repeat_along is itself jit-compiled.
prim_repeat_along(x, times = 2L, axis = 2L)
#> AnvlArray
#>  1 1
#>  2 2
#>  3 3
#> [ CPUf32{3,2} ]
# Traced into the outer graph when composed with another jit-compiled function.
jit(function(x) prim_repeat_along(x, times = 2L, axis = 2L))(x)
#> AnvlArray
#>  1 1
#>  2 2
#>  3 3
#> [ CPUf32{3,2} ]

And compute gradients through it.

f <- function(x) {
  repeated <- prim_repeat_along(x, times = 2L, axis = 2L)
  sum(repeated)
}

grad_f <- jit(gradient(f))
grad_f(x)[[1L]]
#> AnvlArray
#>  2
#>  2
#>  2
#> [ CPUf32{3,1} ]

Contributing to the Package

If you want to contribute your primitive to {anvl}, there are some additional things to be aware of.

File Organization

R/primitives.R: Define the prim_* primitive via new_primitive()
R/rules-inference.R: Add the inference rule infer_<primitive>() (unless an existing one can be reused)
R/rules-stablehlo.R: Add the StableHLO lowering rule
R/rules-reverse.R: Add the reverse rule (if differentiable)
R/rules-quickr.R: Add the quickr lowering rule (optional; only if the primitive should run on the quickr backend)
R/api.R: Add the nv_* wrapper function (or possibly in another api file).

Testing

Tests can go in two places:

inst/extra-tests/: For tests that compare against torch. These live in inst/ to avoid listing torch as a dependency.
tests/testthat/: For tests without a torch counterpart.

Important: the describe() / test_that() label must start with the full primitive name, e.g. describe("prim_repeat_along", { ... }). The meta tests in tests/testthat/test-primitives-meta.R verify that every primitive has corresponding stablehlo and reverse tests, and flag any that are missing.

Since no torch counterpart exists for prim_repeat_along, we would add manual tests in:

tests/testthat/test-primitives-stablehlo.R
tests/testthat/test-primitives-reverse.R

Also, ensure that no linter errors are present, devtools::check() passes, and format the code using make format.

Higher-Order Primitives

Higher-Order Primitives are primitives that are parameterized by an R function or expression. Examples include prim_if and prim_while. These are generally much more complex to handle, so we don’t cover them here in detail (for now). The general idea, however, is that the primitive prim_* function needs to trace the provided function using trace_fn() and then forward this graph to the stablehlo lowering rule and reverse rule.