Blog

Building a PyTorch Backend for a Custom Accelerator

Open-source and open-weight models have made AI easy to prototype. A vision pipeline can start from YOLO, a local language product can start from a Llama model, and the first implementation can run through ordinary framework code on commodity CPUs or GPUs. That is a good starting point, but it is not where efficient products have to stop.

That pressure is especially visible in Edge AI. Cameras, robots, industrial systems, vehicles, and sensing platforms increasingly need to run perception and decision pipelines close to where data is produced. World-model-style systems push that pressure even further: they combine perception, state, prediction, and planning, which makes latency, memory movement, and power budget part of the product architecture rather than late optimization details. FPGA-based solutions are a strong fit for this stage because they let teams shape hardware around the workload while the model is still evolving.

Most AI applications are written in frameworks like PyTorch, not as ad-hoc solutions. The practical question is therefore: how do we keep ordinary PyTorch code while redirecting selected tensor operations into custom hardware?

PyTorch model operations connected to custom AI hardware

Making custom hardware reachable from PyTorch

In PyTorch, one answer is PrivateUse1. It is a reserved backend slot that lets an extension register a custom device without modifying PyTorch itself. The user-facing code can still look like ordinary PyTorch:

a = torch.tensor([1, 2, 3, 4], dtype=torch.float32).to("privateuseone")
b = torch.tensor([10, 20, 30, 40], dtype=torch.float32).to("privateuseone")
c = torch.add(a, b)

Behind that line is a C++ extension that registers kernels for selected ATen operators on the PrivateUse1 dispatch key. In our prototype, the accelerated operators are intentionally small: vector add and vector subtract, each operating on up to four float32 elements. That is not an AI accelerator by itself. It is the smallest useful slice of a full accelerator stack: tensor allocation, operator dispatch, a backend abstraction, a kernel driver, a register map, and hardware.

The goal is not to pretend that a four-element vector unit is the interesting accelerator. The goal is to make the path realistic before making the datapath large: a PyTorch operation crosses into a custom backend, the backend submits a hardware-shaped command, and the result comes back as a tensor the rest of the model can keep using.

Preparing tensors for the custom device

PyTorch does not only need kernels for add and sub. Before user code can do anything interesting, tensors must be allocated, copied, reshaped, printed, and converted to and from CPU. A backend that only implements one arithmetic operator will usually fail much earlier, during tensor construction or transfer.

The extension therefore registers three categories of behavior.

First, it gives PrivateUse1 a device guard implementation and a user-facing backend name:

C10_REGISTER_GUARD_IMPL(PrivateUse1, PrivateUse1Guard);

PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) {
  c10::register_privateuse1_backend("privateuseone");
}

Second, it registers an allocator for PrivateUse1. In this experiment the allocator is deliberately malloc-backed. That means device tensors live in ordinary host memory, which keeps the software path easy to inspect and lets the driver copy directly from user-space pointers. This is not how a high-performance accelerator with device memory would finish, but it is a very good bringup model.

Third, it implements the infrastructure operators that make a tensor usable. The extension handles operators such as empty, empty_strided, copy_, _copy_from, view, resize_, _local_scalar_dense, fill_, and zero_. Many of these wrappers reinterpret the malloc-backed storage as CPU-accessible memory and delegate to existing native CPU behavior.

That choice is pragmatic. We only want to accelerate the operations under test; the rest of the PyTorch surface should remain boring.

Accelerating selected ATen operators

The arithmetic kernels are registered against the aten namespace for the PrivateUse1 dispatch key:

TORCH_LIBRARY_IMPL(aten, PrivateUse1, m) {
  m.impl("add.Tensor", &tensor_add_wrapper);
}

The add kernel checks the contract before it touches the lower layer:

  • both operands must be on PrivateUse1
  • both operands must be float32
  • both operands must be contiguous
  • the tensors must have the same number of elements
  • alpha must be 1
  • the length must fit the hardware vector width

Once those checks pass, the kernel allocates an output tensor and submits one operation to the backend:

int rc = get_backend().submit(
    TENSOR_OP_VADD_F32,
    (uint32_t)a.numel(),
    4,
    a.data_ptr<float>(),
    b.data_ptr<float>(),
    out.data_ptr<float>());

The subtract kernel follows the same pattern with TENSOR_OP_VSUB_F32.

The boundary between PyTorch and custom hardware

The call above is the important boundary. PyTorch has already done the framework-level work: it dispatched the operator, checked the tensor device, and entered our backend kernel. From that point on, the software has to describe a hardware-shaped command without leaking board details back into the model code.

The C++ extension keeps that boundary small:

class HardwareBackend {
public:
    virtual int submit(uint32_t op, uint32_t len, uint32_t elem_bytes,
                       const void* a_data, const void* b_data,
                       void* out_data) = 0;
};

The ATen kernel does not know about register offsets, AXI4-Lite transactions, device-tree nodes, boot scripts, or simulation setup. It sends an opcode, an element count, an element size, two input buffers, and one output buffer. Below that interface, the implementation can evolve from a mock driver to real MMIO hardware or an emulated target. Those are bringup choices. The model-facing operation stays the same.

This is the part that matters for custom AI acceleration: the hardware path can stay narrow and focused on the operations that justify a custom datapath. We do not need to reimplement all of PyTorch before one accelerator block becomes useful.

The extension also registers a boxed CPU fallback for unimplemented operators:

TORCH_LIBRARY_IMPL(_, PrivateUse1, m) {
  m.fallback(torch::CppFunction::makeFromBoxedFunction<&pu1_cpu_fallback>());
}

For a prototype, this is extremely helpful. It means the custom device can participate in more PyTorch code than the accelerator itself implements. Unimplemented operations can be copied to CPU, executed there, and copied back. That fallback is not the performance story; the accelerated operations are. The fallback is what keeps the integration usable while the custom datapath targets the expensive parts of the workload first.

The result is a realistic software shell around a tiny accelerator. The hardware only implements vector add and subtract, but the PyTorch side behaves like a device backend: it can allocate tensors, move data, execute selected operations in custom hardware, and fall back to CPU when needed.

This first step gives us the shape of the integration: ordinary model code can call into a custom device path without knowing how the accelerator is built. From here, the work moves downward. The next pieces are the driver interface that accepts these commands and the register map that exposes the hardware block that performs the operation.