Repository navigation
Commit de9d9fe
committed
[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
- python/cudf/cudf
- core/udf/mlir_backend
- tests/private_objects/mlir_backend
0 commit comments