Skip to content

Commit bab707d

Browse files
ngoldbaumngoldbaum
authored
Merge pull request #78 from SwayamInSync/promoter-fix
[BUG]: Avoiding promoter overrides leading to global state corruption
2 parents c4e39c6 + 673d115 commit bab707d

7 files changed

Lines changed: 175 additions & 15 deletions

File tree

‎src/csrc/umath/binary_ops.cpp‎

Lines changed: 56 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -296,8 +296,7 @@ quad_ldexp_resolve_descriptors(PyObject *self, PyArray_DTypeMeta *const dtypes[]
296296
Py_INCREF(given_descrs[0]);
297297
loop_descrs[0] = given_descrs[0];
298298

299-
// Input 1: Use NPY_INTP (int64 on 64-bit, int32 on 32-bit) to match platform integer size
300-
// This ensures we can handle the full range of PyArray_PyLongDType without data loss
299+
// Input 1: Use NPY_INTP to match the registered PyArray_IntpDType
301300
loop_descrs[1] = PyArray_DescrFromType(NPY_INTP);
302301

303302
// Output: QuadPrecDType with same backend as input
@@ -408,7 +407,7 @@ create_quad_ldexp_ufunc(PyObject *numpy, const char *ufunc_name)
408407
return -1;
409408
}
410409

411-
PyArray_DTypeMeta *dtypes[3] = {&QuadPrecDType, &PyArray_PyLongDType, &QuadPrecDType};
410+
PyArray_DTypeMeta *dtypes[3] = {&QuadPrecDType, &PyArray_IntpDType, &QuadPrecDType};
412411

413412
PyType_Slot slots[] = {
414413
{NPY_METH_resolve_descriptors, (void *)&quad_ldexp_resolve_descriptors},
@@ -429,28 +428,33 @@ create_quad_ldexp_ufunc(PyObject *numpy, const char *ufunc_name)
429428
};
430429

431430
if (PyUFunc_AddLoopFromSpec(ufunc, &Spec) < 0) {
431+
Py_DECREF(ufunc);
432432
return -1;
433433
}
434434

435435
PyObject *promoter_capsule =
436-
PyCapsule_New((void *)&quad_ufunc_promoter, "numpy._ufunc_promoter", NULL);
436+
PyCapsule_New((void *)&quad_ldexp_promoter, "numpy._ufunc_promoter", NULL);
437437
if (promoter_capsule == NULL) {
438+
Py_DECREF(ufunc);
438439
return -1;
439440
}
440441

441-
PyObject *DTypes = PyTuple_Pack(3, &PyArrayDescr_Type, &PyArray_PyLongDType, &PyArrayDescr_Type);
442+
PyObject *DTypes = PyTuple_Pack(3, &QuadPrecDType, &PyArrayDescr_Type, &PyArrayDescr_Type);
442443
if (DTypes == 0) {
443444
Py_DECREF(promoter_capsule);
445+
Py_DECREF(ufunc);
444446
return -1;
445447
}
446448

447449
if (PyUFunc_AddPromoter(ufunc, DTypes, promoter_capsule) < 0) {
448450
Py_DECREF(promoter_capsule);
449451
Py_DECREF(DTypes);
452+
Py_DECREF(ufunc);
450453
return -1;
451454
}
452455
Py_DECREF(promoter_capsule);
453456
Py_DECREF(DTypes);
457+
Py_DECREF(ufunc);
454458
return 0;
455459
}
456460

@@ -486,29 +490,52 @@ create_quad_binary_2out_ufunc(PyObject *numpy, const char *ufunc_name)
486490
};
487491

488492
if (PyUFunc_AddLoopFromSpec(ufunc, &Spec) < 0) {
493+
Py_DECREF(ufunc);
489494
return -1;
490495
}
491496

492497
PyObject *promoter_capsule =
493498
PyCapsule_New((void *)&quad_ufunc_promoter, "numpy._ufunc_promoter", NULL);
494499
if (promoter_capsule == NULL) {
500+
Py_DECREF(ufunc);
495501
return -1;
496502
}
497503

