1212only works inside kernels; device functions must receive it externally.
1313"""
1414
15- import operator
1615from functools import lru_cache
1716
1817from numba_cuda_mlir ._mlir import ir
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-
8173def _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-
152110def 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-
278162def _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
946830def _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