Tutorial: A PyTorch MLP on TinyTPU, Then Change It¶
This tutorial takes a two-layer PyTorch MLP, compiles it onto the TinyTPU with ACT, and runs it on the TinyTPU design itself. Then it changes the machine twice and runs the same model again:
A parameter: a larger systolic array. One environment variable.
An instruction: a new fused
mvoutrelu, added to the ISA, the hardware and the compiler. Four small patches.
After each change ACT recompiles the unchanged model for the new machine,
and the result is checked byte for byte against PyTorch. Nothing on this page
needs Vitis. Every output shown is a real run, regenerated by
make tutorial-outputs.
PyTorch nn.Module workloads/models.py
-> AlloTracer + torch.fx workloads/extract.py shapes and ops
-> one workload spec per layer act/corpus/'s JSON
-> ACT mapspace search act/search.py, act_target.py
-> a TinyTPU program checked on isa_ref
-> the design, simulated ip/, df.build(target="simulator")
-> compared with PyTorch workloads/run.py
Paths are relative to examples/tinytpu/. The flow and its design choices
are explained on The Workload Suite: PyTorch MLPs on TinyTPU, the machine on TinyTPU-isa: A Programmable GEMM Accelerator.
Setup¶
Once per checkout. The allo conda environment, torch (optional
upstream, required here) and this checkout’s own MLIR bindings:
conda activate allo
pip install --index-url https://download.pytorch.org/whl/cpu torch==2.14.0
export LLVM_BUILD_DIR=/path/to/llvm-project/build # the env does not set it
examples/tinytpu/reproduce.sh --no-cosim # builds mlir/build, runs the gates
cd examples/tinytpu
The Makefile exports LLVM_BUILD_DIR and OMP_NUM_THREADS if they
are unset, and falls back to conda run -n allo when the shell’s python
cannot import allo.
The model¶
An ordinary nn.Module, with nothing in it about Allo or the TinyTPU:
class MlpSmall(nn.Module):
def __init__(self):
super().__init__()
self.fc1 = nn.Linear(32, 32, bias=False)
self.fc2 = nn.Linear(32, 16, bias=False)
def forward(self, x):
return self.fc2(torch.relu(self.fc1(x)))
It is registered with its batch size and input width, which is how make
mlp MODEL=<name> finds it. A new model is a class and one line here:
MODELS = {
"mlp_tiny": (MlpTiny, 4, 16),
"mlp_deep": (MlpDeep, 4, 16),
"mlp_small": (MlpSmall, 8, 32),
"mlp_wide": (MlpWide, 64, 64),
"mlp_bias": (MlpBias, 4, 16),
}
models.build fills the weights with integers in [-8, 8), so casting the
model to int8 is exact. The suite measures a mapping, not a quantisation
scheme.
Run it¶
make mlp # ~10 s: compile, build the design, run, compare
make mlp runs workloads/demo.py, which prints one section per stage.
Each stage is a function in that file, so a script can reuse any of them.
Trace and specs. fx gives the graph with shapes. Each nn.Linear
becomes one GEMM spec, and the ReLU that follows it becomes that spec’s
epilogue:
[2] Trace: AlloTracer + torch.fx ShapeProp
------------------------------------------
graph():
%x : [num_users=1] = placeholder[target=x]
%fc1 : [num_users=1] = call_module[target=fc1](args = (%x,), kwargs = {})
%relu : [num_users=1] = call_function[target=torch.relu](args = (%fc1,), kwargs = {})
%fc2 : [num_users=1] = call_module[target=fc2](args = (%relu,), kwargs = {})
return fc2
[3] Workload specs: 2 layers
----------------------------
mlp_small_l0 mk,kn->mn m=8 k=32 n=32 epilogue=relu+saturate
mlp_small_l1 mk,kn->mn m=8 k=32 n=16 epilogue=saturate
Compile. ACT lists every loop nest that tiles the GEMM onto a 4x4 array,
lowers each one to a TinyTPU program, and refuses the nests the instruction
word cannot express, giving the reason. The rest it ranks with its cost
model. The winner is then run on isa_ref (the ISA as numpy) against the
spec’s einsum:
[4] Compile with ACT: search the mapspace onto the TinyTPU
----------------------------------------------------------
mlp_small_l0: 2036 loop nests, 3 legal, 2033 refused
refused 1722 acc-split
refused 242 acc-position
refused 61 ar-distance
refused 8 AGU_TERMS
rank mapping instrs est. cycles
1 N8>K8 14 1186
2 M2>N8>K8 16 1588
3 N8>M2>K8 16 1790
picked N8>K8; the program against the spec's einsum (isa_ref, 4 operand distributions): BIT-EXACT
mlp_small_l1: 1175 loop nests, 3 legal, 1172 refused
refused 990 acc-split
refused 141 acc-position
refused 37 ar-distance
refused 4 AGU_TERMS
rank mapping instrs est. cycles
1 N4>K8 13 682
2 M2>N4>K8 15 884
3 N4>M2>K8 15 984
picked N4>K8; the program against the spec's einsum (isa_ref, 4 operand distributions): BIT-EXACT
The program. disasm.py prints it using the operand names from the
ISA spec. Note the vrelu then mvout at the end: the ReLU is a pass
over the accumulator of its own, and it is what the second change removes.
[5] The program ACT emitted for mlp_small_l0 (14 instructions)
--------------------------------------------------------------
0 loop x8
1 dma_ld rows=8 mode=2 dram_row0=0 col_block=0 dst_row0=0 agu: col_block+=iv0*1, dst_row0+=iv0*64
2 endloop
3 loop x8
4 dma_ld rows=32 mode=1 dram_row0=0 col_block=0 dst_row0=0 agu: col_block+=iv0*1, dst_row0+=iv0*64
5 endloop
6 loop x8
7 mm rows=8 vr_a=0 ar0=0 acc=0 spad_w=0 agu: spad_w+=iv0*64
8 loop x7
9 mm rows=8 vr_a=64 ar0=0 acc=1 spad_w=4 agu: vr_a+=iv1*64, spad_w+=iv0*64, spad_w+=iv1*4
10 endloop
11 vrelu rows=8 ar_d=0 ar_s=0
12 mvout rows=8 ar0=0 dram_row0=0 col_block=0 agu: col_block+=iv0*1
13 endloop
Run it on the design. df.build(tinytpu_isa, target="simulator")
builds the eight-unit dataflow design from ip/ (about 5 s). Each layer’s
program runs on it with the real weights in DRAM, and each layer’s output
becomes the next layer’s input:
[7] Run the model on the design, one program per layer
------------------------------------------------------
mlp_small_l0 8x32x32 design vs isa_ref over all 4096 bytes of C: 0 differ
mlp_small_l1 8x32x16 design vs isa_ref over all 4096 bytes of C: 0 differ
[✓] Against PyTorch
-------------------
0 of 384 output bytes differ
torch [127, -128, 127, -128, -128, -128, -128, -128, 127, -128, -128, 127, 127, -128, -128, -128]
tpu [127, -128, 127, -128, -128, -128, -128, -128, 127, -128, -128, 127, 127, -128, -128, -128]
Summary (cycles are the cost model's estimate; `make mlp-cosim` measures)
layer mapping instrs est. cycles
mlp_small_l0 N8>K8 14 1186
mlp_small_l1 N4>K8 13 682
model 1868
Two checks are made here, and they catch different faults. The design
against isa_ref, over the whole 4096-byte C buffer, catches the
hardware. The model’s output against PyTorch catches the compiler and the
flow. Most output bytes are saturated at 127 or -128, because weights and
inputs in [-8, 8) summed over K=32 overflow int8. So the row comparison is
weak evidence on its own, and the per-layer full-buffer check is the strong
one.
make mlp MODEL=mlp_deep TOP=5 # four layers, five mappings each
make mlp-compile # stop after ACT; no design build
A model the machine cannot run is refused with the reason, not approximated:
REFUSED fc1: nn.Linear with bias: the ISA has vadd but no mapping that broadcasts a bias row into the accumulator, so the baseline cannot express it
REFUSED sigmoid: call_function torch.sigmoid is not an int8 GEMM or a fused ReLU; this build computes int8 x int8 -> int32 with a ReLU epilogue
The ISA has no instruction for the refused nodes, so the model does not compile.
Change 1: a larger array¶
The design’s parameters live in ip/params.py:
class TpuParams:
"""T is the only parameter that changes the *shape* of the region: the
array is T*T instances and the chains are T and T*T stream arrays."""
T: int = 4 # SIMD width == array dimension
MAXDIM: int = 64 # largest M, K, N supported
SPAD_ROWS: int = None # scratchpad rows, each one packed word
NVR: int = None # operand vector registers
NAR: int = None # accumulator vector registers
QD: int = 16 # stream depth (16 since item 24: depth 8 deadlocks legal programs)
IMEM_SIZE: int = 56 # instruction memory words, header included
DMA_WORDS: int = 1 # packed words per operand-burst iteration
T is the array dimension, and the only parameter that changes the
shape of the design: T*T processing elements and T- and
T*T-wide stream arrays. Every size is derived from it. The shipped build
reads TPU_T from the environment, so the change is:
TPU_T=8 make mlp # an 8x8 array; ~25 s, the design is 4x larger
ACT searches again and finds different tilings, with half the loop trips:
[5] The program ACT emitted for mlp_small_l0 (14 instructions)
--------------------------------------------------------------
0 loop x4
1 dma_ld rows=8 mode=2 dram_row0=0 col_block=0 dst_row0=0 agu: col_block+=iv0*1, dst_row0+=iv0*64
2 endloop
3 loop x4
4 dma_ld rows=32 mode=1 dram_row0=0 col_block=0 dst_row0=0 agu: col_block+=iv0*1, dst_row0+=iv0*64
5 endloop
6 loop x4
7 mm rows=8 vr_a=0 ar0=0 acc=0 spad_w=0 agu: spad_w+=iv0*64
8 loop x3
9 mm rows=8 vr_a=64 ar0=0 acc=1 spad_w=8 agu: vr_a+=iv1*64, spad_w+=iv0*64, spad_w+=iv1*8
10 endloop
11 vrelu rows=8 ar_d=0 ar_s=0
12 mvout rows=8 ar0=0 dram_row0=0 col_block=0 agu: col_block+=iv0*1
13 endloop
The design is still bit-exact, and the model still matches PyTorch:
[✓] Against PyTorch
-------------------
0 of 384 output bytes differ
torch [127, -128, 127, -128, -128, -128, -128, -128, 127, -128, -128, 127, 127, -128, -128, -128]
tpu [127, -128, 127, -128, -128, -128, -128, -128, 127, -128, -128, 127, 127, -128, -128, -128]
Summary (cycles are the cost model's estimate; `make mlp-cosim` measures)
layer mapping instrs est. cycles
mlp_small_l0 N4>K4 14 607
mlp_small_l1 N2>K4 13 393
model 1000
against the T=4 build’s estimate:
Summary (cycles are the cost model's estimate; `make mlp-cosim` measures)
layer mapping instrs est. cycles
mlp_small_l0 N8>K8 14 1186
mlp_small_l1 N4>K8 13 682
model 1868
The cost model is fitted to measurements at T=4, so the near-halving is
its extrapolation, not a measurement. TPU_T=8 make mlp-cosim MODEL=mlp_small
measures it (Vitis, minutes).
Note
To change the shipped default rather than one run, edit it in both
isa_spec.json (the T parameter) and microarch_isa.py, then
python gen_isa.py --write. Editing only microarch_isa.py fails:
isa_encoding.py, which isa_ref reads, is generated from the spec
and keeps modelling a 4x4 machine, and gen_isa.py --check reports
“the reference model would be modelling a different machine”. The two
defaults are kept apart on purpose, so the design is checked against an
ISA that was written down independently of it.
Change 2: a fused instruction¶
The vrelu and mvout at the end of every ReLU layer are two passes
over the accumulator rows. mvout already clips each lane to int8 on its
way out, so it can apply the ReLU in the same step: mvoutrelu. The change
is four patches in tutorial/mvoutrelu/, in the order a designer would
write them:
git apply tutorial/mvoutrelu/*.patch
python gen_isa.py --write
make mlp
Step 1: declare it in the ISA¶
One entry in isa_spec.json: an opcode number, the operand fields (the
same as mvout’s) and its actions, the per-unit steps that say what
the instruction does. Compared with mvout it adds one max0 before
the clip:
diff --git a/examples/tinytpu/isa_spec.json b/examples/tinytpu/isa_spec.json
index 59bc4b88..78ad5285 100644
--- a/examples/tinytpu/isa_spec.json
+++ b/examples/tinytpu/isa_spec.json
@@ -983,6 +983,90 @@
}
],
"note": "`vadd` then `vrelu` in one pass: the accumulator already has both the adder and the rectifier, so the composition costs a mux and saves a pass over `ar` and an intermediate row. Two `ar` reads a row, so two accu steps a row -- which nothing here says, because it follows from `accu` having one `ar` read port."
+ },
+ {
+ "name": "mvoutrelu",
+ "value": 11,
+ "software_constant": "OP_MVOUTRELU",
+ "rows": "nr accumulator rows",
+ "operands": [
+ {
+ "field": "f0",
+ "name": "ar0",
+ "meaning": "first accumulator row to retire"
+ },
+ {
+ "field": "f1",
+ "name": "dram_row0",
+ "meaning": "first DRAM row of C"
+ },
+ {
+ "field": "f2",
+ "name": "col_block",
+ "meaning": "packed-word column block of C, 0 .. WPR-1"
+ }
+ ],
+ "note": "`vrelu` then `mvout` in one pass: the rectify happens on the way out, in the step that already clips to int8, so a fused-ReLU layer spends one accumulator pass and no intermediate write instead of two. Two chained lane operations a step (max0, then the clip), within the alu's declared depth. The worked example of adding a fused instruction: docs/source/designs/tinytpu_tutorial.rst.",
+ "rows_expression": "nr",
+ "actions": [
+ {
+ "unit": "accu",
+ "kind": "read",
+ "port": "ar.read",
+ "state": "ar",
+ "base": "ar0",
+ "into": "value",
+ "role": "the value to retire"
+ },
+ {
+ "unit": "accu",
+ "kind": "compute",
+ "port": "alu",
+ "compute": "max0",
+ "args": [
+ "value"
+ ],
+ "into": "rectified"
+ },
+ {
+ "unit": "accu",
+ "kind": "compute",
+ "port": "alu",
+ "compute": "to_operand",
+ "args": [
+ "rectified"
+ ],
+ "into": "clipped"
+ },
+ {
+ "unit": "accu",
+ "kind": "emit",
+ "port": "ac2sp",
+ "args": [
+ "clipped"
+ ]
+ },
+ {
+ "unit": "dma_st",
+ "kind": "receive",
+ "port": "ac2sp",
+ "into": "clipped_in",
+ "args": [
+ "clipped"
+ ]
+ },
+ {
+ "unit": "dma_st",
+ "kind": "write",
+ "port": "dram.write",
+ "state": "C",
+ "base": "dram_row0",
+ "args": [
+ "clipped_in"
+ ],
+ "offset": "col_block * T"
+ }
+ ]
}
],
"dispatch": {
python gen_isa.py --write then regenerates the encoder, the decoder, the
reference model’s meaning for the instruction and the ISA reference page,
none of which is edited by hand:
wrote examples/tinytpu/isa_encoding.py
wrote docs/source/designs/tinytpu_isa_spec.rst
Step 2: build it in the hardware¶
Three Allo units change. The opcode constant goes in ip/isa.py. The
accumulator (ip/units/accumulator.py, an Allo @df.kernel) reads the
source row from f0 as mvout does, and rectifies in the step that
already clips. The sequencer sends the instruction to the same two units as
mvout. dma_st only gets a comment, because it receives rows that
accu has already rectified.
diff --git a/examples/tinytpu/ip/isa.py b/examples/tinytpu/ip/isa.py
index d93cbff9..71c3edab 100644
--- a/examples/tinytpu/ip/isa.py
+++ b/examples/tinytpu/ip/isa.py
@@ -22,11 +22,12 @@ OP_MVOUT = 7 # f0=ar0 f1=dram_row0 f2=col_block nr=rows acc -> DRAM
OP_LOOP = 8 # open a loop, body is the next instruction nr=trip count
OP_ENDLOOP = 9 # close the innermost loop
OP_VADDRELU = 10 # f0=ar_d f1=ar_s1 f2=ar_s2 nr=rows
+OP_MVOUTRELU = 11 # f0=ar0 f1=dram_row0 f2=col_block nr=rows relu, acc -> DRAM
OPCODE_NAMES = {OP_NOP: "nop", OP_DMA_LD: "dma_ld", OP_DMA_ST: "dma_st",
OP_VLD: "vld", OP_MM: "mm", OP_VADD: "vadd", OP_VRELU: "vrelu",
OP_MVOUT: "mvout", OP_LOOP: "loop", OP_ENDLOOP: "endloop",
- OP_VADDRELU: "vaddrelu"}
+ OP_VADDRELU: "vaddrelu", OP_MVOUTRELU: "mvoutrelu"}
# `dma_ld`'s f0: bit 0 is the SOURCE matrix, bit 1 the DESTINATION memory.
# 0: A -> spad 1: B -> spad 2: A -> vr 3: B -> vr
@@ -82,6 +83,7 @@ ISA_NAMESPACE = {
"OP_VLD": OP_VLD, "OP_MM": OP_MM, "OP_VADD": OP_VADD,
"OP_VRELU": OP_VRELU, "OP_MVOUT": OP_MVOUT, "OP_LOOP": OP_LOOP,
"OP_ENDLOOP": OP_ENDLOOP, "OP_VADDRELU": OP_VADDRELU,
+ "OP_MVOUTRELU": OP_MVOUTRELU,
"DMA_SRC_B": DMA_SRC_B, "DMA_TO_VR": DMA_TO_VR,
"AGU_TERMS": AGU_TERMS, "AGU_F0": AGU_F0, "AGU_F1": AGU_F1,
"AGU_F2": AGU_F2, "AGU_F3": AGU_F3,
diff --git a/examples/tinytpu/ip/units/accumulator.py b/examples/tinytpu/ip/units/accumulator.py
index a0092791..bf0c0c47 100644
--- a/examples/tinytpu/ip/units/accumulator.py
+++ b/examples/tinytpu/ip/units/accumulator.py
@@ -4,7 +4,8 @@
"""The accumulator vector registers, their sole owner, and the vector ALU.
`mm` deposits or accumulates one k-tile's partial sums, `vadd` and `vrelu` are
-elementwise, and `mvout` clips to int8 on the way out. Every arm is muxed down
+elementwise, and `mvout` clips to int8 on the way out (`mvoutrelu` also
+rectifies, in the same step). Every arm is muxed down
to ONE `ar` read and ONE `ar` write per iteration; `vadd` needs two reads, so
it takes two steps per row, which is why the sequencer sends this unit twice
its row count for a `vadd`. Its II=1 rests on a dependence claim the assembler
@@ -37,7 +38,8 @@ def accumulator_directives(s, ctx):
reads=("c_acc", "cw"),
writes=("ac2sp",),
parameters=("NAR", "AW", "VW", "T"),
- isa=("OP_MM", "OP_VADD", "OP_VRELU", "OP_VADDRELU", "OP_MVOUT"),
+ isa=("OP_MM", "OP_VADD", "OP_VRELU", "OP_VADDRELU", "OP_MVOUT",
+ "OP_MVOUTRELU"),
directives=accumulator_directives,
)
def accu():
@@ -77,6 +79,8 @@ def accu():
write_row = f1 + row
if op == OP_MVOUT:
read_row = f0 + row
+ if op == OP_MVOUTRELU: # retires from f0, exactly as `mvout`
+ read_row = f0 + row
if two_source == 1:
if phase == 1:
read_row = f2 + row
@@ -124,6 +128,10 @@ def accu():
write_word[32 * relu_lane : 32 * (relu_lane + 1)] = rectified
else:
do_write = 0
+ # `mvoutrelu` rectifies in the same step `mvout` clips in.
+ rectify: int32 = 0
+ if op == OP_MVOUTRELU:
+ rectify = 1
clipped_word: UInt(VW) = 0
with allo.meta_for(T) as clip_lane:
retiring: int32 = read_word[32 * clip_lane : 32 * (clip_lane + 1)]
@@ -131,6 +139,9 @@ def accu():
retiring = 127
if retiring < -128:
retiring = -128
+ if rectify == 1:
+ if retiring < 0:
+ retiring = 0
clipped: int8 = retiring
clipped_word[8 * clip_lane : 8 * (clip_lane + 1)] = clipped
ac2sp.put(clipped_word)
diff --git a/examples/tinytpu/ip/units/dma_store.py b/examples/tinytpu/ip/units/dma_store.py
index 8111c6ab..3b020822 100644
--- a/examples/tinytpu/ip/units/dma_store.py
+++ b/examples/tinytpu/ip/units/dma_store.py
@@ -34,7 +34,8 @@ def dma_st(dram_c: int8[MAXDIM * MAXDIM]):
for work in range(n_row):
row += 1
if row >= instr_rows:
- # Only `mvout` reaches this queue, so there is no opcode test.
+ # Only `mvout` and `mvoutrelu` reach this queue, and `accu` has
+ # already applied the difference, so there is no opcode test.
word: UInt(64) = c_dst.get()
dram_row0 = word[18:30]
col_block = word[30:42]
diff --git a/examples/tinytpu/ip/units/sequencer.py b/examples/tinytpu/ip/units/sequencer.py
index b101676c..9a7adc9a 100644
--- a/examples/tinytpu/ip/units/sequencer.py
+++ b/examples/tinytpu/ip/units/sequencer.py
@@ -33,7 +33,7 @@ def sequencer_directives(s, ctx):
isa=("LOOP_DEPTH", "NHDR", "IWORDS", "AGU_TERMS", "AGU_F0", "AGU_F1",
"AGU_F2", "AGU_F3", "OP_LOOP", "OP_ENDLOOP", "OP_DMA_LD", "OP_VLD",
"OP_MM", "OP_VADD", "OP_VRELU", "OP_VADDRELU", "OP_MVOUT",
- "DMA_TO_VR"),
+ "OP_MVOUTRELU", "DMA_TO_VR"),
directives=sequencer_directives,
)
def sequencer(dram_imem: UInt(64)[IMEM_SIZE]):
@@ -145,6 +145,9 @@ def sequencer(dram_imem: UInt(64)[IMEM_SIZE]):
if op == OP_MVOUT:
c_acc.put(resolved)
c_dst.put(resolved)
+ if op == OP_MVOUTRELU:
+ c_acc.put(resolved)
+ c_dst.put(resolved)
pc += 1
if pc >= n_instr:
Step 3: teach the toolchain¶
The assembler checks the new instruction’s accumulator reads and adds its
rows to the per-unit work counts in the program header. The program DSL gets
an mvoutrelu() method, and the stress fuzzer mixes it into its random
programs half the time:
diff --git a/examples/tinytpu/ip/assembler.py b/examples/tinytpu/ip/assembler.py
index b1712020..706d5f93 100644
--- a/examples/tinytpu/ip/assembler.py
+++ b/examples/tinytpu/ip/assembler.py
@@ -19,7 +19,7 @@ from examples.tinytpu.ip.isa import (
AGU_TERMS, AGU_TERM_BITS, AGU_LEVEL_BITS, DMA_SRC_B, DMA_TO_VR, FIELD_LO,
FIELD_MASK, LOOP_DEPTH, NHDR, NR_LO, NR_MASK, OPCODE_NAMES, OP_DMA_LD,
OP_DMA_ST, OP_ENDLOOP, OP_LOOP, OP_MASK, OP_MM, OP_MVOUT, OP_NOP, OP_VADD,
- OP_VADDRELU, OP_VLD, OP_VRELU)
+ OP_MVOUTRELU, OP_VADDRELU, OP_VLD, OP_VRELU)
def opcode_of(word):
@@ -249,7 +249,7 @@ class Assembler:
ar_read(source[i], accu_step + i, "a source")
ar_write(dest[i], accu_step + i)
accu_step += nr
- elif op == OP_MVOUT:
+ elif op in (OP_MVOUT, OP_MVOUTRELU):
for i, row in enumerate(span("ar", f0, nr)):
ar_read(row, accu_step + i, "the value to retire")
accu_step += nr
@@ -309,9 +309,9 @@ class Assembler:
rows_to_spad + rows(OP_VLD) + n_mm * (p.T + 1),
rows_to_vr + rows(OP_VLD) + mm_rows,
n_mm | (mm_rows << 16),
- rows(OP_MM, OP_VRELU, OP_MVOUT)
+ rows(OP_MM, OP_VRELU, OP_MVOUT, OP_MVOUTRELU)
+ 2 * rows(OP_VADD, OP_VADDRELU),
- rows(OP_MVOUT),
+ rows(OP_MVOUT, OP_MVOUTRELU),
a_span | (b_span << 16)]
assert len(header) == NHDR
# Every count is read back through a 16-bit slice, which used to
diff --git a/examples/tinytpu/isa_dsl.py b/examples/tinytpu/isa_dsl.py
index 7bb6a5cb..18a1d811 100644
--- a/examples/tinytpu/isa_dsl.py
+++ b/examples/tinytpu/isa_dsl.py
@@ -87,7 +87,7 @@ sys.path.insert(0, os.path.abspath(os.path.join(os.path.dirname(__file__),
from examples.tinytpu.microarch_isa import ( # noqa: E402
AGU_F0, AGU_F1, AGU_F2, AGU_F3, AGU_TERMS, LOOP_DEPTH,
OP_DMA_LD, OP_ENDLOOP, OP_LOOP, OP_MM, OP_MVOUT, OP_NOP, OP_VADD, OP_VLD,
- OP_VADDRELU, OP_VRELU, enc, enc_agu,
+ OP_MVOUTRELU, OP_VADDRELU, OP_VRELU, enc, enc_agu,
A_VR, AR_C, B_SP, DMA_SRC_B, DMA_TO_VR, MAXDIM, MAXROWS, T,
SPAD_ROWS, NVR, NAR,
)
@@ -247,6 +247,10 @@ class Program:
"""Accumulator -> DRAM, clipped to int8 on the way."""
self._ins(OP_MVOUT, ar, dram_row, col_block, nr=rows)
+ def mvoutrelu(self, ar, dram_row=0, col_block=0, rows=0):
+ """`vrelu` then `mvout`, in one pass of the accumulator."""
+ self._ins(OP_MVOUTRELU, ar, dram_row, col_block, nr=rows)
+
def emit(self):
"""The `(word, agu_word)` list `assemble()` takes."""
if self._open:
diff --git a/examples/tinytpu/microarch_isa.py b/examples/tinytpu/microarch_isa.py
index a3a5513f..c723fda3 100644
--- a/examples/tinytpu/microarch_isa.py
+++ b/examples/tinytpu/microarch_isa.py
@@ -36,8 +36,8 @@ from examples.tinytpu.ip.assembler import ( # noqa: F401
from examples.tinytpu.ip.isa import ( # noqa: F401
AGU_F0, AGU_F1, AGU_F2, AGU_F3, AGU_TERMS, DMA_SRC_B, DMA_TO_VR, IWORDS,
LOOP_DEPTH, MAXROWS, NHDR, OP_DMA_LD, OP_DMA_ST, OP_ENDLOOP, OP_LOOP,
- OP_MM, OP_MVOUT, OP_NOP, OP_VADD, OP_VADDRELU, OP_VLD, OP_VRELU, enc,
- enc_agu)
+ OP_MM, OP_MVOUT, OP_MVOUTRELU, OP_NOP, OP_VADD, OP_VADDRELU, OP_VLD,
+ OP_VRELU, enc, enc_agu)
# Instruction slots. Sized to the longest program shipped (the stress harness's
# random programs), not to a round number: the sequencer's prefetch is
diff --git a/examples/tinytpu/stress_isa.py b/examples/tinytpu/stress_isa.py
index a859de84..5f4d3be2 100644
--- a/examples/tinytpu/stress_isa.py
+++ b/examples/tinytpu/stress_isa.py
@@ -282,7 +282,8 @@ def random_program(seed):
s = run_of("ar", n)
if s is None or not accu_ok([(s + i, i) for i in range(n)]):
return False
- k.mvout(s, dram_row=ri(0, MAXDIM - n), col_block=ri(0, WPR - 1), rows=n)
+ retire = k.mvoutrelu if ri(0, 1) else k.mvout
+ retire(s, dram_row=ri(0, MAXDIM - n), col_block=ri(0, WPR - 1), rows=n)
acc["it"] += n
return True
Step 4: let the compiler use it¶
ACT’s cost model gets the opcode, with the same unit loads as mvout. The
lowering gets one rule: if a layer’s epilogue ends in an op the retiring
instruction can apply (FUSED_RETIRE), fold it into the retire. Nothing
else in ACT changes. The mapspace search, ranking and checks are the same:
diff --git a/examples/tinytpu/act_machine.py b/examples/tinytpu/act_machine.py
index 8861df60..3ec8bc1e 100644
--- a/examples/tinytpu/act_machine.py
+++ b/examples/tinytpu/act_machine.py
@@ -13,7 +13,8 @@ from allo.act.schedule import Region # noqa: E402
from allo.act.schedule import Step # noqa: E402
from examples.tinytpu.microarch_isa import ( # noqa: E402
AGU_TERMS, DMA_SRC_B, DMA_TO_VR, LOOP_DEPTH, NAR, NVR, OP_DMA_LD,
- OP_ENDLOOP, OP_LOOP, OP_MM, OP_MVOUT, OP_VADD, OP_VLD, OP_VRELU,
+ OP_ENDLOOP, OP_LOOP, OP_MM, OP_MVOUT, OP_MVOUTRELU, OP_VADD, OP_VLD,
+ OP_VRELU,
SPAD_ROWS, T, expand,
)
@@ -85,10 +86,14 @@ MVOUT = Opcode(
reads=(lambda nr, f0, **_: Region(AR, f0, nr),),
writes=(lambda nr, f1, f2, **_: Region(dram("C", f2), f1, nr),))
+MVOUTRELU = Opcode(
+ code=OP_MVOUTRELU, name="mvoutrelu",
+ loads=MVOUT.loads, reads=MVOUT.reads, writes=MVOUT.writes)
+
MACHINE = Machine(
name="tinytpu-isa",
spaces=(Space(SPAD, SPAD_ROWS), Space(VR, NVR), Space(AR, NAR)),
- opcodes=(DMA_LD, VLD, MM, VADD, VRELU, MVOUT))
+ opcodes=(DMA_LD, VLD, MM, VADD, VRELU, MVOUT, MVOUTRELU))
def is_control(prog):
diff --git a/examples/tinytpu/act_target.py b/examples/tinytpu/act_target.py
index 4c882c5e..8029f87d 100644
--- a/examples/tinytpu/act_target.py
+++ b/examples/tinytpu/act_target.py
@@ -32,6 +32,13 @@ EPILOGUE = {
"relu": lambda k, dst, src, rows: k.vrelu(dst, src, rows),
}
+# An epilogue op the retiring instruction can apply on its way out. When the
+# LAST op of the epilogue is one of these, it is folded into the `mvout` and
+# costs no accumulator pass of its own.
+FUSED_RETIRE = {
+ "relu": lambda k: k.mvoutrelu,
+}
+
# Which kind of gap each refusal is. `express` means the instruction word has no
# field or no room for what the nest asks; the design is fine and the encoding
# is not. `refuse` means the machine would accept the program and produce a
@@ -260,10 +267,13 @@ class TinyTpu(Target):
result = into[workload.contractions[0]]
for term in workload.contractions[1:]:
k.vadd(result, result, into[term], rows=rows)
- for op in workload.epilogue:
+ epilogue, retire = list(workload.epilogue), k.mvout
+ if epilogue and epilogue[-1] in FUSED_RETIRE:
+ retire = FUSED_RETIRE[epilogue.pop()](k)
+ for op in epilogue:
EPILOGUE[op](k, result, result, rows)
- k.mvout(result, dram_row=walked(0, row_terms),
- col_block=walked(0, column_terms), rows=rows)
+ retire(result, dram_row=walked(0, row_terms),
+ col_block=walked(0, column_terms), rows=rows)
def walk(loops, open_loops):
if not loops:
Run it again¶
make mlp on the same model. ACT now retires the ReLU layer with the new
instruction, one instruction and one accumulator pass shorter:
[5] The program ACT emitted for mlp_small_l0 (13 instructions)
--------------------------------------------------------------
0 loop x8
1 dma_ld rows=8 mode=2 dram_row0=0 col_block=0 dst_row0=0 agu: col_block+=iv0*1, dst_row0+=iv0*64
2 endloop
3 loop x8
4 dma_ld rows=32 mode=1 dram_row0=0 col_block=0 dst_row0=0 agu: col_block+=iv0*1, dst_row0+=iv0*64
5 endloop
6 loop x8
7 mm rows=8 vr_a=0 ar0=0 acc=0 spad_w=0 agu: spad_w+=iv0*64
8 loop x7
9 mm rows=8 vr_a=64 ar0=0 acc=1 spad_w=4 agu: vr_a+=iv1*64, spad_w+=iv0*64, spad_w+=iv1*4
10 endloop
11 mvoutrelu rows=8 ar0=0 dram_row0=0 col_block=0 agu: col_block+=iv0*1
12 endloop
[7] Run the model on the design, one program per layer
------------------------------------------------------
mlp_small_l0 8x32x32 design vs isa_ref over all 4096 bytes of C: 0 differ
mlp_small_l1 8x32x16 design vs isa_ref over all 4096 bytes of C: 0 differ
[✓] Against PyTorch
-------------------
0 of 384 output bytes differ
torch [127, -128, 127, -128, -128, -128, -128, -128, 127, -128, -128, 127, 127, -128, -128, -128]
tpu [127, -128, 127, -128, -128, -128, -128, -128, 127, -128, -128, 127, 127, -128, -128, -128]
Summary (cycles are the cost model's estimate; `make mlp-cosim` measures)
layer mapping instrs est. cycles
mlp_small_l0 N8>K8 13 1085
mlp_small_l1 N4>K8 13 682
model 1767
mlp_small_l1 has no ReLU and is unchanged. The model’s estimate falls by
what the vrelu pass cost (compare the baseline summary above).
Check it¶
make mlp is one model. Two gates check the instruction itself, and both
passed on this change:
python gen_isa.py --check # the spec, generated files and 9 consumers agree
python stress_isa.py # ~3 min: 640 programs on the design vs isa_ref
Both earned their keep while the change was written. The first version failed each of them:
gen_isa.py --checkfailed twice. The accumulator’s clip no longer read as a literal saturation at -128, which is how the check verifies it. Anddma_st, now reached by two opcodes, looked like it should decode them. It shouldn’t: both do the same thing there. That second failure was a gap in the check, which now lets a unit skip decoding when every opcode that reaches it has identical actions.stress_isa.pydeadlocked inkpn_model, a channel-level model of the design that had its own hand-written “onlymvoutsends a row todma_st”. That rule is now derived from the spec’s actions.stress_isa.pythen found a real hardware bug: the accumulator read the source row fromf1, the DRAM-row field, instead off0.make mlphad passed anyway, because in its programs both fields are 0. Only the random programs exposed it. That is the lineif op == OP_MVOUTRELU: read_row = f0 + rowin step 2.
The first two fixes are part of the tree, so the patches above are just the designer’s change.
Undo it:
git apply -R tutorial/mvoutrelu/*.patch
python gen_isa.py --write
What is and is not established¶
Cycles on this page are estimates from ACT’s cost model, fitted to cosim at
T=4. The fused instruction has not been through RTL cosim, and theT=8figure is an extrapolation.make mlp-cosim MODEL=mlp_smallmeasures either one (Vitis, onecsynth, minutes).The fused instruction is not on
main. With the patches applied, thegemm.relucases intests/act/test_tinytpu.pyfail, and should: they pin ReLU layers to the shippedvrelu; mvoutprograms, which cosim has measured. Mergingmvoutrelumeans measuring the new programs withTPU_TB=stresscosim and re-pinning those tests.Untested combination:
TPU_T=8with the patches applied.
Reproducing this page¶
make tutorial-outputs # re-runs everything above, ~1 min
pytest tests/act/test_tutorial.py # from the repo root, seconds
make tutorial-outputs runs tutorial/capture.sh, which writes the
outputs this page includes to docs/source/designs/tinytpu_tutorial/. It
applies the patches for the last run and reverts them however it exits. The
test checks that the walk-through still compiles and matches PyTorch, that
mlp_bias is still refused, and that the four patches still apply to the
tree. If a file the patches edit moves on, that test fails first.