|
26 | 26 | RefType) |
27 | 27 | from numba.core.typing import signature |
28 | 28 | from numba.core.extending import (register_jitable, overload, overload_method, |
29 | | - intrinsic, overload_attribute) |
| 29 | + intrinsic) |
30 | 30 | from numba.misc import quicksort, mergesort |
31 | 31 | from numba.cpython import slicing |
32 | 32 | from numba.cpython.unsafe.tuple import tuple_setitem, build_full_slice_tuple |
@@ -2618,127 +2618,87 @@ def array_flags_f_contiguous(context, builder, typ, value): |
2618 | 2618 | # ------------------------------------------------------------------------------ |
2619 | 2619 | # .real / .imag |
2620 | 2620 |
|
2621 | | -@overload_attribute(types.Array, 'real') |
2622 | | -def ol_array_real(arr): |
2623 | | - typ = arr |
| 2621 | +@lower_getattr(types.Array, "real") |
| 2622 | +def array_real_part(context, builder, typ, value): |
2624 | 2623 | if typ.dtype in types.complex_domain: |
2625 | | - def impl(arr): |
2626 | | - return _array_real_attr(arr) |
2627 | | - return impl |
| 2624 | + return array_complex_attr(context, builder, typ, value, attr='real') |
2628 | 2625 | elif typ.dtype in types.number_domain: |
2629 | | - def impl(arr): |
2630 | | - return arr |
2631 | | - return impl |
| 2626 | + # as an identity function |
| 2627 | + return impl_ret_borrowed(context, builder, typ, value) |
2632 | 2628 | else: |
2633 | | - msg = f"cannot access .real of array of {arr.dtype}" |
2634 | | - raise errors.TypingError(msg) |
| 2629 | + raise NotImplementedError('unsupported .real for {}'.format(type.dtype)) |
2635 | 2630 |
|
2636 | 2631 |
|
2637 | | -@intrinsic |
2638 | | -def _force_readonly(tyctx, arr): |
2639 | | - """Takes an array, and returns it with the readonly bit set through |
2640 | | - typing""" |
2641 | | - sig = arr.copy(readonly=True)(arr) |
2642 | | - |
2643 | | - def impl(cgctx, builder, sig, llargs): |
2644 | | - arrayty = make_array(arr) |
2645 | | - array = arrayty(cgctx, builder, llargs[0]) |
2646 | | - return impl_ret_borrowed(cgctx, builder, sig.return_type, |
2647 | | - array._getvalue()) |
2648 | | - return sig, impl |
2649 | | - |
2650 | | - |
2651 | | -@overload_attribute(types.Array, 'imag') |
2652 | | -def ol_array_imag(arr): |
2653 | | - typ = arr |
| 2632 | +@lower_getattr(types.Array, "imag") |
| 2633 | +def array_imag_part(context, builder, typ, value): |
2654 | 2634 | if typ.dtype in types.complex_domain: |
2655 | | - def impl(arr): |
2656 | | - return _array_imag_attr(arr) |
2657 | | - return impl |
| 2635 | + return array_complex_attr(context, builder, typ, value, attr='imag') |
2658 | 2636 | elif typ.dtype in types.number_domain: |
2659 | | - def impl(arr): |
2660 | | - # .imag on a real domain dtype is a readonly zeros array |
2661 | | - tmp = np.zeros_like(arr) |
2662 | | - rotmp = _force_readonly(tmp) |
2663 | | - return rotmp |
2664 | | - return impl |
| 2637 | + # return a readonly zero array |
| 2638 | + sig = signature(typ.copy(readonly=True), typ) |
| 2639 | + arrtype, shapes = _parse_empty_like_args(context, builder, sig, [value]) |
| 2640 | + ary = _empty_nd_impl(context, builder, arrtype, shapes) |
| 2641 | + cgutils.memset(builder, ary.data, builder.mul(ary.itemsize, |
| 2642 | + ary.nitems), 0) |
| 2643 | + return impl_ret_new_ref(context, builder, sig.return_type, |
| 2644 | + ary._getvalue()) |
2665 | 2645 | else: |
2666 | | - msg = f"cannot access .imag of array of {arr.dtype}" |
2667 | | - raise errors.TypingError(msg) |
| 2646 | + raise NotImplementedError('unsupported .imag for {}'.format(type.dtype)) |
2668 | 2647 |
|
2669 | 2648 |
|
2670 | | -def _generate_real_imag_attr(attr): |
2671 | | - """ |
2672 | | - Generates intrinsics for obtaining the real/imaginary part of complex array |
| 2649 | +def array_complex_attr(context, builder, typ, value, attr): |
2673 | 2650 | """ |
| 2651 | + Given a complex array, it's memory layout is: |
2674 | 2652 |
|
2675 | | - @intrinsic |
2676 | | - def intrin_typing(tyctx, arr): |
2677 | | - sig = arr.copy(dtype=arr.dtype.underlying_float, layout='A')(arr) |
| 2653 | + R C R C R C |
| 2654 | + ^ ^ ^ |
2678 | 2655 |
|
2679 | | - def codegen(cgctx, builder, sig, llargs): |
2680 | | - value = llargs[0] |
2681 | | - typ = sig.args[0] |
2682 | | - """ |
2683 | | - Given a complex array, it's memory layout is: |
| 2656 | + (`R` indicates a float for the real part; |
| 2657 | + `C` indicates a float for the imaginary part; |
| 2658 | + the `^` indicates the start of each element) |
2684 | 2659 |
|
2685 | | - R C R C R C |
2686 | | - ^ ^ ^ |
| 2660 | + To get the real part, we can simply change the dtype and itemsize to that |
| 2661 | + of the underlying float type. The new layout is: |
2687 | 2662 |
|
2688 | | - (`R` indicates a float for the real part; |
2689 | | - `C` indicates a float for the imaginary part; |
2690 | | - the `^` indicates the start of each element) |
| 2663 | + R x R x R x |
| 2664 | + ^ ^ ^ |
2691 | 2665 |
|
2692 | | - To get the real part, we can simply change the dtype and itemsize to |
2693 | | - that of the underlying float type. The new layout is: |
| 2666 | + (`x` indicates unused) |
2694 | 2667 |
|
2695 | | - R x R x R x |
2696 | | - ^ ^ ^ |
| 2668 | + A load operation will use the dtype to determine the number of bytes to |
| 2669 | + load. |
2697 | 2670 |
|
2698 | | - (`x` indicates unused) |
| 2671 | + To get the imaginary part, we shift the pointer by 1 float offset and |
| 2672 | + change the dtype and itemsize. The new layout is: |
2699 | 2673 |
|
2700 | | - A load operation will use the dtype to determine the number of bytes |
2701 | | - to load. |
| 2674 | + x C x C x C |
| 2675 | + ^ ^ ^ |
| 2676 | + """ |
| 2677 | + if attr not in ['real', 'imag'] or typ.dtype not in types.complex_domain: |
| 2678 | + raise NotImplementedError("cannot get attribute `{}`".format(attr)) |
2702 | 2679 |
|
2703 | | - To get the imaginary part, we shift the pointer by 1 float offset |
2704 | | - and change the dtype and itemsize. The new layout is: |
| 2680 | + arrayty = make_array(typ) |
| 2681 | + array = arrayty(context, builder, value) |
2705 | 2682 |
|
2706 | | - x C x C x C |
2707 | | - ^ ^ ^ |
2708 | | - """ |
2709 | | - if (attr not in ['real', 'imag'] or |
2710 | | - typ.dtype not in types.complex_domain): |
2711 | | - raise NotImplementedError(f"cannot get attribute `{attr}`") |
2712 | | - |
2713 | | - arrayty = make_array(typ) |
2714 | | - array = arrayty(cgctx, builder, value) |
2715 | | - |
2716 | | - # sizeof underlying float type |
2717 | | - flty = typ.dtype.underlying_float |
2718 | | - sizeof_flty = cgctx.get_abi_sizeof(cgctx.get_data_type(flty)) |
2719 | | - itemsize = array.itemsize.type(sizeof_flty) |
2720 | | - |
2721 | | - # cast data pointer to float type |
2722 | | - llfltptrty = cgctx.get_value_type(flty).as_pointer() |
2723 | | - dataptr = builder.bitcast(array.data, llfltptrty) |
2724 | | - |
2725 | | - # add offset |
2726 | | - if attr == 'imag': |
2727 | | - dataptr = builder.gep(dataptr, [ir.IntType(32)(1)]) |
2728 | | - |
2729 | | - # make result |
2730 | | - resultty = typ.copy(dtype=flty, layout='A') |
2731 | | - result = make_array(resultty)(cgctx, builder) |
2732 | | - repl = dict(data=dataptr, itemsize=itemsize) |
2733 | | - cgutils.copy_struct(result, array, repl) |
2734 | | - return impl_ret_borrowed(cgctx, builder, resultty, |
2735 | | - result._getvalue()) |
2736 | | - return sig, codegen |
2737 | | - return intrin_typing |
2738 | | - |
2739 | | - |
2740 | | -_array_real_attr = _generate_real_imag_attr("real") |
2741 | | -_array_imag_attr = _generate_real_imag_attr("imag") |
| 2683 | + # sizeof underlying float type |
| 2684 | + flty = typ.dtype.underlying_float |
| 2685 | + sizeof_flty = context.get_abi_sizeof(context.get_data_type(flty)) |
| 2686 | + itemsize = array.itemsize.type(sizeof_flty) |
| 2687 | + |
| 2688 | + # cast data pointer to float type |
| 2689 | + llfltptrty = context.get_value_type(flty).as_pointer() |
| 2690 | + dataptr = builder.bitcast(array.data, llfltptrty) |
| 2691 | + |
| 2692 | + # add offset |
| 2693 | + if attr == 'imag': |
| 2694 | + dataptr = builder.gep(dataptr, [ir.IntType(32)(1)]) |
| 2695 | + |
| 2696 | + # make result |
| 2697 | + resultty = typ.copy(dtype=flty, layout='A') |
| 2698 | + result = make_array(resultty)(context, builder) |
| 2699 | + repl = dict(data=dataptr, itemsize=itemsize) |
| 2700 | + cgutils.copy_struct(result, array, repl) |
| 2701 | + return impl_ret_borrowed(context, builder, resultty, result._getvalue()) |
2742 | 2702 |
|
2743 | 2703 |
|
2744 | 2704 | @overload_method(types.Array, 'conj') |
|
0 commit comments