已开启

grouped_sum优化 #10

QLiangong创建于 5月18日

5月18日创建

Describe the solution you'd like

1 前言

1.1 背景

在 aarch64 和 X86 平台进行基准测试，测试结果显示 Daft 的部分场景下在 aarch64 架构上的性能相比 X86 劣化。Daft 在一些场景下，依赖了 daft_core::array::ops::sum::grouped_sum 接口且该接口性能在不同的场景下均有性能劣化，本文主要分析 daft 中 grouped_sum 函数性能劣化原因和优化方法。

2 Story概述

2.1 Story需求描述（必填）

SR/AR ID	SR/AR描述	Story标题	Story描述
SRxxxxx	grouped_sum接口优化及效果验证，TPC-H Benchmark性能提升	grouped_sum接口优化	优化daft 聚合算子中 grouped_sum 接口在鲲鹏 920B 上的性能

2.2 Story功能描述（必填）

grouped_sum 函数是 Daft 聚合算子中的核心分组求和函数，负责对按 GroupIndices（Vec<Vec<u64>>）描述的每个分组内的若干行值进行求和。该函数通过 impl_daft_numeric_agg! 宏统一为 5 种数值类型生成实现：Int64Type / UInt64Type / Float32Type / Float64Type / Decimal128Type。

该函数的核心工作流程包括：

通过 as_arrow() 获取底层 Arrow PrimitiveArray
检查 null_count() > 0 选择无 null 快速路径或含 null 慢速路径
对每个 group 通过 g.iter().fold(...) 将该 group 中所有索引对应的值累加
将每个 group 的求和结果通过 Vec::from_iter 收集为输出 DataArray

原始实现中存在以下性能瓶颈：

fold 串行依赖：单累加器 fold 形成 loop-carried dependency chain，CPU 无法并行 issue 多个独立的 add
Vec<Vec<u64>> pointer chase：跨 group 切换时硬件预取器失效，产生 L1 cache miss
as_arrow() 重复 downcast：每个 self.get(idx) 内部都做一次 Any::downcast_ref，在大数据量下累积可观开销
未利用鲲鹏 920B 的 SVE 指令集：LLVM auto-vectorizer 无法为 fold 模式生成 SVE gather + 横向归约代码

2.3 Story用户使用场景分析（必填）

一、用户场景分析

Daft 执行 groupby(...).sum() 或 agg(col(...).sum()) 操作时，grouped_sum 是核心路径。
在 TPC-H 基准测试中，多个查询（Q1、Q13、Q14、Q15、Q22 等）都依赖分组求和操作，grouped_sum 的性能直接影响整体查询性能。Q1 是典型的 sum-heavy group-by 查询，对 lineitem 表按 l_returnflag 和 l_linestatus 分组后对多列做求和。

二、新增/变更的文件、脚本（必填）

本次接口优化需要修改的文件如下：

src/daft-core/src/array/ops/aarch64/mod.rs       (新增)
src/daft-core/src/array/ops/aarch64/prfm.rs      (新增)
src/daft-core/src/array/ops/aarch64/sve_sum.rs   (新增)
src/daft-core/src/array/ops/aarch64/tests.rs     (新增)
src/daft-core/src/array/ops/mod.rs               (修改)
src/daft-core/src/array/ops/sum.rs               (修改)
docs/kunpeng-920b-build-guide.md                 (新增)

2.4 Story约束（必填）

基于 Daft v0.7.5 版本进行优化。

3 Story设计描述

3.1 Story设计思路（必填）

问题分析

原始实现中 grouped_sum 存在以下性能瓶颈：

