ENH: Enhanced scalar fastpaths

numpy · ganesh-k13 · Dec 9, 2020 · Dec 12, 2020 · Dec 12, 2020 · Dec 12, 2020
commit 708dd751976c768d3677d6c91fa4e154ee61457d
diff --git a/numpy/core/src/umath/scalarmath.c.src b/numpy/core/src/umath/scalarmath.c.src
@@ -595,6 +595,10 @@ descr_from_basic_scalar(PyObject *obj, PyArray_Descr **descr)
         *descr = PyArray_DescrFromType(NPY_DOUBLE);
         return 1;
     }
+    else if (PyBool_Check(obj)) {
+        *descr = PyArray_DescrFromType(NPY_BOOL);
+        return 1;
+    }
     else if (PyLong_CheckExact(obj)) {
         if (PyLong_AsLong(obj) == -1 && PyErr_Occurred()) {
             PyErr_Clear();
@@ -610,6 +614,8 @@ descr_from_basic_scalar(PyObject *obj, PyArray_Descr **descr)
     return -1;
 }
 
+#define PyDescr_PythonCRepresentable(d) (d->type_num == NPY_LONGDOUBLE || \
+        d->type_num > NPY_CDOUBLE)
 
 /**
  * This function attempts to compare NumPy or Python scalars
@@ -618,37 +624,51 @@ descr_from_basic_scalar(PyObject *obj, PyArray_Descr **descr)
  *
  * @param self The first object which will be converted to scalar item
  * @param other The second object to compare to
- * @param other The value of comparison operator
+ * @param cmp_op The value of comparison operator
  * @returns Py_False or Py_True if richcompare is successfull
  *          , otherwise NULL.
  */
 static PyObject*
 do_richcompare_on_scalars(PyObject *self, PyObject *other, int cmp_op) {
-    PyObject *item_self, *ret=NULL;
+    PyObject *cmp_item_self, *cmp_item_other, *ret=NULL;
     PyArray_Descr *self_descr=NULL, *other_descr=NULL;
-    void *data_self;
+    void *data_self, *data_other;
     int pyscalar_other, pyscalar_self;
-    int is_complex_operands, is_equality_operator;
+    int is_complex_operands, is_equality_operator, is_python_scalar;
 
     pyscalar_self = descr_from_basic_scalar(self, &self_descr);
     pyscalar_other = descr_from_basic_scalar(other, &other_descr);
 
     if (pyscalar_self >= 0 && pyscalar_other >= 0) {
         /*
-         * If either of the operands are complex AND the operator is not equality,
-         * python's built-in richcompare cannot be used as it is not supported
+         * If either of the operands are complex and operator is not equality,
+         * or the operands are not representable as a native type,
+         * python's built-in richcompare cannot be used as it is not supported.
          */
         is_complex_operands = (PyTypeNum_ISCOMPLEX(self_descr->type_num) ||
             PyTypeNum_ISCOMPLEX(other_descr->type_num));
         is_equality_operator = (cmp_op != Py_EQ && cmp_op != Py_NE);
+        is_python_scalar = PyDescr_PythonCRepresentable(self_descr) ||
+            PyDescr_PythonCRepresentable(other_descr);
 
-        if (!(is_complex_operands && is_equality_operator)) {
+        if (!(is_complex_operands && is_equality_operator) && !is_python_scalar) {
             data_self = scalar_value(self, NULL);
-            item_self = self_descr->f->getitem(data_self, NULL);
-            if (item_self != NULL) {
-                ret = PyObject_RichCompare(item_self, other, cmp_op);
+            data_other = scalar_value(other, NULL);
+            cmp_item_self = pyscalar_self == 0 ? self_descr->f->getitem(data_self, NULL):
+                self;
+            cmp_item_other = pyscalar_other == 0 ? other_descr->f->getitem(data_other, NULL):
+                other;
+
+            if (cmp_item_self != NULL && cmp_item_other != NULL) {
+                ret = PyObject_RichCompare(cmp_item_self, cmp_item_other, cmp_op);
+            }
+
+            if (pyscalar_self == 0) {
+                Py_XDECREF(cmp_item_self);
+            }
+            if (pyscalar_other == 0) {
+                Py_XDECREF(cmp_item_other);
             }
-            Py_XDECREF(item_self);
         }
     }
     Py_XDECREF(self_descr);