如何使用SIMD指令加速有序整数数组匹配场景的哈达玛积计算

我正在尝试实现两个数据集的最快连接运算，当键匹配时对对应值执行乘积运算。目前我已经实现了标量版本，但我认为使用SIMD可以进一步加速该算法。我有两个作为匹配键的Span<int>对象：aKeys和bKeys，当aKeys中的值与bKeys中的值匹配时，需要将aValues和bValues中对应的数值相乘并存储。

我希望使用SIMD指令，将aKeys中的单个键与bKeys中的多个值同时比较。由于数组有序，一旦bKeys中的值大于当前待测试的aKey值，就需要切换到下一个aKey值。
SpanHelpers中的IndexOf<T>算法和需求很接近，但它仅支持查询单个值。我希望利用待查询值全部唯一且按升序排列的特性进一步优化。

标量实现代码如下：

// 注意：aKeys和aValues长度相同
// bKeys和bValues长度相同
// aKeys和bKeys中的值都是唯一且有序的，但aKeys中的值不一定都在bKeys中存在，反之亦然
let hadamardProduct (aKeys: Span<int>, aValues: Span<float>, bKeys: Span<int>, bValues: Span<float>) =
    let maxN = Math.Min (aKeys.Length, bKeys.Length)
    let outKeys = Array.zeroCreate maxN
    let outValues = Array.zeroCreate maxN

    let mutable aIdx = 0
    let mutable bIdx = 0
    let mutable outIdx = 0

    while aIdx < aKeys.Length && bIdx < bKeys.Length do
        
        if aKeys.[aIdx] = bKeys.[bIdx] then
            outKeys.[outIdx] <- aKeys.[aIdx]
            outValues.[outIdx] <- aValues.[aIdx] * bValues.[bIdx]
            outIdx <- outIdx + 1
            aIdx <- aIdx + 1
            bIdx <- bIdx + 1
        elif aKeys.[aIdx] < bKeys.[bIdx] then
            aIdx <- aIdx + 1
        else
            bIdx <- bIdx + 1

    let resultKeys = Memory (outKeys, 0, outIdx)
    let resultValues = Memory (outValues, 0, outIdx)

    resultKeys, resultValues

2021-08-23 09:26 更新 #

经过进一步开发我实现了如下SIMD版本的方案，请问还有什么方法可以进一步提升代码运行速度？

#nowarn "9" "51" "20" // 忽略指针相关的警告
open System
open FSharp.NativeInterop
open System.Numerics
open System.Runtime.Intrinsics.X86
open System.Runtime.Intrinsics

let hadamardProduct (aKeys: Span<int>, aValues: Span<float>, bKeys: Span<int>, bValues: Span<float>) =
    let maxN = Math.Min (aKeys.Length, bKeys.Length)
    let outKeys = Array.zeroCreate maxN
    let outValues = Array.zeroCreate maxN

    let mutable aIdx = 0
    let mutable bIdx = 0
    let mutable outIdx = 0

    if bKeys.Length > 4 then

        let lastBlockIdx = bKeys.Length - (bKeys.Length % 4)
        let bPointer = && (bKeys.GetPinnableReference ())
        let mutable bVector = Sse2.LoadVector128 (NativePtr.add bPointer bIdx)

        while aIdx < aKeys.Length && bIdx < lastBlockIdx do
            let aVector = Vector128.Create aKeys.[aIdx]
            let comparison = Sse2.CompareEqual (aVector, bVector)
            let matches = Sse2.MoveMask (comparison.AsByte ())

            if matches > 0 then
                let bIdxOffset = (BitOperations.TrailingZeroCount matches) / 4 // 将字节偏移转换为索引
                outKeys.[outIdx] <- aKeys.[aIdx]
                outValues.[outIdx] <- aValues.[aIdx] * bValues.[bIdx + bIdxOffset]
                aIdx <- aIdx + 1
                outIdx <- outIdx + 1
                // 注意：bIdx每次是同时检测4个值，不需要每次都跳转

            elif aKeys.[aIdx] > bKeys.[bIdx + 3] then
                // 注意：bIdx需要步长前进，而非单步递增
                bIdx <- bIdx + 4
                // 仅在必要时加载新的向量值
                if bIdx < lastBlockIdx then
                    bVector <- Sse2.LoadVector128 (NativePtr.add bPointer bIdx)
            else
                aIdx <- aIdx + 1

    // 处理剩余不满足向量长度的零散数据
    while aIdx < aKeys.Length && bIdx < bKeys.Length do
        
        if aKeys.[aIdx] = bKeys.[bIdx] then
            outKeys.[outIdx] <- aKeys.[aIdx]
            outValues.[outIdx] <- aValues.[aIdx] * bValues.[bIdx]
            outIdx <- outIdx + 1
            aIdx <- aIdx + 1
            bIdx <- bIdx + 1
        elif aKeys.[aIdx] < bKeys.[bIdx] then
            aIdx <- aIdx + 1
        else
            bIdx <- bIdx + 1

    let resultKeys = Memory (outKeys, 0, outIdx)
    let resultValues = Memory (outValues, 0, outIdx)

    resultKeys, resultValues

内容的提问来源于stack exchange，提问作者Matthew Crews