PyTorchSim: A Comprehensive, Fast, and Accurate NPU Simulation Framework

Wonhyuk Yang^∗, Yunseon Shin^∗, Okkyun Woo^∗, Geonwoo Park^§, Hyunkyu Ham,

Jeehoon Kang^†¶, Jongse Park^¶, Gwangsun Kim

MICRO 2025 Session 7B: Tools and Simulators

^* Equal contribution

^§ Currently with Samsung Electronics

AI Model & Hardware Trends

• Deep Neural Network (DNN) are growing exponentially in size
• Hardware performance improves relatively slowly

2

1.4x / year

4.4x / year

We need high-performance & efficient hardware

Evolution of Neural Processing Units

Neural Processing Units (NPUs) are designed to address this challenge

• Efficient dataflow units (e.g., systolic arrays, adder-tree)

3

NPU simulators play a crucial role in designing NPUs

NPU Simulator Requirements – AI Model

• Both inference and training are important
• (Vector) operations are becoming increasingly diverse
• Sparsity is now widely exploited for efficiency
• Latency varies with input data (i.e., data-dependent timing behavior)

4

NPU Simulator Requirements – NPU Hardware

• Need to accurately model key components of modern NPUs
• NPU core: matrix, vector, and scalar units and scratchpad memory
• Shared resources: interconnect and DRAM
• Detailed simulation often necessary to model inter-core contention

​

5

NPU Simulator Requirements – Compiler

• Necessary to bridge AI models and NPU hardware
• Lowers DNN models into machine code
• Apply various optimizations
• Enables full-model simulation for both inference and training

6

For compiler support, an NPU ISA must first be defined

NPU ISA

• No de-facto standard ISA for NPUs (unlike CPUs and GPUs)

​

• RISC-V can be a strong candidate for an NPU ISA
• Generality:
• Vector-length agnostic design
• A generic interface (VCIX) for dataflow units
• Extensibility: Reserved opcode space for custom instructions
• Openness: Contributions from both academia and industry
• SW Infrastructure: Rich software ecosystem (e.g., compilers, simulators)

​

​

​

7

Existing NPU Models

• Various NPU simulators and analytical models have been proposed

8

None of them meets all the key requirements

We propose PyTorchSim, which can better satisfy the key requirements for NPU simulation

Contents

• Background / Motivation
• PyTorchSim
• PyTorchSim overview
• NPU core modeling
• Compilation flow
• Methodology & Evaluation
• Case study
• Impact of DNN Training Hyperparameter
• Scheduling for Chiplet-based NPUs
• Summary

​

PyTorchSim Framework: Overview

10

Modeling Common Building Blocks in NPU Cores

11

Scalar

unit

Vector

unit

VCIX

VCIX: A Generic Interface for Dataflow Units

12

We adopt VCIX, a generic interface for diverse dataflow units

PyTorchSim Compilation Flow

• NPU backend generates a RISC-V binary
• Applies target-dependent optimizations

13

TorchInductor

CPU backend

GPU backend

FX graphs

Loop-level IR

torch.compile()

• Benefits of integration with PyTorch 2
• Simulate existing PyTorch models without any need for manual conversion
• Supports training simulation through� PyTorch’s auto-differentiation
• Leverage existing target-independent optimizations

TorchDynamo/Autograd

PyTorch models

Instruction-Level Simulation (ILS)

• Simulate every single instruction one-by-one
• Inherently slow

​

14

DMA load

Tile A

DMA load

Tile B

Matrix multiply

​

(~100 instructions)

Key Insight behind NPU Execution

15

Tile Operation Graph (TOG)

16

.insn r 43, 3, 0, zero, t2, t4

.insn r 43, 3, 4, zero, a2, ra

.insn r 43, 3, 5, zero, a2, s8

.insn r 43, 3, 2, zero, a3, s0

add a0, s1, a6

.insn r 43, 3, 0, zero, a7, t4

.insn r 43, 3, 4, zero, a2, ra

.insn r 43, 3, 5, zero, a2, s8

.insn r 43, 3, 2, zero, a0, a5

addi a4, a4, 1

addi s1, s1, 512

addi a5, a5, 1024

bge s3, a4, .LBB0_16

j .LBB0_13

...

​

add a0, a0, s7

vmv1r.v v11, v6

vle32.v v11, (a0), v0.t

sf.vc.iv 1, 0, v8, 0

sf.vc.iv 1, 0, v9, 0

...

sf.vc.iv 1, 0, v8, 0

sf.vc.iv 1, 0, v9, 0

bltz s3, .LBB0_21