fold 单累加器串行依赖：g.iter().fold(0, |acc, idx| acc + self.get(idx).unwrap()) 形成 loop-carried dependency chain，下一次 add 必须等上一次完成。CPU 无法并行 issue 多个独立的 add 指令，IPC 严重受限于 add latency。f64 的 fadd latency 在鲲鹏 920B 上是 4-5 周期，单累加器把 IPC 压到 ~0.2。
Vec<Vec<u64>> 两级 Vec pointer chase：外层 Vec 顺序访问 OK，但每个 inner Vec<u64> 是独立堆分配，group 之间存在 pointer chase。鲲鹏 920B 的硬件预取器对跨 group 的随机布局无效，反映在 perf stat 上是 920B cache-misses（2.71B）比 9654（1.92B）高 41%。
as_arrow() 在 get(idx) 中重复 downcast：get 内部每次调用都做 Any::downcast_ref + Result::unwrap，在 group size 为数千万行的场景下累计上亿次冗余 downcast。
未利用 SVE 指令集：鲲鹏 920B 拥有 SVE 1.0（VL=128），其中 ld1d {z1.d}, p0/z, [base, z0.d, lsl #3]（向量 gather）和 faddv（FP 横向归约）是 x86 没有高效等价物的指令。LLVM auto-vectorizer 无法为 fold 模式生成这些指令，需要手写内联汇编。
null 路径 Option 模式匹配开销：含 null 时 fold 中每次迭代的 match (acc, self.get(idx)) 都要做 4-way pattern match，分支预测压力大。

优化设计

优化一：SVE gather + 谓词累加 + uaddv 横向归约（i64 / u64）

为 i64 / u64 实现 SVE 内核，使用一条 ld1d 指令一次性 gather 整个 SVE 向量长度的值（920B VL=128 时一次 2 个 i64），用谓词受控的 add 累加，最后用 uaddv 做横向归约。

优化前：

g.iter().fold(0 as i64, |acc, index| {
    let idx = *index as usize;
    acc + self.get(idx).unwrap()    // 单累加器串行依赖 + 重复 downcast
})

优化后（核心 inline asm）：

core::arch::asm!(
    "mov     z2.d, #0",                                  // accumulator = 0
    "ptrue   p1.d",                                      // all-true for reduce
    "whilelt p0.d, {pos}, {n}",                          // loop predicate
    "2:",
    "ld1d    {{z0.d}}, p0/z, [{idx}, {pos}, lsl #3]",    // load indices
    "ld1d    {{z1.d}}, p0/z, [{val}, z0.d, lsl #3]",     // SVE gather (920B-specific)
    "add     z2.d, p0/m, z2.d, z1.d",                    // predicated accumulate
    "incd    {pos}",                                     // VL-agnostic step
    "whilelt p0.d, {pos}, {n}",
    "b.first 2b",                                        // canonical SVE loop tail
    "uaddv   d3, p1, z2.d",                              // horizontal reduce
    "fmov    {sum}, d3",
    ...
);

whilelt + incd + b.first 是 SVE 的标准 vector-length-agnostic 循环模式，自动处理尾部不完整向量。u64 内核复用 i64 内核（整数 add 对有无符号一致）。

优化二：SVE gather + 4 路独立累加器 + faddv（f64）

f64 求和最大瓶颈是 fadd 的 4-5 周期 latency × 单累加器串行依赖。SVE 内核用 4 个独立累加器（z3 / z4 / z5 / z6）打破串行，每 4 个 SVE 向量为一个 chunk 主循环，残余元素用单累加器 tail 循环处理：

// Main loop: 4-way unrolled
"ld1d    {{z0.d}}, p1/z, [{idx}, {pos}, lsl #3]",
"ld1d    {{z1.d}}, p1/z, [{val}, z0.d, lsl #3]",
"fadd    z3.d, p1/m, z3.d, z1.d",     // chain 0
"incd    {pos}",
// ... chunks 1, 2, 3 feeding z4, z5, z6
// Tail loop with whilelt predicate
// Tree reduce 4 accumulators -> 1
"fadd    z3.d, p1/m, z3.d, z4.d",
"fadd    z5.d, p1/m, z5.d, z6.d",
"fadd    z3.d, p1/m, z3.d, z5.d",
"faddv   d7, p1, z3.d",               // horizontal reduce

4 路独立累加器让 CPU 每个周期可以同时 dispatch 4 个 fadd。接受 faddv recursive pairwise reduction 顺序与 strict left-to-right scalar 求和存在 ULP 级数值差异。

优化三：标量 4 路累加器 + PRFM 软件预取（f32 / i128）

f32 和 i128 没有高效的 SVE 路径：

f32：920B VL=128 时只有 4 个 .s lane，且 ld1w 配合 .d 谓词的 gather 需要复杂的 uzp1 / fcvt 数据重排，throughput 增益不抵指令开销
i128：SVE 没有 128-bit 整数 lane

退化为标量 4 路独立累加器 + PRFM 软件预取：

while i + 4 <= n {
    // PRFM with distance 32 for i128 (16 bytes per element)
    if i + 32 < n {
        let pf_idx = *group.get_unchecked(i + 32) as usize;
        prfm_l1_strm(values.add(pf_idx));    // pldl1strm: streaming hint
    }
    s0 = s0.wrapping_add(*values.add(group[i  ] as usize));
    s1 = s1.wrapping_add(*values.add(group[i+1] as usize));
    s2 = s2.wrapping_add(*values.add(group[i+2] as usize));
    s3 = s3.wrapping_add(*values.add(group[i+3] as usize));
    i += 4;
}

LLVM 会自动把 i128 的 wrapping_add 编译成 adds + adc 配对加法，并使用 ldp x, x, [ptr] ARM 配对加载指令。

优化四：跨 group `pldl1keep` 预取

在 grouped_sum 的外层 group 循环中，对未来第 4 个 group 的 indices 头部地址做软件预取：

for (gi, g) in groups.iter().enumerate() {
    if gi + 4 < n_groups {
        unsafe {
            sve_sum::prefetch_next_group(&groups[gi + 4]);   // prfm pldl1keep
        }
    }
    // process current group via SVE kernel
}

pldl1keep（temporal hint）用于 indices 数据，因为它会在内层 SVE 循环中被重复读取，应保留在 L1。value 数据数组的预取在 SVE 内核内部由 pldl1strm（streaming hint）处理。

优化五：消除 `as_arrow()` 重复 downcast + `Vec::with_capacity` 预分配

把原 groups.iter().map(|g| g.iter().fold(...)) 模式改写为显式两层循环，外层一次性 as_arrow() + Vec::with_capacity 预分配：

优化前：

DataArray::<$T>::from_field_and_values(
    self.field.clone(),
    groups.iter().map(|g| {
        g.iter().fold(0, |acc, idx| acc + self.get(*idx as usize).unwrap())
        //                            ^^^ 每次都做 Any::downcast_ref
    }),
)

优化后：

let arrow_array = self.as_arrow()?;
let values_ptr = arrow_array.values().as_ptr();
let nulls = if self.null_count() > 0 { arrow_array.nulls() } else { None };
let mut out: Vec<Option<$AggType>> = Vec::with_capacity(n_groups);
for (gi, g) in groups.iter().enumerate() {
    let v = unsafe { sve_kernel(values_ptr, g) };
    out.push(Some(v));
}
DataArray::<$T>::from_iter(self.field.clone(), out.into_iter())

优化六：null 路径 4 路标量累加器 + any 标志

原始 null 路径用 Option<AggType> + 4-way match 累加，每次迭代都要做 Option 模式匹配 + 两个分支预测。改写为独立的 4 路累加器 + any: bool 标志：

let mut any = false;
let (mut a0, mut a1, mut a2, mut a3) = (0, 0, 0, 0);
while i + 4 <= n {
    if nulls.is_valid(i0) { a0 = a0.wrapping_add(*values.add(i0)); any = true; }
    if nulls.is_valid(i1) { a1 = a1.wrapping_add(*values.add(i1)); any = true; }
    if nulls.is_valid(i2) { a2 = a2.wrapping_add(*values.add(i2)); any = true; }
    if nulls.is_valid(i3) { a3 = a3.wrapping_add(*values.add(i3)); any = true; }
    i += 4;
}
if any { Some(a0.wrapping_add(a1).wrapping_add(a2).wrapping_add(a3)) } else { None }

is_valid(idx) 编译为 bit-test，分支预测器对 "non-null" 占主导的 TPC-H 数据高度准确。

优化推广

所有 6 项优化通过 impl_daft_numeric_agg! 宏统一应用到 5 个数值类型（Int64Type / UInt64Type / Float32Type / Float64Type / Decimal128Type）。aarch64 路径用 #[cfg(target_arch = "aarch64")] 与原 fold 实现隔离，x86 上完全保留原实现，零回归风险。

3.2 Story业务交互流程（必填）

无具体业务交互。

3.3 接口设计

外部调用接口保持一致，DaftSumAggable::grouped_sum 公开 trait 签名不变。

新增内部内核：

sve_grouped_sum_i64 / sve_grouped_sum_i64_with_nulls
sve_grouped_sum_u64 / sve_grouped_sum_u64_with_nulls
sve_grouped_sum_f64 / sve_grouped_sum_f64_with_nulls
sve_grouped_sum_f32 / sve_grouped_sum_f32_with_nulls
sve_grouped_sum_i128 / sve_grouped_sum_i128_with_nulls
prfm_l1_keep / prfm_l1_strm / prfm_l2_keep / prfm_l2_strm 软件预取助手

所有新增内核位于 src/daft-core/src/array/ops/aarch64/ 目录，由 #![cfg(target_arch = "aarch64")] 全局门控。

3.4 DFX可定位性设计

无

3.5 升级兼容性

无

3.6 性能分析

在 TPC-H 基准测试上，针对 Daft v0.7.5 进行编译测试，得到下面的测试结果：

测试用例	优化前(9654)	优化前(920B)	优化前性能比	优化后(9654)	优化后(920B)	优化后性能比
Q1	6.7517066	6.0328852	1.119150519	6.7029808	5.4568594	1.228358715
Q2	1.2203616	1.3354886	0.913794098	1.2060022	1.3435246	0.897640579
Q3	4.4868162	5.6250378	0.797650853	4.4974568	5.6718698	0.792940769
Q4	2.6365156	3.086039	0.854336449	2.6349122	3.1073254	0.847967902
Q5	3.7190698	4.3775678	0.84957446	3.6962438	4.3561414	0.848513274
Q6	1.404526	1.4363802	0.977823281	1.4056444	1.4213296	0.988964418
Q7	4.4859308	6.3077034	0.711182901	4.502043	6.2717898	0.71782428
Q8	5.1249948	6.0524074	0.846769634	5.1174206	6.0664146	0.843565918
Q9	14.2546034	23.202077	0.614367559	14.2671736	23.1268386	0.61690981
Q10	4.8905028	6.1928378	0.789703034	4.8684974	6.2233474	0.782295618
Q11	1.5208956	2.381062	0.638746744	1.5320852	2.3692252	0.646660858
Q12	2.4981312	2.9614004	0.843564146	2.5015758	2.977138	0.840261956
Q13	6.9952248	9.7909056	0.714461469	7.0148916	9.3722464	0.748474944
Q14	2.9131622	3.523356	0.826814605	2.9294578	3.5156314	0.83326648
Q15	2.6283794	2.8118046	0.934766022	2.6303622	2.7290944	0.963822358
Q16	2.16825	3.3735444	0.642721643	2.1925162	3.4125154	0.64249269
Q17	2.6805972	3.3542188	0.799171837	2.6820918	3.3478398	0.801141022
Q18	8.9272486	10.2762256	0.868728359	8.8971816	10.2816132	0.865348796
Q19	4.6196828	4.6451978	0.994507231	4.6389292	4.628015	1.00235829
Q20	3.6656908	4.8202398	0.760478929	3.6535482	4.6340676	0.788410639
Q21	45.9487664	49.9232732	0.920387696	44.8899384	50.4389562	0.889985475
Q22	3.0508056	3.8732616	0.78765803	2.997819	3.8204158	0.78468396
GEOMEAN	-	-	0.818838922	-	-	0.825198708

性能提升：跨平台性能比从 0.818838922 提升至 0.825198708，提升约 0.78%。

3.7 SFMEA分析、测试设计

3.7.1 测试设计

本次优化复用现有测试用例进行验证，无需编写新测试，另外针对 SVE 内核新增 Rust 单元测试。

现有测试覆盖：

测试文件	覆盖的聚合类型	覆盖场景
`tests/dataframe/test_aggregations.py`	sum, count, mean, min, max 等全部聚合	无 group / 单列 group / 多列 group / null 处理 / 字面量聚合 / 表达式聚合
`tests/dataframe/test_decimals.py`	Decimal128 sum	`df.sum()` 全表求和 / `groupby(...).sum()` 分组求和
`tests/dataframe/test_morsels.py`	sum	`df.groupby(...).sum()` 多分区分组求和
`tests/dataframe/test_pivot.py`	sum	pivot 聚合
`tests/dataframe/test_monotonically_increasing_id.py`	sum, min, max	分组聚合 + 表达式聚合

新增单元测试：

测试文件	测试用例数	覆盖类型	覆盖场景
`src/daft-core/src/array/ops/aarch64/tests.rs`	12+	i64 / u64 / f64 / f32 / i128	空 group / 单元素 group / VL 不对齐的尾部 / 大数据量随机索引 / 全 null / 部分 null

测试命令：

安装优化后 wheel 进行以下测试：

# DataFrame 级别聚合测试（覆盖 grouped_sum 的端到端调用路径）
pytest tests/dataframe -k "not lance"

# Rust 层级 SVE 内核单元测试（aarch64 主机上执行）
cargo test -p daft-core --release -- aarch64::tests

关键验证点：

所有 test_aggregations.py 测试通过 → 验证 i64 / u64 / f64 / f32 各类型 grouped_sum 正确性
所有 test_decimals.py 测试通过 → 验证 i128 (Decimal128) grouped_sum 正确性
所有 test_morsels.py 测试通过 → 验证多分区下 grouped_sum 行为一致
所有 SVE 内核单元测试通过 → 验证 SVE 汇编内核的正确性（空 / 单元素 / 长尾 / 大数据量 / null 场景）

Describe alternatives you've considered

Additional Context

Would you like to implement a fix?

Yes

QLiangong成员

5月18日关联了pull request：grouped_sum优化