Systolic Array Simulator and Compiler Extensions

A cycle-approximate output-stationary systolic array simulator in C, extended with PyTorch, TVM, MLIR, and a parametric performance model. Built for Georgia Tech ECE 6100 (Advanced Computer Architecture), then extended for hardware-software co-design and compiler work.

This repository doesn't contain the original course source materials (lab5a, lab5b) and files required to run completely, it only contains the extensions to adhere to academic integrity and ensure future students complete the base assignment themselves. Feel free to reach out if you are interested in reproducing and running it properly.

What This Is

• lab5a/: Blocked vs. naive matrix multiply on CPU. Shows cache locality through tiling.
• lab5b/: 16x16 systolic array simulator in C. PEs hold accumulators, FIFOs feed data from the left and top boundaries, and skewed injection pushes partial sums diagonally.
• ext_pytorch/: The C simulator wrapped as a torch.ops.systolic.matmul custom operator via TORCH_LIBRARY and pybind11.
• ext_tvm/: TVM BYOC pattern registration (dense, dense+bias, dense+bias+relu) and C-source codegen for partitioned Relay graphs.
• ext_mlir/: A custom systolic MLIR dialect with ODS ops (matmul, load_tile, store_tile) and a lowering pass to affine loops.
• ext_perf_model/: Parametric Python simulator, XGBoost cost model, and roofline analysis for design-space exploration.

Getting Started

Prerequisites

• C compiler (gcc, clang, or MSVC)
• Python 3.8+ with PyTorch 2.0+ (for ext_pytorch and ext_perf_model)
• GNU Make (for lab5a/ and lab5b/)
• Optional: TVM source tree with LLVM (for ext_tvm/), LLVM/MLIR built from source (for ext_mlir/)

Quick Run

# 1. CPU baseline: blocked vs naive matmul
cd lab5a
make
./sim

# 2. Systolic array simulator (C, fixed 16x16)
cd lab5b
make
./sim
# Expected: PASS, ~47 cycles

# 3. PyTorch custom op
cd ext_pytorch
python setup.py build_ext --inplace
python benchmark.py
# Expected: Correctness PASS, ~6x overhead vs native torch.matmul

# 4. Performance model
cd ext_perf_model
python collect_data.py        # generates 800 data points
python train_model.py         # trains XGBoost model, reports MAPE/RMSE
python explore.py             # design-space exploration
python plot_roofline.py       # generates roofline_systolic.png

# 5. TVM demo (works without TVM installed)
cd ext_tvm
python demo_tvm_byoc.py

# 6. MLIR dialect (requires LLVM/MLIR build tree)
cd ext_mlir
# Place in your LLVM source tree, build with CMake, then:
# mlir-opt test_matmul.mlir --systolic-to-affine ...

Architecture

The systolic array is output-stationary. Each PE at (i, j) accumulates C[i][j] = dot(A[i, :], B[:, j]). Data enters from the left boundary (rows of A) and top boundary (columns of B), with skewed injection so each PE receives its A and B values at the right cycle.

         top[0] ... top[N-1]   (FIFOs for columns of B)
             |         |
             v         v
left[0] -> PE[0,0] -> PE[0,1] -> ... -> PE[0,N-1]
   |          |          |               |
   v          v          v               v
left[1] -> PE[1,0] -> PE[1,1] -> ... -> PE[1,N-1]

Each PE:

• Accepts west_in (A value) and north_in (B value)
• When both valid: acc += A * B, passes inputs to east_out / south_out
• Sets done = true after N MACs

The C simulator uses #define N 16 in systolic.h. The Python simulator in ext_perf_model/systolic_sim.py takes N as a constructor argument and matches the C cycle order exactly.

Components

lab5a/ - Blocked Matrix Multiply

Compares simple_matmul (naive triple loop) vs blocked_matmul (128x128 tiles with transposed B tile) for N = 256 to 4096. The blocked version keeps data in cache by operating on tiles and copying them into local SRAM buffers before computing.

lab5b/ - Systolic Array Simulator

Four files:

• fifo.c: circular FIFO for boundary data
• pe.c: Processing Element with MAC, accumulator, and pass-through
• systolic.c: grid initialization, cycle driver, data shifting, boundary injection
• main.c: random A/B, golden reference, simulator run, PASS/FAIL check

The simulator is cycle-approximate, not cycle-accurate RTL. PEs compute in a single global step, and FIFOs never stall.

