Lower a graph to StableHLO

Converts a traced AnvlGraph into the StableHLO intermediate representation (IR). Each graph operation 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 rules for translating to stablehlo are stored in $rules[["stablehlo"]] of the primitives.

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

Usage

stablehlo(
  graph,
  id = "main",
  constants_as_inputs = TRUE,
  env = NULL,
  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.
constants_as_inputs: (logical(1))
If TRUE (default), constants are registered as inputs to the StableHLO function so they can be passed in at execution time. If FALSE, they are not added as inputs. Set to FALSE for closures. Note that GraphLiterals are always inlined into the StableHLO function.
env: (HloEnv | NULL)
Optional environment for reusing variable mappings across nested function lowerings (e.g. for higher-order primitives like nv_while).
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. 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

A list of length 3:

the stablehlo::Func
The list of GraphValues holding ConcreteArrays.
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

x <- nv_array(c(1, 2))
graph <- trace_fn(function(y) y + x, list(y = nv_aval("f32", shape = c())))
graph
#> <AnvlGraph>
#>   Inputs:
#>     %x1: f32[]
#>   Constants:
#>     %c1: f32[2]
#>   Body:
#>     %1: f32[2] = broadcast_in_axes [shape = 2, broadcast_axes = <any>] (%x1)
#>     %2: f32[2] = add(%1, %c1)
#>   Outputs:
#>     %2: f32[2] 
stablehlo(graph)
#> [[1]]
#> func.func @main (%0: tensor<2xf32>, %1: tensor<f32>) -> tensor<2xf32> {
#> %2 = "stablehlo.broadcast_in_dim" (%1) {
#> broadcast_dimensions = array<i64>
#> }: (tensor<f32>) -> (tensor<2xf32>)
#> %3 = stablehlo.add %2, %0 : tensor<2xf32>
#> return %3 : tensor<2xf32>
#> }
#> 
#> [[2]]
#> [[2]][[1]]
#> GraphValue(ConcreteArray(f32, (2))) 
#> 
#> 
#> [[3]]
#> list()
#>

Usage

Arguments

Value

See also

Examples