ENH: ufunc: Add the main loop selection/type-determination mechanism

6 files changed, 128 insertions, 54 deletions

diff --git a/numpy/core/code_generators/numpy_api.py b/numpy/core/code_generators/numpy_api.py
index 52c5d8bfc..953c671da 100644
--- a/numpy/core/code_generators/numpy_api.py
+++ b/numpy/core/code_generators/numpy_api.py
@@ -299,6 +299,7 @@ multiarray_funcs_api = {
     'PyArray_MinScalarType':                264,
     'PyArray_ResultType':                   265,
     'PyArray_CanCastArrayTo':               266,
+    'PyArray_CanCastTypeTo':                267,
 }
 
 ufunc_types_api = {
diff --git a/numpy/core/src/multiarray/convert_datatype.c b/numpy/core/src/multiarray/convert_datatype.c
index 6cd4e7a79..b2fa70b19 100644
--- a/numpy/core/src/multiarray/convert_datatype.c
+++ b/numpy/core/src/multiarray/convert_datatype.c
@@ -63,7 +63,7 @@ PyArray_CastToType(PyArrayObject *mp, PyArray_Descr *at, int fortran)
     if (out == NULL) {
         return NULL;
     }
-    ret = PyArray_CastTo((PyArrayObject *)out, mp);
+    ret = PyArray_CopyInto((PyArrayObject *)out, mp);
     if (ret != -1) {
         return out;
     }
@@ -219,7 +219,8 @@ PyArray_CanCastSafely(int fromtype, int totype)
 /*NUMPY_API
  * leaves reference count alone --- cannot be NULL
  *
- * TODO: For NumPy 2.0, add a NPY_CASTING parameter (can_cast_to function).
+ * PyArray_CanCastTypeTo is equivalent to this, but adds a 'casting'
+ * parameter.
  */
 NPY_NO_EXPORT npy_bool
 PyArray_CanCastTo(PyArray_Descr *from, PyArray_Descr *to)
@@ -312,8 +313,13 @@ type_num_unsigned_to_signed(int type_num)
     }
 }
 
+/*NUMPY_API
+ * Returns true if data of type 'from' may be cast to data of type
+ * 'to' according to the rule 'casting'.
+ */
 NPY_NO_EXPORT npy_bool
