5684590474

# SPDX-FileCopyrightText: 2020 - 2023 Intel Corporation
#
# SPDX-License-Identifier: Apache-2.0

import operator
from functools import reduce

import dpctl
from llvmlite import binding as ll
from llvmlite import ir as llvmir
from numba.core import cgutils, types
from numba.core.imputils import Registry
from numba.core.typing.npydecl import parse_dtype

from numba_dpex import config, kernel_target
from numba_dpex.core.codegen import SPIR_DATA_LAYOUT
from numba_dpex.core.types import Array
from numba_dpex.ocl.atomics import atomic_helper
from numba_dpex.utils import address_space

from . import stubs
from ._declare_function import _declare_function

registry = Registry()
lower = registry.lower

_void_value = llvmir.Constant(llvmir.IntType(8).as_pointer(), None)


# -----------------------------------------------------------------------------


@lower(stubs.get_global_id, types.uint32)
def get_global_id_impl(context, builder, sig, args):
    [dim] = args
    get_global_id = _declare_function(
        context, builder, "get_global_id", sig, ["unsigned int"]
    )
    res = builder.call(get_global_id, [dim])
    return context.cast(builder, res, types.uintp, types.intp)


@lower(stubs.get_local_id, types.uint32)
def get_local_id_impl(context, builder, sig, args):
    [dim] = args
    get_local_id = _declare_function(
        context, builder, "get_local_id", sig, ["unsigned int"]
    )
    res = builder.call(get_local_id, [dim])
    return context.cast(builder, res, types.uintp, types.intp)


@lower(stubs.get_group_id, types.uint32)
def get_group_id_impl(context, builder, sig, args):
    [dim] = args
    get_group_id = _declare_function(
        context, builder, "get_group_id", sig, ["unsigned int"]
    )
    res = builder.call(get_group_id, [dim])
    return context.cast(builder, res, types.uintp, types.intp)


@lower(stubs.get_num_groups, types.uint32)
def get_num_groups_impl(context, builder, sig, args):
    [dim] = args
    get_num_groups = _declare_function(
        context, builder, "get_num_groups", sig, ["unsigned int"]
    )
    res = builder.call(get_num_groups, [dim])
    return context.cast(builder, res, types.uintp, types.intp)


@lower(stubs.get_work_dim)
def get_work_dim_impl(context, builder, sig, args):
    get_work_dim = _declare_function(
        context, builder, "get_work_dim", sig, ["void"]
    )
    res = builder.call(get_work_dim, [])
    return res


@lower(stubs.get_global_size, types.uint32)
def get_global_size_impl(context, builder, sig, args):
    [dim] = args
    get_global_size = _declare_function(
        context, builder, "get_global_size", sig, ["unsigned int"]
    )
    res = builder.call(get_global_size, [dim])
    return context.cast(builder, res, types.uintp, types.intp)


@lower(stubs.get_local_size, types.uint32)
def get_local_size_impl(context, builder, sig, args):
    [dim] = args
    get_local_size = _declare_function(
        context, builder, "get_local_size", sig, ["unsigned int"]
    )
    res = builder.call(get_local_size, [dim])
    return context.cast(builder, res, types.uintp, types.intp)


@lower(stubs.barrier, types.uint32)
def barrier_one_arg_impl(context, builder, sig, args):
    [flags] = args
    barrier = _declare_function(
        context, builder, "barrier", sig, ["unsigned int"]
    )
    builder.call(barrier, [flags])
    return _void_value


@lower(stubs.barrier)
def barrier_no_arg_impl(context, builder, sig, args):
    assert not args
    sig = types.void(types.uint32)
    barrier = _declare_function(
        context, builder, "barrier", sig, ["unsigned int"]
    )
    flags = context.get_constant(types.uint32, stubs.GLOBAL_MEM_FENCE)
    builder.call(barrier, [flags])
    return _void_value


@lower(stubs.mem_fence, types.uint32)
def mem_fence_impl(context, builder, sig, args):
    [flags] = args
    mem_fence = _declare_function(
        context, builder, "mem_fence", sig, ["unsigned int"]
    )
    builder.call(mem_fence, [flags])
    return _void_value


