AssemblyScript Wasm基数排序比纯JS慢4倍，求优化方案

问题背景 #

已将JS基数排序顶点变换基准测试移植到AssemblyScript WASM，完成了noAssert、运行时存根、optimizeLevel 3、启用SIMD、复用全局变量、乒乓缓冲区、减少批量复制、显式StaticArray计数等常规优化，但Wasm版本运行速度仍远慢于原JS版本，需排查遗漏的优化点。

现有代码实现 #

assembly/radix.ts #

// 8-bit radix parameters
const BUCKETS = 256;

// Temp buffers used by the sort. Stored globally to avoid any allocations in the hot path.
let b_keys: Uint32Array | null = null;
let b_indices: Uint32Array | null = null;

// One single StaticArray for bucket counts — no GC, no reallocations ever.
const b_counts = new StaticArray<i32>(BUCKETS);

/**
 * Sorts two parallel Uint32Array buffers in place by their keys.
 */
export function radixSort(keys: Uint32Array, indices: Uint32Array): void {
    const n = keys.length;

    // Allocate scratch buffers once, or grow them if needed
    if (b_keys === null || b_keys!.length < n) {
        b_keys = new Uint32Array(n);
        b_indices = new Uint32Array(n);
    }
    const tempKeys = b_keys!;
    const tempIndices = b_indices!;

    // Do one pass per byte
    for (let shift = 0; shift < 32; shift += 8) {
        // 1) Zero the counts
        for (let b = 0; b < BUCKETS; ++b) {
            b_counts[b] = 0;
        }

        // 2) Counting pass
        for (let i = 0; i < n; ++i) {
            const k = keys[i];
            b_counts[(k >>> shift) & (BUCKETS - 1)]++;
        }

        // 3) Prefix sum → bucket starting indices
        for (let b = 1; b < BUCKETS; ++b) {
            b_counts[b] += b_counts[b - 1];
        }

        // 4) Placement pass (in reverse) to keep stability
        for (let i = n - 1; i >= 0; --i) {
            const k = keys[i];
            const bucket = (k >>> shift) & (BUCKETS - 1);
            const dst = --b_counts[bucket];
            tempKeys[dst] = k;
            tempIndices[dst] = indices[i];
        }

        // 5) Bulk-copy back into the original arrays
        memory.copy(
            keys.dataStart,       // dest ptr
            tempKeys.dataStart,   // src ptr
            <usize>n << 2         // n * 4 bytes
        );
        memory.copy(
            indices.dataStart,
            tempIndices.dataStart,
            <usize>n << 2
        );
    }
}

/**
 * Float→sortable-u32 conversion (unchanged).
 */
@inline
export function f32ToSortable(val: f32): u32 {
    const iv = reinterpret<u32>(val);
    // flip bits so negatives sort before positives
    return ((iv >> 31) & 1) ? ~iv : (iv | 0x8000_0000);
}

assembly/simd.ts #

// This file now contains both SIMD and Scalar versions of the transform function,
// selected at compile-time by the ASC_SIMD flag in your package.json.

@inline
export function transform_vertex_simd(vertexPtr: usize, matrixPtr: usize): f32 {

    // --- The 'if (ASC_SIMD)' block is evaluated at COMPILE TIME ---
    if (ASC_SIMD) {
        // --- SIMD PATH ---
        // This code will only be included in your 'release-simd.wasm' build.
        // It requires the matrix to be transposed in JavaScript.

        const vert_x = load<f32>(vertexPtr, 0);
        const vert_y = load<f32>(vertexPtr, 4);
        const vert_z = load<f32>(vertexPtr, 8);
        const v_vertex = f32x4(vert_x, vert_y, vert_z, 1.0);

        const v_row2 = v128.load(matrixPtr, 32);

        const products = f32x4.mul(v_vertex, v_row2);

        const sum_xy = f32x4.extract_lane(products, 0) + f32x4.extract_lane(products, 1);
        const sum_zw = f32x4.extract_lane(products, 2) + f32x4.extract_lane(products, 3);
        const res_z = sum_xy + sum_zw;

        return res_z;

    } else {
        // --- SCALAR PATH ---
        // This code will only be included in your 'release-scalar.wasm' build.
        // It uses the original, non-transposed matrix.

        const x = load<f32>(vertexPtr, 0);
        const y = load<f32>(vertexPtr, 4);
        const z = load<f32>(vertexPtr, 8);
        const w: f32 = 1.0; // w is implicitly 1

        // Direct calculation of Z_out = m[8]*x + m[9]*y + m[10]*z + m[11]*w
        // using the original column-major matrix layout.
        const m8 = load<f32>(matrixPtr, 32);
        const m9 = load<f32>(matrixPtr, 36);
        const m10 = load<f32>(matrixPtr, 40);
        const m11 = load<f32>(matrixPtr, 44);

        // This scalar version is simplified to match the SIMD version's logic exactly.
        // It calculates only the Z-depth from the original (non-transposed) matrix.
        const res_z = m8 * x + m9 * y + m10 * z + m11 * w;

        return res_z;
    }
}

assembly/index.ts #

// assembly/index.ts (Corrected Benchmarking Logic)

import { transform_vertex_simd } from "./simd"; // We now only import the SIMD version
import { radixSort, f32ToSortable } from "./radix";
import { lastProfile, ProfileData } from "./profiler";

let inputVerticesPtr: usize = 0;
let inputMatrixPtr: usize = 0;
let outputVerticesPtr: usize = 0;

