torch-xdna

torch-xdna is an experimental downstream PyTorch backend for one AMD XDNA2 NPU. It uses PyTorch's existing PrivateUse1 and out-of-tree Inductor extension contracts, owns real XRT buffer-object storage, and launches validated IRON or Triton-XDNA-generated device code directly through XRT.

The project is independently maintained research. It is not affiliated with, sponsored by, endorsed by, or supported by AMD, Xilinx, or the PyTorch project. It is not production ready and does not claim general PyTorch, torch.compile, model, or XDNA coverage.

Current boundary

The validated host has one RyzenAI-npu4, Ubuntu 26.04, kernel 7.0.0-27-generic, the unchanged in-tree amdxdna driver, and XRT 2.20.0 at 021204355eeaa034ff69aae407ace2265adf047a.

Surface	Current contract	Evidence boundary
Standalone IRON	User-authored INT32 vector-add-one	1,024/1,024 exact on physical XDNA2
PyTorch eager	Copy-backed contiguous INT32, 1–1,024 elements; add-one and same-shape tensor add; contiguous rank-1/rank-2 BF16 storage/copy with 1–16,384 elements; full-extent BF16 view/reshape/contiguous no-op	Real XRT BO storage, explicit copies, authenticated metadata aliases, pinned-core ATen rejection shield, no CPU/arithmetic fallback
Inductor / native IRON	One static one-dimensional contiguous int32[n] + int32[n], 1 <= n <= 1024, alpha=1	Every declared length passed cold and warm physical XRT execution with bit-exact output and no fallback
Inductor / Strategy B	One static contiguous bfloat16[1024] + bfloat16[1024], alpha=1, plus exactly two static BF16 aten.mm.default shapes: [128,64] @ [64,128] -> [128,128] and [64,64] @ [64,128] -> [64,128]	Each graph has its own pinned Inductor-generated Triton body, Triton-XDNA/AIR/AIE/Peano ELF, receipt, direct torch-xdna BO launch, and focused physical proof; this is not general BF16, matmul, or linear support
Compiler cache	Schema-2 device-artifact, native-extension, and generated-wrapper identity domains	Hardware-free canonicalization, separation, atomicity, and fail-closed corruption coverage; prior INT32 range evidence and focused Strategy B evidence remain exact-checkpoint claims

Each compiled route emits one direct synchronous native XRT launch and keeps the result on xdna:0 until explicit copy-back. Unsupported dtype, shape, layout, operator, scalar, stream, autograd, training, dynamic-shape, and general Triton requests fail explicitly. CPU, NumPy, eager redispatch, extern kernels, FakeTensor, and Meta may not perform user arithmetic.

The native INT32 kernel processes 1,024 lanes. A shorter compiled INT32 input uses an XRT BO of that extent; the native runtime zeroes both input tails under the execution lock and returns only the logical [n] result. Logical length and padding policy are part of the cache identity. The Strategy B add accepts only exact BF16 [1024]; the Strategy B matrix product accepts only the independently receipt-bound shapes [128,64] @ [64,128] -> [128,128] and [64,64] @ [64,128] -> [64,128]. All three Strategy B graphs apply no padding.

The broader BF16 contract is storage usability, not broader arithmetic: rank-one or rank-two contiguous zero-offset tensors may contain 1 through 16,384 elements, including the three BO extents selected for the first matmul slice. Full-extent view, reshape, and already-contiguous contiguous() preserve the same authenticated BO without allocation. Partial, offset, non-contiguous, rank-three, and arithmetic requests remain rejected.

Each matrix-product contract admits every BF16 bit pattern without a hidden device-to-host value scan. For finite normal values, the Inductor/compiler IR retains FP32 arithmetic intent, while the physical AIE2P program uses native BF16 operands and its ACCFP accumulator; results are compared with an FP32-dot-then-BF16 oracle using rtol=0.016 and atol=0.03125, not advertised as bitwise IEEE-FP32 accumulation. Input subnormals are explicitly flush-to-zero. Focused qNaN, infinity, and infinity-times-zero cases establish only the recorded output classifications; signaling-NaN behavior, NaN payload preservation, mixed nonfinite inputs, and device-output rounding remain unqualified.

