NxIREE

Mix.install([
  {:nx_iree, "~> 0.0"}
])

デバイスの取得

NxIREE.list_drivers()
|> elem(1)

dev =
  NxIREE.list_devices("metal")
  |> elem(1)
  |> hd()

中間表現のコンパイル

mlir_module = """
module {
  func.func public @main(%arg0: tensor<4xf32>, %arg1: tensor<4xi32>) -> tensor<4xf32> {
    %0 = stablehlo.cosine %arg0 : tensor<4xf32>
    %1 = stablehlo.convert %arg1 : (tensor<4xi32>) -> tensor<4xf32>
    %2 = stablehlo.sine %1 : tensor<4xf32>
    %3 = stablehlo.add %0, %2 : tensor<4xf32>
    return %3 : tensor<4xf32>
  }
}
"""

flags = [
  "--iree-hal-target-backends=metal-spirv",
  "--iree-input-type=stablehlo_xla",
  "--iree-execution-model=async-internal"
]

module = NxIREE.compile(mlir_module, flags, output_container: Nx.template({4}, :f32))

行列演算の実行

arg0 = Nx.tensor([1.0, 2.0, 3.0, 4.0])
arg1 = Nx.tensor([1, -1, 1, -1])

NxIREE.call(module, [arg0, arg1], device: dev)

nx_function = fn a, b ->
  Nx.add(Nx.cos(a), Nx.sin(b))
end

nx_function.(arg0, arg1)