498-
PyObject *DTypes = PyTuple_Pack(4, &PyArrayDescr_Type, &PyArrayDescr_Type,
504+
// Register promoter for (QuadPrecDType, Any, Any, Any)
505+
PyObject *DTypes = PyTuple_Pack(4, &QuadPrecDType, &PyArrayDescr_Type,
499506
&PyArrayDescr_Type, &PyArrayDescr_Type);
500507
if (DTypes == 0) {
501508
Py_DECREF(promoter_capsule);
509+
Py_DECREF(ufunc);
510+
return -1;
511+
}
512+
513+
if (PyUFunc_AddPromoter(ufunc, DTypes, promoter_capsule) < 0) {
514+
Py_DECREF(promoter_capsule);
515+
Py_DECREF(DTypes);
516+
Py_DECREF(ufunc);
517+
return -1;
518+
}
519+
Py_DECREF(DTypes);
520+
521+
// Register promoter for (Any, QuadPrecDType, Any, Any)
522+
DTypes = PyTuple_Pack(4, &PyArrayDescr_Type, &QuadPrecDType,
523+
&PyArrayDescr_Type, &PyArrayDescr_Type);
524+
if (DTypes == 0) {
525+
Py_DECREF(promoter_capsule);
526+
Py_DECREF(ufunc);
502527
return -1;
503528
}
504529

505530
if (PyUFunc_AddPromoter(ufunc, DTypes, promoter_capsule) < 0) {
506531
Py_DECREF(promoter_capsule);
507532
Py_DECREF(DTypes);
533+
Py_DECREF(ufunc);
508534
return -1;
509535
}
510536
Py_DECREF(promoter_capsule);
511537
Py_DECREF(DTypes);
538+
Py_DECREF(ufunc);
512539
return 0;
513540
}
514541

@@ -542,28 +569,50 @@ create_quad_binary_ufunc(PyObject *numpy, const char *ufunc_name)
542569
};
543570

544571
if (PyUFunc_AddLoopFromSpec(ufunc, &Spec) < 0) {
572+
Py_DECREF(ufunc);
545573
return -1;
546574
}
547575

548576
PyObject *promoter_capsule =
549577
PyCapsule_New((void *)&quad_ufunc_promoter, "numpy._ufunc_promoter", NULL);
550578
if (promoter_capsule == NULL) {
579+
Py_DECREF(ufunc);
551580
return -1;
552581
}
553582

554-
PyObject *DTypes = PyTuple_Pack(3, &PyArrayDescr_Type, &PyArrayDescr_Type, &PyArrayDescr_Type);
583+
// Register promoter for (QuadPrecDType, Any, Any)
584+
PyObject *DTypes = PyTuple_Pack(3, &QuadPrecDType, &PyArrayDescr_Type, &PyArrayDescr_Type);
585+
if (DTypes == 0) {
586+
Py_DECREF(promoter_capsule);
587+
Py_DECREF(ufunc);
588+
return -1;
589+
}
590+
591+
if (PyUFunc_AddPromoter(ufunc, DTypes, promoter_capsule) < 0) {
592+
Py_DECREF(promoter_capsule);
593+
Py_DECREF(DTypes);
594+
Py_DECREF(ufunc);
595+
return -1;
596+
}
597+
Py_DECREF(DTypes);
598+
599+
// Register promoter for (Any, QuadPrecDType, Any)
600+
DTypes = PyTuple_Pack(3, &PyArrayDescr_Type, &QuadPrecDType, &PyArrayDescr_Type);
555601
if (DTypes == 0) {
556602
Py_DECREF(promoter_capsule);
603+
Py_DECREF(ufunc);
557604
return -1;
558605
}
559606

560607
if (PyUFunc_AddPromoter(ufunc, DTypes, promoter_capsule) < 0) {
561608
Py_DECREF(promoter_capsule);
562609
Py_DECREF(DTypes);
610+
Py_DECREF(ufunc);
563611
return -1;
564612
}
565613
Py_DECREF(promoter_capsule);
566614
Py_DECREF(DTypes);
615+
Py_DECREF(ufunc);
567616
return 0;
568617
}
569618

‎src/csrc/umath/comparison_ops.cpp‎

Lines changed: 12 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -257,16 +257,17 @@ quad_reduce_comp_strided_loop_unaligned(PyArrayMethod_Context *context, char *co
257257

258258

259259
NPY_NO_EXPORT int
260-
comparison_ufunc_promoter(PyUFuncObject *ufunc, PyArray_DTypeMeta *op_dtypes[],
261-
PyArray_DTypeMeta *signature[], PyArray_DTypeMeta *new_op_dtypes[])
260+
comparison_ufunc_promoter(PyObject *ufunc_obj, PyArray_DTypeMeta *const op_dtypes[],
261+
PyArray_DTypeMeta *const signature[], PyArray_DTypeMeta *new_op_dtypes[])
262262
{
263263
PyArray_DTypeMeta *new_signature[NPY_MAXARGS];
264264
memcpy(new_signature, signature, 3 * sizeof(PyArray_DTypeMeta *));
265265
new_signature[2] = NULL;
266-
int res = quad_ufunc_promoter(ufunc, op_dtypes, new_signature, new_op_dtypes);
266+
int res = quad_ufunc_promoter(ufunc_obj, op_dtypes, new_signature, new_op_dtypes);
267267
if (res < 0) {
268268
return -1;
269269
}
270+
Py_INCREF(&PyArray_BoolDType);
270271
Py_XSETREF(new_op_dtypes[2], &PyArray_BoolDType);
271272
return 0;
272273
}
@@ -301,6 +302,7 @@ create_quad_comparison_ufunc(PyObject *numpy, const char *ufunc_name)
301302
};
302303

