80 lines
2.1 KiB
TypeScript
80 lines
2.1 KiB
TypeScript
|
import { Worker } from "worker_threads";
|
|||
|
|
|
|||
|
|
function createWorker(data: Float32Array): Promise<Float32Array> {
|
|||
|
|
return new Promise((resolve, reject) => {
|
|||
|
|
const worker = new Worker("./vino.worker.js");
|
|||
|
|
worker.on("message", (result) => resolve(result as Float32Array));
|
|||
|
|
worker.on("error", reject);
|
|||
|
|
worker.on("exit", (code) => {
|
|||
|
|
if (code !== 0) {
|
|||
|
|
reject(new Error(`Worker stopped with exit code ${code}`));
|
|||
|
|
}
|
|||
|
|
});
|
|||
|
|
worker.postMessage(data);
|
|||
|
|
});
|
|||
|
|
}
|
|||
|
|
|
|||
|
|
class TaskQueue {
|
|||
|
|
private concurrency: number;
|
|||
|
|
private queue: (() => void)[];
|
|||
|
|
private running: number;
|
|||
|
|
|
|||
|
|
constructor(concurrency: number) {
|
|||
|
|
this.concurrency = concurrency;
|
|||
|
|
this.queue = [];
|
|||
|
|
this.running = 0;
|
|||
|
|
}
|
|||
|
|
|
|||
|
|
private async runTask(task: () => Promise<any>): Promise<any> {
|
|||
|
|
this.running++;
|
|||
|
|
try {
|
|||
|
|
const result = await task();
|
|||
|
|
this.running--;
|
|||
|
|
this.next();
|
|||
|
|
return result;
|
|||
|
|
} catch (error) {
|
|||
|
|
this.running--;
|
|||
|
|
this.next();
|
|||
|
|
throw error;
|
|||
|
|
}
|
|||
|
|
}
|
|||
|
|
|
|||
|
|
public addTask(task: () => Promise<any>): Promise<any> {
|
|||
|
|
return new Promise((resolve, reject) => {
|
|||
|
|
const run = () => this.runTask(task).then(resolve).catch(reject);
|
|||
|
|
this.queue.push(run);
|
|||
|
|
this.next();
|
|||
|
|
});
|
|||
|
|
}
|
|||
|
|
|
|||
|
|
private next(): void {
|
|||
|
|
if (this.running < this.concurrency && this.queue.length > 0) {
|
|||
|
|
const task = this.queue.shift();
|
|||
|
|
task?.();
|
|||
|
|
}
|
|||
|
|
}
|
|||
|
|
}
|
|||
|
|
|
|||
|
|
export async function cpuStartInfer() {
|
|||
|
|
const numTasks = 128;
|
|||
|
|
const patchSize = [1, 1, 80, 160, 160];
|
|||
|
|
const taskQueue = new TaskQueue(2); // 设置并发任务数，例如2
|
|||
|
|
|
|||
|
|
console.time("总推理用时");
|
|||
|
|
|
|||
|
|
try {
|
|||
|
|
const tasks = Array.from({ length: numTasks }, () => {
|
|||
|
|
const inferSample = new Float32Array(
|
|||
|
|
patchSize.reduce((a, b) => a * b)
|
|||
|
|
).map(() => Math.random() * 2.0 - 1.0);
|
|||
|
|
return taskQueue.addTask(() => createWorker(inferSample));
|
|||
|
|
});
|
|||
|
|
|
|||
|
|
const results = await Promise.all(tasks);
|
|||
|
|
console.log("所有任务完成，结果数量：", results.length);
|
|||
|
|
} catch (e) {
|
|||
|
|
console.error("推理失败:", e);
|
|||
|
|
}
|
|||
|
|
console.timeEnd("总推理用时");
|
|||
|
|
}
|