See the developer-preview contract for the exact accepted and rejected surface. Direct torch.arange(..., device="xdna") is not currently supported and fails closed before allocation or launch. Create inputs on CPU and copy them with .to("xdna").

Installation and acceptance

Product Python support is exactly CPython >=3.12,<3.13. The accepted build is pinned more narrowly to CPython 3.12.13, PyTorch Git 85ef666ce184b62346977ae19330cda383b49ab3, GCC/G++ 14, XRT 2.20, and the external xdna2-aie2p-v1 artifact manifest. The PyTorch input is a locally produced repaired CPU manywheel containing the frozen metadata-only SDPA setup experiment; it is not an official downloaded release wheel and is not distributed here.

Start from an object-independent clean checkout at an exact reviewed commit. Keep the new environment, build root, PyTorch wheel, artifacts, and caches outside that checkout:

scripts/bootstrap_dev_environment.sh \
  --environment "$ENV_ROOT" \
  --build-root "$BUILD_ROOT" \
  --torch-wheel "$TORCH_WHEEL" \
  --xrt-root "$XRT_ROOT" \
  --jobs 6

env -u PYTHONPATH -u PYTHONHOME \
  "$ENV_ROOT/bin/python" -I -B scripts/run_acceptance.py \
  --hardware-free \
  --source-checkout "$PWD" \
  --xrt-root "$XRT_ROOT"

env -u PYTHONPATH -u PYTHONHOME \
  "$ENV_ROOT/bin/python" -I -B scripts/run_acceptance.py \
  --physical \
  --source-checkout "$PWD" \
  --xrt-root "$XRT_ROOT" \
  --artifact-dir "$XDNA_ARTIFACT_DIR" \
  --strategy-b-elf "$XDNA_STRATEGY_B_ELF" \
  --elements 1024

The exact matrix-product slices are exercised separately by tests/test_physical_inductor_strategy_b_bfloat16_mm.py with TORCH_XDNA_STRATEGY_B_MATMUL_ELF or TORCH_XDNA_STRATEGY_B_MATMUL_M64_ELF set to the selected private mode-0600 ELF. Select the matching pytest case with -k m128-k64-n128-existing or -k m64-k64-n128-new; the test then selects the corresponding runner profile. Neither shape is yet part of run_acceptance.py --physical, so these focused passes do not complete the object-independent integrated-qualification or physical-CI gate.

Run hardware-free acceptance first. It verifies packaging, installed-source identity, PyTorch and native ABI/linkage, registration, lowering, cache, artifact-validation, and rejection contracts without claiming physical execution. Physical mode validates the artifacts before backend initialization, requires the real device and runtime, records bit-exact output and XRT completion evidence, and rechecks XRT health. A missing NPU is an explicit hardware-unavailable result, not success.

The full object-independent procedure and its original physical proof are in Fresh-checkout acceptance and Fresh-checkout result. The latest exhaustive 1-through-1,024 INT32 proof is in the 1fc6676 downstream-main acceptance result. That receipt predates the consolidated Strategy B product path. The complete alpha-foundation stack was subsequently qualified from an object-independent checkout at 927b7f7cdc284637b32e94a6cbeb3a462fb3a2c1: a clean non-editable wheel passed 704/704 hardware-free tests, representative INT32 cold/warm physical execution, the exact Strategy B BF16 cold/warm/trapped path, explicit rejection checks, and pre/post XRT health validation. Downstream main at 0d8216ac790e7179478d59b31265f97df6a25cbb has the byte-identical qualified tree, and merged-main hosted run 30751703259 passed.

External artifacts and cache

