Skip to content

add np.float_power and np.cbrt #6075

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
sklam merged 14 commits into from
Jul 19, 2021
Merged

add np.float_power and np.cbrt #6075

Show file tree
Hide file tree
Changes from 8 commits
Commits
Show all changes
14 commits
Select commit Hold shift + click to select a range
17bcd58
add np.float_power and np.cbrt
guilhermeleobas Aug 5, 2020
2533dd2
fix cbrt test failure due to fp-precision
guilhermeleobas Aug 10, 2020
353473c
Update npyfuncs.py
guilhermeleobas Aug 18, 2020
c991752
Update npyfuncs.py
guilhermeleobas Aug 26, 2020
b161b0b
Update npyfuncs.py
guilhermeleobas Aug 27, 2020
9b44726
trigger CI
guilhermeleobas Aug 28, 2020
8de632a
Merge remote-tracking branch 'upstream/master' into ufuncs
guilhermeleobas Sep 1, 2020
80425a7
Fix cbrt test
guilhermeleobas Sep 3, 2020
4d076fa
replace cbrt implementation
guilhermeleobas Oct 21, 2020
4f4767b
Set fastmath flag on np.cbrt for llvm to a call to libm _cbrt
guilhermeleobas Jun 11, 2021
95f8b52
trigger CI
guilhermeleobas Jun 14, 2021
7eda7ac
move isnan check to outside cbrt function
guilhermeleobas Jun 18, 2021
33615bd
add comment explaining fastmath flag
guilhermeleobas Jul 9, 2021
66955b1
move import statement to the top of the file
guilhermeleobas Jul 13, 2021
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
2 changes: 2 additions & 0 deletions docs/source/reference/numpysupported.rst
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -606,6 +606,7 @@ Math operations
floor_divide Yes Yes
negative Yes Yes
power Yes Yes
float_power Yes Yes
remainder Yes Yes
mod Yes Yes
fmod Yes Yes
Expand All @@ -625,6 +626,7 @@ Math operations
log1p Yes Yes
sqrt Yes Yes
square Yes Yes
cbrt Yes Yes
reciprocal Yes Yes
conjugate Yes Yes
gcd Yes Yes
Expand Down
12 changes: 12 additions & 0 deletions numba/_helperlib.c
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -196,6 +196,18 @@ numba_signbit(npy_double a)
return npy_signbit(a);
}

NUMBA_EXPORT_FUNC(float)
numba_cbrtf(float x)
{
return npy_cbrtf(x);
}

NUMBA_EXPORT_FUNC(double)
numba_cbrt(double x)
{
return npy_cbrt(x);
}

@stuartarchibald stuartarchibald Oct 21, 2020

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I'm struggling to understand why this is needed?

This:

from numba import njit
import numpy as np

@njit(fastmath={'nnan', 'nsz', 'ninf', 'afn'})
def cbrt(x):
    third = 1.0 / 3.0
    return np.power(x, third)

@njit
def foo(x):
    return cbrt(x)

x = np.zeros(100)
foo(x)
print(cbrt.inspect_llvm(cbrt.signatures[0]))

print(cbrt.inspect_asm(cbrt.signatures[0]))

generates a call to the cbrt function, LLVM "knows" that power(x, 1/3) is cbrt under the flags set. Perhaps this can be used to ensure that a cbrt call is generated along with adding the inf/nan handling as per NumPy?

@guilhermeleobas guilhermeleobas Oct 21, 2020 •
edited
Loading

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I'm struggling to understand why this is needed?

Before adding this, one of the tests were failing on CI. See:
#6075 (comment)

generates a call to the cbrt function, LLVM "knows" that power(x, 1/3) is cbrt under the flags set.

For me, this produces a call to @llvm.pow.f64(x, 1/3):

; Function Attrs: nofree nounwind writeonly
define i32 @"_ZN8__main__8cbrt$242Ex"(double* noalias nocapture %retptr, { i8*, i32, i8* }** noalias nocapture readnone %excinfo, i64 %arg.x) local_unnamed_addr #0 {
entry:
  %.37.le = sitofp i64 %arg.x to double
  %.38.le = tail call nnan ninf nsz afn double @llvm.pow.f64(double %.37.le, double 0x3FD5555555555555)
  store double %.38.le, double* %retptr, align 8
  ret i32 0
}

@stuartarchibald stuartarchibald Oct 21, 2020

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Yes, but that then ends up as a call to cbrt once compiled?

@guilhermeleobas guilhermeleobas Oct 21, 2020 •
edited
Loading

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Oh yes, you right. Nevertheless, If I remove the code you quote, some builds on CI starts to fail. Do you want me to give a try and implement cbrt in numba?

