open-xdna / examples /npu_patch_embed_scaled.py

Source snapshot of github.com/Scottcjn/open-xdna@6c8c680

e629bba verified 12 days ago

2.7 kB

#!/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()

	#!/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=2MPDIM*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()