This repository contains the package-pinned artifact manifest, source recipes, tests, and sanitized evidence. It does not contain XCLBINs, instruction ELFs, raw instruction binaries, compiler wheels, PyTorch wheels, XRT packages, virtual environments, compiler caches, build trees, or raw logs.

Normal import does not download or compile artifacts. Two external boundaries must not be conflated:

• TORCH_XDNA_PREBUILT_ARTIFACT_DIR names the exact nine-file IRON set used by the native eager and INT32 paths. Every filename, size, SHA-256, XCLBIN UUID, MLIR_AIE kernel identity, target, provenance, PyTorch identity, and XRT compatibility field must match the package manifest.
• TORCH_XDNA_STRATEGY_B_ELF optionally names the separately approved direct ELF for the exact BF16 add graph. TORCH_XDNA_STRATEGY_B_MATMUL_ELF names the [128,64] @ [64,128] -> [128,128] direct ELF, while TORCH_XDNA_STRATEGY_B_MATMUL_M64_ELF independently names the [64,64] @ [64,128] -> [64,128] direct ELF. Each is accepted only as a current-user-owned regular file with no group/other permissions, exact size and SHA-256, its pinned compiler lock and graph-specific receipt, and its expected kernel identity. None has an XCLBIN UUID.

A Strategy B source ELF may be removed after atomic cache publication; subsequent processes may replay only its complete, hash-validated private cache entry. When both exact graphs are provisioned, admission validates the complete requested set before publishing any clean miss. Artifact redistribution is unreviewed and not approved by this project.

Native artifact/snapshot configuration is one-shot and compiler registration is process-global. If import fails after either phase begins, correct the configuration and start a new process; retrying in the same process is unsupported and remains fail-closed.

• Artifact acquisition separates the current private external path from incomplete local-build/source-reconstruction and future-download paths, including the exact legal-review blocker.
• Compiler cache defines canonical identity, atomic publication, path independence, corruption handling, and deterministic invalidation across the IRON, Strategy B device-artifact, native-extension, and wrapper domains. A warm hit validates identity; it does not skip Inductor code generation or establish a speedup.
• CI boundaries separates the hosted adjacent-nightly source lane, exact-product acceptance, and the future self-hosted physical job.

Project map

• Developer-preview contract: support, compatibility, install, acceptance, and remaining release gates.
• Strategy B compiler contract: target propagation, compiler lock, receipt, final ABI, and focused physical proof for each exact BF16 graph.
• Strategy B BF16 matrix-product results: focused physical results and numerical boundaries for the two independently authorized aten.mm.default shapes.
• Generated-kernel manifest: native capability authority, typed legacy-receipt projection, and the admission boundary for future generated kernel families.
• Milestones: immutable physical tags and Issue #1 gates.
• CI boundaries: what hosted and self-hosted workflows can and cannot prove.
• Fresh-checkout acceptance: reproducible isolated bootstrap and installed-package verification.
• Integrated acceptance result: historical exact-commit fresh-clone, cache, and exhaustive static INT32 physical evidence; it does not qualify the later Strategy B integration.
• Eager design: storage, copies, dispatch, synchronization, and explicit failure behavior.
• Inductor design and first physical result: the initial [1024] vertical slice from which the current range was generalized.
• System validation, findings, and sanitized validation: Ubuntu 26.04/XRT recovery.
• Contributing, security, and repository instructions: ownership, review, AI-assisted development, and publication policy.

Ownership and release status

Generic PyTorch extension-contract work, this downstream XDNA runtime, and vendor/system XRT-driver work are separate ownership tracks. PyTorch issue #189138 concerns generic out-of-tree compile extensibility; it does not provide or validate this XDNA backend. Upstream PyTorch work is frozen pending maintainer direction.

The package remains 0.1.0.dev0. v0.1.0a1 is reserved until every Issue #1 gate has reviewed evidence and the maintainer explicitly approves the tag and release. A workflow definition, a local-only pass, or broader research does not complete a gate.