diff --git a/src/csrc/umath/binary_ops.cpp b/src/csrc/umath/binary_ops.cpp index 9f27ee6..56aaf55 100644 --- a/src/csrc/umath/binary_ops.cpp +++ b/src/csrc/umath/binary_ops.cpp @@ -438,7 +438,7 @@ create_quad_ldexp_ufunc(PyObject *numpy, const char *ufunc_name) } PyObject *DTypes = PyTuple_Pack(3, &QuadPrecDType, &PyArrayDescr_Type, &PyArrayDescr_Type); - if (DTypes == 0) { + if (DTypes == NULL) { Py_DECREF(promoter_capsule); Py_DECREF(ufunc); return -1; @@ -502,7 +502,7 @@ create_quad_binary_2out_ufunc(PyObject *numpy, const char *ufunc_name) // Register promoter for (QuadPrecDType, Any, Any, Any) PyObject *DTypes = PyTuple_Pack(4, &QuadPrecDType, &PyArrayDescr_Type, &PyArrayDescr_Type, &PyArrayDescr_Type); - if (DTypes == 0) { + if (DTypes == NULL) { Py_DECREF(promoter_capsule); Py_DECREF(ufunc); return -1; @@ -519,7 +519,7 @@ create_quad_binary_2out_ufunc(PyObject *numpy, const char *ufunc_name) // Register promoter for (Any, QuadPrecDType, Any, Any) DTypes = PyTuple_Pack(4, &PyArrayDescr_Type, &QuadPrecDType, &PyArrayDescr_Type, &PyArrayDescr_Type); - if (DTypes == 0) { + if (DTypes == NULL) { Py_DECREF(promoter_capsule); Py_DECREF(ufunc); return -1; @@ -533,6 +533,7 @@ create_quad_binary_2out_ufunc(PyObject *numpy, const char *ufunc_name) } Py_DECREF(promoter_capsule); Py_DECREF(DTypes); + Py_DECREF(ufunc); return 0; } @@ -580,7 +581,7 @@ create_quad_binary_ufunc(PyObject *numpy, const char *ufunc_name) // Register promoter for (QuadPrecDType, Any, Any) PyObject *DTypes = PyTuple_Pack(3, &QuadPrecDType, &PyArrayDescr_Type, &PyArrayDescr_Type); - if (DTypes == 0) { + if (DTypes == NULL) { Py_DECREF(promoter_capsule); Py_DECREF(ufunc); return -1; @@ -596,7 +597,7 @@ create_quad_binary_ufunc(PyObject *numpy, const char *ufunc_name) // Register promoter for (Any, QuadPrecDType, Any) DTypes = PyTuple_Pack(3, &PyArrayDescr_Type, &QuadPrecDType, &PyArrayDescr_Type); - if (DTypes == 0) { + if (DTypes == NULL) { Py_DECREF(promoter_capsule); Py_DECREF(ufunc); return -1; @@ -610,6 +611,7 @@ create_quad_binary_ufunc(PyObject *numpy, const char *ufunc_name) } Py_DECREF(promoter_capsule); Py_DECREF(DTypes); + Py_DECREF(ufunc); return 0; } diff --git a/src/csrc/umath/comparison_ops.cpp b/src/csrc/umath/comparison_ops.cpp index 9924bd3..f9a0539 100644 --- a/src/csrc/umath/comparison_ops.cpp +++ b/src/csrc/umath/comparison_ops.cpp @@ -8,6 +8,7 @@ #include #include +#include #include "numpy/arrayobject.h" #include "numpy/ufuncobject.h" @@ -256,6 +257,33 @@ quad_reduce_comp_strided_loop_unaligned(PyArrayMethod_Context *context, char *co } +// Picks the output dtype NumPy's own object loop would use, so that deferring to +// it does not change the result dtype. The comparison ufuncs register both +// `OO->?` and `OO->O` and default to the bool one; the `logical_*` ufuncs only +// register `OO->O`, so they must stay object. +static bool +comparison_object_output_is_bool(PyUFuncObject *ufunc) +{ + return strcmp(ufunc->name, "logical_and") != 0 && + strcmp(ufunc->name, "logical_or") != 0 && + strcmp(ufunc->name, "logical_xor") != 0; +} + +// Registers `promoter` for a single (in1, in2, out) DType pattern. +static int +add_comparison_promoter(PyObject *ufunc, PyObject *promoter, PyArray_DTypeMeta *in1, + PyArray_DTypeMeta *in2, PyArray_DTypeMeta *out) +{ + PyObject *DTypes = PyTuple_Pack(3, (PyObject *)in1, (PyObject *)in2, (PyObject *)out); + if (DTypes == NULL) { + return -1; + } + + int res = PyUFunc_AddPromoter(ufunc, DTypes, promoter); + Py_DECREF(DTypes); + return res; +} + NPY_NO_EXPORT int comparison_ufunc_promoter(PyObject *ufunc_obj, PyArray_DTypeMeta *const op_dtypes[], PyArray_DTypeMeta *const signature[], PyArray_DTypeMeta *new_op_dtypes[]) @@ -271,20 +299,24 @@ comparison_ufunc_promoter(PyObject *ufunc_obj, PyArray_DTypeMeta *const op_dtype return 0; } + PyUFuncObject *ufunc = (PyUFuncObject *)ufunc_obj; + if (quad_ufunc_has_object_input(ufunc, op_dtypes)) { + for (int i = 0; i < 2; i++) { + quad_set_promoted_dtype(signature[i], &PyArray_ObjectDType, &new_op_dtypes[i]); + } + PyArray_DTypeMeta *output_dtype = comparison_object_output_is_bool(ufunc) + ? &PyArray_BoolDType + : &PyArray_ObjectDType; + quad_set_promoted_dtype(signature[2], output_dtype, &new_op_dtypes[2]); + return 0; + } + // Normal path: promote both inputs to QuadPrecDType, output is Bool for (int i = 0; i < 2; i++) { - if (signature[i]) { - Py_INCREF(signature[i]); - new_op_dtypes[i] = signature[i]; - } - else { - Py_INCREF(&QuadPrecDType); - new_op_dtypes[i] = &QuadPrecDType; - } + quad_set_promoted_dtype(signature[i], &QuadPrecDType, &new_op_dtypes[i]); } - Py_INCREF(&PyArray_BoolDType); - new_op_dtypes[2] = &PyArray_BoolDType; + quad_set_promoted_dtype(signature[2], &PyArray_BoolDType, &new_op_dtypes[2]); return 0; } @@ -355,38 +387,29 @@ create_quad_comparison_ufunc(PyObject *numpy, const char *ufunc_name) return -1; } - // Register promoter for (QuadPrecDType, Any, Bool) - needed for mixed-type comparisons - PyObject *DTypes = PyTuple_Pack(3, &QuadPrecDType, &PyArrayDescr_Type, &PyArray_BoolDType); - if (DTypes == 0) { - Py_DECREF(promoter_capsule); - Py_DECREF(ufunc); - return -1; - } - - if (PyUFunc_AddPromoter(ufunc, DTypes, promoter_capsule) < 0) { - Py_DECREF(promoter_capsule); - Py_DECREF(DTypes); - Py_DECREF(ufunc); - return -1; - } - Py_DECREF(DTypes); - - // Register promoter for (Any, QuadPrecDType, Bool) - needed for reverse mixed-type comparisons - DTypes = PyTuple_Pack(3, &PyArrayDescr_Type, &QuadPrecDType, &PyArray_BoolDType); - if (DTypes == 0) { - Py_DECREF(promoter_capsule); - Py_DECREF(ufunc); - return -1; - } + // (QuadPrecDType, Any, Bool) and its reverse cover mixed-type comparisons, + // whose result is Bool. The Object patterns are registered separately: an + // explicitly requested object output (`dtype=object`) has to match a + // promoter whose output slot is Object, which the Bool patterns do not. + PyArray_DTypeMeta *any_dt = (PyArray_DTypeMeta *)&PyArrayDescr_Type; + + PyArray_DTypeMeta *promoter_patterns[][3] = { + {&QuadPrecDType, any_dt, &PyArray_BoolDType}, + {any_dt, &QuadPrecDType, &PyArray_BoolDType}, + {&QuadPrecDType, &PyArray_ObjectDType, &PyArray_ObjectDType}, + {&PyArray_ObjectDType, &QuadPrecDType, &PyArray_ObjectDType}, + }; - if (PyUFunc_AddPromoter(ufunc, DTypes, promoter_capsule) < 0) { - Py_DECREF(promoter_capsule); - Py_DECREF(DTypes); - Py_DECREF(ufunc); - return -1; + for (size_t i = 0; i < sizeof(promoter_patterns) / sizeof(promoter_patterns[0]); i++) { + if (add_comparison_promoter(ufunc, promoter_capsule, promoter_patterns[i][0], + promoter_patterns[i][1], promoter_patterns[i][2]) < 0) { + Py_DECREF(promoter_capsule); + Py_DECREF(ufunc); + return -1; + } } Py_DECREF(promoter_capsule); - Py_DECREF(DTypes); + Py_DECREF(ufunc); return 0; diff --git a/src/csrc/umath/matmul.cpp b/src/csrc/umath/matmul.cpp index 6d4b5de..821c288 100644 --- a/src/csrc/umath/matmul.cpp +++ b/src/csrc/umath/matmul.cpp @@ -534,7 +534,10 @@ init_matmul_ops(PyObject *numpy) } if (PyUFunc_AddPromoter(ufunc, DTypes, promoter_capsule) < 0) { - PyErr_Clear(); + Py_DECREF(promoter_capsule); + Py_DECREF(DTypes); + Py_DECREF(ufunc); + return -1; } Py_DECREF(DTypes); @@ -547,11 +550,15 @@ init_matmul_ops(PyObject *numpy) } if (PyUFunc_AddPromoter(ufunc, DTypes, promoter_capsule) < 0) { - PyErr_Clear(); + Py_DECREF(promoter_capsule); + Py_DECREF(DTypes); + Py_DECREF(ufunc); + return -1; } Py_DECREF(DTypes); Py_DECREF(promoter_capsule); + Py_DECREF(ufunc); return 0; diff --git a/src/include/umath/promoters.hpp b/src/include/umath/promoters.hpp index d9545bc..2f41bb5 100644 --- a/src/include/umath/promoters.hpp +++ b/src/include/umath/promoters.hpp @@ -11,6 +11,27 @@ #include "../dtype.h" +inline bool +quad_ufunc_has_object_input(PyUFuncObject *ufunc, PyArray_DTypeMeta *const op_dtypes[]) +{ + for (int i = 0; i < ufunc->nin; i++) { + if (op_dtypes[i] == &PyArray_ObjectDType) { + return true; + } + } + return false; +} + +inline void +quad_set_promoted_dtype(PyArray_DTypeMeta *signature_dtype, + PyArray_DTypeMeta *fallback_dtype, + PyArray_DTypeMeta **new_dtype) +{ + PyArray_DTypeMeta *dtype = signature_dtype != NULL ? signature_dtype : fallback_dtype; + Py_INCREF(dtype); + *new_dtype = dtype; +} + inline int quad_ufunc_promoter(PyObject *ufunc_obj, PyArray_DTypeMeta *const op_dtypes[], PyArray_DTypeMeta *const signature[], PyArray_DTypeMeta *new_op_dtypes[]) @@ -28,17 +49,17 @@ quad_ufunc_promoter(PyObject *ufunc_obj, PyArray_DTypeMeta *const op_dtypes[], return 0; } + if (quad_ufunc_has_object_input(ufunc, op_dtypes)) { + for (int i = 0; i < nargs; i++) { + quad_set_promoted_dtype(signature[i], &PyArray_ObjectDType, &new_op_dtypes[i]); + } + return 0; + } + // This promoter is only registered for patterns where at least one // input is QuadPrecDType, so we always promote all args to QuadPrecDType. for (int i = 0; i < nargs; i++) { - if (signature[i]) { - Py_INCREF(signature[i]); - new_op_dtypes[i] = signature[i]; - } - else { - Py_INCREF(&QuadPrecDType); - new_op_dtypes[i] = &QuadPrecDType; - } + quad_set_promoted_dtype(signature[i], &QuadPrecDType, &new_op_dtypes[i]); } return 0; diff --git a/tests/test_quaddtype.py b/tests/test_quaddtype.py index d298a44..0fb597e 100644 --- a/tests/test_quaddtype.py +++ b/tests/test_quaddtype.py @@ -6511,6 +6511,155 @@ def test_divmod_float64_preserves_dtype(self): assert r.dtype == np.float64 +class TestObjectPromotion: + @pytest.fixture + def backend(self, request): + return request.param + + @pytest.fixture + def operands(self, backend): + quad = np.array([1, 2], dtype=QuadPrecDType(backend=backend)) + objects = np.array([3, 4], dtype=object) + return quad, objects + + @pytest.mark.parametrize("backend", ["sleef", "longdouble"], indirect=True) + @pytest.mark.parametrize("op", [np.add, np.multiply]) + @pytest.mark.parametrize("reverse", [False, True]) + def test_arithmetic_uses_object_loop(self, operands, op, reverse): + quad, objects = operands + left, right = (objects, quad) if reverse else (quad, objects) + + result = op(left, right) + expected = op(left.astype(object), right.astype(object)) + + assert result.dtype == np.dtype(object) + np.testing.assert_array_equal(result, expected, strict=True) + + @pytest.mark.parametrize("backend", ["sleef", "longdouble"], indirect=True) + @pytest.mark.parametrize("op", [np.equal, np.less]) + @pytest.mark.parametrize("reverse", [False, True]) + def test_comparison_uses_object_loop(self, operands, op, reverse): + quad, objects = operands + left, right = (objects, quad) if reverse else (quad, objects) + + result = op(left, right) + expected = op(left.astype(object), right.astype(object)) + + assert result.dtype == np.dtype(np.bool_) + np.testing.assert_array_equal(result, expected, strict=True) + + @pytest.mark.parametrize("backend", ["sleef", "longdouble"], indirect=True) + @pytest.mark.parametrize("op", [np.equal, np.less]) + @pytest.mark.parametrize("reverse", [False, True]) + def test_comparison_honours_explicit_object_dtype(self, operands, op, reverse): + # Comparisons default to a bool result, but NumPy also lets an object + # result be requested explicitly. That needs a promoter whose output slot + # is Object; the Bool ones cannot match it. + quad, objects = operands + left, right = (objects, quad) if reverse else (quad, objects) + + result = op(left, right, dtype=object) + expected = op(left.astype(object), right.astype(object), dtype=object) + + assert result.dtype == np.dtype(object) + np.testing.assert_array_equal(result, expected, strict=True) + + @pytest.mark.parametrize("backend", ["sleef", "longdouble"], indirect=True) + @pytest.mark.parametrize("op", [np.logical_and, np.logical_or]) + @pytest.mark.parametrize("reverse", [False, True]) + def test_logical_uses_object_loop(self, operands, op, reverse): + quad, objects = operands + left, right = (objects, quad) if reverse else (quad, objects) + + result = op(left, right) + expected = op(left.astype(object), right.astype(object)) + + assert result.dtype == np.dtype(object) + np.testing.assert_array_equal(result, expected, strict=True) + + @pytest.mark.parametrize("backend", ["sleef", "longdouble"], indirect=True) + @pytest.mark.parametrize("reverse", [False, True]) + def test_logical_xor_matches_unsupported_object_loop(self, operands, reverse): + quad, objects = operands + left, right = (objects, quad) if reverse else (quad, objects) + + # NumPy has no dedicated object loop for logical_xor, so its generic + # fallback calls a method named after the ufunc on each element, i.e. + # `int.logical_xor`, which does not exist. That AttributeError is plain + # NumPy behaviour on two object arrays; we assert we reproduce it rather + # than inventing a quad-specific result. + with pytest.raises(AttributeError): + np.logical_xor(left.astype(object), right.astype(object)) + with pytest.raises(AttributeError): + np.logical_xor(left, right) + + @pytest.mark.parametrize("backend", ["sleef", "longdouble"], indirect=True) + @pytest.mark.parametrize("reverse", [False, True]) + def test_broadcasting_and_masked_out_use_object_loop(self, backend, reverse): + quad = np.array([[1], [2]], dtype=QuadPrecDType(backend=backend)) + objects = np.array([[3, 4, 5]], dtype=object) + left, right = (objects, quad) if reverse else (quad, objects) + where = np.array([[True, False, True], [False, True, False]]) + out = np.full((2, 3), "unchanged", dtype=object) + expected = out.copy() + + np.add(left, right, out=out, where=where) + np.add(left.astype(object), right.astype(object), out=expected, where=where) + + np.testing.assert_array_equal(out, expected, strict=True) + + @pytest.mark.parametrize("backend", ["sleef", "longdouble"], indirect=True) + @pytest.mark.parametrize("reverse", [False, True]) + def test_divmod_without_object_loop_still_raises(self, operands, reverse): + quad, objects = operands + left, right = (objects, quad) if reverse else (quad, objects) + + # np.divmod registers no object loop at all (unlike floor_divide and + # remainder, which both have OO->O), so divmod on two object arrays is a + # TypeError in plain NumPy. Deferring to NumPy means we raise it too. + with pytest.raises(TypeError): + np.divmod(left, right) + + @pytest.mark.parametrize("backend", ["sleef", "longdouble"], indirect=True) + @pytest.mark.parametrize("reverse", [False, True]) + def test_matmul_uses_object_loop(self, backend, reverse): + quad = np.array([[1, 2], [3, 4]], dtype=QuadPrecDType(backend=backend)) + objects = np.array([[5, 6], [7, 8]], dtype=object) + left, right = (objects, quad) if reverse else (quad, objects) + + result = np.matmul(left, right) + expected = np.matmul(left.astype(object), right.astype(object)) + + assert result.dtype == np.dtype(object) + np.testing.assert_array_equal(result, expected, strict=True) + + @pytest.mark.parametrize("backend", ["sleef", "longdouble"], indirect=True) + @pytest.mark.parametrize( + "other", + [ + np.array([3 + 1j, 4 + 2j], dtype=np.complex128), + np.array(["3", "4"], dtype="U1"), + np.array([b"3", b"4"], dtype="S1"), + ], + ) + def test_unsupported_common_dtypes_still_raise(self, operands, other): + quad, _ = operands + + with pytest.raises(np.exceptions.DTypePromotionError): + np.result_type(quad, other) + with pytest.raises(TypeError): + np.add(quad, other) + + def test_builtin_object_dispatch_is_unchanged(self): + left = np.array([1, 2], dtype=object) + right = np.array([3, 4], dtype=object) + + result = np.add(left, right) + + assert result.dtype == np.dtype(object) + np.testing.assert_array_equal(result, np.array([4, 6], dtype=object), strict=True) + + def test_sleef_purecfma_symbols(): """Test that SLEEF PURECFMA symbols are present in the compiled module.