USS_Tensor_Processing_Unit

This repository contains research on the synthesis flow and hardware code required to accelerate a PyTorch neural network using a custom TPU architecture defined in Verilog HDL that runs on an FPGA. This research is endorsed and funded by Miami University of Ohio and conducted under the Electrical and Computer Engineering department in the College of Engineering and Computing.

What is a TPU?

A TPU (Tensor Processing Unit) is Google's proprietary name for its hardware AI accelerator. TPUs are designed to speed up the large tensor operations required to run advanced neural networks by reusing as many data values as possible. This reuse is made possible by a systolic array — an array of multiply-accumulate units (MACs) all wired together. Google's first-generation TPU had a 256x256 systolic array for a total of 65,536 multipliers!

A TPU can be thought of as an extreme case of computational parallelism. The idea of a TPU is best understood by comparing it to other architectures:

• CPU - dozens of powerful, extremely fast cores that can perform a multitude of advanced operations.
• GPU - hundreds or thousands of cores specialized for graphics processing.
• TPU - tens of thousands of "cores" that only perform one operation: multiply-accumulation.

The reuse is what makes the architecture fast. In a CPU or GPU, every multiplication pays to move its operands out of memory or registers; in a systolic array, each weight and activation is loaded once and then flows from MAC unit to MAC unit, participating in a new multiplication at every step. The memory traffic — usually the true bottleneck of neural-network inference — is amortized across an entire row or column of the array, so the silicon spends its time computing instead of waiting on data.

The TPU in this repository is the same idea at a smaller scale: an 8x8 systolic array of MACs operating on 8-bit signed integers, plus the supporting subsystems a real accelerator needs — on-chip memory, per-layer requantization, a custom ISA, and a compiler that lowers PyTorch networks onto it.

Folder Hierarchy

This is a general overview of the project folder hierarchy. Only relevant directories are included; further subdirectories are specific to components of an individual system.

.
├── Benchmarks/ - Example transmissions generated by the Assembler.
├── Demo/ - Demo scripts
├── Resources/ - The resources I found most helpful for this project
├── Specifications/ - ISA, Message Protocol, and Hardware Spec
│   └── Old/ - Old versions of specifications for historical reference
└── System/ - Contains all project code
    ├── Assembler/ - Assembles a neural network from /System/Neural_Networks/ into a transmission binary
    ├── Communication/ - Comm link between PC and FPGA via an Arduino
    │   ├── Arduino_2_FPGA/ - PlatformIO Arduino project that receives serial data from a computer and forwards it to SPI pins
    │   └── PC_2_Arduino/ - Small Python script that sends a binary file byte-by-byte to a connected Arduino via serial port
    ├── Neural_Networks/ - Contains the PyTorch neural networks used in this project
    └── Tensor_Processing_Unit/ - Quartus project with TPU RTL (Verilog 2001)
        └── demo/ - Top-level module that corresponds to /Demo/

Documentation

See /Specifications/ for the documents used to design and build the TPU.

• Functional_TPU_ISA.md -> Custom instruction set used by the TPU.
• Functional_TPU_Message_Protocol.md -> Protocol used to communicate with the TPU over an SPI/Serial connection via an Arduino (see the Synthesis Flow section for more detail).
• Functional_TPU_Hardware_Specification.md -> Abstract description of each Verilog module built as part of the TPU.

These three documents are a great starting point for building your own custom neural network accelerator!

Usage

The system is a four-stage pipeline: a PyTorch network from System/Neural_Networks/ is trained and exported, the Assembler compiles the export into a transmission binary, the Communication link streams that binary from the PC through an Arduino to the FPGA, and the TPU executes the resident program and presents its results over VGA. Each stage hands the next a plain artifact on disk, so every subsystem can also be operated — and understood — on its own. The subsections below cover setting up the toolchain, running the full flow end to end, and driving each subsystem independently; if you only want to program a board, the pre-built transmissions in /Benchmarks let you skip the compilation pipeline entirely.

Environment Setup

Development targets Debian Linux with Python 3.13. From the repository root:

./setup.sh
source USS_TPU.venv/bin/activate

setup.sh creates the USS_TPU.venv virtual environment, installs the exact dependency versions pinned in requirements.txt (PyTorch, torchax, etc.), installs the Assembler's nn_assembler package in editable mode, and warns about any missing optional tools. It is safe to re-run.

The remaining tools are stage-specific — install only what you plan to use:

• stablehlo-opt (Assembler) — required to assemble any network into a transmission. The exact binary used by this project is published in the StableHLO Binaries release (Linux x86-64, glibc ≥ 2.38, i.e. Debian 13 / Ubuntu 24.04 or newer) — download it, chmod +x, and put it on your PATH; alternatively build it from openxla/stablehlo. The release also carries the matching stablehlo Python wheel (CPython 3.13) for anyone who wants the StableHLO Python API — the pipeline itself needs only the binary. Neither is needed if you stream the pre-built /Benchmarks transmissions instead.
• python3-tk (Demo GUI) — sudo apt install python3-tk.
• Quartus Prime Lite 23.1std (RTL) — synthesizes the FPGA bitstream and programs the board; Questa FSE (bundled) runs the testbench regression.
• PlatformIO (Communication) — builds and uploads the Arduino firmware.

Running on hardware additionally requires a Terasic DE1-SoC (Cyclone V 5CSEMA5F31C6), an Arduino Uno, a VGA display, and jumper wires from the Arduino's SPI pins to the DE1-SoC GPIO header per the pin assignments in the Quartus project.

Synthesis Flow

