5742028140

pending completion

Build # 5742028140

Build Type

push

github

Committed by

web-flow

Commit Message

Merge pull request #1103 from IntelPython/feature/remove_atomic_emulation

Remove atomic emulation

Run Details

1079 of 1626 branches covered (66.36%)

Branch coverage included in aggregate %.

5 of 7 new or added lines in 2 files covered. (71.43%)

2 existing lines in 1 file now uncovered.

4474 of 5593 relevant lines covered (79.99%)

0.8 hits per line

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

80.08

/numba_dpex/ocl/oclimpl.py

# 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
    # If atomic is not supported, it will be emulated by the sycl compiler.
    return (
        dtype == types.int32
        or dtype == types.uint32
        or dtype == types.float32
        or dtype == types.int64
        or dtype == types.uint64
        or dtype == types.float64
    )


@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✔
147	aryty, indty, valty = sig.args	1✔
148	ary, inds, val = args	1✔
149	dtype = aryty.dtype	1✔
150
151	if indty == types.intp:	1✔
152	indices = [inds] # just a single integer	1✔
153	indty = [indty]	1✔
154	else:
155	indices = cgutils.unpack_tuple(builder, inds, count=len(indty))	1✔
156	indices = [	1✔
157	context.cast(builder, i, t, types.intp)
158	for t, i in zip(indty, indices)
159	]
160
161	if dtype != valty:	1!
162	raise TypeError("expecting %s but got %s" % (dtype, valty))	×
163
164	if aryty.ndim != len(indty):	1!
165	raise TypeError(	×
166	"indexing %d-D array with %d-D index" % (aryty.ndim, len(indty))
167	)
168
169	lary = context.make_array(aryty)(context, builder, ary)	1✔
170	ptr = cgutils.get_item_pointer(context, builder, aryty, lary, indices)	1✔
171
172	if dtype == types.float32 or dtype == types.float64:	1✔
173	context.extra_compile_options[kernel_target.LLVM_SPIRV_ARGS] = [	1✔
174	"--spirv-ext=+SPV_EXT_shader_atomic_float_add"
175	]
176	name = "__spirv_AtomicFAddEXT"	1✔
177	elif dtype == types.int32 or dtype == types.int64:	1!
178	name = "__spirv_AtomicIAdd"	1✔
179	else:
180	raise TypeError("Unsupported type")	×
181
182	assert name != ""	1✔
183
184	ptr_type = context.get_value_type(dtype).as_pointer()	1✔
185	ptr_type.addrspace = aryty.addrspace	1✔
186
187	retty = context.get_value_type(sig.return_type)	1✔
188	spirv_fn_arg_types = [	1✔
189	ptr_type,
190	llvmir.IntType(32),
191	llvmir.IntType(32),
192	context.get_value_type(sig.args[2]),
193	]
194
195	from numba_dpex.core import itanium_mangler as ext_itanium_mangler	1✔
196
197	numba_ptr_ty = types.CPointer(dtype, addrspace=ptr_type.addrspace)	1✔
198	mangled_fn_name = ext_itanium_mangler.mangle_ext(	1✔
199	name,
200	[
201	numba_ptr_ty,
202	"__spv.Scope.Flag",
203	"__spv.MemorySemanticsMask.Flag",
204	valty,
205	],
206	)
207
208	fnty = llvmir.FunctionType(retty, spirv_fn_arg_types)	1✔
209	fn = cgutils.get_or_insert_function(builder.module, fnty, mangled_fn_name)	1✔
210	fn.calling_convention = kernel_target.CC_SPIR_FUNC	1✔
211
212	sycl_memory_order = atomic_helper.sycl_memory_order.relaxed	1✔
213	sycl_memory_scope = atomic_helper.sycl_memory_scope.device	1✔
214	spirv_scope = atomic_helper.get_scope(sycl_memory_scope)	1✔
215	spirv_memory_semantics_mask = atomic_helper.get_memory_semantics_mask(	1✔
216	sycl_memory_order
217	)
218	fn_args = [	1✔
219	ptr,
220	context.get_constant(types.int32, spirv_scope),
221	context.get_constant(types.int32, spirv_memory_semantics_mask),
222	val,
223	]
224
225	return builder.call(fn, fn_args)	1✔
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
231	# If atomic is not supported, it will be emulated by the sycl compiler.
232	return (	1✔
233	dtype == types.int32
234	or dtype == types.uint32
235	or dtype == types.float32
236	or dtype == types.int64
237	or dtype == types.uint64
238	or dtype == types.float64
239	)
240
241
242	@lower(stubs.atomic.add, types.Array, types.intp, types.Any)	1✔
243	@lower(stubs.atomic.add, types.Array, types.UniTuple, types.Any)	1✔
244	@lower(stubs.atomic.add, types.Array, types.Tuple, types.Any)	1✔
245	def atomic_add_tuple(context, builder, sig, args):	1✔
246	dtype = sig.args[0].dtype	1✔
247	if support_atomic(dtype):	1!
248	return native_atomic_add(context, builder, sig, args)	1✔
249	else:
NEW 250	raise TypeError(f"Atomic operation on unsupported type {dtype}")	×
251
252
253	def atomic_sub_wrapper(context, builder, sig, args):	1✔
254	# dpcpp yet does not support ``__spirv_AtomicFSubEXT``. To support atomic.sub we
255	# reuse atomic.add and negate the value. For example, atomic.add(A, index, -val) is
256	# equivalent to atomic.sub(A, index, val).
257	val = args[2]	1✔
258	new_val = cgutils.alloca_once(	1✔
259	builder,
260	context.get_value_type(sig.args[2]),
261	size=context.get_constant(types.uintp, 1),
262	name="new_val_0",
263	)
264	val_dtype = sig.args[2]	1✔
265	if val_dtype == types.float32 or val_dtype == types.float64:	1✔
266	builder.store(	1✔
267	builder.fmul(val, context.get_constant(sig.args[2], -1)), new_val
268	)
269	elif val_dtype == types.int32 or val_dtype == types.int64:	1!
270	builder.store(	1✔
271	builder.mul(val, context.get_constant(sig.args[2], -1)), new_val
272	)
273	else:
274	raise TypeError("Unsupported type %s" % val_dtype)	×
275
276	args[2] = builder.load(new_val)	1✔
277
278	return native_atomic_add(context, builder, sig, args)	1✔
279
280
281	@lower(stubs.atomic.sub, types.Array, types.intp, types.Any)	1✔
282	@lower(stubs.atomic.sub, types.Array, types.UniTuple, types.Any)	1✔
283	@lower(stubs.atomic.sub, types.Array, types.Tuple, types.Any)	1✔
284	def atomic_sub_tuple(context, builder, sig, args):	1✔
285	dtype = sig.args[0].dtype	1✔
286	if support_atomic(dtype):	1!
287	return atomic_sub_wrapper(context, builder, sig, args)	1✔
288	else:
NEW 289	raise TypeError(f"Atomic operation on unsupported type {dtype}")	×
290
291
292	@lower(stubs.private.array, types.IntegerLiteral, types.Any)	1✔
293	def dpex_private_array_integer(context, builder, sig, args):	1✔
294	length = sig.args[0].literal_value	1✔
295	dtype = parse_dtype(sig.args[1])	1✔
296	return _generic_array(	1✔
297	context,
298	builder,
299	shape=(length,),
300	dtype=dtype,
301	symbol_name="_dpex_pmem",
302	addrspace=address_space.PRIVATE,
303	)
304
305
306	@lower(stubs.private.array, types.Tuple, types.Any)	1✔
307	@lower(stubs.private.array, types.UniTuple, types.Any)	1✔
308	def dpex_private_array_tuple(context, builder, sig, args):	1✔
309	shape = [s.literal_value for s in sig.args[0]]	×
310	dtype = parse_dtype(sig.args[1])	×
311	return _generic_array(	×
312	context,
313	builder,
314	shape=shape,
315	dtype=dtype,
316	symbol_name="_dpex_pmem",
317	addrspace=address_space.PRIVATE,
318	)
319
320
321	@lower(stubs.local.array, types.IntegerLiteral, types.Any)	1✔
322	def dpex_local_array_integer(context, builder, sig, args):	1✔
323	length = sig.args[0].literal_value	1✔
324	dtype = parse_dtype(sig.args[1])	1✔
325	return _generic_array(	1✔
326	context,
327	builder,
328	shape=(length,),
329	dtype=dtype,
330	symbol_name="_dpex_lmem",
331	addrspace=address_space.LOCAL,
332	)
333
334
335	@lower(stubs.local.array, types.Tuple, types.Any)	1✔
336	@lower(stubs.local.array, types.UniTuple, types.Any)	1✔
337	def dpex_local_array_tuple(context, builder, sig, args):	1✔
338	shape = [s.literal_value for s in sig.args[0]]	×
339	dtype = parse_dtype(sig.args[1])	×
340	return _generic_array(	×
341	context,
342	builder,
343	shape=shape,
344	dtype=dtype,
345	symbol_name="_dpex_lmem",
346	addrspace=address_space.LOCAL,
347	)
348
349
350	def _generic_array(context, builder, shape, dtype, symbol_name, addrspace):	1✔
351	"""
352	This function allows us to create generic arrays in different
353	address spaces.
354	"""
355	elemcount = reduce(operator.mul, shape)	1✔
356	lldtype = context.get_data_type(dtype)	1✔
357	laryty = llvmir.ArrayType(lldtype, elemcount)	1✔
358
359	if addrspace == address_space.LOCAL:	1✔
360	lmod = builder.module	1✔
361
362	# Create global variable in the requested address-space
363	gvmem = cgutils.add_global_variable(	1✔
364	lmod, laryty, symbol_name, addrspace
365	)
366
367	if elemcount <= 0:	1!
368	raise ValueError("array length <= 0")	×
369	else:
370	gvmem.linkage = "internal"	1✔
371
372	if dtype not in types.number_domain:	1!
373	raise TypeError("unsupported type: %s" % dtype)	×
374
375	elif addrspace == address_space.PRIVATE:	1✔
376	gvmem = cgutils.alloca_once(builder, laryty, name=symbol_name)	1✔
377	else:
378	raise NotImplementedError("addrspace {addrspace}".format(**locals()))
379
380	# We need to add the addrspace to _make_array() function call as we want
381	# the variable containing the reference of the memory to retain the
382	# original address space of that memory. Before, we were casting the
383	# memories allocated in local address space to global address space. This
384	# approach does not let us identify the original address space of a memory
385	# down the line.
386	return _make_array(	1✔
387	context, builder, gvmem, dtype, shape, addrspace=addrspace
388	)
389
390
391	def _make_array(	1✔
392	context,
393	builder,
394	dataptr,
395	dtype,
396	shape,
397	layout="C",
398	addrspace=address_space.GENERIC,
399	):
400	ndim = len(shape)	1✔
401	# Create array object
402	aryty = Array(dtype=dtype, ndim=ndim, layout="C", addrspace=addrspace)	1✔
403	ary = context.make_array(aryty)(context, builder)	1✔
404
405	targetdata = _get_target_data(context)	1✔
406	lldtype = context.get_data_type(dtype)	1✔
407	itemsize = lldtype.get_abi_size(targetdata)	1✔
408	# Compute strides
409	rstrides = [itemsize]	1✔
410	for i, lastsize in enumerate(reversed(shape[1:])):	1!
411	rstrides.append(lastsize * rstrides[-1])	×
412	strides = [s for s in reversed(rstrides)]	1✔
413
414	kshape = [context.get_constant(types.intp, s) for s in shape]	1✔
415	kstrides = [context.get_constant(types.intp, s) for s in strides]	1✔
416
417	context.populate_array(	1✔
418	ary,
419	data=builder.bitcast(dataptr, ary.data.type),
420	shape=cgutils.pack_array(builder, kshape),
421	strides=cgutils.pack_array(builder, kstrides),
422	itemsize=context.get_constant(types.intp, itemsize),
423	meminfo=None,
424	)
425
426	return ary._getvalue()	1✔
427
428
429	def _get_target_data(context):	1✔
430	return ll.create_target_data(SPIR_DATA_LAYOUT[context.address_size])	1✔

IntelPython / numba-dpex / 5742028140

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