97 lines
2.9 KiB
Python
97 lines
2.9 KiB
Python
#!/usr/bin/env python3
|
|
"""
|
|
Test VitisAI EP with the quantized nomic-embed-text-v1.5 model.
|
|
Checks whether the NPU VAIML pass achieves meaningful GOPs coverage.
|
|
"""
|
|
|
|
import os
|
|
import json
|
|
import time
|
|
import numpy as np
|
|
from pathlib import Path
|
|
|
|
MODEL_PATH = Path.home() / ".cache/breadsearch/models/model_quantized_static.onnx"
|
|
CACHE_DIR = Path.home() / ".cache/breadsearch/npu/nomic-quantized-static"
|
|
VAIP_CONFIG = Path.home() / ".config/breadsearch/vaip_config.json"
|
|
RYZEN_AI_LIB = Path.home() / ".local/share/ryzen-ai-1.7.1/lib"
|
|
|
|
# libvaiml.so must be discoverable
|
|
os.environ["RYZEN_AI_INSTALLATION_PATH"] = str(RYZEN_AI_LIB)
|
|
os.environ["LD_LIBRARY_PATH"] = str(RYZEN_AI_LIB) + ":" + os.environ.get("LD_LIBRARY_PATH", "")
|
|
|
|
CACHE_DIR.mkdir(parents=True, exist_ok=True)
|
|
|
|
if not VAIP_CONFIG.exists():
|
|
print(f"ERROR: vaip_config.json not found at {VAIP_CONFIG}")
|
|
print("Check breadmill embed.rs for the find_vaip_config() paths")
|
|
exit(1)
|
|
|
|
print(f"Testing quantized model: {MODEL_PATH}")
|
|
print(f"Cache dir: {CACHE_DIR}")
|
|
print(f"VAIP config: {VAIP_CONFIG}")
|
|
print(f"RYZEN_AI_INSTALLATION_PATH: {os.environ['RYZEN_AI_INSTALLATION_PATH']}")
|
|
|
|
import onnxruntime as ort
|
|
|
|
providers = [
|
|
("VitisAIExecutionProvider", {
|
|
"config_file": str(VAIP_CONFIG),
|
|
"cacheDir": str(CACHE_DIR),
|
|
"cacheKey": "nomic-quantized-static",
|
|
}),
|
|
"CPUExecutionProvider",
|
|
]
|
|
|
|
print("\nCreating InferenceSession with VitisAI EP...")
|
|
t0 = time.time()
|
|
try:
|
|
sess = ort.InferenceSession(str(MODEL_PATH), providers=providers)
|
|
t1 = time.time()
|
|
print(f"Session created in {t1-t0:.1f}s")
|
|
print(f"Active providers: {sess.get_providers()}")
|
|
except Exception as e:
|
|
print(f"ERROR creating session: {e}")
|
|
exit(1)
|
|
|
|
# Check for the VAIML pass summary
|
|
summary_path = CACHE_DIR / "nomic-quantized" / "preliminary-vaiml-pass-summary.txt"
|
|
if not summary_path.exists():
|
|
# Try variations
|
|
for p in CACHE_DIR.rglob("preliminary-vaiml-pass-summary.txt"):
|
|
summary_path = p
|
|
break
|
|
|
|
if summary_path.exists():
|
|
print(f"\n--- VAIML Pass Summary ---")
|
|
print(summary_path.read_text())
|
|
else:
|
|
print(f"\nNo VAIML summary found at {summary_path}")
|
|
print("Files in cache dir:")
|
|
for f in CACHE_DIR.rglob("*"):
|
|
if f.is_file():
|
|
print(f" {f}")
|
|
|
|
# Run a quick inference test
|
|
print("\nRunning inference test...")
|
|
seq_len = 128
|
|
dummy_ids = np.ones((1, seq_len), dtype=np.int64)
|
|
dummy_mask = np.ones((1, seq_len), dtype=np.int64)
|
|
dummy_types = np.zeros((1, seq_len), dtype=np.int64)
|
|
|
|
input_names = [inp.name for inp in sess.get_inputs()]
|
|
print(f"Input names: {input_names}")
|
|
|
|
feed = {}
|
|
for name in input_names:
|
|
if "type" in name:
|
|
feed[name] = dummy_types
|
|
elif "mask" in name:
|
|
feed[name] = dummy_mask
|
|
else:
|
|
feed[name] = dummy_ids
|
|
|
|
t0 = time.time()
|
|
outputs = sess.run(None, feed)
|
|
t1 = time.time()
|
|
print(f"Inference completed in {(t1-t0)*1000:.1f}ms")
|
|
print(f"Output shape: {outputs[0].shape}")
|