stablehlo

Package website: release | dev

The {stablehlo} R package provides a functional API to create stableHLO programs. These programs can be executed using the R package pjrt.

Installation

From GitHub:

pak::pak("r-xla/stablehlo")

You can also install from r-universe, by adding the code below to your .Rprofile.

options(repos = c(
  rxla = "https://r-xla.r-universe.dev",
  CRAN = "https://cloud.r-project.org/"
))

Quickstart

Below, we create a function that takes two input arguments x and y of type f32 and shape 2x2 and adds them. Passing func to hlo_input() is optional, because it will automatically use the last function created with hlo_func().

library(stablehlo)
func <- hlo_func("main")
func
#> func.func @main () ->  {
#> 
#> }
x <- hlo_input("x", "f32", shape = c(2, 2), func = func)
x
#> Variable %x in:
#> func.func @main (%x: tensor<2x2xf32>) ->  {
#> 
#> }
y <- hlo_input("y", "f32", shape = c(2, 2), func = func)
y
#> Variable %y in:
#> func.func @main (%x: tensor<2x2xf32>, %y: tensor<2x2xf32>) ->  {
#> 
#> }
z <- hlo_add(x, y)
z
#> Variable %0 in:
#> func.func @main (%x: tensor<2x2xf32>, %y: tensor<2x2xf32>) ->  {
#> %0 = stablehlo.add %x, %y : tensor<2x2xf32>
#> }
f <- hlo_return(z)
identical(f, func)
#> [1] TRUE
f
#> func.func @main (%x: tensor<2x2xf32>, %y: tensor<2x2xf32>) -> tensor<2x2xf32> {
#> %0 = stablehlo.add %x, %y : tensor<2x2xf32>
#> return %0 : tensor<2x2xf32>
#> }

Main Features

• Implements type inference to easily build up stableHLO programs.
• Simple, functional API.
• Clear error messages.
• Easy installation because the implementation is in R and does not depend on the stableHLO C++ builder, which depends on LLVM and MLIR.

Important notes

• stableHLO uses 0-based indexing. Wherever operations take dimension indices (e.g., axes, start indices, permutation dimensions), use 0-based values. This differs from R’s 1-based indexing.
• We try to use "bool" for the BooleanType. However, in the generated stableHLO representation this will show up as "i1". Think of these as interchangeable.

Limitations

The R package is a partial implementation of the stableHLO specification. At least initially, it will:

• only support a subset of the available operations, see this issue for an overview.
• not support all data types, e.g. quantized types and complex numbers are not supported.
• not support shape dynamism.

Acknowledgments

• This work is supported by MaRDI.
• This work is built upon the stableHLO specification from the OpenXLA project.