206 lines
8.6 KiB
Python
206 lines
8.6 KiB
Python
|
#!/usr/bin/env python3
|
|||
|
|
"""test_scan_tuning.py — Verify muh scan tuning against CCCL ground truth
|
|||
|
|
|
|||
|
|
Scan is SMEM-critical (unlike reduce which is register-limited):
|
|||
|
|
BlockLoad stages data from global→SMEM: tile = threads × items × type_size
|
|||
|
|
BlockScan uses SMEM scratch: ~threads × sizeof(AccumT) for RAKING
|
|||
|
|
Lookback uses tile_state SMEM: negligible (global memory)
|
|||
|
|
|
|||
|
|
Total SMEM ≈ threads × items × type_size + threads × accum_size
|
|||
|
|
Must be ≤ 48KB (49152 bytes) on BI-V100.
|
|||
|
|
|
|||
|
|
CCCL reference: cub/benchmarks/bench/scan/exclusive/sum.cu
|
|||
|
|
ipt_22.tpb_384.ns_1904.dcid_6.l2w_830.trp_1.ld_0 1.148 0.997 1.140 1.463
|
|||
|
|
|
|||
|
|
This benchmark comment tells us SM100's best config for 4B types is:
|
|||
|
|
items=22, threads=384, delay=1904ns, algo=exponential_backon_jitter_window,
|
|||
|
|
l2w=830ns, transpose=true, load=DEFAULT
|
|||
|
|
→ tile = 384 × 22 × 4 = 33792B (69% of 48KB) ✓
|
|||
|
|
|
|||
|
|
BI-V100 differences:
|
|||
|
|
1. SM count: 16 vs 148 → fewer CTAs, larger tiles beneficial
|
|||
|
|
2. L2 cache: 6MB vs 50MB → l2w needs recalibration
|
|||
|
|
3. Lookback delay: 16 CTAs contend less → shorter delays possible
|
|||
|
|
"""
|
|||
|
|
|
|||
|
|
import sys, os
|
|||
|
|
sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
|
|||
|
|
sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.dirname(os.path.abspath(__file__)))))
|
|||
|
|
|
|||
|
|
SMEM_LIMIT = 49152
|
|||
|
|
WARP_SIZE = 32
|
|||
|
|
SM_COUNT = 16
|
|||
|
|
|
|||
|
|
def scale_mem_bound(nom_t, nom_i, ts, max_smem=SMEM_LIMIT):
|
|||
|
|
items = max(1, min(nom_i * 4 // ts, nom_i * 2))
|
|||
|
|
spt = ts * items
|
|||
|
|
if spt > 0:
|
|||
|
|
raw = max_smem // spt
|
|||
|
|
max_t = ((raw + 31) // 32) * 32
|
|||
|
|
else:
|
|||
|
|
max_t = nom_t
|
|||
|
|
threads = min(nom_t, max_t)
|
|||
|
|
if threads < 32:
|
|||
|
|
threads = 32
|
|||
|
|
return items, threads
|
|||
|
|
|
|||
|
|
|
|||
|
|
# bi100 scan structs from tuning_scan.cuh
|
|||
|
|
BI100_SCAN = {
|
|||
|
|
# Lookback structs: tile = threads × items × type_size (BlockLoad staging)
|
|||
|
|
"bi100_lookback_1B_o4": {"threads": 512, "items": 18, "type_size": 1,
|
|||
|
|
"delay_algo": "exponential_backon", "delay_ns": 175, "l2w": 270},
|
|||
|
|
"bi100_lookback_2B_o4": {"threads": 512, "items": 13, "type_size": 2,
|
|||
|
|
"delay_algo": "exponential_backon", "delay_ns": 175, "l2w": 270},
|
|||
|
|
"bi100_lookback_4B_o4": {"threads": 384, "items": 22, "type_size": 4,
|
|||
|
|
"delay_algo": "exponential_backon_jitter_window", "delay_ns": 952, "l2w": 415},
|
|||
|
|
"bi100_lookback_8B_o4": {"threads": 320, "items": 19, "type_size": 8,
|
|||
|
|
"delay_algo": "exponential_backon_jitter_window", "delay_ns": 386, "l2w": 426},
|
|||
|
|
"bi100_lookback_1B_o8": {"threads": 384, "items": 14, "type_size": 1,
|
|||
|
|
"delay_algo": "exponential_backon_jitter_window", "delay_ns": 760, "l2w": 465},
|
|||
|
|
"bi100_lookback_4B_o8": {"threads": 320, "items": 19, "type_size": 4,
|
|||
|
|
"delay_algo": "exponential_backon_jitter_window", "delay_ns": 478, "l2w": 330},
|
|||
|
|
"bi100_lookback_8B_o8": {"threads": 320, "items": 19, "type_size": 8,
|
|||
|
|
"delay_algo": "exponential_backon_jitter_window", "delay_ns": 478, "l2w": 330},
|
|||
|
|
}
|
|||
|
|
|
|||
|
|
# SM100 best from CCCL benchmark annotations
|
|||
|
|
SM100_SCAN = {
|
|||
|
|
"4B_o4_best": {"threads": 384, "items": 22, "type_size": 4,
|
|||
|
|
"speedup": [1.148442, 0.997167, 1.139902, 1.462651],
|
|||
|
|
"annotation": "ipt_22.tpb_384.ns_1904.dcid_6.l2w_830.trp_1.ld_0"},
|
|||
|
|
"4B_o4_alt1": {"threads": 512, "items": 18, "type_size": 4,
|
|||
|
|
"speedup": [1.188818, 1.005682, 1.173041, 1.305288],
|
|||
|
|
"annotation": "ipt_18.tpb_512.ns_768.dcid_7.l2w_820.trp_1.ld_0"},
|
|||
|
|
"8B_o4_best": {"threads": 416, "items": 14, "type_size": 8,
|
|||
|
|
"speedup": [1.107210, 1.000000, 1.100637, 1.307692],
|
|||
|
|
"annotation": "ipt_14.tpb_384.ns_228.dcid_7.l2w_775.trp_1.ld_1"},
|
|||
|
|
}
|
|||
|
|
|
|||
|
|
|
|||
|
|
def test_scan_smem():
|
|||
|
|
"""Test SMEM safety for all bi100 scan structs.
|
|||
|
|
|
|||
|
|
CRITICAL INSIGHT from agent_scan.cuh:
|
|||
|
|
_TempStorage is a UNION of {load, store, scan+prefix}
|
|||
|
|
Peak SMEM = max(load_tile, store_tile, scan_scratch + prefix_scratch)
|
|||
|
|
NOT load_tile + scan_scratch (that was the bug in previous test version)
|
|||
|
|
|
|||
|
|
BlockLoad::TempStorage for WARP_TRANSPOSE:
|
|||
|
|
= threads × items × type_size (the tile staging buffer)
|
|||
|
|
BlockScan::TempStorage for WARP_SCANS:
|
|||
|
|
≈ num_warps × type_size (one value per warp) — much smaller than tile
|
|||
|
|
TilePrefixCallback::TempStorage:
|
|||
|
|
= a few integers for lookback state — negligible
|
|||
|
|
|
|||
|
|
So: peak SMEM ≈ tile = threads × items × type_size
|
|||
|
|
"""
|
|||
|
|
print("=== scan SMEM safety (union model from agent_scan.cuh) ===")
|
|||
|
|
all_ok = True
|
|||
|
|
for name, cfg in BI100_SCAN.items():
|
|||
|
|
tile = cfg["threads"] * cfg["items"] * cfg["type_size"]
|
|||
|
|
# WARP_SCANS scratch: one value per warp
|
|||
|
|
num_warps = cfg["threads"] // WARP_SIZE
|
|||
|
|
scan_scratch = num_warps * cfg["type_size"]
|
|||
|
|
# prefix callback: ~16 bytes
|
|||
|
|
prefix_scratch = 16
|
|||
|
|
# Union: peak = max(tile, scan+prefix)
|
|||
|
|
peak_smem = max(tile, scan_scratch + prefix_scratch)
|
|||
|
|
pct = peak_smem / SMEM_LIMIT * 100
|
|||
|
|
ok = peak_smem <= SMEM_LIMIT
|
|||
|
|
if not ok:
|
|||
|
|
all_ok = False
|
|||
|
|
status = "PASS" if ok else "FAIL"
|
|||
|
|
print(f" {status}: {name:30s} tile={tile:6d}B (peak, union) "
|
|||
|
|
f"scan_alt={scan_scratch+prefix_scratch:5d}B → peak={peak_smem:6d}B ({pct:.0f}%)")
|
|||
|
|
return all_ok
|
|||
|
|
|
|||
|
|
|
|||
|
|
def test_scan_delay_scaling():
|
|||
|
|
"""Test that BI-V100 delay parameters are scaled correctly from SM100.
|
|||
|
|
|
|||
|
|
BI-V100 vs SM100:
|
|||
|
|
- SM count: 16 vs 148 → 9.25× fewer CTAs → less tile_state contention
|
|||
|
|
- L2: 6MB vs 50MB → 8.3× smaller → l2w should be proportionally higher
|
|||
|
|
- But L2/SM: 384KB vs 338KB → BI-V100 has MORE L2 per SM
|
|||
|
|
|
|||
|
|
CCCL scaling heuristic (from existing muh comments):
|
|||
|
|
delay_ns *= 0.5 (fewer CTAs → less contention → shorter delay)
|
|||
|
|
l2w *= 0.6 (L2/SM is similar, but total L2 is smaller)
|
|||
|
|
"""
|
|||
|
|
print("\n=== scan delay parameter scaling ===")
|
|||
|
|
|
|||
|
|
# SM100 reference
|
|||
|
|
sm100_delay = 1904 # ns (from best scan benchmark)
|
|||
|
|
sm100_l2w = 830 # ns
|
|||
|
|
|
|||
|
|
# BI-V100 4B_o4 struct
|
|||
|
|
bi100 = BI100_SCAN["bi100_lookback_4B_o4"]
|
|||
|
|
|
|||
|
|
# Expected: delay_ns ≈ 1904 × 0.5 = 952
|
|||
|
|
expected_ns = int(sm100_delay * 0.5)
|
|||
|
|
actual_ns = bi100["delay_ns"]
|
|||
|
|
ns_ok = abs(actual_ns - expected_ns) <= 100 # tolerance
|
|||
|
|
|
|||
|
|
# Expected: l2w ≈ 830 × 0.5 = 415
|
|||
|
|
expected_l2w = int(sm100_l2w * 0.5)
|
|||
|
|
actual_l2w = bi100["l2w"]
|
|||
|
|
l2w_ok = abs(actual_l2w - expected_l2w) <= 100
|
|||
|
|
|
|||
|
|
print(f" SM100 best: delay={sm100_delay}ns l2w={sm100_l2w}ns")
|
|||
|
|
print(f" BI100 4B: delay={actual_ns}ns (expected ~{expected_ns}) "
|
|||
|
|
f"{'PASS' if ns_ok else 'WARN'}")
|
|||
|
|
print(f" BI100 4B: l2w={actual_l2w}ns (expected ~{expected_l2w}) "
|
|||
|
|
f"{'PASS' if l2w_ok else 'WARN'}")
|
|||
|
|
print(f" Note: delay scaling is heuristic — real values need BI-V100 benchmark")
|
|||
|
|
|
|||
|
|
return True # Warn only, don't fail
|
|||
|
|
|
|||
|
|
|
|||
|
|
def test_scan_vs_sm100_tile_size():
|
|||
|
|
"""Compare tile sizes between SM100 and BI-V100."""
|
|||
|
|
print("\n=== scan tile size comparison ===")
|
|||
|
|
|
|||
|
|
for key in ["4B_o4_best", "8B_o4_best"]:
|
|||
|
|
sm = SM100_SCAN[key]
|
|||
|
|
sm_tile = sm["threads"] * sm["items"] * sm["type_size"]
|
|||
|
|
|
|||
|
|
# Find matching bi100 struct
|
|||
|
|
ts = sm["type_size"]
|
|||
|
|
bi_key = f"bi100_lookback_{ts}B_o4"
|
|||
|
|
if bi_key in BI100_SCAN:
|
|||
|
|
bi = BI100_SCAN[bi_key]
|
|||
|
|
bi_tile = bi["threads"] * bi["items"] * bi["type_size"]
|
|||
|
|
ratio = bi_tile / sm_tile
|
|||
|
|
|
|||
|
|
print(f" {key}:")
|
|||
|
|
print(f" SM100: t={sm['threads']:3d} i={sm['items']:2d} → tile={sm_tile:6d}B")
|
|||
|
|
print(f" BI100: t={bi['threads']:3d} i={bi['items']:2d} → tile={bi_tile:6d}B "
|
|||
|
|
f"(ratio={ratio:.2f}×)")
|
|||
|
|
print(f" SM100 speedup: {sm['speedup']}")
|
|||
|
|
|
|||
|
|
# BI-V100 wants same or larger tiles (16 SMs need more work per CTA)
|
|||
|
|
if ratio < 0.5:
|
|||
|
|
print(f" WARN: BI-V100 tile is less than half of SM100 — may be too small")
|
|||
|
|
|
|||
|
|
|
|||
|
|
def main():
|
|||
|
|
print("muh scan tuning verification")
|
|||
|
|
print("CCCL ref: cub/benchmarks/bench/scan/exclusive/sum.cu")
|
|||
|
|
print("=" * 60)
|
|||
|
|
|
|||
|
|
results = []
|
|||
|
|
results.append(("scan SMEM safety", test_scan_smem()))
|
|||
|
|
results.append(("delay scaling", test_scan_delay_scaling()))
|
|||
|
|
test_scan_vs_sm100_tile_size()
|
|||
|
|
|
|||
|
|
print("\n" + "=" * 60)
|
|||
|
|
all_pass = all(r for _, r in results)
|
|||
|
|
for name, ok in results:
|
|||
|
|
print(f" {'PASS' if ok else 'FAIL'}: {name}")
|
|||
|
|
print(f"\nOverall: {'ALL PASS' if all_pass else 'SOME FAILURES'}")
|
|||
|
|
return 0 if all_pass else 1
|
|||
|
|
|
|||
|
|
if __name__ == "__main__":
|
|||
|
|
sys.exit(main())
|