@lower(stubs.sub_group_barrier)
def sub_group_barrier_impl(context, builder, sig, args):
    assert not args
    sig = types.void(types.uint32)
    barrier = _declare_function(
        context, builder, "barrier", sig, ["unsigned int"]
    )
    flags = context.get_constant(types.uint32, stubs.LOCAL_MEM_FENCE)
    builder.call(barrier, [flags])
    return _void_value


def native_atomic_add(context, builder, sig, args):
    aryty, indty, valty = sig.args
    ary, inds, val = args
    dtype = aryty.dtype

    if indty == types.intp:
        indices = [inds]  # just a single integer
        indty = [indty]
    else:
        indices = cgutils.unpack_tuple(builder, inds, count=len(indty))
        indices = [
            context.cast(builder, i, t, types.intp)
            for t, i in zip(indty, indices)
        ]

    if dtype != valty:
        raise TypeError("expecting %s but got %s" % (dtype, valty))

    if aryty.ndim != len(indty):
        raise TypeError(
            "indexing %d-D array with %d-D index" % (aryty.ndim, len(indty))
        )

    lary = context.make_array(aryty)(context, builder, ary)
    ptr = cgutils.get_item_pointer(context, builder, aryty, lary, indices)

    if dtype == types.float32 or dtype == types.float64:
        context.extra_compile_options[kernel_target.LLVM_SPIRV_ARGS] = [
            "--spirv-ext=+SPV_EXT_shader_atomic_float_add"
        ]
        name = "__spirv_AtomicFAddEXT"
    elif dtype == types.int32 or dtype == types.int64:
        name = "__spirv_AtomicIAdd"
    else:
        raise TypeError("Unsupported type")

    assert name != ""

    ptr_type = context.get_value_type(dtype).as_pointer()
    ptr_type.addrspace = aryty.addrspace

    retty = context.get_value_type(sig.return_type)
    spirv_fn_arg_types = [
        ptr_type,
        llvmir.IntType(32),
        llvmir.IntType(32),
        context.get_value_type(sig.args[2]),
    ]

    from numba_dpex.core import itanium_mangler as ext_itanium_mangler

    numba_ptr_ty = types.CPointer(dtype, addrspace=ptr_type.addrspace)
    mangled_fn_name = ext_itanium_mangler.mangle_ext(
        name,
        [
            numba_ptr_ty,
            "__spv.Scope.Flag",
            "__spv.MemorySemanticsMask.Flag",
            valty,
        ],
    )

    fnty = llvmir.FunctionType(retty, spirv_fn_arg_types)
    fn = cgutils.get_or_insert_function(builder.module, fnty, mangled_fn_name)
    fn.calling_convention = kernel_target.CC_SPIR_FUNC

    sycl_memory_order = atomic_helper.sycl_memory_order.relaxed
    sycl_memory_scope = atomic_helper.sycl_memory_scope.device
    spirv_scope = atomic_helper.get_scope(sycl_memory_scope)
    spirv_memory_semantics_mask = atomic_helper.get_memory_semantics_mask(
        sycl_memory_order
    )
    fn_args = [
        ptr,
        context.get_constant(types.int32, spirv_scope),
        context.get_constant(types.int32, spirv_memory_semantics_mask),
        val,
    ]

    return builder.call(fn, fn_args)


def support_atomic(dtype: types.Type) -> bool:
    # This check should be the same as described in sycl documentation:
    # https://registry.khronos.org/SYCL/specs/sycl-2020/html/sycl-2020.html#sec:atomic-references
    return (
        dtype == types.int32
        or dtype == types.uint32
        or dtype == types.float32
        or (
            dtype == types.int64
            or dtype == types.uint64
            or dtype == types.float64
        )
        and dpctl.get_current_queue().sycl_device.has_aspect_atomic64
    )


@lower(stubs.atomic.add, types.Array, types.intp, types.Any)
@lower(stubs.atomic.add, types.Array, types.UniTuple, types.Any)
@lower(stubs.atomic.add, types.Array, types.Tuple, types.Any)
def atomic_add_tuple(context, builder, sig, args):
    dtype = sig.args[0].dtype
    if support_atomic(dtype):
        return native_atomic_add(context, builder, sig, args)
    else:
        raise TypeError(f"Atomic operation on unsupported type {dtype}")


