Skip to content

Commit de9d9fe

Browse files
[MLIR] GroupBy JIT cooperative block reductions
Add the pure-MLIR cooperative reduction backend used by the JIT GroupBy.apply engine, as a standalone building block (no GroupBy.apply wiring yet): - cooperative_types: CooperativeArrayType {meminfo, data, size, index} + model. - cooperative_typing: reduction attributes (min/max/sum/mean/var/std/idxmin/ idxmax/corr), unsupported-dtype columns modelled as Poison, and rejection of operators applied directly to a grouped column. - cooperative_device_funcs: per-dtype MLIR reduction modules (shared-memory tree reductions) with NaN skipping (pandas skipna) for min/max/sum, int->int64 sum promotion, float32 mean/var/std, ddof=1 var/std, Pearson corr (int-only), and (value,index) idxmin/idxmax with smallest-index tie-break. - cooperative_lowering: thin lowerings that call the outlined device functions, plus the linker hook that adds the needed PTX at complete(). Self-contained device tests (test_cooperative_reductions.py) launch one cooperative block per reduction and verify the scalar result, covering all four numeric dtypes, NaN skipping, idxmin/idxmax, corr (incl. the float->UDFError case), and operator rejection. compute-sanitizer memcheck clean.
1 parent dac5249 commit de9d9fe

5 files changed

Lines changed: 2799 additions & 0 deletions

File tree

0 commit comments

Comments
 (0)