303304
if (PyUFunc_AddLoopFromSpec(ufunc, &Spec) < 0) {
305+
Py_DECREF(ufunc);
304306
return -1;
305307
}
306308

@@ -326,25 +328,29 @@ create_quad_comparison_ufunc(PyObject *numpy, const char *ufunc_name)
326328
};
327329

328330
if (PyUFunc_AddLoopFromSpec(ufunc, &Spec_reduce) < 0) {
331+
Py_DECREF(ufunc);
329332
return -1;
330333
}
331334

332335
PyObject *promoter_capsule =
333336
PyCapsule_New((void *)&comparison_ufunc_promoter, "numpy._ufunc_promoter", NULL);
334337
if (promoter_capsule == NULL) {
338+
Py_DECREF(ufunc);
335339
return -1;
336340
}
337341

338342
// Register promoter for (QuadPrecDType, Any, Bool) - needed for mixed-type comparisons
339343
PyObject *DTypes = PyTuple_Pack(3, &QuadPrecDType, &PyArrayDescr_Type, &PyArray_BoolDType);
340344
if (DTypes == 0) {
341345
Py_DECREF(promoter_capsule);
346+
Py_DECREF(ufunc);
342347
return -1;
343348
}
344349

345350
if (PyUFunc_AddPromoter(ufunc, DTypes, promoter_capsule) < 0) {
346351
Py_DECREF(promoter_capsule);
347352
Py_DECREF(DTypes);
353+
Py_DECREF(ufunc);
348354
return -1;
349355
}
350356
Py_DECREF(DTypes);
@@ -353,16 +359,19 @@ create_quad_comparison_ufunc(PyObject *numpy, const char *ufunc_name)
353359
DTypes = PyTuple_Pack(3, &PyArrayDescr_Type, &QuadPrecDType, &PyArray_BoolDType);
354360
if (DTypes == 0) {
355361
Py_DECREF(promoter_capsule);
362+
Py_DECREF(ufunc);
356363
return -1;
357364
}
358365

359366
if (PyUFunc_AddPromoter(ufunc, DTypes, promoter_capsule) < 0) {
360367
Py_DECREF(promoter_capsule);
361368
Py_DECREF(DTypes);
369+
Py_DECREF(ufunc);
362370
return -1;
363371
}
364372
Py_DECREF(promoter_capsule);
365373
Py_DECREF(DTypes);
374+
Py_DECREF(ufunc);
366375

367376
return 0;
368377
}

‎src/csrc/umath/matmul.cpp‎

Lines changed: 14 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -431,7 +431,8 @@ init_matmul_ops(PyObject *numpy)
431431
return -1;
432432
}
433433

434-
PyObject *DTypes = PyTuple_Pack(3, &PyArrayDescr_Type, &PyArrayDescr_Type, &PyArrayDescr_Type);
434+
// Register promoter for (QuadPrecDType, Any, Any)
435+
PyObject *DTypes = PyTuple_Pack(3, &QuadPrecDType, &PyArrayDescr_Type, &PyArrayDescr_Type);
435436
if (DTypes == NULL) {
436437
Py_DECREF(promoter_capsule);
437438
Py_DECREF(ufunc);
@@ -441,10 +442,21 @@ init_matmul_ops(PyObject *numpy)
441442
if (PyUFunc_AddPromoter(ufunc, DTypes, promoter_capsule) < 0) {
442443
PyErr_Clear();
443444
}
444-
else {
445+
Py_DECREF(DTypes);
446+
447+
// Register promoter for (Any, QuadPrecDType, Any)
448+
DTypes = PyTuple_Pack(3, &PyArrayDescr_Type, &QuadPrecDType, &PyArrayDescr_Type);
449+
if (DTypes == NULL) {
450+
Py_DECREF(promoter_capsule);
451+
Py_DECREF(ufunc);
452+
return -1;
445453
}
446454

455+
if (PyUFunc_AddPromoter(ufunc, DTypes, promoter_capsule) < 0) {
456+
PyErr_Clear();
457+
}
447458
Py_DECREF(DTypes);
459+
448460
Py_DECREF(promoter_capsule);
449461
Py_DECREF(ufunc);
450462

‎src/csrc/umath/unary_ops.cpp‎

Lines changed: 8 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -138,9 +138,11 @@ create_quad_unary_ufunc(PyObject *numpy, const char *ufunc_name)
138138
};
139139

140140
if (PyUFunc_AddLoopFromSpec(ufunc, &Spec) < 0) {
141+
Py_DECREF(ufunc);
141142
return -1;
142143
}
143144

145+
Py_DECREF(ufunc);
144146
return 0;
145147
}
146148

