| |
| """Operand build of the K-concatenated split matmul (``SPLIT_KCAT=2``, OPT round 2 item 2) as one ``generic_op``. |
| |
| ``kcat_operand(x)``: fp32 TILE ``[..., M, K]`` (K a multiple of 32) -> fp32 ``[..., M, 3K + 32]`` = |
| ``[x_hi | x_hi | x_lo | ones]`` with ``x_hi`` = bf16(x) (the ``ttnn.typecast`` LLK) and ``x_lo = x - x_hi``; the |
| last tile has columns 0 and 1 = 1 (the exact bias rows of the concatenated weight). It replaces the 5 stock programs |
| of ``SPLIT_KCAT=1`` (typecast, typecast, subtract, concat over 4 tensors). x tiles are split in contiguous ranges |
| over the grid (``kernels/kcat_*.cpp``). |
| """ |
| from __future__ import annotations |
|
|
| import os |
| from typing import Any |
|
|
| __all__ = ["kcat_operand", "kcat_tiles", "supported"] |
|
|
| _KDIR = os.path.join(os.path.dirname(os.path.abspath(__file__)), "kernels") |
| TB = 4096 |
| N_RT = 2 |
|
|
|
|
| def supported(x: Any) -> bool: |
| """True for a device fp32 TILE interleaved tensor with tile-aligned rows and columns (not the host fake).""" |
| 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 hasattr(ttnn, "generic_op")) |
| except Exception: |
| return False |
|
|
|
|
| def kcat_tiles(k: int, pad: int) -> int: |
| """Row tiles of X' for an input of K = ``k`` columns: 3 Kt + 1, rounded up to a multiple of ``pad``.""" |
| return -(-(3 * (k // 32) + 1) // pad) * pad |
|
|
|
|
| ACTS = {None: 0, "gelu": 1, "gelu_tanh": 2} |
|
|
|
|
| def kcat_operand(x: Any, pad: int = 1, act: Any = None, memory_config: Any = None, act_once: bool = False): |
| """``pad``: X' row tiles rounded up to a multiple of ``pad`` with zero tiles (the matmul's K block must divide |
| them; 3 Kt + 1 is odd, e.g. 97 for K = 1024). ``act`` (``KCAT_ACT``): ``"gelu"`` / ``"gelu_tanh"`` applied to |
| x first (the previous linear's activation, the same LLK as the stock ``ttnn.gelu`` program). ``memory_config``: |
| of X' (default DRAM interleaved; ``KCAT_L1``: the L1 block-sharded layout the consumer matmul reads in place, the |
| writer's TensorAccessor resolves the shard of each tile). ``act_once`` (``KCAT_ACT_ONCE``): the activation on one |
| DEST copy of the tile, copied to the second by ``copy_dest_values`` (instead of on both copies).""" |
| import ttnn |
|
|
| dev = x.device() |
| shp = list(x.padded_shape) |
| K = shp[-1] |
| assert K % 32 == 0 and shp[-2] % 32 == 0 and x.dtype == ttnn.float32 and x.layout == ttnn.TILE_LAYOUT |
| Kt = K // 32 |
| rows = 1 |
| for d in shp[:-1]: |
| rows *= d |
| rows //= 32 |
| total = rows * Kt |
| Ktp = kcat_tiles(K, pad) |
| oshape = list(x.shape)[:-1] + [32 * Ktp] |
| out = ttnn.allocate_tensor_on_device(ttnn.Shape(oshape), ttnn.float32, ttnn.TILE_LAYOUT, dev, |
| memory_config or ttnn.DRAM_MEMORY_CONFIG) |
| g = dev.compute_with_storage_grid_size() |
| n = min(total, g.x * g.y) |
| 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))) |
| crs = ttnn.CoreRangeSet(set(rs)) |
| base, extra = divmod(total, n) |
| rd, wr, cp = ttnn.RuntimeArgs(), ttnn.RuntimeArgs(), ttnn.RuntimeArgs() |
| t0 = 0 |
| for i, (cx, cy) in enumerate(cs): |
| k = base + (1 if i < extra else 0) |
| rd[cx][cy] = [t0, k] |
| wr[cx][cy] = [t0, k] |
| cp[cx][cy] = [k, 0] |
| t0 += k |
|
|
| 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(16, 2), cb(17, 2), cb(18, 1)] + ([cb(19, 1)] if Ktp > 3 * Kt + 1 else []) |
| um = [ttnn.UnpackToDestMode.Default] * 64 |
| um[0] = 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 |
| SRC = ttnn.KernelDescriptor.SourceType.FILE_PATH |
|
|
| def acc(t): |
| return list(ttnn.TensorAccessorArgs(t).get_compile_time_args()) |
|
|
| 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) |
|
|
| ks = [kd("kcat_reader.cpp", [N_RT] + acc(x), rd, [x.buffer_address()], ttnn.ReaderConfigDescriptor()), |
| kd("kcat_writer.cpp", [Kt, N_RT, Ktp] + acc(out), wr, [out.buffer_address()], ttnn.WriterConfigDescriptor()), |
| kd("kcat_compute.cpp", [N_RT, ACTS[act], int(bool(act_once))], cp, [], ccfg)] |
| ttnn.generic_op([x, out], ttnn.ProgramDescriptor(kernels=ks, semaphores=[], cbs=cbs)) |
| return out |
|
|