-can_cast_to(PyArray_Descr *from, PyArray_Descr *to, NPY_CASTING casting)
+PyArray_CanCastTypeTo(PyArray_Descr *from, PyArray_Descr *to,
+                                                    NPY_CASTING casting)
 {
     /* If unsafe casts are allowed */
     if (casting == NPY_UNSAFE_CASTING) {
@@ -405,7 +411,7 @@ PyArray_CanCastArrayTo(PyArrayObject *arr, PyArray_Descr *to,
 
     /* If it's not a scalar, use the standard rules */
     if (PyArray_NDIM(arr) > 0 || !PyTypeNum_ISNUMBER(from->type_num)) {
-        return can_cast_to(from, to, casting);
+        return PyArray_CanCastTypeTo(from, to, casting);
     }
     /* Otherwise, check the value */
     else {
@@ -446,7 +452,7 @@ PyArray_CanCastArrayTo(PyArrayObject *arr, PyArray_Descr *to,
         PyObject_Print(to, stdout, 0);
         printf("\n");
 #endif
-        ret = can_cast_to(dtype, to, casting);
+        ret = PyArray_CanCastTypeTo(dtype, to, casting);
         Py_DECREF(dtype);
         return ret;
     }
diff --git a/numpy/core/src/multiarray/convert_datatype.h b/numpy/core/src/multiarray/convert_datatype.h
index b45d0dfbd..844cce0c9 100644
--- a/numpy/core/src/multiarray/convert_datatype.h
+++ b/numpy/core/src/multiarray/convert_datatype.h
@@ -1,26 +1,10 @@
 #ifndef _NPY_ARRAY_CONVERT_DATATYPE_H_
 #define _NPY_ARRAY_CONVERT_DATATYPE_H_
 
-NPY_NO_EXPORT PyObject *
-PyArray_CastToType(PyArrayObject *mp, PyArray_Descr *at, int fortran);
-
-NPY_NO_EXPORT int
-PyArray_CastTo(PyArrayObject *out, PyArrayObject *mp);
-
 NPY_NO_EXPORT PyArray_VectorUnaryFunc *
 PyArray_GetCastFunc(PyArray_Descr *descr, int type_num);
 
 NPY_NO_EXPORT int
-PyArray_CanCastSafely(int fromtype, int totype);
-
-/* This can replace PyArray_CanCastTo for NumPy 2.0 (ABI change) */
-NPY_NO_EXPORT npy_bool
-can_cast_to(PyArray_Descr *from, PyArray_Descr *to, NPY_CASTING casting);
-
-NPY_NO_EXPORT npy_bool
-PyArray_CanCastTo(PyArray_Descr *from, PyArray_Descr *to);
-
-NPY_NO_EXPORT int
 PyArray_ObjectType(PyObject *op, int minimum_type);
 
 NPY_NO_EXPORT PyArrayObject **
diff --git a/numpy/core/src/multiarray/multiarraymodule.c b/numpy/core/src/multiarray/multiarraymodule.c
index 96118acaf..f5184966c 100644
--- a/numpy/core/src/multiarray/multiarraymodule.c
+++ b/numpy/core/src/multiarray/multiarraymodule.c
@@ -2229,14 +2229,14 @@ array_can_cast_safely(PyObject *NPY_UNUSED(self), PyObject *args,
         ret = PyArray_CanCastArrayTo(arr, d2, casting);
         Py_DECREF(arr);
     }
-    /* Otherwise use CanCastTo */
+    /* Otherwise use CanCastTypeTo */
     else {
         if (!PyArray_DescrConverter2(from_obj, &d1) || d1 == NULL) {
             PyErr_SetString(PyExc_TypeError,
                     "did not understand one of the types; 'None' not accepted");
             goto finish;
         }
-        ret = can_cast_to(d1, d2, casting);
+        ret = PyArray_CanCastTypeTo(d1, d2, casting);
     }
 
     retobj = ret ? Py_True : Py_False;
diff --git a/numpy/core/src/multiarray/new_iterator.c.src b/numpy/core/src/multiarray/new_iterator.c.src
index b7b928051..f2bc9b29d 100644
--- a/numpy/core/src/multiarray/new_iterator.c.src
+++ b/numpy/core/src/multiarray/new_iterator.c.src
@@ -2671,7 +2671,7 @@ npyiter_check_casting(npy_intp niter, PyArrayObject **op,
             }
             /* Check write (temp -> op) casting */
             if ((op_itflags[iiter]&NPY_OP_ITFLAG_WRITE) &&
-                        !can_cast_to(op_dtype[iiter],
+                        !PyArray_CanCastTypeTo(op_dtype[iiter],
                                           PyArray_DESCR(op[iiter]),
                                           casting)) {
                 PyErr_Format(PyExc_TypeError,
diff --git a/numpy/core/src/umath/ufunc_object.c b/numpy/core/src/umath/ufunc_object.c
index 6a25e3a4c..f3415f723 100644
--- a/numpy/core/src/umath/ufunc_object.c
+++ b/numpy/core/src/umath/ufunc_object.c
@@ -3503,7 +3503,7 @@ static PyObject *
 ufunc_generic_call_iter(PyUFuncObject *self, PyObject *args, PyObject *kwds)
 {
     npy_intp nargs, nin = self->nin, nout = self->nout;
-    npy_intp i, niter = nin + nout;
+    npy_intp i, j, niter = nin + nout;
     PyObject *obj, *context;
     char *ufunc_name;
     /* This contains the all the inputs and outputs */
@@ -3511,9 +3511,11 @@ ufunc_generic_call_iter(PyUFuncObject *self, PyObject *args, PyObject *kwds)
     PyArray_Descr *dtype[NPY_MAXARGS];
     PyArray_Descr *result_type;
 
+    int no_castable_output, all_inputs_scalar;
+
     /* TODO: For 1.6, the default should probably be NPY_CORDER */
     NPY_ORDER order = NPY_KEEPORDER;
-    NPY_CASTING casting = NPY_SAFE_CASTING;
+    NPY_CASTING casting = NPY_SAFE_CASTING, input_casting, output_casting;
     PyObject *extobj = NULL, *typetup = NULL;
 
     ufunc_name = self->name ? self->name : "";
@@ -3539,6 +3541,7 @@ ufunc_generic_call_iter(PyUFuncObject *self, PyObject *args, PyObject *kwds)
     result_type = NULL;
 
     /* Get input arguments */
+    all_inputs_scalar = 1;
     for(i = 0; i < nin; ++i) {
         obj = PyTuple_GET_ITEM(args, i);
         if (!PyArray_Check(obj) && !PyArray_IsScalar(obj, Generic)) {
@@ -3555,18 +3558,14 @@ ufunc_generic_call_iter(PyUFuncObject *self, PyObject *args, PyObject *kwds)
             context = NULL;
         }
         op[i] = (PyArrayObject *)PyArray_FromAny(obj, NULL, 0, 0, 0, context);
-
-        /* Start with a native byte-order data type */
-        /*
-        if (PyArray_ISNBO(PyArray_DESCR(op[i])->byteorder)) {
-            dtype[i] = PyArray_DESCR(op[i]);
-            Py_INCREF(dtype[i]);
+        Py_XDECREF(context);
+        if (op[i] == NULL) {
+            goto fail;
         }
-        else {
-            dtype[i] = PyArray_DescrNewByteorder(PyArray_DESCR(op[i]),
-                                                            NPY_NATIVE);
+        /* If there's a non-scalar input, indicate so */
+        if (PyArray_NDIM(op[i]) > 0) {
+            all_inputs_scalar = 0;
         }
-        */
     }
 
     /* Get positional output arguments */
@@ -3650,6 +3649,7 @@ ufunc_generic_call_iter(PyUFuncObject *self, PyObject *args, PyObject *kwds)
                                                 "of ArrayType");
                                 goto fail;
                             }
+                        
                         }
                     }
                     else if (strncmp(str,"order",5) == 0) {
@@ -3676,31 +3676,114 @@ ufunc_generic_call_iter(PyUFuncObject *self, PyObject *args, PyObject *kwds)
     }
 
     /*
-     * Determine the result type.  A better approach would be for
-     * the ufunc to provide a function which gives back the result
-     * type and inner loop function.  The default would work as follows.
+     * Now that we have the casting parameter, decide on the casting rules
+     * for inputs and outputs.  For inputs, we want NPY_SAFE_CASTING or
+     * stricter, and for outputs, anything is ok.
      */
-    result_type = PyArray_ResultType(nin, op, 0, NULL);
-    if (result_type == NULL) {
-        goto fail;
+    if (casting > NPY_SAFE_CASTING) {
+        input_casting = NPY_SAFE_CASTING;
     }
-    /* Take into account the types of any output parameters */
-    for (i = nin; i < nargs; ++i) {
-        PyArray_Descr *tmp;
-        if (op[i] != NULL) {
-            tmp = PyArray_PromoteTypes(result_type, PyArray_DESCR(op[i]));
-            if (tmp == NULL) {
-                goto fail;
+    else {
+        input_casting = casting;
+    }
+    output_casting = casting;
+
+    /*
+     * Determine the UFunc loop.  This could in general be *much* faster,
+     * and a better way to implement it might be for the ufunc to
+     * provide a function which gives back the result type and inner
+     * loop function.
+     *
+     * The method for finding the loop in the previous code did not
+     * appear consistent (as noted by some asymmetry in the generated
+     * coercion tables for np.add).
+     */
+    no_castable_output = 0;
+    for (i = 0; i < self->ntypes; ++i) {
+        char *types = self->types + i*self->nargs;
+        /*
+         * First check if all the inputs can be safely cast
+         * to the types for this function
+         */
+        for (j = 0; j < nin; ++j) {
+            PyArray_Descr *tmp = PyArray_DescrFromType(types[j]);
+            /*
+             * If all the inputs are scalars, use the regular
+             * promotion rules, not the special value-checking ones.
+             */
+            if (all_inputs_scalar) {
+                if (!PyArray_CanCastTypeTo(PyArray_DESCR(op[j]), tmp,
+                                                    input_casting)) {
+                    Py_DECREF(tmp);
+                    break;
+                }
+            }
+            else {
+                if (!PyArray_CanCastArrayTo(op[j], tmp, input_casting)) {
+                    Py_DECREF(tmp);
+                    break;
+                }
+            }
+            Py_DECREF(tmp);
+        }
+        /*
+         * If all the inputs were ok, then check casting back to the
+         * outputs.
+         */
+        if (j == nin) {
+            for (j = nin; j < niter; ++j) {
+                if (op[j] != NULL) {
+                    PyArray_Descr *tmp = PyArray_DescrFromType(types[j]);
+                    if (!PyArray_CanCastTypeTo(tmp, PyArray_DESCR(op[j]),
+                                                    output_casting)) {
+                        Py_DECREF(tmp);
+                        no_castable_output = 1;
+                        break;
+                    }
+                    Py_DECREF(tmp);
+                }
+            }
+        }
+        /* If all the tests passed, we found our function */
+        if (j == niter) {
+            /* Fill the dtypes array */
+            for (j = 0; j < niter; ++j) {
+                dtype[j] = PyArray_DescrFromType(types[j]);
             }
-            Py_DECREF(result_type);
-            result_type = tmp;
+            break;
         }
     }
+    /* If no function was found, throw an error */
+    if (i == self->ntypes) {
+        if (no_castable_output) {
+            PyErr_Format(PyExc_TypeError,
+                    "function output could not be coerced to a provided "
+                    "output parameter according to the specified casting "
+                    "rule");
+        }
+        else {
+            PyErr_Format(PyExc_TypeError,
+                    "function not supported for the input types, and the "
+                    "inputs could not be safely coerced to any supported "
+                    "types");
+        }
+        return NULL;
+    }
+    else {
+        
+    }
 
-    /* Find the function loop */
 
-    printf("result type: ");
-    PyObject_Print((PyObject *)result_type, stdout, 0);
+    printf("input types:\n");
+    for (i = 0; i < nin; ++i) {
+        PyObject_Print((PyObject *)dtype[i], stdout, 0);
+        printf(" ");
+    }
+    printf("\noutput types:\n");
+    for (i = nin; i < niter; ++i) {
+        PyObject_Print((PyObject *)dtype[i], stdout, 0);
+        printf(" ");
+    }
     printf("\n");
 
     PyErr_SetString(PyExc_RuntimeError, "implementation is not finished!");