def atomic_sub_wrapper(context, builder, sig, args):
    # dpcpp yet does not support ``__spirv_AtomicFSubEXT``. To support atomic.sub we
    # reuse atomic.add and negate the value. For example, atomic.add(A, index, -val) is
    # equivalent to atomic.sub(A, index, val).
    val = args[2]
    new_val = cgutils.alloca_once(
        builder,
        context.get_value_type(sig.args[2]),
        size=context.get_constant(types.uintp, 1),
        name="new_val_0",
    )
    val_dtype = sig.args[2]
    if val_dtype == types.float32 or val_dtype == types.float64:
        builder.store(
            builder.fmul(val, context.get_constant(sig.args[2], -1)), new_val
        )
    elif val_dtype == types.int32 or val_dtype == types.int64:
        builder.store(
            builder.mul(val, context.get_constant(sig.args[2], -1)), new_val
        )
    else:
        raise TypeError("Unsupported type %s" % val_dtype)

    args[2] = builder.load(new_val)

    return native_atomic_add(context, builder, sig, args)


@lower(stubs.atomic.sub, types.Array, types.intp, types.Any)
@lower(stubs.atomic.sub, types.Array, types.UniTuple, types.Any)
@lower(stubs.atomic.sub, types.Array, types.Tuple, types.Any)
def atomic_sub_tuple(context, builder, sig, args):
    dtype = sig.args[0].dtype
    if support_atomic(dtype):
        return atomic_sub_wrapper(context, builder, sig, args)
    else:
        raise TypeError(f"Atomic operation on unsupported type {dtype}")


@lower(stubs.private.array, types.IntegerLiteral, types.Any)
def dpex_private_array_integer(context, builder, sig, args):
    length = sig.args[0].literal_value
    dtype = parse_dtype(sig.args[1])
    return _generic_array(
        context,
        builder,
        shape=(length,),
        dtype=dtype,
        symbol_name="_dpex_pmem",
        addrspace=address_space.PRIVATE,
    )


@lower(stubs.private.array, types.Tuple, types.Any)
@lower(stubs.private.array, types.UniTuple, types.Any)
def dpex_private_array_tuple(context, builder, sig, args):
    shape = [s.literal_value for s in sig.args[0]]
    dtype = parse_dtype(sig.args[1])
    return _generic_array(
        context,
        builder,
        shape=shape,
        dtype=dtype,
        symbol_name="_dpex_pmem",
        addrspace=address_space.PRIVATE,
    )


@lower(stubs.local.array, types.IntegerLiteral, types.Any)
def dpex_local_array_integer(context, builder, sig, args):
    length = sig.args[0].literal_value
    dtype = parse_dtype(sig.args[1])
    return _generic_array(
        context,
        builder,
        shape=(length,),
        dtype=dtype,
        symbol_name="_dpex_lmem",
        addrspace=address_space.LOCAL,
    )


@lower(stubs.local.array, types.Tuple, types.Any)
@lower(stubs.local.array, types.UniTuple, types.Any)
def dpex_local_array_tuple(context, builder, sig, args):
    shape = [s.literal_value for s in sig.args[0]]
    dtype = parse_dtype(sig.args[1])
    return _generic_array(
        context,
        builder,
        shape=shape,
        dtype=dtype,
        symbol_name="_dpex_lmem",
        addrspace=address_space.LOCAL,
    )


