File size: 13,537 Bytes
be62f78 | 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 258 259 260 261 262 263 264 265 266 267 268 269 270 271 272 273 274 275 276 277 278 279 280 281 282 283 284 285 286 | # SPDX-License-Identifier: Apache-2.0
"""Fused fp32 LayerNorm as one ``ttnn.generic_op`` program (``LN_KERNEL``, OPT round 2 item 3).
:func:`layer_norm_fp32_fused` computes exactly what :func:`tt.layers.layer_norm_fp32` computes with 7-9 stock
programs (mean, subtract, square, mean, add eps, rsqrt, multiply, gamma, beta), in one program: the compute kernel
(``kernels/ln32_compute.cpp``) issues the same SFPU LLK calls in the same order as the stock kernels, and every
intermediate the stock graph writes to an fp32 DRAM tensor stays fp32 in L1 / DST, so the output is meant to be
bit-identical (checked on the device: ``code/scripts/ln_kernel_check.py``).
Layout: x is an fp32 TILE DRAM-interleaved tensor ``[..., R, W]`` (R a multiple of 32, W = 32 * Wt); the tile rows
are split in contiguous blocks over ``min(rows, grid)`` cores (one core per tile row up to the grid). gamma / beta
are ``[1, 1, 1, W]`` fp32 TILE rows (row 0 valid) or None.
"""
from __future__ import annotations
import os
import struct
from typing import Any, Optional
__all__ = ["layer_norm_fp32_fused", "supported"]
_KDIR = os.path.join(os.path.dirname(os.path.abspath(__file__)), "kernels")
TB = 4096 # fp32 tile bytes
N_RT = 2 # per-core RT args of every kernel ([row0, n_rows] / [n_rows] + pad): a CT arg (probe P3)
def _bits(v: float) -> int:
return int.from_bytes(struct.pack("<f", float(v)), "little")
def supported(x: Any) -> bool:
"""True when ``x`` is a device fp32 TILE interleaved tensor with tile-aligned rows and columns."""
import ttnn
try:
shp = list(x.padded_shape)
return (x.dtype == ttnn.float32 and x.layout == ttnn.TILE_LAYOUT and not x.is_sharded()
and shp[-1] % 32 == 0 and shp[-2] % 32 == 0 and int(x.shape[-1]) == shp[-1]
and hasattr(ttnn, "generic_op"))
except Exception: # noqa: BLE001 - the host fake ttnn
return False
def _cores(n: int, g):
"""The first ``n`` cores of the grid in row-major order: a CoreRangeSet and the coordinate list."""
import ttnn
cs = [(i % g.x, i // g.x) for i in range(n)]
full, rem = divmod(n, g.x)
rs = []
if full:
rs.append(ttnn.CoreRange(ttnn.CoreCoord(0, 0), ttnn.CoreCoord(g.x - 1, full - 1)))
if rem:
rs.append(ttnn.CoreRange(ttnn.CoreCoord(0, full), ttnn.CoreCoord(rem - 1, full)))
return ttnn.CoreRangeSet(set(rs)), cs
def layer_norm_fp32_fused(x, gamma=None, beta=None, *, eps: float, residual=None, rgate=None, write_h: bool = True,
lean: bool = True, sfpu_bcast: bool = False, res_t: bool = False, out_t: bool = False,
memory_config=None):
"""LayerNorm over the last dim of fp32 ``x`` -> fp32 tensor of x's shape (see the module docstring).
With ``residual`` (``LN_RESID``): ``h = x + residual (* rgate)`` first (the stock ``ttnn.add(x,
ttnn.multiply(residual, rgate))``, fp32 SFPU ops), then ``LN(h)``; returns ``(h, LN(h))`` (``h`` is None when
``write_h`` is False). ``residual``: fp32 like ``x``; ``rgate``: a ``[1, 1, 1, W]`` fp32 row.
``lean``: one unpacker / SFPU-binary init per phase instead of one per tile (same LLK math calls).
``sfpu_bcast`` (``LN_SFPU_BCAST``): the SFPU row reduce writes the row statistics to every column
(``kernels/ln32_sfpu.h``) instead of the writer's RISC-V column fill (same values).
``res_t`` / ``out_t`` (``LN_TR``, x ``[.., E, T, W]``): the residual comes as ``[.., E, W, T]`` (its per-entity
transpose) / the output is written as ``[.., E, W, T]``, the tiles transposed in the kernel with the stock
``ttnn.transpose`` LLK (exact): the two transposes around the mixer's token-mixing MLP. ``memory_config``: of
the outputs (default DRAM; interleaved L1 for ``ENC_L1``)."""
import ttnn
dev = x.device()
shp = list(x.padded_shape)
W = shp[-1]
Wt = W // 32
rows = 1
for d in shp[:-1]:
rows *= d
rows //= 32
hr = int(residual is not None)
hrg = int(hr and rgate is not None)
wh = int(hr and write_h)
Tt = shp[-2] // 32
assert not (res_t and hrg), "LN_TR: no residual gate"
oshape = x.shape if not out_t else ttnn.Shape(list(x.shape)[:-2] + [shp[-1], shp[-2]])
omem = memory_config or ttnn.DRAM_MEMORY_CONFIG
out = ttnn.allocate_tensor_on_device(oshape, ttnn.float32, ttnn.TILE_LAYOUT, dev, omem)
h = (ttnn.allocate_tensor_on_device(x.shape, ttnn.float32, ttnn.TILE_LAYOUT, dev, omem)
if wh else None)
g = dev.compute_with_storage_grid_size()
n = min(rows, g.x * g.y)
crs, cs = _cores(n, g)
base, extra = divmod(rows, n)
rd, wr, cp = ttnn.RuntimeArgs(), ttnn.RuntimeArgs(), ttnn.RuntimeArgs()
r0 = 0
for i, (cx, cy) in enumerate(cs):
k = base + (1 if i < extra else 0)
rd[cx][cy] = [r0, k]
wr[cx][cy] = [r0, k]
cp[cx][cy] = [k, 0]
r0 += k
hg, hb = int(gamma is not None), int(beta is not None)
def acc(t):
return list(ttnn.TensorAccessorArgs(t).get_compile_time_args())
def cb(idx, pages):
return ttnn.CBDescriptor(total_size=pages * TB, core_ranges=crs, format_descriptors=[
ttnn.CBFormatDescriptor(buffer_index=idx, data_format=ttnn.float32, page_size=TB)])
cbs = [cb(0, 2 * Wt), cb(3, 1), cb(4, 2), cb(5, 2), cb(6, Wt), cb(16, 2 * Wt if Wt <= 4 else Wt)]
if hg:
cbs.append(cb(1, Wt))
if hb:
cbs.append(cb(2, Wt))
if hr:
cbs += [cb(7, Wt), cb(9, Wt)]
if hrg:
cbs.append(cb(8, Wt))
if wh:
cbs.append(cb(17, 2))
if out_t:
cbs.append(cb(18, 1))
um = [ttnn.UnpackToDestMode.Default] * 64
for i in (0, 1, 2, 3, 5, 6, 7, 8, 9, 18):
um[i] = ttnn.UnpackToDestMode.UnpackToDestFp32
ccfg = ttnn.ComputeConfigDescriptor(math_fidelity=ttnn.MathFidelity.HiFi4, fp32_dest_acc_en=True,
math_approx_mode=False)
ccfg.unpack_to_dest_mode = um
g_t = gamma if hg else x
b_t = beta if hb else x
r_t = residual if hr else x
rg_t = rgate if hrg else x
h_t = h if wh else out
reader_ct = [Wt, hg, hb, _bits(eps), N_RT, hr, hrg, int(bool(res_t)), Tt] + acc(x) + acc(g_t) + acc(b_t) + acc(
r_t) + acc(rg_t)
writer_ct = [Wt, N_RT, wh, int(bool(sfpu_bcast)), int(bool(out_t)), Tt] + acc(out) + acc(h_t)
compute_ct = [Wt, hg, hb, _bits(1.0 / W), N_RT, hr, hrg, wh, int(bool(lean)), int(bool(sfpu_bcast)),
int(bool(res_t)), int(bool(out_t))]
SRC = ttnn.KernelDescriptor.SourceType.FILE_PATH
def kd(name, ct, rt, common, config):
return ttnn.KernelDescriptor(kernel_source=os.path.join(_KDIR, name), source_type=SRC, core_ranges=crs,
compile_time_args=ct, defines=[], runtime_args=rt, common_runtime_args=common,
config=config, compiler_include_paths=[_KDIR])
ks = [kd("ln32_reader.cpp", reader_ct, rd, [x.buffer_address(), g_t.buffer_address(), b_t.buffer_address(),
r_t.buffer_address(), rg_t.buffer_address()],
ttnn.ReaderConfigDescriptor()),
kd("ln32_writer.cpp", writer_ct, wr, [out.buffer_address(), h_t.buffer_address()],
ttnn.WriterConfigDescriptor()),
kd("ln32_compute.cpp", compute_ct, cp, [], ccfg)]
ins = [x] + ([gamma] if hg else []) + ([beta] if hb else []) + ([residual] if hr else []) + (
[rgate] if hrg else []) + ([h] if wh else [])
ttnn.generic_op(ins + [out], ttnn.ProgramDescriptor(kernels=ks, semaphores=[], cbs=cbs))
if hr:
return h, out
return out
def reference_decomposition(x, gamma: Optional[Any] = None, beta: Optional[Any] = None, *, eps: float):
"""The stock decomposition (``tt.layers.layer_norm_fp32``), for the device check."""
from .layers import layer_norm_fp32
return layer_norm_fp32(x, gamma, beta, eps=eps)
def split_supported(x: Any) -> bool:
"""The split-row form needs every tile row's Wt column tiles on distinct cores (rows * Wt <= grid cores)."""
if not supported(x):
return False
shp = list(x.padded_shape)
rows = 1
for d in shp[:-1]:
rows *= d
rows //= 32
g = x.device().compute_with_storage_grid_size()
wt = shp[-1] // 32
return 2 <= wt <= g.x * g.y and rows * wt <= g.x * g.y
def layer_norm_fp32_split(x, gamma=None, beta=None, *, eps: float, residual=None, rgate=None, write_h: bool = True,
sfpu_bcast: bool = False, kcat_ktp: int = 0, memory_config=None):
""":func:`layer_norm_fp32_fused` with each tile row spread over Wt cores (``LN_SPLIT``, ``kernels/ln32s_*.cpp``):
member j owns column tile j; the root (member 0) folds the gathered tiles in order and broadcasts the mean and
rstd, so the result is the same bit for bit. For few rows (the decoder's 11 tile rows: 88 cores instead of 11).
``kcat_ktp`` (``KCAT_EMIT``): instead of y, write the split operand ``[y_hi | y_hi | y_lo | 1 | 0..]`` of the next
K-concatenated linear (``[..., R, 32 * kcat_ktp]``, the ``kcat_operand`` layout and LLK calls).
``memory_config``: of the outputs (default DRAM; interleaved L1 for ``DEC_L1``)."""
import ttnn
dev = x.device()
shp = list(x.padded_shape)
W = shp[-1]
Wt = W // 32
omem = memory_config or ttnn.DRAM_MEMORY_CONFIG
rows = 1
for d in shp[:-1]:
rows *= d
rows //= 32
hg, hb = int(gamma is not None), int(beta is not None)
hr = int(residual is not None)
hrg = int(hr and rgate is not None)
wh = int(hr and write_h)
oshape = x.shape if not kcat_ktp else ttnn.Shape(list(x.shape)[:-1] + [32 * kcat_ktp])
out = ttnn.allocate_tensor_on_device(oshape, ttnn.float32, ttnn.TILE_LAYOUT, dev, omem)
h = (ttnn.allocate_tensor_on_device(x.shape, ttnn.float32, ttnn.TILE_LAYOUT, dev, omem)
if wh else None)
g = dev.compute_with_storage_grid_size()
G = min(rows, (g.x * g.y) // Wt)
n = G * Wt
crs, cs = _cores(n, g)
phys = [dev.worker_core_from_logical_core(ttnn.CoreCoord(cx, cy)) for cx, cy in cs]
rd, wr, cp = ttnn.RuntimeArgs(), ttnn.RuntimeArgs(), ttnn.RuntimeArgs()
n_rt = 6 + 2 * Wt
for gi in range(G):
n_rows = len(range(gi, rows, G))
root = phys[gi * Wt]
members = []
for m in range(Wt):
members += [phys[gi * Wt + m].x, phys[gi * Wt + m].y]
for j in range(Wt):
cx, cy = cs[gi * Wt + j]
args = [gi, n_rows, G, j, root.x, root.y] + members
rd[cx][cy] = args
wr[cx][cy] = args
cp[cx][cy] = [n_rows, int(j == 0)]
def acc(t):
return list(ttnn.TensorAccessorArgs(t).get_compile_time_args())
def cb(idx, pages):
return ttnn.CBDescriptor(total_size=pages * TB, core_ranges=crs, format_descriptors=[
ttnn.CBFormatDescriptor(buffer_index=idx, data_format=ttnn.float32, page_size=TB)])
cbs = [cb(0, 2), cb(3, 1), cb(4, 2), cb(5, 2), cb(6, 2), cb(10, Wt), cb(11, 2), cb(12, 1), cb(16, 2)]
if hg:
cbs.append(cb(1, 1))
if hb:
cbs.append(cb(2, 1))
if hr:
cbs += [cb(7, 2), cb(9, 2)]
if hrg:
cbs.append(cb(8, 1))
if wh:
cbs.append(cb(17, 2))
if kcat_ktp:
cbs += [cb(18, 1), cb(19, 2), cb(20, 1)] + ([cb(21, 1)] if kcat_ktp > 3 * Wt + 1 else [])
um = [ttnn.UnpackToDestMode.Default] * 64
for i in (0, 1, 2, 3, 5, 6, 7, 8, 9, 10, 18):
um[i] = ttnn.UnpackToDestMode.UnpackToDestFp32
ccfg = ttnn.ComputeConfigDescriptor(math_fidelity=ttnn.MathFidelity.HiFi4, fp32_dest_acc_en=True,
math_approx_mode=False)
ccfg.unpack_to_dest_mode = um
g_t = gamma if hg else x
b_t = beta if hb else x
r_t = residual if hr else x
rg_t = rgate if hrg else x
h_t = h if wh else out
reader_ct = [Wt, hg, hb, _bits(eps), n_rt, hr, hrg] + acc(x) + acc(g_t) + acc(b_t) + acc(r_t) + acc(rg_t)
writer_ct = [Wt, n_rt, wh, int(bool(sfpu_bcast)), int(bool(kcat_ktp)), int(kcat_ktp)] + acc(out) + acc(h_t)
compute_ct = [Wt, hg, hb, _bits(1.0 / W), 2, hr, hrg, wh, int(bool(sfpu_bcast)), int(bool(kcat_ktp))]
SRC = ttnn.KernelDescriptor.SourceType.FILE_PATH
def kd(name, ct, rt, common, config):
return ttnn.KernelDescriptor(kernel_source=os.path.join(_KDIR, name), source_type=SRC, core_ranges=crs,
compile_time_args=ct, defines=[], runtime_args=rt, common_runtime_args=common,
config=config, compiler_include_paths=[_KDIR])
ks = [kd("ln32s_reader.cpp", reader_ct, rd, [x.buffer_address(), g_t.buffer_address(), b_t.buffer_address(),
r_t.buffer_address(), rg_t.buffer_address()],
ttnn.ReaderConfigDescriptor()),
kd("ln32s_writer.cpp", writer_ct, wr, [out.buffer_address(), h_t.buffer_address()],
ttnn.WriterConfigDescriptor()),
kd("ln32s_compute.cpp", compute_ct, cp, [], ccfg)]
sems = [ttnn.SemaphoreDescriptor(id=i, core_ranges=crs, initial_value=0) for i in range(2)]
ins = [x] + ([gamma] if hg else []) + ([beta] if hb else []) + ([residual] if hr else []) + (
[rgate] if hrg else []) + ([h] if wh else [])
ttnn.generic_op(ins + [out], ttnn.ProgramDescriptor(kernels=ks, semaphores=sems, cbs=cbs))
if hr:
return h, out
return out
|