def impl(context, builder, sig, args):
[tyinp, tyout] = sig.args
[inp, out] = args
ndim = tyinp.ndim
iary = context.make_array(tyinp)(context, builder, inp)
oary = context.make_array(tyout)(context, builder, out)
if asfloat:
sig = typing.signature(types.float64, types.float64)
else:
sig = typing.signature(tyout.dtype, tyinp.dtype)
fnwork = context.get_function(funckey, sig)
intpty = context.get_value_type(types.intp)
# TODO handle differing shape by mimicking broadcasting
shape = cgutils.unpack_tuple(builder, iary.shape, ndim)
with cgutils.loop_nest(builder, shape, intp=intpty) as indices:
pi = cgutils.get_item_pointer(builder, tyinp, iary, indices)
po = cgutils.get_item_pointer(builder, tyout, oary, indices)
ival = builder.load(pi)
if asfloat:
dval = context.cast(builder, ival, tyinp.dtype, types.float64)
dres = fnwork(builder, [dval])
res = context.cast(builder, dres, types.float64, tyout.dtype)
elif tyinp.dtype != tyout.dtype:
tempres = fnwork(builder, [ival])
res = context.cast(builder, tempres, tyinp.dtype, tyout.dtype)
else:
res = fnwork(builder, [ival])
builder.store(res, po)
return out
def getitem_array_tuple(context, builder, sig, args):
aryty, idxty = sig.args
ary, idx = args
arystty = make_array(aryty)
ary = arystty(context, builder, ary)
ndim = aryty.ndim
if isinstance(sig.return_type, types.Array):
# Slicing
raw_indices = cgutils.unpack_tuple(builder, idx, aryty.ndim)
start = []
shapes = []
strides = []
oshapes = cgutils.unpack_tuple(builder, ary.shape, ndim)
for ax, (indexval, idxty) in enumerate(zip(raw_indices, idxty)):
if idxty == types.slice3_type:
slice = Slice(context, builder, value=indexval)
cgutils.normalize_slice(builder, slice, oshapes[ax])
start.append(slice.start)
shapes.append(cgutils.get_range_from_slice(builder, slice))
strides.append(
cgutils.get_strides_from_slice(builder, ndim, ary.strides,
slice, ax))
else:
ind = context.cast(builder, indexval, idxty, types.intp)
start.append(ind)
dataptr = cgutils.get_item_pointer(
builder,
aryty,
ary,
start,
wraparound=context.metadata['wraparound'])
# Build array
retstty = make_array(sig.return_type)
retary = retstty(context, builder)
retary.data = dataptr
retary.shape = cgutils.pack_array(builder, shapes)
retary.strides = cgutils.pack_array(builder, strides)
return retary._getvalue()
else:
# Indexing
indices = cgutils.unpack_tuple(builder, idx, count=len(idxty))
indices = [
context.cast(builder, i, t, types.intp)
for t, i in zip(idxty, indices)
]
ptr = cgutils.get_item_pointer(
builder,
aryty,
ary,
indices,
wraparound=context.metadata['wraparound'])
return context.unpack_value(builder, aryty.dtype, ptr)
def getitem_array_unituple(context, builder, sig, args):
aryty, idxty = sig.args
ary, idx = args
ndim = aryty.ndim
arystty = make_array(aryty)
ary = arystty(context, builder, ary)
if idxty.dtype == types.slice3_type:
# Slicing
raw_slices = cgutils.unpack_tuple(builder, idx, aryty.ndim)
slices = [Slice(context, builder, value=sl) for sl in raw_slices]
for sl, sh in zip(slices, cgutils.unpack_tuple(builder, ary.shape,
ndim)):
cgutils.normalize_slice(builder, sl, sh)
indices = [sl.start for sl in slices]
dataptr = cgutils.get_item_pointer(
builder,
aryty,
ary,
indices,
wraparound=context.metadata['wraparound'])
# Build array
retstty = make_array(sig.return_type)
retary = retstty(context, builder)
retary.data = dataptr
shapes = [cgutils.get_range_from_slice(builder, sl) for sl in slices]
retary.shape = cgutils.pack_array(builder, shapes)
strides = [
cgutils.get_strides_from_slice(builder, ndim, ary.strides, sl, i)
for i, sl in enumerate(slices)
]
retary.strides = cgutils.pack_array(builder, strides)
return retary._getvalue()
else:
# Indexing
assert idxty.dtype == types.intp
indices = cgutils.unpack_tuple(builder, idx, count=len(idxty))
indices = [
context.cast(builder, i, t, types.intp)
for t, i in zip(idxty, indices)
]
ptr = cgutils.get_item_pointer(
builder,
aryty,
ary,
indices,
wraparound=context.metadata['wraparound'])
return context.unpack_value(builder, aryty.dtype, ptr)
def getitem_array_tuple(context, builder, sig, args):
aryty, idxty = sig.args
ary, idx = args
arystty = make_array(aryty)
ary = arystty(context, builder, ary)
ndim = aryty.ndim
if isinstance(sig.return_type, types.Array):
# Slicing
raw_indices = cgutils.unpack_tuple(builder, idx, aryty.ndim)
start = []
shapes = []
strides = []
oshapes = cgutils.unpack_tuple(builder, ary.shape, ndim)
for ax, (indexval, idxty) in enumerate(zip(raw_indices, idxty)):
if idxty == types.slice3_type:
slice = Slice(context, builder, value=indexval)
cgutils.normalize_slice(builder, slice, oshapes[ax])
start.append(slice.start)
shapes.append(cgutils.get_range_from_slice(builder, slice))
strides.append(cgutils.get_strides_from_slice(builder, ndim,
ary.strides,
slice, ax))
else:
ind = context.cast(builder, indexval, idxty, types.intp)
start.append(ind)
dataptr = cgutils.get_item_pointer(builder, aryty, ary, start,
wraparound=True)
# Build array
retstty = make_array(sig.return_type)
retary = retstty(context, builder)
populate_array(retary,
data=dataptr,
shape=cgutils.pack_array(builder, shapes),
strides=cgutils.pack_array(builder, strides),
itemsize=ary.itemsize,
meminfo=ary.meminfo,
parent=ary.parent)
return retary._getvalue()
else:
# Indexing
indices = cgutils.unpack_tuple(builder, idx, count=len(idxty))
indices = [context.cast(builder, i, t, types.intp)
for t, i in zip(idxty, indices)]
ptr = cgutils.get_item_pointer(builder, aryty, ary, indices,
wraparound=True)
return context.unpack_value(builder, aryty.dtype, ptr)
def ptx_atomic_max_tuple(context, builder, sig, args):
aryty, indty, valty = sig.args
ary, inds, val = args
dtype = aryty.dtype
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("expect %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(builder, aryty, lary, indices)
if aryty.dtype == types.float64:
lmod = builder.module
return builder.call(nvvmutils.declare_atomic_max_float64(lmod),
(ptr, val))
else:
raise TypeError('Unimplemented atomic max with %s array' % dtype)
def hsail_atomic_add_tuple(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(builder, aryty, lary, indices)
return builder.atomic_rmw("add", ptr, val, ordering='monotonic')
def ptx_atomic_add_tuple(context, builder, sig, args):
aryty, indty, valty = sig.args
ary, inds, val = args
dtype = aryty.dtype
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("expect %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(builder, aryty, lary, indices)
if aryty.dtype == types.float32:
lmod = cgutils.get_module(builder)
return builder.call(nvvmutils.declare_atomic_add_float32(lmod),
(ptr, val))
else:
return builder.atomic_rmw('add', ptr, val, 'monotonic')
def getitem_array1d_slice(context, builder, sig, args):
aryty, _ = sig.args
if aryty.ndim != 1:
# TODO
raise NotImplementedError("1D indexing into %dD array" % aryty.ndim)
ary, idx = args
arystty = make_array(aryty)
ary = arystty(context, builder, value=ary)
shapes = cgutils.unpack_tuple(builder, ary.shape, aryty.ndim)
slicestruct = Slice(context, builder, value=idx)
cgutils.normalize_slice(builder, slicestruct, shapes[0])
dataptr = cgutils.get_item_pointer(builder, aryty, ary,
[slicestruct.start],
wraparound=True)
retstty = make_array(sig.return_type)
retary = retstty(context, builder)
shape = cgutils.get_range_from_slice(builder, slicestruct)
retary.shape = cgutils.pack_array(builder, [shape])
stride = cgutils.get_strides_from_slice(builder, aryty.ndim, ary.strides,
slicestruct, 0)
retary.strides = cgutils.pack_array(builder, [stride])
retary.data = dataptr
return retary._getvalue()
def ptx_atomic_add_tuple(context, builder, sig, args):
aryty, indty, valty = sig.args
ary, inds, val = args
dtype = aryty.dtype
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("expect %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(builder, aryty, lary, indices)
if aryty.dtype == types.float32:
lmod = builder.module
return builder.call(nvvmutils.declare_atomic_add_float32(lmod), (ptr, val))
elif aryty.dtype == types.float64:
lmod = builder.module
return builder.call(nvvmutils.declare_atomic_add_float64(lmod), (ptr, val))
else:
return builder.atomic_rmw("add", ptr, val, "monotonic")
def getitem_array1d_slice(context, builder, sig, args):
aryty, _ = sig.args
if aryty.ndim != 1:
# TODO
raise NotImplementedError("1D indexing into %dD array" % aryty.ndim)
ary, idx = args
arystty = make_array(aryty)
ary = arystty(context, builder, value=ary)
shapes = cgutils.unpack_tuple(builder, ary.shape, aryty.ndim)
slicestruct = Slice(context, builder, value=idx)
cgutils.normalize_slice(builder, slicestruct, shapes[0])
dataptr = cgutils.get_item_pointer(builder,
aryty,
ary, [slicestruct.start],
wraparound=True)
retstty = make_array(sig.return_type)
retary = retstty(context, builder)
shape = cgutils.get_range_from_slice(builder, slicestruct)
retary.shape = cgutils.pack_array(builder, [shape])
stride = cgutils.get_strides_from_slice(builder, aryty.ndim, ary.strides,
slicestruct, 0)
retary.strides = cgutils.pack_array(builder, [stride])
retary.data = dataptr
return retary._getvalue()
def getitem_array_unituple(context, builder, sig, args):
aryty, idxty = sig.args
ary, idx = args
ndim = aryty.ndim
arystty = make_array(aryty)
ary = arystty(context, builder, ary)
if idxty.dtype == types.slice3_type:
# Slicing
raw_slices = cgutils.unpack_tuple(builder, idx, aryty.ndim)
slices = [Slice(context, builder, value=sl) for sl in raw_slices]
for sl, sh in zip(slices,
cgutils.unpack_tuple(builder, ary.shape, ndim)):
cgutils.normalize_slice(builder, sl, sh)
indices = [sl.start for sl in slices]
dataptr = cgutils.get_item_pointer(builder, aryty, ary, indices,
wraparound=True)
# Build array
retstty = make_array(sig.return_type)
retary = retstty(context, builder)
shapes = [cgutils.get_range_from_slice(builder, sl)
for sl in slices]
strides = [cgutils.get_strides_from_slice(builder, ndim, ary.strides,
sl, i)
for i, sl in enumerate(slices)]
populate_array(retary,
data=dataptr,
shape=cgutils.pack_array(builder, shapes),
strides=cgutils.pack_array(builder, strides),
itemsize=ary.itemsize,
meminfo=ary.meminfo,
parent=ary.parent)
return retary._getvalue()
else:
# Indexing
assert isinstance(idxty.dtype, types.Integer)
indices = cgutils.unpack_tuple(builder, idx, count=len(idxty))
indices = [context.cast(builder, i, t, types.intp)
for t, i in zip(idxty, indices)]
ptr = cgutils.get_item_pointer(builder, aryty, ary, indices,
wraparound=idxty.dtype.signed)
return context.unpack_value(builder, aryty.dtype, ptr)
def iternext_numpy_flatiter(context, builder, sig, args, result):
[flatiterty] = sig.args
[flatiter] = args
flatitercls = make_array_flat_cls(flatiterty)
flatiter = flatitercls(context, builder, value=flatiter)
arrty = flatiterty.array_type
arrcls = context.make_array(arrty)
arr = arrcls(context, builder, value=builder.load(flatiter.array))
ndim = arrty.ndim
shapes = cgutils.unpack_tuple(builder, arr.shape, ndim)
indptr = flatiter.iters
# Load indices and check if they are valid
indices = []
is_valid = cgutils.true_bit
zero = context.get_constant(types.intp, 0)
one = context.get_constant(types.intp, 1)
for ax in range(ndim):
axsize = shapes[ax]
idxptr = builder.gep(indptr, [context.get_constant(types.intp, ax)])
idx = builder.load(idxptr)
ax_valid = builder.icmp(lc.ICMP_SLT, idx, axsize)
indices.append(idx)
is_valid = builder.and_(is_valid, ax_valid)
result.set_valid(is_valid)
with cgutils.if_likely(builder, is_valid):
# Get yielded value
valptr = cgutils.get_item_pointer(builder, arrty, arr, indices)
yield_value = builder.load(valptr)
result.yield_(yield_value)
# Increment iterator indices
carry_flags = [cgutils.true_bit]
for ax, (idx, axsize) in reversed(list(enumerate(zip(indices,
shapes)))):
idxptr = builder.gep(indptr, [context.get_constant(types.intp, ax)])
lastcarry = carry_flags[-1]
idxp1 = builder.add(idx, one)
carry = builder.icmp(lc.ICMP_SGE, idxp1, axsize)
idxfinal = builder.select(lastcarry,
builder.select(carry, zero, idxp1),
idx)
builder.store(idxfinal, idxptr)
carry_flags.append(builder.and_(carry, lastcarry))
with cgutils.if_unlikely(builder, carry_flags[-1]):
# If we have iterated all elements,
# Set first index to out-of-bound
idxptr = builder.gep(indptr, [context.get_constant(types.intp, 0)])
builder.store(shapes[0], idxptr)
def iternext_numpy_flatiter(context, builder, sig, args, result):
[flatiterty] = sig.args
[flatiter] = args
flatitercls = make_array_flat_cls(flatiterty)
flatiter = flatitercls(context, builder, value=flatiter)
arrty = flatiterty.array_type
arrcls = context.make_array(arrty)
arr = arrcls(context, builder, value=builder.load(flatiter.array))
ndim = arrty.ndim
shapes = cgutils.unpack_tuple(builder, arr.shape, ndim)
indptr = flatiter.iters
# Load indices and check if they are valid
indices = []
is_valid = cgutils.true_bit
zero = context.get_constant(types.intp, 0)
one = context.get_constant(types.intp, 1)
for ax in range(ndim):
axsize = shapes[ax]
idxptr = builder.gep(indptr, [context.get_constant(types.intp, ax)])
idx = builder.load(idxptr)
ax_valid = builder.icmp(lc.ICMP_SLT, idx, axsize)
indices.append(idx)
is_valid = builder.and_(is_valid, ax_valid)
result.set_valid(is_valid)
with cgutils.if_likely(builder, is_valid):
# Get yielded value
valptr = cgutils.get_item_pointer(builder, arrty, arr, indices)
yield_value = builder.load(valptr)
result.yield_(yield_value)
# Increment iterator indices
carry_flags = [cgutils.true_bit]
for ax, (idx, axsize) in reversed(list(enumerate(zip(indices,
shapes)))):
idxptr = builder.gep(indptr,
[context.get_constant(types.intp, ax)])
lastcarry = carry_flags[-1]
idxp1 = builder.add(idx, one)
carry = builder.icmp(lc.ICMP_SGE, idxp1, axsize)
idxfinal = builder.select(lastcarry,
builder.select(carry, zero, idxp1), idx)
builder.store(idxfinal, idxptr)
carry_flags.append(builder.and_(carry, lastcarry))
with cgutils.if_unlikely(builder, carry_flags[-1]):
# If we have iterated all elements,
# Set first index to out-of-bound
idxptr = builder.gep(indptr, [context.get_constant(types.intp, 0)])
builder.store(shapes[0], idxptr)
def setitem_array1d(context, builder, sig, args):
aryty, _, valty = sig.args
ary, idx, val = args
arystty = make_array(aryty)
ary = arystty(context, builder, ary)
ptr = cgutils.get_item_pointer(builder, aryty, ary, [idx], wraparound=True)
val = context.cast(builder, val, valty, aryty.dtype)
context.pack_value(builder, aryty.dtype, val, ptr)
def getitem_array1d_intp(context, builder, sig, args):
aryty, _ = sig.args
if aryty.ndim != 1:
# TODO
raise NotImplementedError("1D indexing into %dD array" % aryty.ndim)
ary, idx = args
arystty = make_array(aryty)
ary = arystty(context, builder, ary)
ptr = cgutils.get_item_pointer(builder, aryty, ary, [idx], wraparound=True)
return context.unpack_value(builder, aryty.dtype, ptr)
def setitem_array1d(context, builder, sig, args):
aryty, _, valty = sig.args
ary, idx, val = args
arystty = make_array(aryty)
ary = arystty(context, builder, ary)
ptr = cgutils.get_item_pointer(builder, aryty, ary, [idx],
wraparound=True)
val = context.cast(builder, val, valty, aryty.dtype)
context.pack_value(builder, aryty.dtype, val, ptr)
def ptx_atomic_cas_tuple(context, builder, sig, args):
aryty, oldty, valty = sig.args
ary, old, val = args
dtype = aryty.dtype
lary = context.make_array(aryty)(context, builder, ary)
zero = context.get_constant(types.intp, 0)
ptr = cgutils.get_item_pointer(builder, aryty, lary, (zero,))
if aryty.dtype == types.int32:
lmod = builder.module
return builder.call(nvvmutils.declare_atomic_cas_int32(lmod), (ptr, old, val))
else:
raise TypeError('Unimplemented atomic compare_and_swap with %s array' % dtype)
def ptx_atomic_add_intp(context, builder, sig, args):
aryty, indty, valty = sig.args
ary, ind, val = args
dtype = aryty.dtype
if dtype != valty:
raise TypeError("expect %s but got %s" % (dtype, valty))
if aryty.ndim != 1:
raise TypeError("indexing %d-D array with 1-D index" % (aryty.ndim,))
lary = context.make_array(aryty)(context, builder, ary)
ptr = cgutils.get_item_pointer(builder, aryty, lary, [ind])
return builder.atomic_rmw('add', ptr, val, 'monotonic')
def setitem_array_tuple(context, builder, sig, args):
aryty, idxty, valty = sig.args
ary, idx, val = args
arystty = make_array(aryty)
ary = arystty(context, builder, ary)
# TODO: other than layout
indices = cgutils.unpack_tuple(builder, idx, count=len(idxty))
indices = [context.cast(builder, i, t, types.intp)
for t, i in zip(idxty, indices)]
ptr = cgutils.get_item_pointer(builder, aryty, ary, indices,
wraparound=True)
context.pack_value(builder, aryty.dtype, val, ptr)
def setitem_array_unituple(context, builder, sig, args):
aryty, idxty, valty = sig.args
ary, idx, val = args
arystty = make_array(aryty)
ary = arystty(context, builder, ary)
# TODO: other than layout
indices = cgutils.unpack_tuple(builder, idx, count=len(idxty))
ptr = cgutils.get_item_pointer(builder, aryty, ary, indices,
wraparound=True)
if context.is_struct_type(aryty.dtype):
stval = builder.load(val)
else:
stval = val
builder.store(stval, ptr)
def getitem_array_unituple(context, builder, sig, args):
aryty, idxty = sig.args
ary, idx = args
arystty = make_array(aryty)
ary = arystty(context, builder, ary)
# TODO: other layout
indices = cgutils.unpack_tuple(builder, idx, count=len(idxty))
# TODO warparound flag
ptr = cgutils.get_item_pointer(builder, aryty, ary, indices,
wraparound=True)
if context.is_struct_type(aryty.dtype):
return ptr
else:
return builder.load(ptr)
def ptx_atomic_max_intp(context, builder, sig, args):
aryty, indty, valty = sig.args
ary, ind, val = args
dtype = aryty.dtype
if dtype != valty:
raise TypeError("expect %s but got %s" % (dtype, valty))
if aryty.ndim != 1:
raise TypeError("indexing %d-D array with 1-D index" % (aryty.ndim,))
lary = context.make_array(aryty)(context, builder, ary)
ptr = cgutils.get_item_pointer(builder, aryty, lary, [ind])
if dtype == types.float64:
lmod = builder.module
return builder.call(nvvmutils.declare_atomic_max_float64(lmod), (ptr, val))
else:
raise TypeError('Unimplemented atomic max with %s array' % dtype)
def ptx_atomic_max_intp(context, builder, sig, args):
aryty, indty, valty = sig.args
ary, ind, val = args
dtype = aryty.dtype
if dtype != valty:
raise TypeError("expect %s but got %s" % (dtype, valty))
if aryty.ndim != 1:
raise TypeError("indexing %d-D array with 1-D index" % (aryty.ndim,))
lary = context.make_array(aryty)(context, builder, ary)
ptr = cgutils.get_item_pointer(builder, aryty, lary, [ind])
if dtype == types.float64:
lmod = builder.module
return builder.call(nvvmutils.declare_atomic_max_float64(lmod), (ptr, val))
else:
raise TypeError("Unimplemented atomic max with %s array" % dtype)
def imp(context, builder, sig, args):
# The common argument handling code
aryty, indty, valty = sig.args
ary, inds, val = args
dtype = aryty.dtype
indty, indices = _normalize_indices(context, builder, indty, inds)
if dtype != valty:
raise TypeError("expect %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(builder, aryty, lary, indices)
# dispatcher to implementation base on dtype
return dispatch_fn(context, builder, dtype, ptr, val)
def ptx_atomic_add_tuple(context, builder, sig, args):
aryty, indty, valty = sig.args
ary, inds, val = args
dtype = aryty.dtype
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("expect %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(builder, aryty, lary, indices)
return builder.atomic_rmw('add', ptr, val, 'monotonic')
def setitem_array_tuple(context, builder, sig, args):
aryty, idxty, valty = sig.args
ary, idx, val = args
arystty = make_array(aryty)
ary = arystty(context, builder, ary)
# TODO: other than layout
indices = cgutils.unpack_tuple(builder, idx, count=len(idxty))
indices = [
context.cast(builder, i, t, types.intp)
for t, i in zip(idxty, indices)
]
ptr = cgutils.get_item_pointer(builder,
aryty,
ary,
indices,
wraparound=True)
context.pack_value(builder, aryty.dtype, val, ptr)
def imp(context, builder, sig, args):
# The common argument handling code
aryty, indty, valty = sig.args
ary, inds, val = args
dtype = aryty.dtype
indty, indices = _normalize_indices(context, builder, indty, inds)
if dtype != valty:
raise TypeError("expect %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(builder, aryty, lary, indices)
# dispatcher to implementation base on dtype
return dispatch_fn(context, builder, dtype, ptr, val)
def ptx_atomic_add_intp(context, builder, sig, args):
aryty, indty, valty = sig.args
ary, ind, val = args
dtype = aryty.dtype
if dtype != valty:
raise TypeError("expect %s but got %s" % (dtype, valty))
if aryty.ndim != 1:
raise TypeError("indexing %d-D array with 1-D index" % (aryty.ndim,))
lary = context.make_array(aryty)(context, builder, ary)
ptr = cgutils.get_item_pointer(builder, aryty, lary, [ind])
if aryty.dtype == types.float32:
lmod = cgutils.get_module(builder)
return builder.call(nvvmutils.declare_atomic_add_float32(lmod), (ptr, val))
elif aryty.dtype == types.float64:
lmod = cgutils.get_module(builder)
return builder.call(nvvmutils.declare_atomic_add_float64(lmod), (ptr, val))
else:
return builder.atomic_rmw('add', ptr, val, 'monotonic')
def ptx_atomic_add_intp(context, builder, sig, args):
aryty, indty, valty = sig.args
ary, ind, val = args
dtype = aryty.dtype
if dtype != valty:
raise TypeError("expect %s but got %s" % (dtype, valty))
if aryty.ndim != 1:
raise TypeError("indexing %d-D array with 1-D index" % (aryty.ndim,))
lary = context.make_array(aryty)(context, builder, ary)
ptr = cgutils.get_item_pointer(builder, aryty, lary, [ind])
if aryty.dtype == types.float32:
lmod = cgutils.get_module(builder)
return builder.call(nvvmutils.declare_atomic_add_float32(lmod), (ptr, val))
elif aryty.dtype == types.float64:
lmod = cgutils.get_module(builder)
return builder.call(nvvmutils.declare_atomic_add_float64(lmod), (ptr, val))
else:
return builder.atomic_rmw('add', ptr, val, 'monotonic')
def hsail_atomic_add_tuple(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(builder, aryty, lary, indices)
return builder.atomic_rmw("add", ptr, val, ordering="monotonic")
def setitem_array1d_slice(context, builder, sig, args):
aryty, idxty, valty = sig.args
ary, idx, val = args
arystty = make_array(aryty)
ary = arystty(context, builder, ary)
shapes = cgutils.unpack_tuple(builder, ary.shape, aryty.ndim)
slicestruct = Slice(context, builder, value=idx)
# the logic here follows that of Python's Objects/sliceobject.c
# in particular PySlice_GetIndicesEx function
ZERO = Constant.int(slicestruct.step.type, 0)
NEG_ONE = Constant.int(slicestruct.start.type, -1)
b_step_eq_zero = builder.icmp(lc.ICMP_EQ, slicestruct.step, ZERO)
# bail if step is 0
with cgutils.ifthen(builder, b_step_eq_zero):
context.return_errcode(builder, errcode.ASSERTION_ERROR)
# adjust for negative indices for start
start = cgutils.alloca_once_value(builder, slicestruct.start)
b_start_lt_zero = builder.icmp(lc.ICMP_SLT, builder.load(start), ZERO)
with cgutils.ifthen(builder, b_start_lt_zero):
add = builder.add(builder.load(start), shapes[0])
builder.store(add, start)
b_start_lt_zero = builder.icmp(lc.ICMP_SLT, builder.load(start), ZERO)
with cgutils.ifthen(builder, b_start_lt_zero):
b_step_lt_zero = builder.icmp(lc.ICMP_SLT, slicestruct.step, ZERO)
cond = builder.select(b_step_lt_zero, NEG_ONE, ZERO)
builder.store(cond, start)
b_start_geq_len = builder.icmp(lc.ICMP_SGE, builder.load(start), shapes[0])
ONE = Constant.int(shapes[0].type, 1)
with cgutils.ifthen(builder, b_start_geq_len):
b_step_lt_zero = builder.icmp(lc.ICMP_SLT, slicestruct.step, ZERO)
cond = builder.select(b_step_lt_zero, builder.sub(shapes[0], ONE), shapes[0])
builder.store(cond, start)
# adjust stop for negative value
stop = cgutils.alloca_once_value(builder, slicestruct.stop)
b_stop_lt_zero = builder.icmp(lc.ICMP_SLT, builder.load(stop), ZERO)
with cgutils.ifthen(builder, b_stop_lt_zero):
add = builder.add(builder.load(stop), shapes[0])
builder.store(add, stop)
b_stop_lt_zero = builder.icmp(lc.ICMP_SLT, builder.load(stop), ZERO)
with cgutils.ifthen(builder, b_stop_lt_zero):
b_step_lt_zero = builder.icmp(lc.ICMP_SLT, slicestruct.step, ZERO)
cond = builder.select(b_step_lt_zero, NEG_ONE, ZERO)
builder.store(cond, start)
b_stop_geq_len = builder.icmp(lc.ICMP_SGE, builder.load(stop), shapes[0])
ONE = Constant.int(shapes[0].type, 1)
with cgutils.ifthen(builder, b_stop_geq_len):
b_step_lt_zero = builder.icmp(lc.ICMP_SLT, slicestruct.step, ZERO)
cond = builder.select(b_step_lt_zero, builder.sub(shapes[0], ONE), shapes[0])
builder.store(cond, stop)
b_step_gt_zero = builder.icmp(lc.ICMP_SGT, slicestruct.step, ZERO)
with cgutils.ifelse(builder, b_step_gt_zero) as (then0, otherwise0):
with then0:
with cgutils.for_range_slice(builder, builder.load(start), builder.load(stop), slicestruct.step, slicestruct.start.type) as loop_idx1:
ptr = cgutils.get_item_pointer(builder, aryty, ary,
[loop_idx1],
wraparound=True)
context.pack_value(builder, aryty.dtype, val, ptr)
with otherwise0:
with cgutils.for_range_slice(builder, builder.load(start), builder.load(stop), slicestruct.step, slicestruct.start.type, inc=False) as loop_idx2:
ptr = cgutils.get_item_pointer(builder, aryty, ary,
[loop_idx2],
wraparound=True)
context.pack_value(builder, aryty.dtype, val, ptr)
def _getitem_array1d(context, builder, arrayty, array, idx, wraparound):
ptr = cgutils.get_item_pointer(builder,
arrayty,
array, [idx],
wraparound=wraparound)
return context.unpack_value(builder, arrayty.dtype, ptr)
def _getitem_array1d(context, builder, arrayty, array, idx):
ptr = cgutils.get_item_pointer(builder, arrayty, array, [idx],
wraparound=True)
return context.unpack_value(builder, arrayty.dtype, ptr)
def setitem_array1d_slice(context, builder, sig, args):
aryty, idxty, valty = sig.args
ary, idx, val = args
arystty = make_array(aryty)
ary = arystty(context, builder, ary)
shapes = cgutils.unpack_tuple(builder, ary.shape, aryty.ndim)
slicestruct = Slice(context, builder, value=idx)
# the logic here follows that of Python's Objects/sliceobject.c
# in particular PySlice_GetIndicesEx function
ZERO = Constant.int(slicestruct.step.type, 0)
NEG_ONE = Constant.int(slicestruct.start.type, -1)
b_step_eq_zero = builder.icmp(lc.ICMP_EQ, slicestruct.step, ZERO)
# bail if step is 0
with cgutils.ifthen(builder, b_step_eq_zero):
context.call_conv.return_user_exc(builder, ValueError,
("slice step cannot be zero", ))
# adjust for negative indices for start
start = cgutils.alloca_once_value(builder, slicestruct.start)
b_start_lt_zero = builder.icmp(lc.ICMP_SLT, builder.load(start), ZERO)
with cgutils.ifthen(builder, b_start_lt_zero):
add = builder.add(builder.load(start), shapes[0])
builder.store(add, start)
b_start_lt_zero = builder.icmp(lc.ICMP_SLT, builder.load(start), ZERO)
with cgutils.ifthen(builder, b_start_lt_zero):
b_step_lt_zero = builder.icmp(lc.ICMP_SLT, slicestruct.step, ZERO)
cond = builder.select(b_step_lt_zero, NEG_ONE, ZERO)
builder.store(cond, start)
b_start_geq_len = builder.icmp(lc.ICMP_SGE, builder.load(start), shapes[0])
ONE = Constant.int(shapes[0].type, 1)
with cgutils.ifthen(builder, b_start_geq_len):
b_step_lt_zero = builder.icmp(lc.ICMP_SLT, slicestruct.step, ZERO)
cond = builder.select(b_step_lt_zero, builder.sub(shapes[0], ONE),
shapes[0])
builder.store(cond, start)
# adjust stop for negative value
stop = cgutils.alloca_once_value(builder, slicestruct.stop)
b_stop_lt_zero = builder.icmp(lc.ICMP_SLT, builder.load(stop), ZERO)
with cgutils.ifthen(builder, b_stop_lt_zero):
add = builder.add(builder.load(stop), shapes[0])
builder.store(add, stop)
b_stop_lt_zero = builder.icmp(lc.ICMP_SLT, builder.load(stop), ZERO)
with cgutils.ifthen(builder, b_stop_lt_zero):
b_step_lt_zero = builder.icmp(lc.ICMP_SLT, slicestruct.step, ZERO)
cond = builder.select(b_step_lt_zero, NEG_ONE, ZERO)
builder.store(cond, start)
b_stop_geq_len = builder.icmp(lc.ICMP_SGE, builder.load(stop), shapes[0])
ONE = Constant.int(shapes[0].type, 1)
with cgutils.ifthen(builder, b_stop_geq_len):
b_step_lt_zero = builder.icmp(lc.ICMP_SLT, slicestruct.step, ZERO)
cond = builder.select(b_step_lt_zero, builder.sub(shapes[0], ONE),
shapes[0])
builder.store(cond, stop)
b_step_gt_zero = builder.icmp(lc.ICMP_SGT, slicestruct.step, ZERO)
with cgutils.ifelse(builder, b_step_gt_zero) as (then0, otherwise0):
with then0:
with cgutils.for_range_slice(builder, builder.load(start),
builder.load(stop), slicestruct.step,
slicestruct.start.type) as loop_idx1:
ptr = cgutils.get_item_pointer(builder,
aryty,
ary, [loop_idx1],
wraparound=True)
context.pack_value(builder, aryty.dtype, val, ptr)
with otherwise0:
with cgutils.for_range_slice(builder,
builder.load(start),
builder.load(stop),
slicestruct.step,
slicestruct.start.type,
inc=False) as loop_idx2:
ptr = cgutils.get_item_pointer(builder,
aryty,
ary, [loop_idx2],
wraparound=True)
context.pack_value(builder, aryty.dtype, val, ptr)
