Skip to content

Commit 1b16a2d

Browse files
Drop NRT-requiring array ops from cooperative groupby backend
The legacy numba-cuda groupby backend was reductions-only. The array- returning cooperative ops (element-wise binops, .exp(), scalar broadcast binops, and cooperative_array_from_ptr) allocated NRT-managed device memory (NRT_MemInfo_new_varsize), a capability the legacy backend never had. Remove those ops entirely. Keep the CooperativeArray struct's meminfo field (null for the non-owning group-column views we build), so the struct layout need not change when a future owning array lands. group + group and friends remain rejected at typing via the UDFError operator layer (see test_operator_on_group_rejected). Removed: _build_{binop,exp,scalar_binop,scalar_rbinop}_module + their func-name helpers + dispatch entries; _lower_binop/_lower_exp/ _lower_scalar_binop + registrations; cooperative_array_from_ptr stub/ lowering; now-orphaned _call_nrt/_arith_op_emit/_shm_ptr_at helpers and broadcast shared-memory slots.
1 parent de9d9fe commit 1b16a2d

3 files changed

Lines changed: 10 additions & 600 deletions

File tree

‎python/cudf/cudf/core/udf/mlir_backend/cooperative_device_funcs.py‎

Lines changed: 2 additions & 325 deletions
Original file line numberDiff line numberDiff line change
@@ -12,7 +12,6 @@
1212
only works inside kernels; device functions must receive it externally.
1313
"""
1414

15-
import operator
1615
from functools import lru_cache
1716

1817
from numba_cuda_mlir._mlir import ir
@@ -27,9 +26,7 @@
2726
# Shared-memory layout (must match cooperative.py)
2827
# ---------------------------------------------------------------------------
2928

30-
_SHM_BROADCAST_OFFSET = 0 # bytes 0..7: data ptr slot
31-
_SHM_MI_SLOT = 8 # bytes 8..15: meminfo ptr slot
32-
_SHM_REDUCTION_OFFSET = 16 # bytes 16+: reduction scratch
29+
_SHM_REDUCTION_OFFSET = 16 # bytes 16+: reduction scratch (0..15 reserved)
3330
_MAX_BLOCK_SIZE = 1024
3431

3532

@@ -73,11 +70,6 @@ def _dtype_tag(dtype):
7370
nb_types.int64: "i64", nb_types.int32: "i32"}[dtype]
7471

7572

76-
def _op_tag(op):
77-
return {operator.add: "add", operator.sub: "sub", operator.mul: "mul",
78-
operator.truediv: "div"}[op]
79-
80-
8173
def _zero_val(mlir_ty):
8274
if isinstance(mlir_ty, ir.F64Type): return arith.constant(mlir_ty, 0.0)
8375
if isinstance(mlir_ty, ir.F32Type): return arith.constant(mlir_ty, 0.0)
@@ -111,44 +103,10 @@ def _bdim_x():
111103
return arith.index_castui(T.i64(), gpu.block_dim("x"))
112104

113105

114-
def _call_nrt(gpu_mod, name, arg_types, result_types, args):
115-
ft = ir.FunctionType.get(arg_types, result_types)
116-
callee = get_or_insert_function(name, ft, gpu_mod)
117-
return func.call(
118-
result=result_types, callee=callee.name.value, operands_=args,
119-
)
120-
121-
122-
def _shm_ptr_at(base_ptr, byte_offset):
123-
"""GEP from raw shm base to a specific byte offset (generic addr space)."""
124-
return llvm.getelementptr(
125-
_ptr(), base_ptr, [_const_i64(byte_offset)],
126-
[GEP_DYNAMIC_INDEX], T.i8(), None,
127-
)
128-
129-
130-
def _arith_op_emit(op, lhs, rhs, mlir_ty):
131-
if isinstance(mlir_ty, (ir.F64Type, ir.F32Type)):
132-
if op is operator.add: return arith.addf(lhs, rhs)
133-
if op is operator.sub: return arith.subf(lhs, rhs)
134-
if op is operator.mul: return arith.mulf(lhs, rhs)
135-
if op is operator.truediv: return arith.divf(lhs, rhs)
136-
elif isinstance(mlir_ty, ir.IntegerType):
137-
if op is operator.add: return arith.addi(lhs, rhs)
138-
if op is operator.sub: return arith.subi(lhs, rhs)
139-
if op is operator.mul: return arith.muli(lhs, rhs)
140-
if op is operator.truediv: return arith.divsi(lhs, rhs)
141-
raise NotImplementedError
142-
143-
144106
# ---------------------------------------------------------------------------
145107
# Function name convention
146108
# ---------------------------------------------------------------------------
147109

148-
def binop_func_name(op, dtype):
149-
return f"_coop_{_op_tag(op)}_{_dtype_tag(dtype)}"
150-
151-
152110
def sum_func_name(dtype):
153111
return f"_coop_sum_{_dtype_tag(dtype)}"
154112

@@ -185,18 +143,6 @@ def idxmin_func_name(dtype):
185143
return f"_coop_idxmin_{_dtype_tag(dtype)}"
186144

187145

188-
def exp_func_name(dtype):
189-
return f"_coop_exp_{_dtype_tag(dtype)}"
190-
191-
192-
def scalar_binop_func_name(op, dtype):
193-
return f"_coop_s{_op_tag(op)}_{_dtype_tag(dtype)}"
194-
195-
196-
def scalar_rbinop_func_name(op, dtype):
197-
return f"_coop_rs{_op_tag(op)}_{_dtype_tag(dtype)}"
198-
199-
200146
# ---------------------------------------------------------------------------
201147
# Shared-memory helpers for address-space-3 access in device functions
202148
# ---------------------------------------------------------------------------
@@ -213,68 +159,6 @@ def _shm3_base(shm_generic_ptr, byte_offset):
213159
# ---------------------------------------------------------------------------
214160
# Module builders: each returns a complete MLIR module string
215161
# ---------------------------------------------------------------------------
216-
217-
def _build_binop_module(op, dtype):
218-
"""Build MLIR module for _coop_{op}_{dtype}(ptr lhs, ptr rhs, i64 size, ptr shm) -> void.
219-
220-
Results written to shm:
221-
shm[0..7] = data pointer (broadcast from thread 0)
222-
shm[8..15] = meminfo pointer (thread 0 only)
223-
"""
224-
elem_bytes = _dtype_byte_size(dtype)
225-
name = binop_func_name(op, dtype)
226-
227-
with context.get_context(), ir.Location.unknown():
228-
elem_ty = _dtype_mlir_type(dtype)
229-
module = ir.Module.create()
230-
with ir.InsertionPoint(module.body):
231-
gpu_mod = gpu.GPUModuleOp(sym_name=f"{name}_module")
232-
233-
body = gpu_mod.bodyRegion.blocks.append()
234-
ft = ir.FunctionType.get([_ptr(), _ptr(), T.i64(), _ptr()], [])
235-
with ir.InsertionPoint(body):
236-
fn = func.FuncOp(name=name, type=ft, visibility="public")
237-
entry = fn.add_entry_block()
238-
239-
with ir.InsertionPoint(entry):
240-
lhs_data, rhs_data, size, shm_base = fn.arguments
241-
242-
data_slot = _shm_ptr_at(shm_base, _SHM_BROADCAST_OFFSET)
243-
mi_slot = _shm_ptr_at(shm_base, _SHM_MI_SLOT)
244-
245-
tid = _tid_x()
246-
is_t0 = arith.cmpi(arith.CmpIPredicate.eq, tid, _const_i64(0))
247-
alloc_bytes = arith.muli(size, _const_i64(elem_bytes))
248-
alloc_bytes = arith.maxui(alloc_bytes, _const_i64(1))
249-
250-
t0_if = scf.IfOp(is_t0, results_=[], has_else=False)
251-
with ir.InsertionPoint(t0_if.then_block):
252-
mi_val = _call_nrt(gpu_mod, "NRT_MemInfo_new_varsize",
253-
[T.i64()], [_ptr()], [alloc_bytes])
254-
data_val = _call_nrt(gpu_mod, "NRT_MemInfo_data_fast",
255-
[_ptr()], [_ptr()], [mi_val])
256-
llvm.store(data_val, data_slot)
257-
llvm.store(mi_val, mi_slot)
258-
scf.YieldOp([])
259-
gpu.barrier()
260-
261-
out_data = llvm.load(_ptr(), data_slot)
262-
bdim = _bdim_x()
263-
loop = scf.ForOp(tid, size, bdim, [])
264-
with ir.InsertionPoint(loop.body):
265-
i = loop.induction_variable
266-
a = llvm.load(elem_ty, _gep_typed(lhs_data, i, elem_ty))
267-
b = llvm.load(elem_ty, _gep_typed(rhs_data, i, elem_ty))
268-
r = _arith_op_emit(op, a, b, elem_ty)
269-
llvm.store(r, _gep_typed(out_data, i, elem_ty))
270-
scf.YieldOp([])
271-
gpu.barrier()
272-
273-
func.ReturnOp([])
274-
275-
return str(module)
276-
277-
278162
def _build_sum_module(dtype):
279163
"""Build MLIR module for _coop_sum_{dtype}(ptr data, i64 size, ptr shm) -> scalar.
280164
@@ -945,191 +829,6 @@ def _build_idxmax_module(dtype):
945829