@@ -261,9 +263,11 @@ create_quad_logical_not_ufunc(PyObject *numpy, const char *ufunc_name)
261263
};
262264

263265
if (PyUFunc_AddLoopFromSpec(ufunc, &Spec) < 0) {
266+
Py_DECREF(ufunc);
264267
return -1;
265268
}
266269

270+
Py_DECREF(ufunc);
267271
return 0;
268272
}
269273

@@ -402,9 +406,11 @@ create_quad_unary_2out_ufunc(PyObject *numpy, const char *ufunc_name)
402406
};
403407

404408
if (PyUFunc_AddLoopFromSpec(ufunc, &Spec) < 0) {
409+
Py_DECREF(ufunc);
405410
return -1;
406411
}
407412

413+
Py_DECREF(ufunc);
408414
return 0;
409415
}
410416

@@ -561,9 +567,11 @@ create_quad_frexp_ufunc(PyObject *numpy, const char *ufunc_name)
561567
};
562568

563569
if (PyUFunc_AddLoopFromSpec(ufunc, &Spec) < 0) {
570+
Py_DECREF(ufunc);
564571
return -1;
565572
}
566573

574+
Py_DECREF(ufunc);
567575
return 0;
568576
}
569577

‎src/csrc/umath/unary_props.cpp‎

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -127,9 +127,11 @@ create_quad_unary_prop_ufunc(PyObject *numpy, const char *ufunc_name)
127127
};
128128

129129
if (PyUFunc_AddLoopFromSpec(ufunc, &Spec) < 0) {
130+
Py_DECREF(ufunc);
130131
return -1;
131132
}
132133

134+
Py_DECREF(ufunc);
133135
return 0;
134136
}
135137

‎src/include/umath/promoters.hpp‎

Lines changed: 27 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -12,9 +12,10 @@
1212
#include "../dtype.h"
1313

1414
inline int
15-
quad_ufunc_promoter(PyUFuncObject *ufunc, PyArray_DTypeMeta *op_dtypes[],
16-
PyArray_DTypeMeta *signature[], PyArray_DTypeMeta *new_op_dtypes[])
15+
quad_ufunc_promoter(PyObject *ufunc_obj, PyArray_DTypeMeta *const op_dtypes[],
16+
PyArray_DTypeMeta *const signature[], PyArray_DTypeMeta *new_op_dtypes[])
1717
{
18+
PyUFuncObject *ufunc = (PyUFuncObject *)ufunc_obj;
1819
int nin = ufunc->nin;
1920
int nargs = ufunc->nargs;
2021
PyArray_DTypeMeta *common = NULL;
@@ -56,7 +57,7 @@ quad_ufunc_promoter(PyUFuncObject *ufunc, PyArray_DTypeMeta *op_dtypes[],
5657
}
5758
// If no common output dtype, use standard promotion for inputs
5859
if (common == NULL) {
59-
common = PyArray_PromoteDTypeSequence(nin, op_dtypes);
60+
common = PyArray_PromoteDTypeSequence(nin, const_cast<PyArray_DTypeMeta **>(op_dtypes));
6061
if (common == NULL) {
6162
if (PyErr_ExceptionMatches(PyExc_TypeError)) {
6263
PyErr_Clear(); // Do not propagate normal promotion errors
@@ -87,4 +88,27 @@ quad_ufunc_promoter(PyUFuncObject *ufunc, PyArray_DTypeMeta *op_dtypes[],
8788
}
8889

8990

91+
inline int
92+
quad_ldexp_promoter(PyObject *ufunc_obj, PyArray_DTypeMeta *const op_dtypes[],
93+
PyArray_DTypeMeta *const signature[], PyArray_DTypeMeta *new_op_dtypes[])
94+
{
95+
Py_INCREF(&QuadPrecDType);
96+
new_op_dtypes[0] = &QuadPrecDType;
97+
98+
// Promote the exponent to PyArray_IntpDType (unless signature specifies otherwise)
99+
if (signature[1] != NULL) {
100+
Py_INCREF(signature[1]);
101+
new_op_dtypes[1] = signature[1];
102+
}
103+
else {
104+
Py_INCREF(&PyArray_IntpDType);
105+
new_op_dtypes[1] = &PyArray_IntpDType;
106+
}
107+
108+
Py_INCREF(&QuadPrecDType);
109+
new_op_dtypes[2] = &QuadPrecDType;
110+
111+
return 0;
112+
}
113+
90114
#endif

0 commit comments

Comments
 (0)