RISC-V binary (GEMM)

( X cycles )

Tile-Level Simulation (TLS) Execution Flow

Tile-Level Simulation works in a two-step flow

​

Offline (compile time)

• Obtain the deterministic compute latency in the TOG

​

Online (simulation time)

• Reuse the obtained compute latency
• Model DMAs with detailed NoC & DRAM�simulators (bookism, ramulator2)

17

TLS can achieve both high simulation speed and accuracy

PyTorchSim TLS Compilation Flow

• Gem5 (timing simulator) models the compute latency for TLS

18

NPU backend

(Ours)

PyTorch

Runtime

Loop-level IR

Challenge: Tile operation latency can depend on input data!

Data-Dependent Timing Behavior Example

19

For each specific pair of input tiles,

the compute latency is still deterministic

Data-dependent attributes

PyTorchSim TLS Compilation Flow

• Spike (functional simulator) executes the binary with PyTorch input:
• To obtain data-dependent attributes
• To validate correctness of compiled binary

20

TOG Simulator

PyTorch

Runtime

Loop-level IR

TOG

RISC-V

binary

Putting It All Together

21

Data-dep.

attributes

Evaluation Methodology

Accuracy validation

• Target: Real Google TPU v3
• Baselines: SCALE-Sim v3, mNPUSim, Timeloop, MAESTRO

​

​

Simulation speed

• Baselines: Accel-Sim, mNPUsim
• We selected Accel-Sim for its rich features and GPUs’ dominance in deep learning

Workloads

• Kernels: GEMM, convolution, layer normalization, softmax, attention
• Full models: ResNet18/50, Bert-base/large

​

​

​

​

​

22

Evaluation: Validation against Real TPU v3

23

860%

GEMM

(NxNxN)

PyTorchSim can accurately simulate full models end-to-end

Evaluation: Speedup

24

Speedup

48

PyTorchSim achieves significant simulation speedup

Contents

• Background / Motivation
• PyTorchSim
• Overview
• NPU core modeling
• Compiler workflow
• Methodology & Evaluation
• Case Studies (CS)
• Impact of DNN training hyperparameter
• Impact of data placement for chiplet-based NPUs
• Summary

​

25

CS1: Impact of DNN Training Hyperparameter�

26

MLP

Training loss

PyTorchSim enables studying training time behaviors of systems

CS2: Impact of Data Placement for Chiplet-based NPUs

27

Chiplet interconnect

(64 GB/s, 20 ns)

CS2: Impact of Data Placement for Chiplet-based NPUs

28

CS2: Impact of Data Placement for Chiplet-based NPUs

29

• High remote traffic becomes the severe bottleneck

Norm. Runtime

PyTorchSim enables the study of diverse NPU architectures

More Results and Discussions in the Paper

Case studies

• Heterogeneous dense-sparse NPU
• DNN inference with multi-model tenancy
• Compiler optimization impact

Discussions

• Extension for other deep learning frameworks
• Extension for modeling GPUs and large-scale systems

​

​

​

30

Summary

• AI workloads require a comprehensive, fast, and accurate simulator

​

• To address this, we propose PyTorchSim
• High-speed, high-accuracy through our Tile-Level Simulation (TLS)
• Integration with PyTorch 2 compilation flow
• General and extensible architecture through a RISC-V-based NPU ISA

​

• Compared to prior NPU simulators, PyTorchSim:
• Enables various case studies not supported by prior simulators
• Significantly improves accuracy and simulation speed

31

PyTorchSim is Open Source

• Contributions are welcome!

​

32

Backup

33

Comparison of the features of different NPU sim.

34

Compilation flow with simplified example IR

• PyTorchSim backend generates MLIR code
• MLIR’s affine, arith, linalg, memref, vector, llvm, etc., which are widely used in AI compilers, are adopted to represent the LLIR.

35

Tile Operation Graph

36

( X cycles )

Custom Instructions

37

CS3: Heterogeneous Dense-Sparse NPU

37

Providing core abstraction interfaces

CS3: Heterogeneous Dense-Sparse NPU

37

• Dense access patterns are favored

​

• Takeaway: Memory scheduling must be carefully designed to ensure fairness

System-driven Requirements

Requirements for modern NPU architecture design

40

Source: Source Hot Chips 2025, “Ironwood,” Aug 26, 2025, PDF (slide 15)

Why AI compiler support matters

• Modern deep learning frameworks now integrate AI compilers
• Compiler optimizations have a huge impact on hardware performance

​

​

41

Source: https://pytorch.org/get-started/pytorch-2-x/

​

Understanding Tile Operation Graph

42