Extension point for lowering selected Nx.block/4 blocks to XLA custom calls (stablehlo.custom_call in StableHLO MLIR).

Other blocks (for example gather-based take or FFT) are lowered inline in EXLA.Defn and do not use this protocol.

When EXLA.Defn calls it

During compilation with compiler: EXLA, each Nx.block(block, inputs, outputs, fn ... end) is processed by this protocol. EXLA invokes call/4 once per block.

If call/4 returns :skip, EXLA compiles the block's default callback (the anonymous function body) instead of emitting a custom call. Default lowerings are provided for Nx.Block.LinAlg.QR and Nx.Block.LinAlg.Eigh.

call/4 arguments

Callback arity is call(struct, out, args, client), matching Nx.block(block, inputs, outputs, fn ... end) (block, outputs, inputs, then client).

• struct — the block passed as the first argument to Nx.block/4 (your own defstruct or an existing block such as %Nx.Block.LinAlg.QR{}).

• out — the output template tuple passed to Nx.block/4 (expression metadata for shapes and types, not runtime tensors).

• args — list of input templates, in the same order as inputs in Nx.block/4.

• client — the active EXLA.Client (use e.g. client.platform to gate host-only lowerings).

call/4 return value

• :skip — this implementation does not apply (unsupported type, non-host platform, wrong arity, etc.). The default block implementation is used instead.

• {:ok, %EXLA.CustomCall.Spec{}} — emit a StableHLO custom call; see EXLA.CustomCall.Spec for call_target_name; optional backend_config ([{name, attr}] string pairs for the stablehlo.custom_call backend_config dictionary); optional attributes (string pairs emitted directly on stablehlo.custom_call); and optional operand_element_types (operand converts when they differ from the lowered inputs).

Dispatch

The protocol uses @fallback_to_any true. Built-in lowerings for known blocks live in defimpl EXLA.CustomCall, for: Any. Your application or dependency can add defimpl EXLA.CustomCall, for: YourStruct; that implementation is chosen whenever the block is %YourStruct{}, instead of the Any fallback.

Native handlers

For a custom call backed by an external native library, load the library that registers the corresponding native handler for the active XLA platform before compiling the block.

Example

defmodule MyApp.CustomQrBlock do
  defstruct []
end

defimpl EXLA.CustomCall, for: MyApp.CustomQrBlock do
  def call(_block, {%{type: {kind, size}}, _r_expr}, [_input], %{platform: :host})
      when kind != :c and kind in [:f, :bf] and size in [16, 32, 64] do
    {:ok, %EXLA.CustomCall.Spec{call_target_name: "my_custom_qr_target"}}
  end

  def call(_, _, _, _), do: :skip
end

Then use Nx.block(%MyApp.CustomQrBlock{}, ...) inside a defn compiled with compiler: EXLA.