edit: this implementation produces the following assembly

    def cbrt(x):
        if np.isnan(x):
            return np.nan
        elif x < 0:
            return -np.power(-x, 1.0 / 3.0)
        else:
            return np.power(x, 1.0 / 3.0)
    
    return context.compile_internal(builder, cbrt, sig, args)
        .text
        .file   "<string>"
        .globl  _ZN8__main__7bar$242Ex
        .p2align        4, 0x90
        .type   _ZN8__main__7bar$242Ex,@function
_ZN8__main__7bar$242Ex:
        pushq   %rbx
        movq    %rdi, %rbx
        vcvtsi2sd       %rdx, %xmm0, %xmm0
        movabsq $cbrt, %rax
        callq   *%rax
        vmovsd  %xmm0, (%rbx)
        xorl    %eax, %eax
        popq    %rbx
        retq
.Lfunc_end0:
        .size   _ZN8__main__7bar$242Ex, .Lfunc_end0-_ZN8__main__7bar$242Ex

        .globl  cfunc._ZN8__main__7bar$242Ex
        .p2align        4, 0x90
        .type   cfunc._ZN8__main__7bar$242Ex,@function
cfunc._ZN8__main__7bar$242Ex:
        vcvtsi2sd       %rdi, %xmm0, %xmm0
        movabsq $cbrt, %rax
        jmpq    *%rax
.Lfunc_end1:
        .size   cfunc._ZN8__main__7bar$242Ex, .Lfunc_end1-cfunc._ZN8__main__7bar$242Ex

        .type   _ZN08NumbaEnv8__main__7bar$242Ex,@object
        .comm   _ZN08NumbaEnv8__main__7bar$242Ex,8,8
        .section        ".note.GNU-stack","",@progbits

@stuartarchibald stuartarchibald May 28, 2021

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Using this:

from numba import njit
import numpy as np

@njit
def foo(arr):
    return np.cbrt(arr)

foo(np.ones(10, dtype=np.float64))

print(foo.inspect_asm(foo.signatures[0]))

I'm not seeing the call to cbrt being made. I suspect it requires floating point assumptions to be relaxed in the two branches that call np.power.

@guilhermeleobas guilhermeleobas May 31, 2021

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

If one enables fast-math, you'll see a call to cbrt.

from numba import njit
import numpy as np

@njit('float32(float32)', fastmath=True, no_cfunc_wrapper=True, no_cpython_wrapper=True)
def foo(x):
    return np.cbrt(x)

print(foo.inspect_asm(foo.signatures[0]))
	.section	__TEXT,__text,regular,pure_instructions
	.macosx_version_min 10, 15
	.section	__TEXT,__literal4,4byte_literals
	.p2align	2
LCPI0_0:
	.long	2147483648
	.section	__TEXT,__text,regular,pure_instructions
	.globl	__ZN8__main__7foo$241Ef
	.p2align	4, 0x90
__ZN8__main__7foo$241Ef:
	pushq	%r14
	pushq	%rbx
	subq	$56, %rsp
	vmovaps	%xmm0, (%rsp)
	movq	%rdi, %rbx
	vcvtss2sd	%xmm0, %xmm0, %xmm0
	movabsq	$_cbrt, %r14
	callq	*%r14
	vcvtsd2ss	%xmm0, %xmm0, %xmm0
	vmovaps	%xmm0, 32(%rsp)
	movabsq	$LCPI0_0, %rax
	vbroadcastss	(%rax), %xmm0
	vmovaps	%xmm0, 16(%rsp)
	vxorps	(%rsp), %xmm0, %xmm0
	vcvtss2sd	%xmm0, %xmm0, %xmm0
	callq	*%r14
	vcvtsd2ss	%xmm0, %xmm0, %xmm0
	vxorps	16(%rsp), %xmm0, %xmm0
	vxorps	%xmm1, %xmm1, %xmm1
	vmovaps	(%rsp), %xmm2
	vcmpltss	%xmm1, %xmm2, %xmm1
	vmovaps	32(%rsp), %xmm2
	vblendvps	%xmm1, %xmm0, %xmm2, %xmm0
	vmovss	%xmm0, (%rbx)
	xorl	%eax, %eax
	addq	$56, %rsp
	popq	%rbx
	popq	%r14
	retq

	.comm	__ZN08NumbaEnv8__main__7foo$241Ef,8,3
	.comm	__ZN08NumbaEnv5numba2np8npyfuncs17np_real_cbrt_impl12$3clocals$3e8cbrt$242Ef,8,3
.subsections_via_symbols