Steps 1–3 are one-time board setup; steps 4–5 are the iteration loop you repeat as you change the network.

1. Synthesize and flash the FPGA. Open System/Tensor_Processing_Unit/Tensor_Processing_Unit.qpf in Quartus, compile, and program the DE1-SoC over USB-Blaster. The default top level (Tensor_Processing_Unit.v) shows a raw-hex debug view of TPU memory on the VGA output; for the LeNet-5 presentation view (ranked class percentages), set demo/demo.v as the top-level entity instead.

2. Flash the Arduino. From System/Communication/Arduino_2_FPGA/, upload with PlatformIO (pio run -t upload). The firmware forwards each verified serial frame to the FPGA over SPI (mode 0, 125 kHz) in lock-step with the PC.

3. Wire the SPI link. The Arduino Uno is 5 V and the DE1-SoC GPIO is 3.3 V — level-shift every Arduino output (this project uses resistive dividers).

   Arduino Uno	DE1-SoC	Signal
   D13 (SCK)	GPIO_0[7]	SPI clock
   D11 (MOSI)	GPIO_0[3]	data to FPGA
   D10 (SS)	GPIO_0[1]	chip select, active LOW
   GND	GND	common ground

   MISO is not wired: results leave the TPU through the VGA display, not back over SPI.

4. Train, assemble, and program. System/run.sh drives the three software stages; each flag is independent, so run any subset:

   ./System/run.sh -m LeNet_5 -r Recent -t -a -p

   • -t trains the model and exports <model>_<run>.mlir + <model>_<run>.weights.npz into System/Neural_Networks/<model>/
   • -a assembles that export into System/Assembler/out/TRANSMISSION.bin
   • -p streams the transmission to the FPGA through the Arduino

   To skip training and assembly entirely, copy a pre-built transmission into place and program it:

   cp Benchmarks/LeNet_5_TRANSMISSION.bin System/Assembler/out/TRANSMISSION.bin
   ./System/run.sh -p

5. Run inference. KEY[0] is reset. SW[9] selects the input source:

   • Device mode (SW[9] LOW): set an input byte on SW[7:0] and pulse KEY[1] to feed it to the network; the HEX displays echo the value. Suits the single-value networks (Tiny_NN, Bigger_NN).
   • External mode (SW[9] HIGH): the input is streamed from the PC as an INPUT transmission to address 0x1 (this is what the Demo does — see /Demo). Required for image-sized inputs like LeNet-5.

   Results appear on the VGA output in both modes.

Running Individual Subsystems

Every stage reads plain artifacts from disk, and the artifacts each stage needs are committed — so any subsystem can be run, studied, or modified on its own without setting up the ones upstream of it.

Neural Networks (PyTorch → StableHLO export)

python System/Neural_Networks/Start.py <model> <run> [-t] [-e] [-r]

<model> is LeNet_5, Tiny_NN, or Bigger_NN; <run> names the weight set (the committed artifacts use Recent). -t trains (MNIST downloads automatically on first use), -e exports <model>_<run>.mlir + <model>_<run>.weights.npz into the model's folder, -r evaluates on the test set.

Assembler (StableHLO → TPU transmission)

nn-assemble <model> <run>       # writes System/Assembler/out/TRANSMISSION.bin
python -m pytest System/Assembler/test/   # unit + end-to-end suite

Consumes the committed .mlir/.npz exports, so it runs with no training step. Two software models of the TPU live in System/Assembler/Interpreter/ and are the fastest way to understand the integer math without hardware:

python System/Assembler/Interpreter/Emulator.py 5     # golden int8 LeNet-5
                                                      # vs float, MNIST test[5]
nn-assemble LeNet_5 Recent                            # then execute the real
python System/Assembler/Interpreter/Interpreter.py 5  # binary instruction-by-
                                                      # instruction vs golden

Communication (PC → Arduino → FPGA)

cd System/Communication/Arduino_2_FPGA && pio run -t upload   # firmware
python System/Communication/PC_2_Arduino/Send_2_Arduino.py   # stream the
                                                              # transmission

The sender streams System/Assembler/out/TRANSMISSION.bin — assemble one or copy a /Benchmarks file into place first. Requires the hardware chain from the Synthesis Flow section.

Tensor Processing Unit (Verilog RTL)

System/Tensor_Processing_Unit/tests/run_regression.sh         # all testbenches
System/Tensor_Processing_Unit/tests/run_regression.sh 3 4 8   # a subset

Synthesis uses the Quartus project (step 1 of the Synthesis Flow). The regression compiles each testbench into its own fresh Questa library and runs it in batch mode; on a machine without Questa it prints a deferral notice and exits without simulating (edit QUESTA_DIR in the script to match your install).

Demo (LeNet-5 live drawing demo)

Demo/demo.sh                            # fullscreen GUI: program, draw, send
python Demo/LeNet5_Demo.py --mnist 0    # headless: send MNIST test[0]

Requires the full hardware chain plus python3-tk for the GUI.

AI Usage

This research is made possible by heavy usage of AI to write code for areas unrelated to the research goal and to expedite the process of writing Verilog code; the goal of the research is to understand how to design an AI accelerator, not to practice coding skills. I chose to use Claude Code for the majority of the AI work.

AI Usage Breakdown

• Specifications: Written and managed by me.
• Assembler: Full dark factory approach (Claude writes code AND manages main.md, CLAUDE.md, and log.md).
• Communication: Mostly dark factory, no spec documents since design is simple.
• Neural_Networks: Mostly handwritten.
• Tensor_Processing_Unit: Dark factory, code tested and validated manually.