ext_pytorch/ - Custom C++ Operator

systolic_op.cpp registers torch.ops.systolic.matmul via TORCH_LIBRARY. It accepts int32 CPU tensors, runs the C simulator, and returns the result. The benchmark shows ~6x overhead vs native torch.matmul because the simulator is cycle-level C code, not optimized BLAS.

ext_tvm/ - TVM BYOC Backend

• systolic_patterns.py: registers nn.dense, nn.bias_add, nn.relu for the systolic target, and defines composite patterns for dense+bias+relu fusion.
• systolic_codegen.py: generates C source with extern declarations for the simulator functions. The partition_and_codegen function runs the full pipeline when TVM is installed.
• demo_tvm_byoc.py: prints the patterns and a C codegen template without needing TVM. The imports are conditionally guarded so it works standalone.

To use in a real TVM build: copy the two .py files into the TVM tree under python/tvm/relay/op/contrib/ and python/tvm/relay/backend/contrib/, then call partition_and_codegen(mod, params) from your Relay workflow.

ext_mlir/ - Custom MLIR Dialect

• include/systolic/SystolicOps.td: ODS definitions for systolic.matmul, systolic.load_tile, systolic.store_tile
• SystolicDialect.cpp: dialect registration and type inference
• SystolicToAffine.cpp: MatmulOpLowering replaces systolic.matmul with a 3-deep affine loop nest (i, j, k) using affine.load, arith.muli, arith.addi, affine.store
• test_matmul.mlir: example IR with systolic ops
• CMakeLists.txt: build config for an LLVM/MLIR source tree

To build: place this directory in your LLVM source tree, add it to the CMake configuration, and build with mlir-tblgen for ODS code generation.

ext_perf_model/ - Performance Modeling and DSE

Simulator: systolic_sim.py replicates the C cycle order in Python and accepts any N at runtime. Run it standalone with python -m systolic_sim -N 16.

Data collection: collect_data.py sweeps 800 configurations across:

• Array size N: 4, 8, 12, 16, 20, 24, 28, 32
• PE latency: 1, 2, 3, 5, 10 cycles per MAC
• Memory bandwidth: 1, 2, 4, 8, 16 bytes/cycle
• Tile size: 4, 8, 16, 32

Cycle counts include Gaussian noise on the memory component to model real hardware variability (memory contention, temperature effects). This prevents the model from simply memorizing a deterministic formula.

Training: train_model.py uses group-based train/test splitting by configuration tuple (N, pe_clock, mem_bw, tile_size) so no config appears in both sets. This avoids data leakage. The XGBoost model achieves:

• MAPE: 5.47%
• RMSE: 272.75 cycles

These are honest metrics on held-out configurations, not held-out seeds of the same config.

Exploration: explore.py loads the trained model and searches the design space for the lowest predicted cycle count.

Roofline: plot_roofline.py draws a roofline plot for a chosen configuration (e.g., N=16, pe_clock=1, tile_size=16) with peak bandwidth taken from the actual dataset max, not hardcoded. The plot shows achieved performance vs. operational intensity for varying memory bandwidths.

Limitations and Honest Notes

• The C simulator is fixed at N=16. Changing it requires editing systolic.h and recompiling. The Python simulator is parametric.
• The simulator is cycle-approximate, not a gate-level or RTL model. There is no contention, no stall modeling, and FIFOs have infinite source/sink abstraction.
• The TVM and MLIR extensions are drop-in code for existing TVM/LLVM build trees. They are not standalone buildable projects. The TVM demo works without TVM installed, but the full integration requires a TVM build.
• The XGBoost cost model is a surrogate for a cycle-approximate simulator, not a real RTL or silicon measurement. The 5.47% MAPE is good for a surrogate but should not be compared to cycle-accurate RTL simulation.
• The MLIR lowering pass replaces systolic.matmul with scalar affine loops. It does not generate vector or systolic-specific instructions. It is a dialect definition and lowering proof-of-concept.

Attribution

The original cycle-approximate systolic array simulator and blocked matrix multiply benchmark were developed as coursework for Georgia Tech ECE 6100 — Advanced Computer Architecture.

This repository contains only the independent extensions: the PyTorch custom operator, TVM BYOC backend, custom MLIR dialect, and parametric performance model. The original assignment materials, starter code, and solutions are not included here to maintain academic integrity.

2026 — Apuroop Mutyala