946830
def _build_idxmin_module(dtype):
947831
return _build_idx_module(dtype, is_max=False)
948-
949-
950-
def _build_exp_module(dtype):
951-
"""Build MLIR module for _coop_exp_{dtype}(ptr data, i64 size, ptr shm) -> void.
952-
953-
Element-wise: out[i] = exp(in[i]).
954-
Results written to shm (same layout as binop).
955-
"""
956-
elem_bytes = _dtype_byte_size(dtype)
957-
name = exp_func_name(dtype)
958-
959-
with context.get_context(), ir.Location.unknown():
960-
from numba_cuda_mlir._mlir.dialects import math as mlir_math
961-
962-
elem_ty = _dtype_mlir_type(dtype)
963-
module = ir.Module.create()
964-
with ir.InsertionPoint(module.body):
965-
gpu_mod = gpu.GPUModuleOp(sym_name=f"{name}_module")
966-
967-
body = gpu_mod.bodyRegion.blocks.append()
968-
ft = ir.FunctionType.get([_ptr(), T.i64(), _ptr()], [])
969-
with ir.InsertionPoint(body):
970-
fn = func.FuncOp(name=name, type=ft, visibility="public")
971-
entry = fn.add_entry_block()
972-
973-
with ir.InsertionPoint(entry):
974-
in_data, size, shm_base = fn.arguments
975-
976-
data_slot = _shm_ptr_at(shm_base, _SHM_BROADCAST_OFFSET)
977-
mi_slot = _shm_ptr_at(shm_base, _SHM_MI_SLOT)
978-
979-
tid = _tid_x()
980-
is_t0 = arith.cmpi(arith.CmpIPredicate.eq, tid, _const_i64(0))
981-
alloc_bytes = arith.muli(size, _const_i64(elem_bytes))
982-
alloc_bytes = arith.maxui(alloc_bytes, _const_i64(1))
983-
984-
t0_if = scf.IfOp(is_t0, results_=[], has_else=False)
985-
with ir.InsertionPoint(t0_if.then_block):
986-
mi_val = _call_nrt(gpu_mod, "NRT_MemInfo_new_varsize",
987-
[T.i64()], [_ptr()], [alloc_bytes])
988-
data_val = _call_nrt(gpu_mod, "NRT_MemInfo_data_fast",
989-
[_ptr()], [_ptr()], [mi_val])
990-
llvm.store(data_val, data_slot)
991-
llvm.store(mi_val, mi_slot)
992-
scf.YieldOp([])
993-
gpu.barrier()
994-
995-
out_data = llvm.load(_ptr(), data_slot)
996-
bdim = _bdim_x()
997-
loop = scf.ForOp(tid, size, bdim, [])
998-
with ir.InsertionPoint(loop.body):
999-
i = loop.induction_variable
1000-
a = llvm.load(elem_ty, _gep_typed(in_data, i, elem_ty))
1001-
if isinstance(elem_ty, (ir.F64Type, ir.F32Type)):
1002-
r = mlir_math.exp(a)
1003-
else:
1004-
a_f = arith.sitofp(elem_ty, a)
1005-
r = mlir_math.exp(a_f)
1006-
llvm.store(r, _gep_typed(out_data, i, elem_ty))
1007-
scf.YieldOp([])
1008-
gpu.barrier()
1009-
1010-
func.ReturnOp([])
1011-
1012-
return str(module)
1013-
1014-
1015-
def _build_scalar_binop_module(op, dtype):
1016-
"""Build MLIR module for _coop_s{op}_{dtype}(ptr arr, scalar, i64 size, ptr shm) -> void.
1017-
1018-
Element-wise: out[i] = arr[i] op scalar.
1019-
Results written to shm (same layout as binop).
1020-
"""
1021-
elem_bytes = _dtype_byte_size(dtype)
1022-
name = scalar_binop_func_name(op, dtype)
1023-
1024-
with context.get_context(), ir.Location.unknown():
1025-
elem_ty = _dtype_mlir_type(dtype)
1026-
module = ir.Module.create()
1027-
with ir.InsertionPoint(module.body):
1028-
gpu_mod = gpu.GPUModuleOp(sym_name=f"{name}_module")
1029-
1030-
body = gpu_mod.bodyRegion.blocks.append()
1031-
ft = ir.FunctionType.get([_ptr(), elem_ty, T.i64(), _ptr()], [])
1032-
with ir.InsertionPoint(body):
1033-
fn = func.FuncOp(name=name, type=ft, visibility="public")
1034-
entry = fn.add_entry_block()
1035-
1036-
with ir.InsertionPoint(entry):
1037-
arr_data, scalar_val, size, shm_base = fn.arguments
1038-
1039-
data_slot = _shm_ptr_at(shm_base, _SHM_BROADCAST_OFFSET)
1040-
mi_slot = _shm_ptr_at(shm_base, _SHM_MI_SLOT)
1041-
1042-
tid = _tid_x()
1043-
is_t0 = arith.cmpi(arith.CmpIPredicate.eq, tid, _const_i64(0))
1044-
alloc_bytes = arith.muli(size, _const_i64(elem_bytes))
1045-
alloc_bytes = arith.maxui(alloc_bytes, _const_i64(1))
1046-
1047-
t0_if = scf.IfOp(is_t0, results_=[], has_else=False)
1048-
with ir.InsertionPoint(t0_if.then_block):
1049-
mi_val = _call_nrt(gpu_mod, "NRT_MemInfo_new_varsize",
1050-
[T.i64()], [_ptr()], [alloc_bytes])
1051-
data_val = _call_nrt(gpu_mod, "NRT_MemInfo_data_fast",
1052-
[_ptr()], [_ptr()], [mi_val])
1053-
llvm.store(data_val, data_slot)
1054-
llvm.store(mi_val, mi_slot)
1055-
scf.YieldOp([])
1056-
gpu.barrier()
1057-
1058-
out_data = llvm.load(_ptr(), data_slot)
1059-
bdim = _bdim_x()
1060-
loop = scf.ForOp(tid, size, bdim, [])
1061-
with ir.InsertionPoint(loop.body):
1062-
i = loop.induction_variable
1063-
a = llvm.load(elem_ty, _gep_typed(arr_data, i, elem_ty))
1064-
r = _arith_op_emit(op, a, scalar_val, elem_ty)
1065-
llvm.store(r, _gep_typed(out_data, i, elem_ty))
1066-
scf.YieldOp([])
1067-
gpu.barrier()
1068-
1069-
func.ReturnOp([])
1070-
1071-
return str(module)
1072-
1073-
1074-
def _build_scalar_rbinop_module(op, dtype):
1075-
"""Build MLIR module for _coop_rs{op}_{dtype}(ptr arr, scalar, i64 size, ptr shm) -> void.
1076-
1077-
Reversed element-wise: out[i] = scalar op arr[i].
1078-
Needed for non-commutative ops (sub, div) when scalar is on the LHS.
1079-
"""
1080-
elem_bytes = _dtype_byte_size(dtype)
1081-
name = scalar_rbinop_func_name(op, dtype)
1082-
1083-
with context.get_context(), ir.Location.unknown():
1084-
elem_ty = _dtype_mlir_type(dtype)
1085-
module = ir.Module.create()
1086-
with ir.InsertionPoint(module.body):
1087-
gpu_mod = gpu.GPUModuleOp(sym_name=f"{name}_module")
1088-
1089-
body = gpu_mod.bodyRegion.blocks.append()
1090-
ft = ir.FunctionType.get([_ptr(), elem_ty, T.i64(), _ptr()], [])
1091-
with ir.InsertionPoint(body):
1092-
fn = func.FuncOp(name=name, type=ft, visibility="public")
1093-
entry = fn.add_entry_block()
1094-
1095-
with ir.InsertionPoint(entry):
1096-
arr_data, scalar_val, size, shm_base = fn.arguments
1097-
1098-
data_slot = _shm_ptr_at(shm_base, _SHM_BROADCAST_OFFSET)
1099-
mi_slot = _shm_ptr_at(shm_base, _SHM_MI_SLOT)
1100-
1101-
tid = _tid_x()
1102-
is_t0 = arith.cmpi(arith.CmpIPredicate.eq, tid, _const_i64(0))
1103-
alloc_bytes = arith.muli(size, _const_i64(elem_bytes))
1104-
alloc_bytes = arith.maxui(alloc_bytes, _const_i64(1))
1105-
1106-
t0_if = scf.IfOp(is_t0, results_=[], has_else=False)
1107-
with ir.InsertionPoint(t0_if.then_block):
1108-
mi_val = _call_nrt(gpu_mod, "NRT_MemInfo_new_varsize",
1109-
[T.i64()], [_ptr()], [alloc_bytes])
1110-
data_val = _call_nrt(gpu_mod, "NRT_MemInfo_data_fast",
1111-
[_ptr()], [_ptr()], [mi_val])
1112-
llvm.store(data_val, data_slot)
1113-
llvm.store(mi_val, mi_slot)
1114-
scf.YieldOp([])
1115-
gpu.barrier()
1116-
1117-
out_data = llvm.load(_ptr(), data_slot)
1118-
bdim = _bdim_x()
1119-
loop = scf.ForOp(tid, size, bdim, [])
1120-
with ir.InsertionPoint(loop.body):
1121-
i = loop.induction_variable
1122-
a = llvm.load(elem_ty, _gep_typed(arr_data, i, elem_ty))
1123-
r = _arith_op_emit(op, scalar_val, a, elem_ty)
1124-
llvm.store(r, _gep_typed(out_data, i, elem_ty))
1125-
scf.YieldOp([])
1126-
gpu.barrier()
1127-
1128-
func.ReturnOp([])
1129-
1130-
return str(module)
1131-
1132-
1133832
# ---------------------------------------------------------------------------
1134833
# Compilation: MLIR module string -> PTX bytes
1135834
# ---------------------------------------------------------------------------
@@ -1182,13 +881,7 @@ def _compile_cooperative_func(func_key, cc):
1182881
Returns PTX bytes.
1183882
"""
1184883
kind = func_key[0]
1185-
if kind == "binop":
1186-
_, op_name, dtype_name = func_key
1187-
op = {"add": operator.add, "sub": operator.sub, "mul": operator.mul,
1188-
"div": operator.truediv}[op_name]
1189-
dtype = _TAG_TO_DTYPE[dtype_name]
1190-
mlir_str = _build_binop_module(op, dtype)
1191-
elif kind == "sum":
884+
if kind == "sum":
1192885
_, dtype_name = func_key
1193886
dtype = _TAG_TO_DTYPE[dtype_name]
1194887
mlir_str = _build_sum_module(dtype)
@@ -1224,22 +917,6 @@ def _compile_cooperative_func(func_key, cc):
1224917
_, dtype_name = func_key
1225918
dtype = _TAG_TO_DTYPE[dtype_name]
1226919
mlir_str = _build_idxmin_module(dtype)
1227-
elif kind == "exp":
1228-
_, dtype_name = func_key
1229-
dtype = _TAG_TO_DTYPE[dtype_name]
1230-
mlir_str = _build_exp_module(dtype)
1231-
elif kind == "scalar_binop":
1232-
_, op_name, dtype_name = func_key
1233-
op = {"add": operator.add, "sub": operator.sub, "mul": operator.mul,
1234-
"div": operator.truediv}[op_name]
1235-
dtype = _TAG_TO_DTYPE[dtype_name]
1236-
mlir_str = _build_scalar_binop_module(op, dtype)
1237-
elif kind == "scalar_rbinop":
1238-
_, op_name, dtype_name = func_key
1239-
op = {"add": operator.add, "sub": operator.sub, "mul": operator.mul,
1240-
"div": operator.truediv}[op_name]
1241-
dtype = _TAG_TO_DTYPE[dtype_name]
1242-
mlir_str = _build_scalar_rbinop_module(op, dtype)
1243920
else:
1244921
raise ValueError(f"Unknown func_key kind: {kind}")
1245922

0 commit comments

Comments
 (0)