3838#define _MULTIARRAYMODULE
3939#define _UMATHMODULE
4040
41+ #include < mutex>
42+ #include < shared_mutex>
43+
4144#define PY_SSIZE_T_CLEAN
4245#include < Python.h>
4346#include < convert_datatype.h>
@@ -504,8 +507,9 @@ call_promoter_and_recurse(PyUFuncObject *ufunc, PyObject *info,
504507 PyObject *promoter = PyTuple_GET_ITEM (info, 1 );
505508 if (PyCapsule_CheckExact (promoter)) {
506509 /* We could also go the other way and wrap up the python function... */
507- PyArrayMethod_PromoterFunction * promoter_function = PyCapsule_GetPointer (
508- promoter , "numpy._ufunc_promoter" );
510+ PyArrayMethod_PromoterFunction *promoter_function =
511+ (PyArrayMethod_PromoterFunction *)PyCapsule_GetPointer (
512+ promoter, " numpy._ufunc_promoter"
509513 if (promoter_function == NULL ) {
510514 return NULL ;
511515 }
@@ -770,8 +774,9 @@ promote_and_get_info_and_ufuncimpl(PyUFuncObject *ufunc,
770774 * 2. Check all registered loops/promoters to find the best match.
771775 * 3. Fall back to the legacy implementation if no match was found.
772776 */
773- PyObject * info = PyArrayIdentityHash_GetItem (ufunc -> _dispatch_cache ,
774- (PyObject * * )op_dtypes );
777+ PyObject *info = PyArrayIdentityHash_GetItem (
778+ (PyArrayIdentityHash *)ufunc->_dispatch_cache ,
779+ (PyObject **)op_dtypes);
775780 if (info != NULL && PyObject_TypeCheck (
776781 PyTuple_GET_ITEM (info, 1 ), &PyArrayMethod_Type)) {
777782 /* Found the ArrayMethod and NOT a promoter: return it */
@@ -793,8 +798,9 @@ promote_and_get_info_and_ufuncimpl(PyUFuncObject *ufunc,
793798 * Found the ArrayMethod and NOT promoter. Before returning it
794799 * add it to the cache for faster lookup in the future.
795800 */
796- if (PyArrayIdentityHash_SetItem (ufunc -> _dispatch_cache ,
797- (PyObject * * )op_dtypes , info , 0 ) < 0 ) {
801+ if (PyArrayIdentityHash_SetItem (
802+ (PyArrayIdentityHash *)ufunc->_dispatch_cache ,
803+ (PyObject **)op_dtypes, info, 0 ) < 0 ) {
798804 return NULL ;
799805 }
800806 return info;
@@ -815,8 +821,9 @@ promote_and_get_info_and_ufuncimpl(PyUFuncObject *ufunc,
815821 }
816822 else if (info != NULL ) {
817823 /* Add result to the cache using the original types: */
818- if (PyArrayIdentityHash_SetItem (ufunc -> _dispatch_cache ,
819- (PyObject * * )op_dtypes , info , 0 ) < 0 ) {
824+ if (PyArrayIdentityHash_SetItem (
825+ (PyArrayIdentityHash *)ufunc->_dispatch_cache ,
826+ (PyObject **)op_dtypes, info, 0 ) < 0 ) {
820827 return NULL ;
821828 }
822829 return info;
@@ -882,13 +889,51 @@ promote_and_get_info_and_ufuncimpl(PyUFuncObject *ufunc,
882889 }
883890
884891 /* Add this to the cache using the original types: */
885- if (cacheable && PyArrayIdentityHash_SetItem (ufunc -> _dispatch_cache ,
886- (PyObject * * )op_dtypes , info , 0 ) < 0 ) {
892+ if (cacheable && PyArrayIdentityHash_SetItem (
893+ (PyArrayIdentityHash *)ufunc->_dispatch_cache ,
894+ (PyObject **)op_dtypes, info, 0 ) < 0 ) {
887895 return NULL ;
888896 }
889897 return info;
890898}
891899
900+ #ifdef Py_GIL_DISABLED
901+ /*
902+ * Fast path for promote_and_get_info_and_ufuncimpl.
903+ * Acquires a read lock to check for a cache hit and then
904+ * only acquires a write lock on a cache miss to fill the cache
905+ */
906+ static inline PyObject *
907+ promote_and_get_info_and_ufuncimpl_with_locking (<
BCA1
/div>
908+ PyUFuncObject *ufunc,
909+ PyArrayObject *const ops[],
910+ PyArray_DTypeMeta *signature[],
911+ PyArray_DTypeMeta *op_dtypes[],
912+ npy_bool legacy_promotion_is_possible)
913+ {
914+ std::shared_mutex *mutex = ((std::shared_mutex *)((PyArrayIdentityHash *)ufunc->_dispatch_cache )->mutex );
915+ mutex->lock_shared ();
916+ PyObject *info = PyArrayIdentityHash_GetItem (
917+ (PyArrayIdentityHash *)ufunc->_dispatch_cache ,
918+ (PyObject **)op_dtypes);
919+ mutex->unlock_shared ();
920+
921+ if (info != NULL && PyObject_TypeCheck (
922+ PyTuple_GET_ITEM (info, 1 ), &PyArrayMethod_Type)) {
923+ /* Found the ArrayMethod and NOT a promoter: return it */
924+ return info;
925+ }
926+
927+ // cache miss, need to acquire a write lock and recursively calculate the
928+ // correct dispatch resolution
929+ mutex->lock ();
930+ info = promote_and_get_info_and_ufuncimpl (ufunc,
931+ ops, signature, op_dtypes, legacy_promotion_is_possible);
932+ mutex->unlock ();
933+
934+ return info;
935+ }
936+ #endif
892937
893938/* *
894939 * The central entry-point for the promotion and dispatching machinery.
@@ -941,6 +986,8 @@ promote_and_get_ufuncimpl(PyUFuncObject *ufunc,
941986{
942987 int nin = ufunc->nin , nargs = ufunc->nargs ;
943988 npy_bool legacy_promotion_is_possible = NPY_TRUE;
989+ PyObject *all_dtypes = NULL ;
990+ PyArrayMethodObject *method = NULL ;
944991
945992 /*
946993 * Get the actual DTypes we operate with by setting op_dtypes[i] from
@@ -976,18 +1023,20 @@ promote_and_get_ufuncimpl(PyUFuncObject *ufunc,
9761023 }
9771024 }
9781025
979- PyObject * info ;
980- Py_BEGIN_CRITICAL_SECTION((PyObject * )ufunc );
981- info = promote_and_get_info_and_ufuncimpl (ufunc ,
1026+ #ifdef Py_GIL_DISABLED
1027+ PyObject *info = promote_and_get_info_and_ufuncimpl_with_locking (ufunc,
1028+ ops, signature, op_dtypes, legacy_promotion_is_possible);
1029+ #else
1030+ PyObject *info = promote_and_get_info_and_ufuncimpl (ufunc,
9821031 ops, signature, op_dtypes, legacy_promotion_is_possible);
983- Py_END_CRITICAL_SECTION ();
1032+ # endif
9841033
9851034 if (info == NULL ) {
9861035 goto handle_error;
9871036 }
9881037
989- PyArrayMethodObject * method = (PyArrayMethodObject * )PyTuple_GET_ITEM (info , 1 );
990- PyObject * all_dtypes = PyTuple_GET_ITEM (info , 0 );
1038+ method = (PyArrayMethodObject *)PyTuple_GET_ITEM (info, 1 );
1039+ all_dtypes = PyTuple_GET_ITEM (info, 0 );
9911040
9921041 /*
9931042 * In certain cases (only the logical ufuncs really), the loop we found may
@@ -1218,7 +1267,7 @@ install_logical_ufunc_promoter(PyObject *ufunc)
12181267 if (dtype_tuple == NULL ) {
12191268 return -1 ;
12201269 }
1221- PyObject * promoter = PyCapsule_New (& logical_ufunc_promoter ,
1270+ PyObject *promoter = PyCapsule_New (( void *) &logical_ufunc_promoter,
12221271 " numpy._ufunc_promoter" NULL );
12231272 if (promoter == NULL ) {
12241273 Py_DECREF (dtype_tuple);
0 commit comments