| | """Portable KernelDMA 1 MiB prefetch-notch profiler (M1 / M3 / M4 / M5). | | | | | | Article: https://eiln.github.io/posts/ane-dma.html | | | | | | Single-cluster (16 ANE cores) — default: | | | | | | python3.9 profile_kernel_dma_mib.py --host-only --label m4p | | | python3.9 profile_kernel_dma_mib.py --out ./kernel_dma_mib_m4p --label m4p | | | | | | Dual-cluster Ultra / 32-core lattice (M3 Ultra, M1/M2 Ultra): | | | |
| | python3.9 profile_kernel_dma_mib.py --host-only --preset dual --label m3u |
| | python3.9 profile_kernel_dma_mib.py --out ./kernel_dma_mib_m3u --preset dual --label m3u |
| | |
| | Wall-clock Core ML predict() is host-dominated and invalid. This script |
| | times _ANEClient evaluateWithModel (eval_us). Do not dump layers 0-4. |
| | """ |
| | from future import annotations |
| | |
| | import argparse |
| | import copy |
| | import json |
| | import os |
| | import platform |
| | import statistics |
| | import subprocess |
| | import sys |
| | from pathlib import Path |
| | |
| | try: |
| | import numpy as np |
| | except ImportError: |
| | np = None |
| | |
| | try: |
| | import coremltools as ct |
| | from coremltools.converters.mil import Builder as mb |
| | from coremltools.converters.mil.mil import types |
| | from coremltools.models.utils import compile_model |
| | except ImportError: |
| | ct = None |
| | mb = None |
| | types = None |
| | compile_model = None |
| | |
| | try: |
| | from dspark_runtime import acquire_native_lock |
| | except ImportError: |
| | acquire_native_lock = None |
| | |
| | HERE = Path(file).resolve().parent |
| | HOST_M = HERE / 'kernel_dma_mib_host.m' |
| | MIB = 1 << 20 |
| | LINE = 64 |
| | CORES = 16 # single ANE cluster |
| | DEFAULT_COUT = 4096 |
| | DEFAULT_CINS = (2016, 2048) |
| | ARTICLE = 'https://eiln.github.io/posts/ane-dma.html' |
| | |
| | SERIALIZED_SPLIT_RELATED = ( |
| | '--enable-global-channel-splitting', |
| | '--enable-forced-maximal-bonded-split=true', |
| | '--fspatial-split-in-x', |
| | '--fkernel-rewind=enabled', |
| | '--split-kernel-section=false', |
| | '--disable-cache-prefetch-mask=0', |
| | '--global-refinement-in-spatial-split=true', |
| | '--enable-segment-aware-kernel-section-split=true', |
| | '--enable-l2-batch-splitting=true', |
| | '--enable-kernel-split-for-multi-palette-lut=true', |
| | ) | | | DROPPED_DMA_SPLIT_CANDIDATES = ( | | | 'SplitKernelDMA', | | | 'EnableKernelDMASplit', | | | 'MaxKernelDMASize', | | | 'AvoidOneMiBKernelDMA', | | | 'SplitCoeffBuffer', | | | ) | | | NOT_ONE_MIB_DMA_SPLIT = ( | | | 'SplitKernelSection', | | | 'SpatialSplit', | | | 'EnableSpatialSplitInX', | | | 'EnableKernelSplitForMultiPaletteLUT', | | | 'EnableSegmentAwareKernelSectionSplit', | | | 'EnableGlobalChannelSplitting', | | | 'EnableL2BatchSplitting', | | | 'EnableForcedMaximalBondedSplit', | | | 'DisableCachePrefetchMask', | | | ) | | | | | | # Dual-cluster Ultra: 2x16 cores. Per-core bytes use cores=32 if Cout is | | | # striped across both clusters. Stock 16-core 1 MiB pair is then 0.5 MiB |
| | # and is the wrong probe. Measured 2026-09-11 on M3 Ultra (Mac15,14). |
| | DUAL_LATTICE = ( |
| | dict(name='16c_stock_1mib', cores=16, cout=4096, cins=(2016, 2048), |
| | note='single-cluster article pair; Ultra off-lattice (0.5 MiB/core)'), |
| | dict(name='32c_1mib_cin2048', cores=32, cout=8192, cins=(2016, 2048), |
| | note='32-core 1 MiB; Ultra milder_1mib_slowdown ~1.61x'), |
| | dict(name='32c_1mib_cin4096', cores=32, cout=4096, cins=(4032, 4096), |
| | note='same 1 MiB/32-core via Cin=4096; Ultra miss'), |
| | dict(name='32c_2mib_cin8192', cores=32, cout=4096, cins=(8128, 8192), |
| | note='32-core 2 MiB; Ultra milder_1mib_slowdown ~1.77x'), |
| | dict(name='32c_2mib_cin4096', cores=32, cout=8192, cins=(4032, 4096), |
| | note='same 2 MiB/32-core via Cin=4096; Ultra miss'), |
| | ) | | | | | | |
| | def fp16_bytes_per_core(cin, cout, cores=CORES): |
| | if cin <= 0 or cout <= 0 or cores <= 0 or cout % cores: |
| | raise ValueError('positive Cin and Cout divisible by core count required') | | | return (cout // cores) * cin * 2 | | | | | | | | | def total_fp16_bytes(cin, cout): | | | return cin * cout * 2 | | | | | | |
| | def on_prefetch_notch(nbytes): |
| | if nbytes <= 0: |
| | raise ValueError('empty transfer') |
| | k = max(1, int(round(nbytes / MIB))) |
| | return abs(nbytes - k * MIB) == 0 |
| | | | | | | | def classify_ratio(ratio, median_us_on=None): | | | """Map on/off eval_us ratio to a notch class. Wall-clock ms is invalid.""" | | | if median_us_on is not None and median_us_on >= 5000: | | | return 'host_dominated_invalid' | | | if ratio >= 2.0: | | | return 'm3_class_notch' | | | if ratio >= 1.3: | | | return 'milder_1mib_slowdown' | | | return 'no_notch' | | | | | | |
| | def sysctl(name): |
| | run = subprocess.run(['sysctl', '-n', name], capture_output=True, text=True, timeout=5) |
| | return run.stdout.strip() if run.returncode == 0 else '' |
| | | | | |
| | def guess_chip(hw_model, brand): |
| | brand_l = (brand or '').lower() |
| | model = hw_model or '' |
| | text = f'{brand_l} {model.lower()}' |
| | for name in ('m5', 'm4', 'm3', 'm2', 'm1'): |
| | if name in text: | | | extra = '' | | | for suf in ('ultra', 'max', 'pro'): | | | if suf in brand_l: | | | extra = ' ' + suf.title() | | | break |
| | return ('Apple ' + name.upper() + extra).strip() |
| | if hw_model.startswith('Mac17'): |
| | return 'Apple M5-class (hw.model)' |
| | if hw_model.startswith('Mac16'): |
| | return 'Apple M4-class (hw.model)' |
| | if hw_model.startswith('Mac15'): |
| | return 'Apple M3-class (hw.model)' |
| | if hw_model.startswith('Mac14'): |
| | return 'Apple M2-class (hw.model)' |
| | if hw_model.startswith('Mac13') or hw_model.startswith('MacBookPro18'): |
| | return 'Apple M1-class (hw.model)' |
| | return 'unknown' | | | | | | |
| | def guess_ane_cores(chip_name): |
| | text = (chip_name or '').lower() |
| | if 'ultra' in text: | | | return 32 | | | return 16 | | | | | | |
| | def chip_info(label=None): |
| | brand = sysctl('machdep.cpu.brand_string') |
| | model = sysctl('hw.model') |
| | guessed = guess_chip(model, brand) |
| | cores = guess_ane_cores(guessed) |
| | return dict( | | | hw_model=model, | | | brand=brand, | | | platform=platform.platform(), | | | guessed_chip=guessed, | | | guessed_ane_cores=cores, | | | dual_cluster=cores >= 32, | | | label=label or guessed, | | | python=sys.version.split()[0], | | | ) | | | | | | |
| | def dense_hits(cin, cout, cores=CORES): |
| | bpc = fp16_bytes_per_core(cin, cout, cores) |
| | row = dict( |
| | cin=cin, cout=cout, cores=cores, |
| | bytes_per_core=bpc, mib_per_core=bpc / MIB, |
| | on_prefetch_notch=on_prefetch_notch(bpc), |
| | total_fp16_bytes=total_fp16_bytes(cin, cout), |
| | ) |
| | if cout % 16 == 0: |
| | bpc16 = fp16_bytes_per_core(cin, cout, 16) |
| | row['mib_per_core_16'] = bpc16 / MIB |
| | row['on_prefetch_notch_16'] = on_prefetch_notch(bpc16) |
| | if cout % 32 == 0: |
| | bpc32 = fp16_bytes_per_core(cin, cout, 32) |
| | row['mib_per_core_32'] = bpc32 / MIB |
| | row['on_prefetch_notch_32'] = on_prefetch_notch(bpc32) |
| | return row | | | | | | |
| | def default_preset(label=None): |
| | chip = chip_info(label) |
| | if chip['dual_cluster'] or (label or '').lower() in ('m3u', 'm1u', 'm2u', 'ultra'): |
| | return 'dual' | | | return 'stock' | | | | | | |
| | def host_report(cins=DEFAULT_CINS, cout=DEFAULT_COUT, label=None, cores=None): |
| | chip = chip_info(label) |
| | cores = int(cores or chip['guessed_ane_cores']) |
| | points = [dense_hits(cin, cout, cores) for cin in cins] |
| | return dict( | | | article=ARTICLE, | | | chip=chip, | | | cores=cores, | | | cout=cout, | | | points=points, |
| | dual_lattice=list(DUAL_LATTICE), |
| | compiler=dict( |
| | has_1mib_kernel_dma_split_flag=False, | | | serialized_split_related=list(SERIALIZED_SPLIT_RELATED), | | | dropped_dma_split_candidates=list(DROPPED_DMA_SPLIT_CANDIDATES), | | | not_one_mib_dma_split=list(NOT_ONE_MIB_DMA_SPLIT), | | | note=('ANEC split flags cover spatial/batch/channel/section/LUT. ' | | | 'Fabricated KernelDMA 1 MiB keys are dropped by ' | | | 'ANECCreateCompilerOptionsCFString. Workaround is pad/split ' | | | 'the GEMM so each cluster core is not k*1 MiB. Ultra dual-cluster ' |
| | 'needs the 32-core lattice; stock 16-core 4096x2048 is 0.5 MiB.')), |
| | reference_m5_evalus=dict( |
| | d2016_us=353.8, d2048_us=785.5, ratio=2.22, gbps=(46.7, 21.4), |
| | class_name='m3_class_notch', cores=16, |
| | host='kernel_dma_mib_host.m'), |
| | reference_m3_article=dict( |
| | d2016_gbps=44.5, d2048_gbps=16.93, ratio=2.63, |
| | class_name='m3_class_notch', cores=16), |
| | reference_m3_ultra_20260911=dict( |
| | note=('Mac15,14 2x16 ANE. Do not cite stock 4096x2048 (0.5 MiB/core). ' |
| | 'Hits: 8192x2048 ~1.61x (66 vs 42 GB/s); 4096x8192 ~1.77x (89 vs 51 GB/s). ' |
| | 'Same byte-count with Cin=4096 missed. M5 Max (Mac17,6) flat on dual lattice.')), |
| | ) | | | | | | |
| | def build_host(dest): |
| | dest = Path(dest) |
| | run = subprocess.run( |
| | ['clang', '-O2', '-fobjc-arc', str(HOST_M), '-o', str(dest), |
| | '-framework', 'Foundation', '-framework', 'CoreVideo', | | | '-framework', 'IOSurface', | | | '-F/System/Library/PrivateFrameworks', '-framework', 'AppleNeuralEngine'], | | | capture_output=True, text=True, timeout=30) | | | if run.returncode: | | | raise RuntimeError(run.stderr[-4000:]) | | | return dest | | | | | | | | | def convert_conv(cin, cout, dest): | | | if ct is None or np is None: | | | raise RuntimeError('coremltools and numpy required for native compile') |
| | dest = Path(dest) |
| | dest.mkdir(parents=True, exist_ok=True) |
| | pkg, mlc = dest / 'model.mlpackage', dest / 'model.mlmodelc' |
| | weight = np.full((cout, cin, 1, 1), np.float16(0.001)) |
| | specs = [mb.TensorSpec(shape=(1, cin, 1, 1), dtype=types.fp16)] |
| | |
| | def conv_main(x): |
| | y = mb.conv(x=x, weight=weight, pad_type='valid', strides=[1, 1]) |
| | return mb.identity(x=y, name='y') |
| | |
| | prog = mb.program(input_specs=specs, opset_version=ct.target.iOS18)(conv_main) |
| | pipeline = copy.deepcopy(ct.PassPipeline.DEFAULT) |
| | pipeline.remove_passes(['common::fuse_conv_scale', 'common::fuse_conv_bias']) |
| | model = ct.convert( |
| | prog, convert_to='mlprogram', compute_precision=ct.precision.FLOAT16, | | | minimum_deployment_target=ct.target.iOS18, pass_pipeline=pipeline, |
| | skip_model_load=True) |
| | model.save(str(pkg)) |
| | compile_model(str(pkg), destination_path=str(mlc)) |
| | return mlc | | | | | | |
| | def time_mlmodelc(host, mlc, warm, timed): |
| | env = dict(os.environ, ANE_IDENTITY='kernel_dma_mib_probe') |
| | for name in ('ANE_INMEM_HWX', 'ANE_INPUT_SPLIT'): |
| | env.pop(name, None) |
| | run = subprocess.run( |
| | [str(host), str(mlc), str(warm), str(timed)], |
| | capture_output=True, text=True, timeout=120, env=env) |
| | if run.returncode: |
| | raise RuntimeError((run.stderr or run.stdout)[-4000:]) |
| | line = [row for row in run.stdout.splitlines() if row.startswith('DMA_JSON ')] |
| | if not line: | | | raise RuntimeError('native host printed no DMA_JSON') |
| | payload = json.loads(line[-1][9:]) |
| | samples = [float(x) for x in payload['eval_us']] |
| | return dict(eval_us_samples=samples, median_eval_us=statistics.median(samples)) |
| | | | | | | | def probe_anec(out): | | | """Re-serialize split-related keys on this machine's ANECompiler.""" |
| | out = Path(out) |
| | out.mkdir(parents=True, exist_ok=True) |
| | src = HERE / 'anec_options_probe.m' | | | binary = out / 'anec_options_probe' | | | flags = { | | | 'SplitKernelSection': False, | | | 'EnableSpatialSplitInX': True, | | | 'EnableKernelSplitForMultiPaletteLUT': True, | | | 'GlobalRefinementInSpatialSplit': True, | | | 'EnableSegmentAwareKernelSectionSplit': True, | | | 'EnableGlobalChannelSplitting': True, | | | 'EnableL2BatchSplitting': True, | | | 'EnableForcedMaximalBondedSplit': True, | | | 'DisableCachePrefetchMask': 0, | | | 'EnableKernelRewind': True, | | | 'SplitKernelDMA': True, | | | 'EnableKernelDMASplit': True, | | | 'MaxKernelDMASize': 1048575, | | | 'AvoidOneMiBKernelDMA': True, | | | 'SplitCoeffBuffer': True, | | | } | | | flags_path = out / 'flags.json' |
| | flags_path.write_text(json.dumps(flags) + '\n') |
| | if not src.is_file(): |
| | return dict(status='skipped_no_probe_source', path=str(src)) |
| | build = subprocess.run( |
| | ['clang', '-O2', '-fobjc-arc', str(src), '-o', str(binary), |
| | '-framework', 'Foundation', | | | '-F/System/Library/PrivateFrameworks', '-framework', 'ANECompiler'], | | | capture_output=True, text=True, timeout=30) | | | if build.returncode: |
| | return dict(status='probe_build_failed', stderr=build.stderr[-2000:]) |
| | run = subprocess.run( |
| | [str(binary), str(flags_path)], capture_output=True, text=True, timeout=30) |
| | serialized = (run.stdout or '') + (run.stderr or '') |
| | has_dma_split = any( |
| | token in serialized.lower() for token in ( |
| | 'split-kernel-dma', 'kernel-dma-split', 'max-kernel-dma', 'one-mib')) |
| | return dict( | | | status='probed' if run.returncode == 0 else 'probe_failed', | | | returncode=run.returncode, | | | serialized=serialized.strip()[-2000:], | | | fabricated_keys_dropped=not has_dma_split, | | | has_1mib_kernel_dma_split_flag=bool(has_dma_split), | | | ) | | | | | | |
| | def _payload_gbps(cin, cout, median_us): |
| | return (total_fp16_bytes(cin, cout) / 1e9) / (median_us / 1e6) |
| | | | | |
| | def run_native(out, cins=DEFAULT_CINS, cout=DEFAULT_COUT, warm=3, timed=12, |
| | label=None, cores=None, host=None): |
| | out = Path(out).resolve() |
| | if out.exists(): |
| | raise FileExistsError('fresh kernel-DMA profile directory required: %s' % out) |
| | out.mkdir(parents=True) |
| | report = host_report(cins, cout, label, cores) |
| | cores = report['cores'] |
| | report.update(status='preparing', passed=False, metric='native_eval_us') |
| | (out / 'report.json').write_text(json.dumps(report, indent=2) + '\n') |
| | lock = None | | | try: | | | if acquire_native_lock is not None: | | | lock = acquire_native_lock(timeout=1) | | | if host is None: |
| | host = build_host(out / 'kernel_dma_mib_host') |
| | measured = [] |
| | for cin in cins: | | | dest = out / ('d%d' % cin) | | | print('DMA_NOTCH_COMPILE', cin, 'cout', cout, 'cores', cores, flush=True) |
| | mlc = convert_conv(cin, cout, dest) |
| | timed_row = time_mlmodelc(host, mlc, warm, timed) |
| | gbps = _payload_gbps(cin, cout, timed_row['median_eval_us']) |
| | row = dense_hits(cin, cout, cores) |
| | row.update(gbps_weight_payload=gbps, **timed_row) |
| | measured.append(row) |
| | print('DMA_NOTCH_TIME', cin, row['median_eval_us'], gbps, flush=True) |
| | off, on = measured[0], measured[-1] |
| | slow = on['median_eval_us'] / off['median_eval_us'] |
| | class_name = classify_ratio(slow, on['median_eval_us']) |
| | report.update( | | | status='passed_kernel_dma_profile', passed=True, points=measured, | | | ratio_on_over_off=slow, | | | ratio_2048_over_2016=slow if set(cins) >= {2016, 2048} else None, | | | class_name=class_name, |
| | m3_class_notch=class_name == 'm3_class_notch', |
| | milder_1mib_slowdown=class_name == 'milder_1mib_slowdown') |
| | (out / 'report.json').write_text(json.dumps(report, indent=2) + '\n') |
| | print('DMA_NOTCH_RESULT', report['status'], slow, class_name, flush=True) |
| | return report | | | except Exception as exc: |
| | report.update(status='stopped_no_retry', error='%s: %s' % (type(exc).__name__, exc)) |
| | (out / 'report.json').write_text(json.dumps(report, indent=2) + '\n') |
| | raise | | | finally: | | | if lock is not None: | | | lock.close() | | | | | | |
| | def run_preset(out, preset, warm, timed, label, cores_override=None): |
| | out = Path(out).resolve() |
| | if out.exists(): |
| | raise FileExistsError('fresh kernel-DMA profile directory required: %s' % out) |
| | out.mkdir(parents=True) |
| | lattices = DUAL_LATTICE if preset == 'dual' else (DUAL_LATTICE[0],) |
| | host = build_host(out / 'kernel_dma_mib_host') |
| | summary = [] |
| | for spec in lattices: |
| | sub = out / spec['name'] |
| | cores = cores_override or spec['cores'] |
| | try: |
| | report = run_native( |
| | sub, spec['cins'], spec['cout'], warm, timed, label, cores, host=host) |
| | summary.append(dict( |
| | name=spec['name'], note=spec['note'], cores=cores, |
| | cout=spec['cout'], cins=list(spec['cins']), |
| | class_name=report.get('class_name'), |
| | ratio_on_over_off=report.get('ratio_on_over_off'), |
| | points=[dict( |
| | cin=p['cin'], median_eval_us=p.get('median_eval_us'), |
| | gbps_weight_payload=p.get('gbps_weight_payload'), |
| | mib_per_core=p.get('mib_per_core'), |
| | on_prefetch_notch=p.get('on_prefetch_notch')) |
| | for p in report.get('points', [])], |
| | )) |
| | except Exception as exc: |
| | summary.append(dict(name=spec['name'], error='%s: %s' % (type(exc).__name__, exc))) |
| | blob = dict(preset=preset, chip=chip_info(label), scans=summary) |
| | (out / 'summary.json').write_text(json.dumps(blob, indent=2) + '\n') |
| | print('DMA_NOTCH_SUMMARY', json.dumps(blob['scans'], indent=2), flush=True) |
| | return blob | | | | | | |
| | def main(argv=None): |
| | parser = argparse.ArgumentParser(description=__doc__) |
| | parser.add_argument('--host-only', action='store_true') |
| | parser.add_argument('--out', type=Path) |
| | parser.add_argument('--probe-anec', action='store_true') |
| | parser.add_argument('--label', help='override chip label, e.g. m1m / m3u / m4 / m5m') |
| | parser.add_argument('--cout', type=int, default=DEFAULT_COUT) |
| | parser.add_argument('--cins', default='2016,2048') |
| | parser.add_argument('--cores', type=int, default=0, |
| | help='ANE cores for MiB/core math (0=auto: 32 if Ultra else 16)') |
| | parser.add_argument('--preset', choices=('stock', 'dual', 'auto'), default='auto', |
| | help='stock=16-core article pair; dual=16c+32c lattice; auto=dual on Ultra') |
| | parser.add_argument('--warm', type=int, default=3) |
| | parser.add_argument('--timed', type=int, default=12) |
| | args = parser.parse_args(argv) |
| | cins = tuple(int(x) for x in args.cins.split(',')) |
| | cores = args.cores or None | | | preset = args.preset |
| | if preset == 'auto': |
| | preset = default_preset(args.label) |
| | if args.host_only or args.out is None: | | | report = host_report(cins, args.cout, args.label, cores) |
| | report['preset'] = preset |
| | print(json.dumps(report, indent=2)) |
| | if args.out is None and not args.probe_anec: | | | return 0 | | | if args.probe_anec: | | | dest = args.out / 'anec_probe' if args.out else Path('anec_dma_split_probe') | | | print(json.dumps(probe_anec(dest), indent=2)) | | | if args.out is not None and not args.host_only: | | | if preset == 'dual': | | | run_preset(args.out, 'dual', args.warm, args.timed, args.label, cores) | | | else: | | | run_native(args.out, cins, args.cout, args.warm, args.timed, args.label, cores) | | | return 0 | | | | | | |
| | if __name__ == '__main__': |
| | raise SystemExit(main()) |