def _generic_array(context, builder, shape, dtype, symbol_name, addrspace):
    """
    This function allows us to create generic arrays in different
    address spaces.
    """
    elemcount = reduce(operator.mul, shape)
    lldtype = context.get_data_type(dtype)
    laryty = llvmir.ArrayType(lldtype, elemcount)

    if addrspace == address_space.LOCAL:
        lmod = builder.module

        # Create global variable in the requested address-space
        gvmem = cgutils.add_global_variable(
            lmod, laryty, symbol_name, addrspace
        )

        if elemcount <= 0:
            raise ValueError("array length <= 0")
        else:
            gvmem.linkage = "internal"

        if dtype not in types.number_domain:
            raise TypeError("unsupported type: %s" % dtype)

    elif addrspace == address_space.PRIVATE:
        gvmem = cgutils.alloca_once(builder, laryty, name=symbol_name)
    else:
        raise NotImplementedError("addrspace {addrspace}".format(**locals()))

    # We need to add the addrspace to _make_array() function call as we want
    # the variable containing the reference of the memory to retain the
    # original address space of that memory. Before, we were casting the
    # memories allocated in local address space to global address space. This
    # approach does not let us identify the original address space of a memory
    # down the line.
    return _make_array(
        context, builder, gvmem, dtype, shape, addrspace=addrspace
    )


def _make_array(
    context,
    builder,
    dataptr,
    dtype,
    shape,
    layout="C",
    addrspace=address_space.GENERIC,
):
    ndim = len(shape)
    # Create array object
    aryty = Array(dtype=dtype, ndim=ndim, layout="C", addrspace=addrspace)
    ary = context.make_array(aryty)(context, builder)

    targetdata = _get_target_data(context)
    lldtype = context.get_data_type(dtype)
    itemsize = lldtype.get_abi_size(targetdata)
    # Compute strides
    rstrides = [itemsize]
    for i, lastsize in enumerate(reversed(shape[1:])):
        rstrides.append(lastsize * rstrides[-1])
    strides = [s for s in reversed(rstrides)]

    kshape = [context.get_constant(types.intp, s) for s in shape]
    kstrides = [context.get_constant(types.intp, s) for s in strides]

    context.populate_array(
        ary,
        data=builder.bitcast(dataptr, ary.data.type),
        shape=cgutils.pack_array(builder, kshape),
        strides=cgutils.pack_array(builder, kstrides),
        itemsize=context.get_constant(types.intp, itemsize),
        meminfo=None,
    )

    return ary._getvalue()


def _get_target_data(context):
    return ll.create_target_data(SPIR_DATA_LAYOUT[context.address_size])

