File size: 7,076 Bytes
1d30035 | 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 | # SPDX-License-Identifier: Apache-2.0
"""MB0: prove the ``ttnn.generic_op`` + ``UnifiedKernelDescriptor`` plumbing works on the full 11x10 grid.
Device probe for docs/research/device-stack-smoke.md (target 5). It
1. opens the device and records the compute grid, DRAM grid, the runtime
``get_optimal_dram_bank_to_logical_worker_assignment`` list (the harvest-pattern question of
unified-kernels-11x10-audit.md section 1.3) and any harvesting info exposed by the bindings;
2. builds a HEIGHT_SHARDED uint32 output tensor with one [32, 32] shard per compute core;
3. compiles ``kernels/hello_kernel.cpp`` through ``UnifiedKernelDescriptor`` (the same descriptor path the
deepseek_v3_b1 micro-ops use) for every core of the grid and runs it via ``ttnn.generic_op``;
4. checks that each core wrote (logical_x, logical_y, row-major linear id, magic) into its own shard;
5. reports cold (JIT compile) and warm (program cache) wall times.
Run through the device lock::
bin/with-device.sh python models/experimental/gr00t/benchmarks/mb0_generic_op_hello.py
"""
from __future__ import annotations
import os
import sys
import time
import torch
import ttnn
from models.demos.deepseek_v3_b1.unified_kernel_descriptor import (
UnifiedCompileTimeCoreDescriptor,
UnifiedKernelDescriptor,
)
KERNEL = "models/experimental/gr00t/benchmarks/kernels/hello_kernel.cpp"
MAGIC = 0x00C0FFEE
SHARD = (32, 32) # 1024 uint32 words per core; we only use the first 4
def _probe_device_facts(device) -> None:
grid = device.compute_with_storage_grid_size()
dram = device.dram_grid_size()
print(
f"FACT arch={ttnn.get_arch_name()} compute_grid={grid.x}x{grid.y} ({grid.x * grid.y} cores) dram_grid={dram.x}x{dram.y}"
)
workers = device.get_optimal_dram_bank_to_logical_worker_assignment(ttnn.NOC.NOC_0)
print(f"FACT optimal_dram_bank_workers(NOC_0)={[(c.x, c.y) for c in workers]}")
workers1 = device.get_optimal_dram_bank_to_logical_worker_assignment(ttnn.NOC.NOC_1)
print(f"FACT optimal_dram_bank_workers(NOC_1)={[(c.x, c.y) for c in workers1]}")
pinned = [(0, 9), (0, 0), (0, 7), (0, 3), (7, 9), (7, 1), (7, 6), (7, 4)]
print(f"FACT deepseek_pinned_list_matches_runtime={[(c.x, c.y) for c in workers] == pinned}")
for name in dir(device):
if "harvest" in name.lower():
try:
print(f"FACT device.{name}={getattr(device, name)()}")
except Exception as exc: # noqa: BLE001 - probe only
print(f"FACT device.{name} -> {exc!r}")
for x in range(grid.x):
for y in (0, grid.y - 1):
phys = device.worker_core_from_logical_core(ttnn.CoreCoord(x, y))
print(f"FACT logical({x},{y}) -> noc0/virtual ({phys.x},{phys.y})")
def run_hello(device, grid_end: ttnn.CoreCoord, tag: str) -> float:
"""Compile+run the hello kernel on the rectangle (0,0)-(grid_end); return the wall time in seconds."""
grid = ttnn.CoreRangeSet([ttnn.CoreRange(ttnn.CoreCoord(0, 0), grid_end)])
num_cores = (grid_end.x + 1) * (grid_end.y + 1)
torch_out = torch.zeros((num_cores * SHARD[0], SHARD[1]), dtype=torch.int32)
shard_spec = ttnn.ShardSpec(grid, list(SHARD), ttnn.ShardOrientation.ROW_MAJOR)
mem = ttnn.MemoryConfig(ttnn.TensorMemoryLayout.HEIGHT_SHARDED, ttnn.BufferType.L1, shard_spec)
out_t = ttnn.from_torch(torch_out, dtype=ttnn.uint32, layout=ttnn.TILE_LAYOUT, device=device, memory_config=mem)
# generic_op insists on >= 1 input and >= 1 output tensor (generic_op_device_operation.cpp:135); the kernel
# never reads this dummy input, it only exists to satisfy that precondition.
in_t = ttnn.from_torch(torch.zeros(32, 32, dtype=torch.int32), dtype=ttnn.uint32, layout=ttnn.TILE_LAYOUT, device=device)
out_cb = 0
cb = ttnn.cb_descriptor_from_sharded_tensor(out_cb, out_t)
num_pages = 1
cb.format_descriptors[0].page_size = SHARD[0] * SHARD[1] * 4
ncrisc_args = [
("hello_out_cb", out_cb),
("hello_num_pages", num_pages),
("hello_grid_start_x", 0),
("hello_grid_start_y", 0),
("hello_grid_end_x", grid_end.x),
("hello_grid_end_y", grid_end.y),
]
ukd = UnifiedKernelDescriptor(
kernel_source=KERNEL,
core_ranges=grid,
ncrisc_named_compile_time_args=ncrisc_args,
brisc_named_compile_time_args=[],
trisc_named_compile_time_args=[],
ncrisc_common_runtime_args=[MAGIC],
trisc_compute_config=ttnn.ComputeConfigDescriptor(
math_fidelity=ttnn.MathFidelity.LoFi, math_approx_mode=False, fp32_dest_acc_en=False, dst_full_sync_en=False
),
unified_compile_time_core_descriptors=[
UnifiedCompileTimeCoreDescriptor(
named_compile_time_arg="is_active_core", core_range=grid, value=1, other_value=0
)
],
)
program = ttnn.ProgramDescriptor(kernels=ukd.get_kernel_descriptors().kernels, cbs=[cb])
t0 = time.perf_counter()
ttnn.generic_op([in_t, out_t], program)
ttnn.synchronize_device(device)
dt = time.perf_counter() - t0
got = ttnn.to_torch(out_t).view(num_cores, SHARD[0] * SHARD[1])
bad = []
for x in range(grid_end.x + 1):
for y in range(grid_end.y + 1):
lin = y * (grid_end.x + 1) + x # row-major id, matches linear_id_in_grid<true>
# HEIGHT_SHARDED shards are laid out in ROW_MAJOR core order -> shard index == lin
words = got[lin, :4].tolist()
if words != [x, y, lin, MAGIC]:
bad.append(((x, y), words))
status = "OK" if not bad else f"MISMATCH on {len(bad)} cores, first {bad[:4]}"
print(
f"RESULT hello[{tag}] cores={num_cores} grid=(0,0)-({grid_end.x},{grid_end.y}) wall={dt * 1e3:.2f} ms {status}"
)
if bad:
raise SystemExit(f"hello kernel verification failed: {status}")
ttnn.deallocate(out_t)
ttnn.deallocate(in_t)
return dt
def main() -> int:
print(
f"ENV TT_METAL_CACHE={os.environ.get('TT_METAL_CACHE', '<default>')} slow_dispatch={os.environ.get('TT_METAL_SLOW_DISPATCH_MODE')}"
)
t0 = time.perf_counter()
device = ttnn.open_device(device_id=0)
print(f"RESULT open_device {time.perf_counter() - t0:.2f} s")
try:
_probe_device_facts(device)
grid = device.compute_with_storage_grid_size()
full = ttnn.CoreCoord(grid.x - 1, grid.y - 1)
cold = run_hello(device, full, "cold")
warm = run_hello(device, full, "warm")
warm2 = run_hello(device, full, "warm2")
# A second grid shape forces a second JIT (different CT args) - measures compile alone (firmware cached)
cold_small = run_hello(device, ttnn.CoreCoord(grid.x - 1, 0), "cold_1row")
print(
f"RESULT compile_estimate full_grid={cold - warm:.3f} s (cold {cold:.3f} s, warm {warm:.4f}/{warm2:.4f} s); "
f"second-variant cold={cold_small:.3f} s"
)
finally:
ttnn.close_device(device)
return 0
if __name__ == "__main__":
sys.exit(main())
|