const F32_SIZE: u32 = 4;
const VERTEX_SZ: u32 = 3 * F32_SIZE;

export function getProfileReport(): ProfileData {
    return lastProfile;
}

export function alloc(numVertices: i32): StaticArray<u32> {
    const vc = <u32>numVertices;
    inputVerticesPtr = heap.alloc(vc * VERTEX_SZ);
    inputMatrixPtr = heap.alloc(16 * F32_SIZE);
    outputVerticesPtr = heap.alloc(vc * VERTEX_SZ);
    const out = new StaticArray<u32>(2);
    out[0] = <u32>inputVerticesPtr;
    out[1] = <u32>inputMatrixPtr;
    return out;
}

// The sort function now runs ONLY ONCE. The repetition is handled in JavaScript.
export function sort(numVertices: i32): u32 {
    const vc = <u32>numVertices;
    const keys = new Uint32Array(vc);
    const indices = new Uint32Array(vc);

    // Stage 1: Key Calculation
    const t1 = Date.now();
    for (let i: u32 = 0; i < vc; ++i) {
        const ptr = inputVerticesPtr + i * VERTEX_SZ;
        // The matrix pointer is now passed directly to the transform function
        const depth = transform_vertex_simd(ptr, inputMatrixPtr);
        keys[i] = f32ToSortable(depth);
        indices[i] = i;
    }
    lastProfile.keyCalculation = Date.now() - t1;

    // Stage 2: Radix Sort
    const t2 = Date.now();
    radixSort(keys, indices);
    lastProfile.radixSort = Date.now() - t2;

    // Stage 3: Reordering
    const t3 = Date.now();
    for (let i: u32 = 0; i < vc; ++i) {
        const orig = indices[i];
        const src = inputVerticesPtr + orig * VERTEX_SZ;
        const dst = outputVerticesPtr + i * VERTEX_SZ;
        memory.copy(dst, src, VERTEX_SZ);
    }
    lastProfile.reordering = Date.now() - t3;

    lastProfile.total = lastProfile.keyCalculation + lastProfile.radixSort + lastProfile.reordering;
    return <u32>outputVerticesPtr;
}

assembly/profiler.ts #

// This class will hold our detailed performance report.
// The ESM bindings will automatically convert this to a JS object.
export class ProfileData {
  keyCalculation: i64 = 0;
  radixSort: i64 = 0;
  reordering: i64 = 0;
  total: i64 = 0;
}

// A global instance to store the data from the last run.
export const lastProfile = new ProfileData();

关键优化遗漏点及修复方案 #

1. 消除数组边界检查开销 #

当前使用的Uint32Array即使开启noAssert仍可能保留隐式边界检查，改用StaticArray<u32>直接操作内存，彻底消除检查：

// 修改radixSort的参数与全局buffer类型
let b_keys: StaticArray<u32> | null = null;
let b_indices: StaticArray<u32> | null = null;

export function radixSort(keys: StaticArray<u32>, indices: StaticArray<u32>): void {
    const n = keys.length;
    // 后续数组访问直接用下标，无边界检查
}

2. 充分利用SIMD求和指令 #

当前SIMD路径手动提取lane相加浪费性能，改用SIMD水平求和指令：

// 替换SIMD路径求和代码
const products = f32x4.mul(v_vertex, v_row2);
const sum1 = f32x4.hadd(products, products);
const sum2 = f32x4.hadd(sum1, sum1);
const res_z = f32x4.extract_lane(sum2, 0);

3. 减少内存拷贝次数 #

基数排序每轮的双向拷贝是性能瓶颈，改用双缓冲指针切换策略：

// radixSort中引入源/目标指针切换
let srcKeys = keys;
let srcIndices = indices;
let dstKeys = tempKeys;
let dstIndices = tempIndices;

for (let shift = 0; shift < 32; shift += 8) {
    // 计数、前缀和、放置阶段写入dstKeys/dstIndices
    // 交换源和目标，下一轮反向处理
    [srcKeys, dstKeys] = [dstKeys, srcKeys];
    [srcIndices, dstIndices] = [dstIndices, srcIndices];
}

// 若为奇数轮，最后一次拷贝回原数组
if (srcKeys !== keys) {
    memory.copy(keys.dataStart, srcKeys.dataStart, <usize>n << 2);
    memory.copy(indices.dataStart, srcIndices.dataStart, <usize>n << 2);
}

4. 优化内存访问模式 #

• 将顶点数据从结构数组（AoS）转为数组结构（SoA），提升缓存命中率；
• 重排阶段分块处理，利用CPU缓存行特性减少缓存失效。

5. 强化编译优化选项 #

确保编译时添加以下参数：

asc assembly/index.ts --target release --optimizeLevel 3 --shrinkLevel 2 --noAssert --runtime stub --enable simd --disable bulk-memory

• --shrinkLevel 2：进一步压缩代码消除冗余；
• --disable bulk-memory：关闭未使用的特性，提升编译器优化空间；
• 确认ASC_SIMD=1环境变量已设置，确保SIMD代码被编译。

6. 修正基准测试精度 #

• 确保JS与Wasm使用完全相同的输入数据；
• 改用performance.now()替代Date.now()提升计时精度；
• 排除Wasm模块加载、内存初始化时间，仅统计核心逻辑运行时长。

内容的提问来源于stack exchange，提问作者Sidhant Kumar