File size: 7,552 Bytes
d160c10 | 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 | # SPDX-License-Identifier: Apache-2.0
"""Probe (OPT round 6, R6.1): grid / shard orientation of the large block-sharded 3x3 convs at their model shapes.
ttnn block-shards e.g. the eight 512 -> 512 3x3 convs of the CustomFPN / DepthNet ([6, 16, 44, 512]) ROW_MAJOR on
8 x 10 cores (16 channel tiles do not split over 12 columns; 132 tile rows -> 14 per grid row). Candidates
``T<x>x<y>``: ``transpose_shards`` (COL_MAJOR: rows over the x columns, channels over the y grid rows) with
``override_sharding_config`` on an x-by-y grid; ``Tauto``: ttnn's own grid with ``transpose_shards``. Each candidate runs
(a) from a DRAM input with a DRAM output (``slicing="l1"``: I2S / S2I inside; the model's form for most convs) and
(b) from an input already in the candidate's own L1 shards with the output left there (the DN_L1 chain case). Per
candidate: output bits against the base (same input), max |diff|, device time per call (a trace of ``reps`` calls,
best of 3 blocking replays).
bin/devrun -t 900 -- python -u code/scripts/probe_conv512.py --json logs/bevdet/opt_r6/probe_c512.json \
--conv img_depth_net... --shape 6,16,44,512 --act relu --cands base,T12x8
"""
from __future__ import annotations
import argparse
import json
import sys
import time
from pathlib import Path
import numpy as np
def timed(device, fn, reps):
"""Device time per call: a trace of ``reps`` calls, best of 3 blocking replays. The trace capture is always
ended (an error inside the capture would otherwise leave the device capturing)."""
import ttnn
fn().deallocate()
ttnn.synchronize_device(device)
outs = []
tid = ttnn.begin_trace_capture(device, cq_id=0)
try:
for _ in range(reps):
outs.append(fn())
finally:
ttnn.end_trace_capture(device, tid, cq_id=0)
best = 1e9
try:
for _ in range(3):
ttnn.synchronize_device(device)
t0 = time.perf_counter()
ttnn.execute_trace(device, tid, cq_id=0, blocking=True)
best = min(best, (time.perf_counter() - t0) * 1e3 / reps)
finally:
ttnn.release_trace(device, tid)
for o in outs:
o.deallocate()
return best
def main():
ap = argparse.ArgumentParser()
ap.add_argument("--json")
ap.add_argument("--reps", type=int, default=8)
ap.add_argument("--conv", default="img_bev_encoder_neck.conv.3",
help="module name for params.provider (DN-relative names start with 'dn:')")
ap.add_argument("--shape", default="6,16,44,512", help="n,h,w,cin")
ap.add_argument("--act", default="relu")
ap.add_argument("--cands", default="base,T12x8,T11x8,Tauto")
ap.add_argument("--abh", type=int, default=0, help="act_block_h_override for every candidate")
ap.add_argument("--nodb", action="store_true", help="no double buffering")
a = ap.parse_args()
import torch
import ttnn
from tt_bevdet.api import KNOBS
from tt_bevdet.device import close_device, open_device
from tt_bevdet.reference.config import default_config
from tt_bevdet.reference.weights import BEVDetWeights
from tt_bevdet.tests import paths
from tt_bevdet.tt.conv_tune import DB
from tt_bevdet.tt.params import DN, BEVDetParams
from tt_bevdet.ttaw.ops.conv import Conv2d, FeatureMap
cfg = default_config()
params = BEVDetParams(BEVDetWeights(paths.weights_onnx()), cfg)
prec = str(KNOBS.read().PRECISION)
name = DN + a.conv[3:] if a.conv.startswith("dn:") else a.conv
p = params.provider(name)
n, h, w, c = (int(v) for v in a.shape.split(","))
BS = ttnn.TensorMemoryLayout.BLOCK_SHARDED
base = {} if a.nodb else dict(DB)
if a.abh:
base["act_block_h_override"] = a.abh
act = None if a.act in ("", "none") else a.act
def grid(x, y):
return ttnn.CoreRangeSet({ttnn.CoreRange(ttnn.CoreCoord(0, 0), ttnn.CoreCoord(x - 1, y - 1))})
def cand(label):
if label in ("base", "auto"):
return dict(base)
if label in ("autoT", "h4T", "w2T"):
return {**base, "shard_layout": BS, "transpose_shards": True}
if label == "autoBS":
return {**base, "shard_layout": BS}
if label == "h4":
return dict(base)
if label == "Tauto":
return {**base, "shard_layout": BS, "transpose_shards": True}
x, y = (int(v) for v in label[1:].split("x"))
return {**base, "shard_layout": BS, "override_sharding_config": True, "core_grid": grid(x, y),
"transpose_shards": label[0] == "T"}
SLICING = {"auto": "auto", "autoT": "auto", "autoBS": "auto", "h4": ("height", 4), "h4T": ("height", 4),
"w2T": ("width", 2)}
dev = open_device(trace_region_size=512 << 20)
print(f"grid {dev.compute_with_storage_grid_size()}, conv {name} {tuple(p.weight.shape)} in {a.shape} prec {prec}",
flush=True)
rng = np.random.default_rng(0)
rows = n * h * w
host = torch.from_numpy(rng.standard_normal((1, 1, rows, c)).astype(np.float32)).relu()
xin = ttnn.from_torch(host, dtype=ttnn.bfloat16, layout=ttnn.TILE_LAYOUT, device=dev,
memory_config=ttnn.DRAM_MEMORY_CONFIG)
fm = FeatureMap(xin, n, h, w, c)
ref = None
res = []
try:
for label in [s for s in a.cands.split(",") if s]:
ov = cand(label)
cc = Conv2d(p.weight, p.bias, stride=p.stride, padding=p.padding, dilation=p.dilation, groups=p.groups,
activation=act, precision=prec, slicing=SLICING.get(label, "l1"), conv_config=ov,
name="probe." + label)
row = {"label": label}
try:
cc.output_memory_config = None if label not in SLICING else "dram" # output left in its L1 shards
o = cc(fm)
mc_out = o.tensor.memory_config()
got = ttnn.to_torch(ttnn.from_device(o.tensor)).float()
o.deallocate()
if ref is None:
ref = got
d = (got - ref).abs()
row.update(bad=int((d != 0).sum()), maxdiff=float(d.max()),
out_shard=(str(mc_out.shard_spec.shape) + " " + str(mc_out.shard_spec.grid.bounding_box()))
if mc_out.is_sharded() else "dram")
cc.output_memory_config = "dram"
row["ms_dram_io"] = round(timed(dev, lambda: cc(fm), a.reps), 4)
cc.output_memory_config = None
if mc_out.memory_layout == BS and label not in SLICING:
xs = ttnn.to_memory_config(xin, mc_out) # input in this candidate's own sharding (chain)
fs = FeatureMap(xs, n, h, w, c)
o2 = cc(fs)
row["chain_same_bits"] = bool((ttnn.to_torch(ttnn.from_device(o2.tensor)).float() == got).all())
o2.deallocate()
row["ms_l1_chain"] = round(timed(dev, lambda: cc(fs), a.reps), 4)
ttnn.deallocate(xs)
except Exception as e: # noqa: BLE001
row["error"] = str(e).splitlines()[0][:200]
finally:
cc.release()
dev.clear_program_cache()
print(json.dumps(row), flush=True)
res.append(row)
if a.json:
Path(a.json).write_text(json.dumps(res, indent=1))
finally:
close_device(dev)
return 0
if __name__ == "__main__":
sys.exit(main())
|