tiny.simd.dotBFloat16Assume

1pub fn computeBFloatAssume(2    comptime D: type,3    a: []const bfloat.BFloat16,4    b: []const bfloat.BFloat16,5    count: usize,6    comptime assumptions: Assumptions,7) f32 {8    requireBFloatTag(D);9    validateInputs(D.lane_count, a.len, b.len, count, assumptions);10    if (D.lane_count < 2) return computeBFloatScalar(a, b, count);11    const DF = D.repartition(f32);12    var sum0: DF.Vector = @splat(0);13    var sum1: DF.Vector = @splat(0);14    var sum2: DF.Vector = @splat(0);15    var sum3: DF.Vector = @splat(0);16    var index: usize = 0;17    while (index + 2 * D.lane_count <= count) : (index += 2 * D.lane_count) {18        const a0 = bfloat.load(D, a[index..]);19        const b0 = bfloat.load(D, b[index..]);20        sum0 = bfloat.reorderWidenMulAccumulate(DF, a0, b0, sum0, &sum1);21        const a1 = bfloat.load(D, a[index + D.lane_count ..]);22        const b1 = bfloat.load(D, b[index + D.lane_count ..]);23        sum2 = bfloat.reorderWidenMulAccumulate(DF, a1, b1, sum2, &sum3);24    }25    if (index + D.lane_count <= count) {26        const av = bfloat.load(D, a[index..]);27        const bv = bfloat.load(D, b[index..]);28        sum0 = bfloat.reorderWidenMulAccumulate(DF, av, bv, sum0, &sum1);29        index += D.lane_count;30    }31    if (!assumptions.multiple_of_vector and index != count) {32        const remaining = count - index;33        const av = loadBFloatTail(D, a[index..], remaining, assumptions.padded_to_vector);34        const bv = loadBFloatTail(D, b[index..], remaining, assumptions.padded_to_vector);35        sum2 = bfloat.reorderWidenMulAccumulate(DF, av, bv, sum2, &sum3);36    }37    return reduce.sum(DF, (sum0 + sum1) + (sum2 + sum3));38}

1pub const dotBFloat16Assume = dot.computeBFloatAssume;