# SPDX-License-Identifier: Apache-2.0 """mb2 part B (7) — ``ttnn.embedding`` gather from a resident [151680, 2048] bf16 table vs host gather + upload. 01-architecture-synthesis §2.1 "Embedding gather": weight-loader-spec §4.4 recommends a host gather of S rows over a resident 621 MB table. This measures both on the p150a: * device: table ROW_MAJOR bf16 [1,1,151680,2048] in DRAM (``embedding_device_operation.cpp``: weights ROW_MAJOR bf16, indices uint32/bf16 ``[B,1,1,L]`` with L % 32 == 0 for a TILE output); ``ttnn.embedding(ids, table, layout=TILE_LAYOUT)`` eager (program cache warm) and inside a 1-op trace, for L in {128, 512}; indices written with ``copy_host_to_device_tensor`` between replays. * host: ``table[ids]`` in torch (bf16) + ``ttnn.from_torch(..., TILE, device)`` (tilize + PCIe upload) and, separately, tilize on host + ``copy_host_to_device_tensor`` into a preallocated device tensor. Correctness of the device gather is checked exactly against the torch gather. """ from __future__ import annotations import argparse import os import sys import time from typing import Any, Dict import torch sys.path.insert(0, os.path.dirname(os.path.abspath(__file__))) import mb2_common as C # noqa: E402 VOCAB = 151680 E = 2048 def main() -> int: ap = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter) ap.add_argument("--out", type=str, default=C.default_out("mb2_embedding")) ap.add_argument("--reps", type=int, default=20) args = ap.parse_args() import ttnn # lazy rep = C.Report("mb2_embedding") rep.env = C.env_info(ttnn) device = C.open_device(ttnn) try: g = torch.Generator().manual_seed(9) t0 = time.perf_counter() table_t = torch.randn(VOCAB, E, generator=g, dtype=torch.float32).to(torch.bfloat16) rep.results["host_table_randn_s"] = round(time.perf_counter() - t0, 2) t1 = time.perf_counter() table = ttnn.from_torch( table_t.reshape(1, 1, VOCAB, E), dtype=ttnn.bfloat16, layout=ttnn.ROW_MAJOR_LAYOUT, device=device, memory_config=ttnn.DRAM_MEMORY_CONFIG, ) ttnn.synchronize_device(device) rep.results["table_upload"] = { "table_upload_s": round(time.perf_counter() - t1, 2), "table_bytes": VOCAB * E * 2, } print("table_upload", rep.results["table_upload"], flush=True) for L in (128, 512): ids_t = torch.randint(0, VOCAB, (1, 1, 1, L), generator=g, dtype=torch.int64) ref = table_t[ids_t.reshape(-1)].float().reshape(1, 1, L, E) def t_device(L=L, ids_t=ids_t, ref=ref): out: Dict[str, Any] = {} ids = ttnn.from_torch( ids_t.to(torch.int32), dtype=ttnn.uint32, layout=ttnn.ROW_MAJOR_LAYOUT, device=device ) # eager, TILE output emb = ttnn.embedding(ids, table, layout=ttnn.TILE_LAYOUT) o = C.to_torch(ttnn, emb) out["tile_out_shape"] = list(emb.shape) out["tile_out_vs_torch_gather"] = C.compare(ref.reshape(o.shape), o) out["eager_tile_us"] = C.time_calls( ttnn, device, lambda: ttnn.embedding(ids, table, layout=ttnn.TILE_LAYOUT), reps=args.reps ) # eager, ROW_MAJOR output (default) emb_rm = ttnn.embedding(ids, table) out["rowmajor_out_shape"] = list(emb_rm.shape) out["eager_rowmajor_us"] = C.time_calls( ttnn, device, lambda: ttnn.embedding(ids, table), reps=args.reps ) # traced (1 op), indices refreshed via copy_host_to_device_tensor tid = ttnn.begin_trace_capture(device, cq_id=0) emb_tr = ttnn.embedding(ids, table, layout=ttnn.TILE_LAYOUT) ttnn.end_trace_capture(device, tid, cq_id=0) out["trace_replay_us"] = C.time_calls( ttnn, device, lambda: ttnn.execute_trace(device, tid, cq_id=0, blocking=True), reps=args.reps ) ids2_t = torch.randint(0, VOCAB, (1, 1, 1, L), generator=g, dtype=torch.int64) ttnn.copy_host_to_device_tensor( ttnn.from_torch(ids2_t.to(torch.int32), dtype=ttnn.uint32, layout=ttnn.ROW_MAJOR_LAYOUT), ids ) ttnn.execute_trace(device, tid, cq_id=0, blocking=True) o2 = C.to_torch(ttnn, emb_tr) ref2 = table_t[ids2_t.reshape(-1)].float().reshape(o2.shape) out["trace_after_new_ids_vs_torch"] = C.compare(ref2, o2) # write indices + replay + read back, end to end def e2e(): ttnn.copy_host_to_device_tensor( ttnn.from_torch(ids2_t.to(torch.int32), dtype=ttnn.uint32, layout=ttnn.ROW_MAJOR_LAYOUT), ids ) ttnn.execute_trace(device, tid, cq_id=0, blocking=True) out["trace_e2e_write_ids_plus_replay_us"] = C.time_calls(ttnn, device, e2e, reps=args.reps) ttnn.release_trace(device, tid) return out rep.run(f"device_gather_L{L}", t_device) def t_host(L=L, ids_t=ids_t): out: Dict[str, Any] = {} flat = ids_t.reshape(-1) # (a) torch gather only t1 = time.perf_counter() for _ in range(args.reps): rows = table_t[flat] out["torch_gather_us"] = (time.perf_counter() - t1) / args.reps * 1e6 # (b) gather + from_torch(TILE, device) (tilize + upload) samples = [] for _ in range(args.reps): t1 = time.perf_counter() rows = table_t[flat].reshape(1, 1, L, E) dev = ttnn.from_torch(rows, dtype=ttnn.bfloat16, layout=ttnn.TILE_LAYOUT, device=device) ttnn.synchronize_device(device) samples.append((time.perf_counter() - t1) * 1e6) ttnn.deallocate(dev) out["gather_plus_from_torch_tile_device_us"] = { "min": min(samples), "median": sorted(samples)[len(samples) // 2], } # (c) gather + host tilize + copy_host_to_device_tensor into a preallocated buffer pre = ttnn.from_torch( torch.zeros(1, 1, L, E, dtype=torch.bfloat16), dtype=ttnn.bfloat16, layout=ttnn.TILE_LAYOUT, device=device, ) samples = [] for _ in range(args.reps): t1 = time.perf_counter() rows = table_t[flat].reshape(1, 1, L, E) host = ttnn.from_torch(rows, dtype=ttnn.bfloat16, layout=ttnn.TILE_LAYOUT) ttnn.copy_host_to_device_tensor(host, pre) ttnn.synchronize_device(device) samples.append((time.perf_counter() - t1) * 1e6) out["gather_plus_tilize_plus_copy_host_to_device_us"] = { "min": min(samples), "median": sorted(samples)[len(samples) // 2], } # (d) upload only (ROW_MAJOR host tensor, no tilize) for reference samples = [] pre_rm = ttnn.from_torch( torch.zeros(1, 1, L, E, dtype=torch.bfloat16), dtype=ttnn.bfloat16, layout=ttnn.ROW_MAJOR_LAYOUT, device=device, ) for _ in range(args.reps): rows = table_t[flat].reshape(1, 1, L, E) host = ttnn.from_torch(rows, dtype=ttnn.bfloat16, layout=ttnn.ROW_MAJOR_LAYOUT) t1 = time.perf_counter() ttnn.copy_host_to_device_tensor(host, pre_rm) ttnn.synchronize_device(device) samples.append((time.perf_counter() - t1) * 1e6) out["copy_host_to_device_rowmajor_only_us"] = { "min": min(samples), "median": sorted(samples)[len(samples) // 2], } out["bytes_moved"] = L * E * 2 return out rep.run(f"host_gather_L{L}", t_host) finally: C.close_device(ttnn, device) rep.write(args.out) return 0 if __name__ == "__main__": sys.exit(main())