1	# SPDX-FileCopyrightText: 2020 - 2023 Intel Corporation
2	#
3	# SPDX-License-Identifier: Apache-2.0
4
5	import operator	1✔
6	from functools import reduce	1✔
7
8	import dpctl	1✔
9	from llvmlite import binding as ll	1✔
10	from llvmlite import ir as llvmir	1✔
11	from numba.core import cgutils, types	1✔
12	from numba.core.imputils import Registry	1✔
13	from numba.core.typing.npydecl import parse_dtype	1✔
14
15	from numba_dpex import config, kernel_target	1✔
16	from numba_dpex.core.codegen import SPIR_DATA_LAYOUT	1✔
17	from numba_dpex.core.types import Array	1✔
18	from numba_dpex.ocl.atomics import atomic_helper	1✔
19	from numba_dpex.utils import address_space	1✔
20
21	from . import stubs	1✔
22	from ._declare_function import _declare_function	1✔
23
24	registry = Registry()	1✔
25	lower = registry.lower	1✔
26
27	_void_value = llvmir.Constant(llvmir.IntType(8).as_pointer(), None)	1✔
28
29
30	# -----------------------------------------------------------------------------
31
32
33	@lower(stubs.get_global_id, types.uint32)	1✔
34	def get_global_id_impl(context, builder, sig, args):	1✔
35	[dim] = args	1✔
36	get_global_id = _declare_function(	1✔
37	context, builder, "get_global_id", sig, ["unsigned int"]
38	)
39	res = builder.call(get_global_id, [dim])	1✔
40	return context.cast(builder, res, types.uintp, types.intp)	1✔
41
42
43	@lower(stubs.get_local_id, types.uint32)	1✔
44	def get_local_id_impl(context, builder, sig, args):	1✔
45	[dim] = args	1✔
46	get_local_id = _declare_function(	1✔
47	context, builder, "get_local_id", sig, ["unsigned int"]
48	)
49	res = builder.call(get_local_id, [dim])	1✔
50	return context.cast(builder, res, types.uintp, types.intp)	1✔
51
52
53	@lower(stubs.get_group_id, types.uint32)	1✔
54	def get_group_id_impl(context, builder, sig, args):	1✔
55	[dim] = args	1✔
56	get_group_id = _declare_function(	1✔
57	context, builder, "get_group_id", sig, ["unsigned int"]
58	)
59	res = builder.call(get_group_id, [dim])	1✔
60	return context.cast(builder, res, types.uintp, types.intp)	1✔
61
62
63	@lower(stubs.get_num_groups, types.uint32)	1✔
64	def get_num_groups_impl(context, builder, sig, args):	1✔
65	[dim] = args	×
66	get_num_groups = _declare_function(	×
67	context, builder, "get_num_groups", sig, ["unsigned int"]
68	)
69	res = builder.call(get_num_groups, [dim])	×
70	return context.cast(builder, res, types.uintp, types.intp)	×
71
72
73	@lower(stubs.get_work_dim)	1✔
74	def get_work_dim_impl(context, builder, sig, args):	1✔
75	get_work_dim = _declare_function(	×
76	context, builder, "get_work_dim", sig, ["void"]
77	)
78	res = builder.call(get_work_dim, [])	×
79	return res	×
80
81
82	@lower(stubs.get_global_size, types.uint32)	1✔
83	def get_global_size_impl(context, builder, sig, args):	1✔
84	[dim] = args	×
85	get_global_size = _declare_function(	×
86	context, builder, "get_global_size", sig, ["unsigned int"]
87	)
88	res = builder.call(get_global_size, [dim])	×
89	return context.cast(builder, res, types.uintp, types.intp)	×
90
91
92	@lower(stubs.get_local_size, types.uint32)	1✔
93	def get_local_size_impl(context, builder, sig, args):	1✔
94	[dim] = args	1✔
95	get_local_size = _declare_function(	1✔
96	context, builder, "get_local_size", sig, ["unsigned int"]
97	)
98	res = builder.call(get_local_size, [dim])	1✔
99	return context.cast(builder, res, types.uintp, types.intp)	1✔
100
101
102	@lower(stubs.barrier, types.uint32)	1✔
103	def barrier_one_arg_impl(context, builder, sig, args):	1✔
104	[flags] = args	1✔
105	barrier = _declare_function(	1✔
106	context, builder, "barrier", sig, ["unsigned int"]
107	)
108	builder.call(barrier, [flags])	1✔
109	return _void_value	1✔
110
111
112	@lower(stubs.barrier)	1✔
113	def barrier_no_arg_impl(context, builder, sig, args):	1✔
114	assert not args	1✔
115	sig = types.void(types.uint32)	1✔
116	barrier = _declare_function(	1✔
117	context, builder, "barrier", sig, ["unsigned int"]
118	)
119	flags = context.get_constant(types.uint32, stubs.GLOBAL_MEM_FENCE)	1✔
120	builder.call(barrier, [flags])	1✔
121	return _void_value	1✔
122
123
124	@lower(stubs.mem_fence, types.uint32)	1✔
125	def mem_fence_impl(context, builder, sig, args):	1✔
126	[flags] = args	×
127	mem_fence = _declare_function(	×
128	context, builder, "mem_fence", sig, ["unsigned int"]
129	)
130	builder.call(mem_fence, [flags])	×
131	return _void_value	×
132
133
134	@lower(stubs.sub_group_barrier)	1✔
135	def sub_group_barrier_impl(context, builder, sig, args):	1✔
136	assert not args	×
137	sig = types.void(types.uint32)	×
138	barrier = _declare_function(	×
139	context, builder, "barrier", sig, ["unsigned int"]
140	)
141	flags = context.get_constant(types.uint32, stubs.LOCAL_MEM_FENCE)	×
142	builder.call(barrier, [flags])	×
143	return _void_value	×
144
145
146	def native_atomic_add(context, builder, sig, args):	1✔
UNCOV 147	aryty, indty, valty = sig.args	×
UNCOV 148	ary, inds, val = args	×
UNCOV 149	dtype = aryty.dtype	×
150
UNCOV 151	if indty == types.intp:	×
UNCOV 152	indices = [inds] # just a single integer	×
UNCOV 153	indty = [indty]	×
154	else:
UNCOV 155	indices = cgutils.unpack_tuple(builder, inds, count=len(indty))	×
UNCOV 156	indices = [	×
157	context.cast(builder, i, t, types.intp)
158	for t, i in zip(indty, indices)
159	]
160
UNCOV 161	if dtype != valty:	×
162	raise TypeError("expecting %s but got %s" % (dtype, valty))	×
163
UNCOV 164	if aryty.ndim != len(indty):	×
165	raise TypeError(	×
166	"indexing %d-D array with %d-D index" % (aryty.ndim, len(indty))
167	)
168
UNCOV 169	lary = context.make_array(aryty)(context, builder, ary)	×
UNCOV 170	ptr = cgutils.get_item_pointer(context, builder, aryty, lary, indices)	×
171
UNCOV 172	if dtype == types.float32 or dtype == types.float64:	×
173	context.extra_compile_options[kernel_target.LLVM_SPIRV_ARGS] = [	×
174	"--spirv-ext=+SPV_EXT_shader_atomic_float_add"
175	]
176	name = "__spirv_AtomicFAddEXT"	×
UNCOV 177	elif dtype == types.int32 or dtype == types.int64:	×
UNCOV 178	name = "__spirv_AtomicIAdd"	×
179	else:
180	raise TypeError("Unsupported type")	×
181
UNCOV 182	assert name != ""	×
183
UNCOV 184	ptr_type = context.get_value_type(dtype).as_pointer()	×
UNCOV 185	ptr_type.addrspace = aryty.addrspace	×
186
UNCOV 187	retty = context.get_value_type(sig.return_type)	×
UNCOV 188	spirv_fn_arg_types = [	×
189	ptr_type,
190	llvmir.IntType(32),
191	llvmir.IntType(32),
192	context.get_value_type(sig.args[2]),
193	]
194
UNCOV 195	from numba_dpex.core import itanium_mangler as ext_itanium_mangler	×
196
UNCOV 197	numba_ptr_ty = types.CPointer(dtype, addrspace=ptr_type.addrspace)	×
UNCOV 198	mangled_fn_name = ext_itanium_mangler.mangle_ext(	×
199	name,
200	[
201	numba_ptr_ty,
202	"__spv.Scope.Flag",
203	"__spv.MemorySemanticsMask.Flag",
204	valty,
205	],
206	)
207
UNCOV 208	fnty = llvmir.FunctionType(retty, spirv_fn_arg_types)	×
UNCOV 209	fn = cgutils.get_or_insert_function(builder.module, fnty, mangled_fn_name)	×
UNCOV 210	fn.calling_convention = kernel_target.CC_SPIR_FUNC	×
211
UNCOV 212	sycl_memory_order = atomic_helper.sycl_memory_order.relaxed	×
UNCOV 213	sycl_memory_scope = atomic_helper.sycl_memory_scope.device	×
UNCOV 214	spirv_scope = atomic_helper.get_scope(sycl_memory_scope)	×
UNCOV 215	spirv_memory_semantics_mask = atomic_helper.get_memory_semantics_mask(	×
216	sycl_memory_order
217	)
UNCOV 218	fn_args = [	×
219	ptr,
220	context.get_constant(types.int32, spirv_scope),
221	context.get_constant(types.int32, spirv_memory_semantics_mask),
222	val,
223	]
224
UNCOV 225	return builder.call(fn, fn_args)	×
226
227
228	def support_atomic(dtype: types.Type) -> bool:	1✔
229	# This check should be the same as described in sycl documentation:
230	# https://registry.khronos.org/SYCL/specs/sycl-2020/html/sycl-2020.html#sec:atomic-references
NEW 231	return (	×
232	dtype == types.int32
233	or dtype == types.uint32
234	or dtype == types.float32
235	or (
236	dtype == types.int64
237	or dtype == types.uint64
238	or dtype == types.float64
239	)
240	and dpctl.get_current_queue().sycl_device.has_aspect_atomic64
241	)
242
243
244	@lower(stubs.atomic.add, types.Array, types.intp, types.Any)	1✔
245	@lower(stubs.atomic.add, types.Array, types.UniTuple, types.Any)	1✔
246	@lower(stubs.atomic.add, types.Array, types.Tuple, types.Any)	1✔
247	def atomic_add_tuple(context, builder, sig, args):	1✔
UNCOV 248	dtype = sig.args[0].dtype	×
NEW 249	if support_atomic(dtype):	×
UNCOV 250	return native_atomic_add(context, builder, sig, args)	×
251	else:
NEW 252	raise TypeError(f"Atomic operation on unsupported type {dtype}")	×
253
254
255	def atomic_sub_wrapper(context, builder, sig, args):	1✔
256	# dpcpp yet does not support ``__spirv_AtomicFSubEXT``. To support atomic.sub we
257	# reuse atomic.add and negate the value. For example, atomic.add(A, index, -val) is
258	# equivalent to atomic.sub(A, index, val).
UNCOV 259	val = args[2]	×
UNCOV 260	new_val = cgutils.alloca_once(	×
261	builder,
262	context.get_value_type(sig.args[2]),
263	size=context.get_constant(types.uintp, 1),
264	name="new_val_0",
265	)
UNCOV 266	val_dtype = sig.args[2]	×
UNCOV 267	if val_dtype == types.float32 or val_dtype == types.float64:	×
268	builder.store(	×
269	builder.fmul(val, context.get_constant(sig.args[2], -1)), new_val
270	)
UNCOV 271	elif val_dtype == types.int32 or val_dtype == types.int64:	×
UNCOV 272	builder.store(	×
273	builder.mul(val, context.get_constant(sig.args[2], -1)), new_val
274	)
275	else:
276	raise TypeError("Unsupported type %s" % val_dtype)	×
277
UNCOV 278	args[2] = builder.load(new_val)	×
279
UNCOV 280	return native_atomic_add(context, builder, sig, args)	×
281
282
283	@lower(stubs.atomic.sub, types.Array, types.intp, types.Any)	1✔
284	@lower(stubs.atomic.sub, types.Array, types.UniTuple, types.Any)	1✔
285	@lower(stubs.atomic.sub, types.Array, types.Tuple, types.Any)	1✔
286	def atomic_sub_tuple(context, builder, sig, args):	1✔
UNCOV 287	dtype = sig.args[0].dtype	×
NEW 288	if support_atomic(dtype):	×
UNCOV 289	return atomic_sub_wrapper(context, builder, sig, args)	×
290	else:
NEW 291	raise TypeError(f"Atomic operation on unsupported type {dtype}")	×
292
293
294	@lower(stubs.private.array, types.IntegerLiteral, types.Any)	1✔
295	def dpex_private_array_integer(context, builder, sig, args):	1✔
296	length = sig.args[0].literal_value	1✔
297	dtype = parse_dtype(sig.args[1])	1✔
298	return _generic_array(	1✔
299	context,
300	builder,
301	shape=(length,),
302	dtype=dtype,
303	symbol_name="_dpex_pmem",
304	addrspace=address_space.PRIVATE,
305	)
306
307
308	@lower(stubs.private.array, types.Tuple, types.Any)	1✔
309	@lower(stubs.private.array, types.UniTuple, types.Any)	1✔
310	def dpex_private_array_tuple(context, builder, sig, args):	1✔
311	shape = [s.literal_value for s in sig.args[0]]	×
312	dtype = parse_dtype(sig.args[1])	×
313	return _generic_array(	×
314	context,
315	builder,
316	shape=shape,
317	dtype=dtype,
318	symbol_name="_dpex_pmem",
319	addrspace=address_space.PRIVATE,
320	)
321
322
323	@lower(stubs.local.array, types.IntegerLiteral, types.Any)	1✔
324	def dpex_local_array_integer(context, builder, sig, args):	1✔
325	length = sig.args[0].literal_value	1✔
326	dtype = parse_dtype(sig.args[1])	1✔
327	return _generic_array(	1✔
328	context,
329	builder,
330	shape=(length,),
331	dtype=dtype,
332	symbol_name="_dpex_lmem",
333	addrspace=address_space.LOCAL,
334	)
335
336
337	@lower(stubs.local.array, types.Tuple, types.Any)	1✔
338	@lower(stubs.local.array, types.UniTuple, types.Any)	1✔
339	def dpex_local_array_tuple(context, builder, sig, args):	1✔
340	shape = [s.literal_value for s in sig.args[0]]	×
341	dtype = parse_dtype(sig.args[1])	×
342	return _generic_array(	×
343	context,
344	builder,
345	shape=shape,
346	dtype=dtype,
347	symbol_name="_dpex_lmem",
348	addrspace=address_space.LOCAL,
349	)
350
351
352	def _generic_array(context, builder, shape, dtype, symbol_name, addrspace):	1✔
353	"""
354	This function allows us to create generic arrays in different
355	address spaces.
356	"""
357	elemcount = reduce(operator.mul, shape)	1✔
358	lldtype = context.get_data_type(dtype)	1✔
359	laryty = llvmir.ArrayType(lldtype, elemcount)	1✔
360
361	if addrspace == address_space.LOCAL:	1✔
362	lmod = builder.module	1✔
363
364	# Create global variable in the requested address-space
365	gvmem = cgutils.add_global_variable(	1✔
366	lmod, laryty, symbol_name, addrspace
367	)
368
369	if elemcount <= 0:	1!
370	raise ValueError("array length <= 0")	×
371	else:
372	gvmem.linkage = "internal"	1✔
373
374	if dtype not in types.number_domain:	1!
375	raise TypeError("unsupported type: %s" % dtype)	×
376
377	elif addrspace == address_space.PRIVATE:	1✔
378	gvmem = cgutils.alloca_once(builder, laryty, name=symbol_name)	1✔
379	else:
380	raise NotImplementedError("addrspace {addrspace}".format(**locals()))
381
382	# We need to add the addrspace to _make_array() function call as we want
383	# the variable containing the reference of the memory to retain the
384	# original address space of that memory. Before, we were casting the
385	# memories allocated in local address space to global address space. This
386	# approach does not let us identify the original address space of a memory
387	# down the line.
388	return _make_array(	1✔
389	context, builder, gvmem, dtype, shape, addrspace=addrspace
390	)
391
392
393	def _make_array(	1✔
394	context,
395	builder,
396	dataptr,
397	dtype,
398	shape,
399	layout="C",
400	addrspace=address_space.GENERIC,
401	):
402	ndim = len(shape)	1✔
403	# Create array object
404	aryty = Array(dtype=dtype, ndim=ndim, layout="C", addrspace=addrspace)	1✔
405	ary = context.make_array(aryty)(context, builder)	1✔
406
407	targetdata = _get_target_data(context)	1✔
408	lldtype = context.get_data_type(dtype)	1✔
409	itemsize = lldtype.get_abi_size(targetdata)	1✔
410	# Compute strides
411	rstrides = [itemsize]	1✔
412	for i, lastsize in enumerate(reversed(shape[1:])):	1!
413	rstrides.append(lastsize * rstrides[-1])	×
414	strides = [s for s in reversed(rstrides)]	1✔
415
416	kshape = [context.get_constant(types.intp, s) for s in shape]	1✔
417	kstrides = [context.get_constant(types.intp, s) for s in strides]	1✔
418
419	context.populate_array(	1✔
420	ary,
421	data=builder.bitcast(dataptr, ary.data.type),
422	shape=cgutils.pack_array(builder, kshape),
423	strides=cgutils.pack_array(builder, kstrides),
424	itemsize=context.get_constant(types.intp, itemsize),
425	meminfo=None,
426	)
427
428	return ary._getvalue()	1✔
429
430
431	def _get_target_data(context):	1✔
432	return ll.create_target_data(SPIR_DATA_LAYOUT[context.address_size])	1✔

IntelPython / numba-dpex / 5684590474

Source File Press 'n' to go to next uncovered line, 'b' for previous

Source File
Press 'n' to go to next uncovered line, 'b' for previous