File size: 2,699 Bytes
e629bba | 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 | #!/usr/bin/env python3
# SPDX-License-Identifier: AGPL-3.0-or-later
# Copyright (C) 2026 Scott Boudreaux / Elyan Labs. Commercial license: see ../COMMERCIAL.md
#
# open-xdna :: SigLIP patch-embed on the NPU, SCALED by image size / batch — measured.
# Tests the hypothesis: bigger M (more patches × batch) amortizes NPU dispatch overhead, so the
# NPU's effective GFLOP/s climbs toward its ~68 peak. Compared against fair CPU f32 BLAS.
import os, sys, time
import numpy as np
MLIR_AIE = os.environ.get("MLIR_AIE_DIR", os.path.expanduser("~/open-xdna/mlir-aie"))
sys.path.insert(0, os.path.join(MLIR_AIE, "programming_examples/basic/matrix_multiplication/single_core"))
import aie.iron as iron
from single_core import single_core
PDIM, EMBED = 768, 768 # SigLIP-base/16 patch_dim, embed
TILE = dict(m=32, k=32, n=32, dtype_in_str="i16", dtype_out_str="i32")
def npu_mm(A_i16, B_i16, M):
a = iron.tensor(A_i16.reshape(-1), dtype=np.int16, device="npu")
b = iron.tensor(B_i16.reshape(-1), dtype=np.int16, device="npu")
c = iron.tensor(np.zeros(M*EMBED, dtype=np.int32), dtype=np.int32, device="npu")
single_core(a, b, c, M=M, K=PDIM, N=EMBED, **TILE)
return c.numpy().reshape(M, EMBED).astype(np.int32)
def main():
rng = np.random.default_rng(0)
Wproj = rng.integers(-4, 5, size=(PDIM, EMBED), dtype=np.int16)
Wf = Wproj.astype(np.float32)
print("open-xdna :: SigLIP patch-embed SCALED (NPU int16 vs fair CPU f32 BLAS)")
print(f" {'~config':>26} {'M (patches)':>11} {'NPU ms':>8} {'NPU GF/s':>9} {'CPU ms':>8} {'CPU GF/s':>9} {'NPU/CPU':>8}")
# tile-friendly M (single_core needs M divisible into tile groups); label by ~equivalent config
for label, M in [("224² ×1 (~196)", 256), ("448² ×1 (~784)", 1024),
("224² ×~21 / 448²×~5", 4096), ("batch ~42 × 224²", 8192)]:
A = rng.integers(-8, 8, size=(M, PDIM), dtype=np.int16)
npu_mm(A, Wproj, M) # warm/compile
t0=time.perf_counter(); npu_mm(A, Wproj, M); t_npu=(time.perf_counter()-t0)*1e3
Af=A.astype(np.float32)
for _ in range(3): Af@Wf
t0=time.perf_counter()
for _ in range(10): Af@Wf
t_cpu=(time.perf_counter()-t0)/10*1e3
flop=2*M*PDIM*EMBED
print(f" {label:>26} {M:>11} {t_npu:>8.2f} {flop/(t_npu/1e3)/1e9:>9.1f} {t_cpu:>8.2f} {flop/(t_cpu/1e3)/1e9:>9.1f} {t_npu/t_cpu:>7.2f}x")
print(" NPU GF/s climbs with M (dispatch amortized) toward its ~68 peak, but stays below CPU f32 BLAS.")
print(" -> raw-speed crossover does NOT happen; NPU's case is offload + ~6.6W power + concurrency.")
if __name__ == "__main__":
main()
|