Loom Frontends

PyTorch

Register the backend once, then torch.compile. The graph is captured, lowered through the shared Triton importer into Tile IR, and execution hands back a native torch.Tensor:

import torch
import loom.experimental.torch as torch_loom

torch_loom.register_torch_compile_backend(name="loom")
compiled = torch.compile(model, backend="loom")
out = compiled(*inputs)

torch.device("loom") also exists, as a staging and integration surface for tensor placement and dispatch bring-up. Graph execution goes through torch.compile.

JAX

JAX uses Loom through the PJRT plugin, the same role torch.compile plays for Torch. Register it with register_loom_pjrt_plugin(), and an ordinary jax.jit returns a native jax.Array:

import os
os.environ.setdefault("JAX_PLATFORMS", "loom,cpu")

from loom.experimental.jax.native_plugin import register_loom_pjrt_plugin
register_loom_pjrt_plugin()

import jax
out = jax.jit(fn, backend="loom")(*args)

A supported upstream pallas_call imports through the same path. Input specs become DMA prefetch staging, scratch becomes local SRAM windows, and the grid becomes bounded ACP-loop metadata.

ONNX

A portable interchange from any training framework. The compiler handles lowering, fusion, and scheduling. No custom kernel authoring required.

Loom provides native kernel authoring through two compatible frontends: Palladium/LPL for Pallas-style tiled programming, and LTL for Triton-style kernel migration and direct Loom hardware controls. Both lower into the same Loom compiler IRs and object/runtime path. Upstream Triton source is also ingested directly, no rewrite required.

Torch/Triton

Loom is not a general CUDA-style GPU, but its Triton path is broader than only block pointers. It imports structured TTIR/TTGIR from TorchInductor and native Triton-style sources when the memory, layout, control-flow, and dtype semantics can be mapped onto Loom's Tile/CFG IR and hardware commands. Structured block pointers, tensor descriptors, affine and static address trees, legal gather, scatter, and scatter-add forms, reductions, dots, selected layout movement, and witnessed CFG loops and joins are compiler-lowered.

Forms that depend on GPU-specific SIMT execution semantics, arbitrary dynamic pointer manipulation, unclassified layout conversion, warp or block synchronization, or unconstrained async and barrier behavior fail closed with explicit diagnostics.

JAX/Pallas

Loom is not a general TPU/GPU Mosaic backend, but its JAX/Pallas path is broader than plain elementwise StableHLO. It uses the real JAX/Pallas call site as the semantic unit: JAX owns tracing, shapes, static arguments, grids, BlockSpecs, refs, and lowering, then Loom imports the closed StableHLO/Pallas program when its memory, layout, control-flow, and dtype semantics can be mapped onto Loom's Tile/CFG IR and hardware commands. Static grids, program-id-derived affine indexing, legal block specs, masked ref loads and stores, dynamic-update-style writes, gather, scatter, scatter-add, reductions, dots, selected layout movement, and witnessed CFG branches, loops, joins, and carried state are compiler-lowered.

Forms that depend on TPU/GPU-specific Mosaic execution semantics, unsupported memory-space colors, arbitrary dynamic ref or pointer aliasing, unclassified layout conversion, backend-specific async copy, semaphore, or barrier behavior, unsupported dtype or precision contracts, or Python behavior that is not present in the lowered program fail closed with explicit diagnostics.

Palladium (Native DSL)

For maximum control, Loom exposes Palladium, a Pallas-like domain-specific language embedded in Python and imported as from loom.experimental import palladium as lpl. It maps directly to Loom's physical execution model: explicit asynchronous DMA, hardware-managed SRAM allocation, and user-defined block grids. A separate loom.experimental.pallas shim keeps the upstream pl.* spelling (BlockSpec, pallas_call, program_id, dot, softmax) for authors porting existing Pallas code.

The dividing line is simple: run stock JAX or Pallas through PJRT to port an existing model with zero friction; reach for Palladium when you want to hand-tune the data movement and engine scheduling that a general-purpose compiler would decide for you.

LTL (Loom Triton Language)

Triton-shaped Loom-native authoring, imported as import loom.triton.language as ltl. It keeps Triton authors in a familiar tl-like mental model while exposing Loom-native primitives and explicit hardware controls: windows, descriptors, tokens, and DMA, MXU, and SCE primitives. The natural landing point for Triton migration.

From Torch

capture = torch_loom.capture_backend(...)
compiled = torch.compile(..., backend=capture)
compiled(*inputs)
capture.write_loom("build/model.loom")

From JAX

capture = jax_loom.capture_jit(fn, backend="loom")
compiled = capture.lower(*inputs)
capture.write_loom("build/model.loom")

From Palladium and Pallas

loomc compile-palladium builds objects from Palladium source. loomc migrate-pallas and loomc migrate-triton report what lowered, what did not, and the equivalent Loom, Palladium, or LTL surface. Compatibility reporting, not the execution path.

To compile a checkpoint-backed autoregressive model into a host-free Loom object, first canonicalize the weights with loomc ingest-checkpoint, then build the executable nested ACP model object with loomc build-checkpoint-block --kind model-loop:

loomc ingest-checkpoint /models/qwen3-0.6b \
    --out build/model_weights.loom \
    --dst-fmt lut4b_fp16

loomc build-checkpoint-block \
    --weights build/model_weights.loom \
    --out build/model_loop \
    --kind model-loop \
    --model-loop-control acp-large-body \
    --decode-steps 1

That writes build/model_loop.loom: the whole decode loop, driven by the on-chip command processor, with the host out of the loop.

Ingest accepts checkpoint directories, .safetensors files, sharded safetensors indexes, and GGUF. Destination formats include bf16, lut2b_fp16, lut3b_fp16, lut4b_fp16, lut8_fp16, bin1_packed16, packed_trit5, packed_trit5_fp16_lut3, packed_int3, int8_in_16lane, int8_packed2, int4_packed4, fp8_e4m3, and fp8_e5m2. Lossy low-bit requantization is intentionally gated behind --allow-lossy-requant.

Inspect and Verify

Every .loom object is a deterministic record of what the chip will do, and loomc reads it back:

loomc manifest build/model.loom
loomc validate build/model.loom --strict
loomc explain --object-file build/model.loom
loomc profile --object-file build/model.loom
loomc visualize --object-file build/model.loom --dot

Golden execution is the bit-exact simulator path; device execution uses the same object-file contract.

Run an Object

The host runtime is libloom. Open the device and submit the object:

from libloom.python.loom_runtime import open_devnode
from libloom.python.object_file_runtime import submit_object_file

with open_devnode(devnode_path="/dev/loom0") as dev:
    report = submit_object_file(dev, "build/model.loom")

For deterministic file-backed bring-up, open_files(bar0_path=..., bar1_path=...) runs the same submission against files instead of a device.