Lower a Graph to StableHLO

Converts a traced AnvlGraph into the StableHLO intermediate representation (IR), building a stablehlo::Func with the stablehlo package. Each statement of the graph is translated to its corresponding StableHLO op. The result can be serialized to MLIR text via stablehlo::repr() and subsequently compiled to an XLA executable with pjrt::pjrt_compile().

The rule for translating a primitive to stablehlo is prim_<name>[["stablehlo"]].

The arguments of the resulting function are, in this order: the graph's constants (when constants_as_inputs = TRUE), the graph's inputs, and the phantom donated inputs (when donate_unaliased_outputs = TRUE), one per output that is not already aliased to a donated input, in the order of the outputs.

This is a low-level function; most users should use jit() instead.

Usage

stablehlo(
  graph,
  id = "main",
  captured = list(),
  donate = character(),
  donate_unaliased_outputs = FALSE,
  platform = NULL
)

Arguments

graph: (AnvlGraph)
The graph to lower (e.g. produced by trace_fn()).
id: (character(1))
The id of the resulting StableHLO function. Use "main" (the default) for a top-level lowering (returning from the main function finalizes the module) and "" for a closure/region lowering (e.g. a while body or a scatter update computation) that builds an anonymous nested function inside an enclosing build.
captured: (list(FuncValue))
Values of the enclosing function to bind the graph's last length(captured) inputs to. The lowered function reads them the way an MLIR region captures a value from above, instead of declaring them as inputs. Used for the sub-graphs of higher-order primitives (e.g. nv_while), whose trailing inputs are the values they close over. The graph's constants always become inputs, and GraphLiterals are always inlined.
donate: (character())
Names of the arguments whose buffers should be donated. Donated buffers can be aliased with outputs of the same type, enabling in-place operations.
donate_unaliased_outputs: (logical(1))
If TRUE and the current target platform is "cpu", append a phantom donated input for every output that isn't already aliased to a user-donated input. They come after all other arguments. This is needed internally so R keeps track of the CPU buffers memory in order to know when to garbage collect.
platform: (NULL | character(1))
Target platform name (e.g. "cpu", "cuda"). Stored on a process-wide global during the call so that platform-aware lowering rules (queried via current_platform()) can branch on it. NULL (the default) leaves the current value untouched — recursive calls from higher-order primitives inherit the platform of the enclosing call.

Value

(list)
Of length 3:

the stablehlo::Func
The graph's constants: the GraphValues holding ConcreteArrays, whose data must be passed as the leading inputs at execution time when constants_as_inputs = TRUE.
A list of phantom-output specs, one per phantom donated input appended when donate_unaliased_outputs = TRUE. Each entry is a list(dtype, shape) describing the buffer the executor must allocate. Empty when no phantoms were added.

Examples

# the closed-over array `x` becomes the constant input %0, before the
# graph's own input `y` (%1)
x <- nv_array(c(1, 2))
graph <- trace_fn(function(y) y + x, list(y = nv_aval("f32", shape = 2)))
out <- stablehlo(graph)
out[[1L]]
#> func.func @main (%0: tensor<2xf32>, %1: tensor<2xf32>) -> tensor<2xf32> {
#> %2 = stablehlo.add %1, %0 : tensor<2xf32>
#> return %2 : tensor<2xf32>
#> }
out[[2L]]
#> [[1]]
#> GraphValue(ConcreteArray(f32, (2))) 
#> 

# a donated input is aliased with an output of the same type
graph <- trace_fn(
  function(a, b) list(a + b, a * b),
  list(a = nv_aval("f32", 3), b = nv_aval("f32", 3))
)
stablehlo(graph, donate = "a")[[1L]]
#> func.func @main (%0: tensor<3xf32> {tf.aliasing_output = 0 : i32}, %1: tensor<3xf32>) -> (tensor<3xf32>, tensor<3xf32>) {
#> %2 = stablehlo.add %0, %1 : tensor<3xf32>
#> %3 = stablehlo.multiply %0, %1 : tensor<3xf32>
#> return %2, %3 : tensor<3xf32>, tensor<3xf32>
#> }

# on CPU, a phantom input is appended for every output not aliased yet, and
# the third element describes the buffers to allocate for them
out <- stablehlo(graph, donate_unaliased_outputs = TRUE, platform = "cpu")
out[[1L]]
#> func.func @main (%0: tensor<3xf32>, %1: tensor<3xf32>, %2: tensor<3xf32> {tf.aliasing_output = 0 : i32}, %3: tensor<3xf32> {tf.aliasing_output = 1 : i32}) -> (tensor<3xf32>, tensor<3xf32>) {
#> %4 = stablehlo.add %0, %1 : tensor<3xf32>
#> %5 = stablehlo.multiply %0, %1 : tensor<3xf32>
#> return %4, %5 : tensor<3xf32>, tensor<3xf32>
#> }
out[[3L]]
#> [[1]]
#> [[1]]$dtype
#> <f32>
#> 
#> [[1]]$shape
#> [1] 3
#> 
#> 
#> [[2]]
#> [[2]]$dtype
#> <f32>
#> 
#> [[2]]$shape
#> [1] 3
#> 
#>

Usage

Arguments

Value

See also

Examples