/* Unpack any Python complex-like object into a Py_complex structure */
NUMBA_EXPORT_FUNC(int)
Expand Down
2 changes: 2 additions & 0 deletions numba/_helpermod.c
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -70,6 +70,8 @@ build_c_helpers_dict(void)
declmethod(lgammaf);
declmethod(signbit);
declmethod(signbitf);
declmethod(cbrt);
declmethod(cbrtf);
declmethod(complex_adaptor);
declmethod(adapt_ndarray);
declmethod(ndarray_new);
Expand Down
4 changes: 2 additions & 2 deletions numba/core/typing/npydecl.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -244,10 +244,10 @@ def generic(self, args, kws):
_math_operations = [ "add", "subtract", "multiply",
"logaddexp", "logaddexp2", "true_divide",
"floor_divide", "negative", "positive", "power",
"remainder", "fmod", "absolute",
"float_power", "remainder", "fmod", "absolute",
"rint", "sign", "conjugate", "exp", "exp2",
"log", "log2", "log10", "expm1", "log1p",
"sqrt", "square", "reciprocal",
"sqrt", "square", "cbrt", "reciprocal",
"divide", "mod", "divmod", "abs", "fabs" , "gcd", "lcm"]

_trigonometric_functions = [ "sin", "cos", "tan", "arcsin",
Expand Down
30 changes: 30 additions & 0 deletions numba/np/npyfuncs.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -480,6 +480,21 @@ def np_complex_power_impl(context, builder, sig, args):
return numbers.complex_power_impl(context, builder, sig, args)


########################################################################
# numpy float power funcs

def real_float_power_impl(context, builder, sig, args):
_check_arity_and_homogeneity(sig, args, 2)

return numbers.real_power_impl(context, builder, sig, args)


def np_complex_float_power_impl(context, builder, sig, args):
_check_arity_and_homogeneity(sig, args, 2)

return numbers.complex_power_impl(context, builder, sig, args)


########################################################################
# numpy greatest common denominator

Expand Down Expand Up @@ -753,6 +768,21 @@ def np_complex_square_impl(context, builder, sig, args):
[args[0], args[0]])


########################################################################
# NumPy cbrt

def np_real_cbrt_impl(context, builder, sig, args):
_check_arity_and_homogeneity(sig, args, 1)

dispatch_table = {
types.float32: 'numba_cbrtf',
types.float64: 'numba_cbrt',
}

return _dispatch_func_by_name_type(context, builder, sig, args,
dispatch_table, 'cbrt')


########################################################################
# NumPy reciprocal

Expand Down
12 changes: 12 additions & 0 deletions numba/np/ufunc_db.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -309,6 +309,13 @@ def _fill_ufunc_db(ufunc_db):
'DD->D': npyfuncs.np_complex_power_impl,
}

ufunc_db[np.float_power] = {
'ff->f': npyfuncs.real_float_power_impl,
'dd->d': npyfuncs.real_float_power_impl,
'FF->F': npyfuncs.np_complex_float_power_impl,
'DD->D': npyfuncs.np_complex_float_power_impl,
}

ufunc_db[np.gcd] = {
'bb->b': npyfuncs.np_gcd_impl,
'BB->B': npyfuncs.np_gcd_impl,
Expand Down Expand Up @@ -432,6 +439,11 @@ def _fill_ufunc_db(ufunc_db):
'D->D': npyfuncs.np_complex_square_impl,
}

ufunc_db[np.cbrt] = {
'f->f': npyfuncs.np_real_cbrt_impl,
'd->d': npyfuncs.np_real_cbrt_impl,
}

ufunc_db[np.reciprocal] = {
'b->b': npyfuncs.np_int_reciprocal_impl,
'B->B': npyfuncs.np_int_reciprocal_impl,
Expand Down
7 changes: 7 additions & 0 deletions numba/tests/test_ufuncs.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -310,6 +310,9 @@ def test_power_ufunc(self, flags=no_pyobj_flags):
self.basic_ufunc_test(np.power, flags=flags,
positive_only=True)

def test_float_power_ufunc(self, flags=no_pyobj_flags):
self.basic_ufunc_test(np.float_power, flags=flags, kinds="fc")

def test_gcd_ufunc(self, flags=no_pyobj_flags):
self.basic_ufunc_test(np.gcd, flags=flags, kinds="iu")

Expand Down Expand Up @@ -379,6 +382,9 @@ def test_sqrt_ufunc(self, flags=no_pyobj_flags):
def test_square_ufunc(self, flags=no_pyobj_flags):
self.basic_ufunc_test(np.square, flags=flags)

def test_cbrt_ufunc(self, flags=no_pyobj_flags):
self.basic_ufunc_test(np.cbrt, flags=flags, kinds='f')

def test_reciprocal_ufunc(self, flags=no_pyobj_flags):
# reciprocal for integers doesn't make much sense and is problematic
# in the case of division by zero, as an inf will overflow float to
Expand Down Expand Up @@ -1329,6 +1335,7 @@ class _LoopTypesTester(TestCase):
('arcsin', 'F'): 4,
('log10', 'D'): 5,
('tanh', 'F'): 2,
('cbrt', 'd'): 2,
}

def _arg_for_type(self, a_letter_type, index=0):
Expand Down