File size: 14,190 Bytes
c699c4c | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 185 186 187 188 189 190 191 192 193 194 195 196 197 198 199 200 201 202 203 204 205 206 207 208 209 210 211 212 213 214 215 216 217 218 219 220 221 222 223 224 225 226 227 228 229 230 231 232 233 234 235 236 237 238 239 240 241 242 243 244 245 246 247 248 249 250 251 252 253 254 255 256 257 | # SPDX-FileCopyrightText: © 2026 Tenstorrent USA, Inc.
# SPDX-License-Identifier: Apache-2.0
"""ABANDONED (round 7): with SP_RC_WPROC=1 the kernel still hung chip 6 after ~700 back-to-back launches
(chipstate/chip6/FAULT, OPT_REPORT.md "Round 7"); do not run it. Kept for reference only.
EXPERIMENTAL, NOT WIRED INTO THE MODEL (round 5): one eager call per Cout=256 was bit-identical to
ttnn.conv2d on chip 10; the next test run (Cout=128 eager + 10 calls in one trace) hung the chip
(chipstate/chip10/FAULT, OPT_REPORT.md "Round 5"). Round 6 (chip 4, end-of-op 'done' barrier on semaphore 2
added): Cout 128 and 256 bit-identical to ttnn.conv2d, 1 and 2 eager calls, traces of 2 and 10 calls (3 replays)
all passed; the timing stage (~500 back-to-back launches) hung chip 4 (chipstate/chip4/FAULT). Leading
hypothesis: the two groups' weight multicasts run concurrently on NOC 1 (PROC 0, start/end swapped), so each
path from its sender in row 0 to the far corner crosses the other group's rectangle (path-reservation
deadlock, rare per launch). Next: SP_RC_WPROC=1 (sender on NOC 0, multicast starts at the sender's own
corner, the two rectangles share no router), staged on a fresh chip. Validate step by step before use.
3x3 / pad-1 conv (+ bias + ReLU) for the 60x80 layers (block 3, heads) as ONE custom generic_op
without halo / im2col / tilize, bit-identical to ttnn.conv2d's height-sharded program.
ttnn's conv for these layers is a halo op (7.7-8.9 us) + a conv on 75 of the 120 cores whose
per-core work is dominated by the im2col reader and tilize (the math is 2 M tiles x 36 K x N).
Here each image row r (80 px, padded to 96 = 3 tile rows m) is computed by two cores, one per half
of the output channels (group g): core (r, g) = logical (x = 6 g + r % 6, y = r // 6).
Input / output: the TILE [H*W, C] layout ttnn's conv uses, height-sharded [64, C] over the grid
(row-major; 75 shards hold data). The data-movement RISCs build the in0 tiles of the core directly
from the (neighbour) shards with a host-built NoC copy schedule (kernels/sp_conv/rc_dm.cpp):
T(ky, kx, m, c) row i = pixel 32 m + i + kx - 1 of image row r + ky - 1, channels 32 c .. 32 c + 31
(zero outside the image), i.e. exactly ttnn's im2col row for the output pixel, K block ky, K tiles
(kx, c) -- 12 tiles per K block for Cin = 128. Each RISC builds half of the channel tiles c.
Compute (kernels/sp_conv/rc_compute.cpp) reproduces ttnn's conv_bmm_tilize.cpp order for
height sharding with packer_l1_acc + fused bias: per K block ky (= ttnn's in0_block_w_i, 12 tiles)
the fp32 DST accumulates the 12 tile products in K order, the block result is packed to fp32
partials (ky = 0 overwrite, ky > 0 packer L1 accumulate), then partials + bias
(add_tiles_bcast_rows) with ReLU on pack -> bf16. Same values, same order => same bits.
Weights: per group the [1152, Cout/2] bf16 matrix (K order ky, kx, ci as ttnn's prepared weights)
is read from DRAM by one sender core (logical (6 g, 0)) in three K-block chunks and multicast to the
group's 6 x 10 rectangle (the two groups use disjoint rectangles); a semaphore carries the number
of chunks delivered. Receivers first signal the sender that they run (no multicast into a core
that has not started the program).
Output rows px 0..79 of tile rows m are written (1 KB = 16-row halves) to the output shards.
"""
from __future__ import annotations
import os
import numpy as np
import torch
import ttnn
_KDIR = os.path.join(os.path.dirname(os.path.dirname(os.path.dirname(os.path.abspath(__file__)))), "kernels", "sp_conv")
H, W = 60, 80
WP = 96 # padded row: 3 tile rows
SHARD = 64 # pixels per input / output shard (ttnn conv layout for 4800 px on the 12 x 10 grid)
NG = 2 # output-channel groups (cores per image row)
GX = 6 # grid columns per group
ZSEL = 7 # schedule source index of the local zero page
NSRC = 6 # source core slots per core (input)
NDST = 3 # destination core slots per core (output)
def _src(name):
with open(os.path.join(_KDIR, name)) as f:
return f.read()
def _unit(row: int, fc: int) -> int:
"""32-byte unit of tile row `row` (0..31), column half fc, inside a tile (face layout)."""
return ((row >> 4) * 2 + fc) * 16 + (row & 15)
def _merge(ent):
"""[(sel, su, du, L)] sorted by du -> merged runs contiguous in source and destination."""
out = []
for e in ent:
if out:
sel, su, du, L = out[-1]
if e[0] == sel and e[2] == du + L and (sel == ZSEL or e[1] == su + L) and (L + e[3] <= 32 if sel == ZSEL else True):
out[-1] = (sel, su, du, L + e[3])
continue
out.append(e)
return out
def core_of(r: int, g: int):
return GX * g + r % GX, r // GX
def input_schedule(r: int, ct: int, ch: int):
"""Copy entries of core row r per K block ky: [3][(sel, src_unit, dst_unit, n_units)] for channel
tile 0 (the kernel adds 64 units per channel tile on both sides), plus the source shard list."""
shards = []
blocks = []
blk_h = 9 * ch
for ky in range(3):
R = r + ky - 1
ent = []
for m in range(3):
for kx in range(3):
dst_tile = (m * 3 + kx) * ch # within this K block's CB_A reservation
for i in range(32):
q = 32 * m + i + kx - 1
valid = 0 <= R < H and 0 <= q < W
for fc in (0, 1):
du = dst_tile * 64 + _unit(i, fc)
if not valid:
ent.append((ZSEL, 0, du, 1))
continue
p = W * R + q
s, lr = divmod(p, SHARD)
if s not in shards:
shards.append(s)
tr, ri = divmod(lr, 32)
su = tr * ct * 64 + _unit(ri, fc)
ent.append((shards.index(s), su, du, 1))
ent.sort(key=lambda e: e[2])
blocks.append(_merge(ent))
assert len(shards) <= NSRC, shards
return blocks, shards
def output_schedule(r: int, g: int, nl: int, cto: int):
"""Output entries of core (r, g): [(sel, src_unit, dst_unit, n_units)] for local n = 0 (the kernel adds
64 units per output channel tile on both sides) + destination shard list."""
shards = []
ent = []
for m in range(3):
for h in range(2):
q0 = 32 * m + 16 * h
if q0 >= W:
continue
p = W * r + q0
s, lr = divmod(p, SHARD)
assert lr % 16 == 0
if s not in shards:
shards.append(s)
tr, ri = divmod(lr, 32)
su = m * nl * 64 + h * 32 # rows 16 h .. 16 h + 15 of out tile (m, n = 0) = units 32 h .. 32 h + 31
du = (tr * cto + g * nl) * 64 + (ri // 16) * 32
ent.append((shards.index(s), su, du, 32))
assert len(shards) <= NDST
return ent, shards
class RowConv:
"""conv3x3(pad 1) + bias (+ ReLU) on [1, 1, 4800, Cin] TILE height-sharded [64, Cin] (60 x 80)."""
def __init__(self, device, weight: torch.Tensor, bias: torch.Tensor, relu: bool = True):
self.device = device
co, ci = int(weight.shape[0]), int(weight.shape[1])
assert tuple(weight.shape[2:]) == (3, 3) and ci % 64 == 0 and co % (32 * NG) == 0
self.ci, self.co, self.relu = ci, co, relu
self.ct = ci // 32 # input channel tiles
self.ch = self.ct // 2 # channel tiles per data-movement RISC
self.nl = co // 32 // NG # output channel tiles per core
self.cto = co // 32
# K order (ky, kx, ci) as ttnn's prepared conv weights; group g = output channels g * co/2 ..
wk = weight.detach().float().permute(2, 3, 1, 0).reshape(9 * ci, co) # [ky, kx, ci] x co
wg = torch.cat([wk[:, g * co // NG:(g + 1) * co // NG] for g in range(NG)], 0) # [NG * 9ci, co/NG]
self.w = ttnn.from_torch(wg.to(torch.bfloat16), dtype=ttnn.bfloat16, layout=ttnn.TILE_LAYOUT, device=device,
memory_config=ttnn.DRAM_MEMORY_CONFIG)
bm = torch.zeros(32, co)
bm[0] = bias.detach().float()
self.b = ttnn.from_torch(bm.to(torch.bfloat16), dtype=ttnn.bfloat16, layout=ttnn.TILE_LAYOUT, device=device,
memory_config=ttnn.DRAM_MEMORY_CONFIG)
# per-core schedule pages (core order = row-major logical grid of the 120 cores)
self.cores = [ttnn.CoreCoord(x, y) for y in range(10) for x in range(12)]
pages = []
self.src_shards, self.dst_shards = {}, {}
for c in self.cores:
g, r = c.x // GX, c.y * GX + c.x % GX
blocks, ss = input_schedule(r, self.ct, self.ch)
oent, ds = output_schedule(r, g, self.nl, self.cto)
self.src_shards[(c.x, c.y)], self.dst_shards[(c.x, c.y)] = ss, ds
words = [len(b) for b in blocks] + [len(oent)]
for b in blocks + [oent]:
for sel, su, du, L in b:
assert su < 1 << 16 and du < 1 << 16 and L < 1 << 16
words += [(sel << 16) | su, (du << 16) | L]
pages.append(words)
self.page_words = (max(len(p) for p in pages) + 7) // 8 * 8
tab = np.zeros((len(pages), self.page_words), np.uint32)
for i, p in enumerate(pages):
tab[i, :len(p)] = p
self.sched = ttnn.from_torch(torch.from_numpy(tab.view(np.int32)), dtype=ttnn.uint32, layout=ttnn.ROW_MAJOR_LAYOUT,
device=device, memory_config=ttnn.DRAM_MEMORY_CONFIG)
self._dm = _src("rc_dm.cpp")
self._cc = _src("rc_compute.cpp")
self.wproc = int(os.environ.get("SP_RC_WPROC", "0"))
def supports(self, x: ttnn.Tensor) -> bool:
if not x.is_sharded() or x.layout != ttnn.TILE_LAYOUT or x.dtype != ttnn.bfloat16:
return False
mc = x.memory_config()
return (mc.memory_layout == ttnn.TensorMemoryLayout.HEIGHT_SHARDED and mc.buffer_type == ttnn.BufferType.L1
and list(mc.shard_spec.shape) == [SHARD, self.ci] and list(x.shape)[-2:] == [H * W, self.ci]
and mc.shard_spec.orientation == ttnn.ShardOrientation.ROW_MAJOR)
def __call__(self, x: ttnn.Tensor) -> ttnn.Tensor:
dev = self.device
mc = x.memory_config()
in_cores = ttnn.corerange_to_cores(mc.shard_spec.grid, row_wise=True)
omc = ttnn.MemoryConfig(ttnn.TensorMemoryLayout.HEIGHT_SHARDED, ttnn.BufferType.L1,
ttnn.ShardSpec(mc.shard_spec.grid, [SHARD, self.co], ttnn.ShardOrientation.ROW_MAJOR))
out = ttnn.allocate_tensor_on_device(ttnn.Shape([1, 1, H * W, self.co]), ttnn.bfloat16, ttnn.TILE_LAYOUT, dev, omc)
grid = ttnn.CoreRangeSet([ttnn.CoreRange(ttnn.CoreCoord(0, 0), ttnn.CoreCoord(11, 9))])
CB_A0, CB_A1, CB_W, CB_B, CB_P, CB_O, CB_T = 0, 1, 2, 3, 4, 5, 6
bf = ttnn.bfloat16
blk_h = 9 * self.ch
def cb(i, pages, page=2048, fmt=bf):
f = ttnn.CBFormatDescriptor(buffer_index=i, data_format=fmt, page_size=page)
return ttnn.CBDescriptor(total_size=page * pages, core_ranges=grid, format_descriptors=[f])
page_bytes = self.page_words * 4
cbs = [cb(CB_A0, 3 * blk_h), cb(CB_A1, 3 * blk_h), cb(CB_W, 36 * self.nl), cb(CB_B, self.nl),
cb(CB_P, 3 * self.nl, 4096, ttnn.float32), cb(CB_O, 3 * self.nl), cb(CB_T, 1, 2 * (page_bytes + 1024))]
acc = (ttnn.TensorAccessorArgs(self.w).get_compile_time_args() + ttnn.TensorAccessorArgs(self.b).get_compile_time_args()
+ ttnn.TensorAccessorArgs(self.sched).get_compile_time_args())
xa, oa, wa, ba, sa = x.buffer_address(), out.buffer_address(), self.w.buffer_address(), self.b.buffer_address(), self.sched.buffer_address()
SC = ttnn.KernelDescriptor.SourceType.SOURCE_CODE
defs = [tuple((d + "=1").split("=")[:2]) for d in os.environ.get("SP_RC_DEFS", "").split(",") if d]
noc = lambda c: dev.worker_core_from_logical_core(c)
ks = []
for proc, cfg, cba in ((0, ttnn.ReaderConfigDescriptor(), CB_A0), (1, ttnn.WriterConfigDescriptor(), CB_A1)):
rt = ttnn.RuntimeArgs()
for k, c in enumerate(self.cores):
g = c.x // GX
snd = noc(ttnn.CoreCoord(GX * g, 0))
m0, m1 = noc(ttnn.CoreCoord(GX * g, 0)), noc(ttnn.CoreCoord(GX * g + GX - 1, 9))
is_snd = int(c.x == GX * g and c.y == 0)
src = []
for s in self.src_shards[(c.x, c.y)]:
n = noc(in_cores[s])
src += [n.x, n.y]
src += [0, 0] * (NSRC - len(src) // 2)
dst = []
for s in self.dst_shards[(c.x, c.y)]:
n = noc(in_cores[s])
dst += [n.x, n.y]
dst += [0, 0] * (NDST - len(dst) // 2)
rt[c.x][c.y] = [xa, oa, wa, ba, sa, k, is_snd, snd.x, snd.y, m0.x, m0.y, m1.x, m1.y, GX * 10 - 1,
g * 36 * self.nl, g * self.nl] + src + dst
ks.append(ttnn.KernelDescriptor(
kernel_source=self._dm, source_type=SC, core_ranges=grid,
compile_time_args=[proc, cba, CB_W, CB_B, CB_O, CB_T, self.nl, self.ch, self.ct, page_bytes, blk_h, self.wproc] + acc,
runtime_args=rt, config=cfg, defines=defs))
ccfg = ttnn.ComputeConfigDescriptor(math_fidelity=ttnn.MathFidelity.HiFi2, fp32_dest_acc_en=True, dst_full_sync_en=False)
ks.append(ttnn.KernelDescriptor(kernel_source=self._cc, source_type=SC, core_ranges=grid,
compile_time_args=[CB_A0, CB_A1, CB_W, CB_B, CB_P, CB_O, self.nl, self.ch, self.ct, blk_h, int(self.relu)],
runtime_args=[], config=ccfg, defines=defs))
sems = [ttnn.SemaphoreDescriptor(id=i, core_ranges=grid, initial_value=0) for i in (0, 1, 2)]
return ttnn.generic_op([x, self.w, self.b, self.sched, out], ttnn.ProgramDescriptor(kernels=ks, semaphores=sems, cbs=cbs))
|