def _argminmax(sdfg: SDFG,
state: SDFGState,
a: str,
axis,
func,
result_type=dace.int32,
return_both=False):
nest = NestedCall(sdfg, state)
assert func in ['min', 'max']
if axis is None or type(axis) is not int:
raise SyntaxError('Axis must be an int')
a_arr = sdfg.arrays[a]
if not 0 <= axis < len(a_arr.shape):
raise SyntaxError("Expected 0 <= axis < len({}.shape), got {}".format(
a, axis))
reduced_shape = list(copy.deepcopy(a_arr.shape))
reduced_shape.pop(axis)
val_and_idx = dace.struct('_val_and_idx', val=a_arr.dtype, idx=result_type)
# HACK: since identity cannot be specified for structs, we have to init the output array
reduced_structs, reduced_struct_arr = sdfg.add_temp_transient(
reduced_shape, val_and_idx)
code = "__init = _val_and_idx(val={}, idx=-1)".format(
dtypes.min_value(a_arr.dtype) if func ==
'max' else dtypes.max_value(a_arr.dtype))
nest.add_state().add_mapped_tasklet(
name="_arg{}_convert_".format(func),
map_ranges={
'__i%d' % i: '0:%s' % n
for i, n in enumerate(a_arr.shape) if i != axis
},
inputs={},
code=code,
outputs={
'__init':
Memlet.simple(
reduced_structs, ','.join('__i%d' % i
for i in range(len(a_arr.shape))
if i != axis))
},
external_edges=True)
nest.add_state().add_mapped_tasklet(
name="_arg{}_reduce_".format(func),
map_ranges={'__i%d' % i: '0:%s' % n
for i, n in enumerate(a_arr.shape)},
inputs={
'__in':
Memlet.simple(
a, ','.join('__i%d' % i for i in range(len(a_arr.shape))))
},
code="__out = _val_and_idx(idx={}, val=__in)".format("__i%d" % axis),
outputs={
'__out':
Memlet.simple(
reduced_structs,
','.join('__i%d' % i for i in range(len(a_arr.shape))
if i != axis),
wcr_str=("lambda x, y:"
"_val_and_idx(val={}(x.val, y.val), "
"idx=(y.idx if x.val {} y.val else x.idx))").format(
func, '<' if func == 'max' else '>'))
},
external_edges=True)
if return_both:
outidx, outidxarr = sdfg.add_temp_transient(
sdfg.arrays[reduced_structs].shape, result_type)
outval, outvalarr = sdfg.add_temp_transient(
sdfg.arrays[reduced_structs].shape, a_arr.dtype)
nest.add_state().add_mapped_tasklet(
name="_arg{}_extract_".format(func),
map_ranges={
'__i%d' % i: '0:%s' % n
for i, n in enumerate(a_arr.shape) if i != axis
},
inputs={
'__in':
Memlet.simple(
reduced_structs, ','.join('__i%d' % i
for i in range(len(a_arr.shape))
if i != axis))
},
code="__out_val = __in.val\n__out_idx = __in.idx",
outputs={
'__out_val':
Memlet.simple(
outval, ','.join('__i%d' % i
for i in range(len(a_arr.shape))
if i != axis)),
'__out_idx':
Memlet.simple(
outidx, ','.join('__i%d' % i
for i in range(len(a_arr.shape))
if i != axis))
},
external_edges=True)
return nest, (outval, outidx)
else:
# map to result_type
out, outarr = sdfg.add_temp_transient(
sdfg.arrays[reduced_structs].shape, result_type)
nest(_elementwise)("lambda x: x.idx", reduced_structs, out_array=out)
return nest, out
import dace
import numpy as np
pair = dace.struct('pair', idx=dace.int32, val=dace.float64)
@dace.program
def argmax(x: dace.float64[1024]):
result = np.ndarray([1], dtype=pair)
with dace.tasklet:
init >> result[0]
init.idx = -1
init.val = -1e38
for i in dace.map[0:1024]:
with dace.tasklet:
inp << x[i]
out >> result(
1, lambda x, y: pair(val=max(x.val, y.val),
idx=(x.idx if x.val > y.val else y.idx)))
out = pair(idx=i, val=inp)
return result
def test_argmax():
A = np.random.rand(1024)
result = argmax(A)
assert result[0][0] == np.argmax(A)
import ctypes
import dace
import numpy as np
csrmatrix = dace.struct(
'csr', # CSR Matrix definition type
rows=dace.int32,
cols=dace.int32,
nnz=dace.int32,
data=(dace.pointer(dace.float32), 'nnz'),
rowsp1=dace.int32,
rowptr=(dace.pointer(dace.int32), 'rowsp1'),
colind=(dace.pointer(dace.int32), 'nnz'))
sdfg = dace.SDFG('addone')
state = sdfg.add_state()
sdfg.add_array('sparsemats_in', [5], dtype=csrmatrix)
sdfg.add_array('sparsemats_out', [5], dtype=csrmatrix)
ome, omx = state.add_map('matrices', dict(i='0:5'))
tasklet = state.add_tasklet('addone', {'mat_in'}, {'mat_out'},
'''
for (int j = 0; j < mat_in.nnz; ++j) {
mat_out.data[j] = mat_in.data[j] + 1.0f;
}
''',
language=dace.Language.CPP)
matr = state.add_read('sparsemats_in')
matw = state.add_write('sparsemats_out')
state.add_memlet_path(matr,
ome,
@pytest.mark.skip
def test_starred_target():
a = np.zeros((1, ), dtype=np.float32)
a[0] = np.pi
b, c, d, e = starred_target(a=a)
assert (b[0] == np.float32(np.pi))
assert (c[0] == np.float32(2) * np.float32(np.pi))
assert (c[1] == np.float32(3) * np.float32(np.pi))
assert (c[2] == np.float32(4) * np.float32(np.pi))
assert (d[0] == np.float32(5) * np.float32(np.pi))
assert (e[0] == np.float32(6) * np.float32(np.pi))
mystruct = dace.struct('mystruct', a=dace.int32, b=dace.float32)
@dace.program
def attribute_reference(a: mystruct[1]):
a.a[0] = 5
a.b[0] = 6
@pytest.mark.skip
def test_attribute_reference():
a = np.ndarray((1, ), dtype=np.dtype(mystruct.as_ctypes()))
attribute_reference(a=a)
assert (a[0]['a'] == np.int32(5))
assert (a[0]['b'] == np.float32(6))
# Copyright 2019-2021 ETH Zurich and the DaCe authors. All rights reserved.
import dace
import numpy as np
vec3d = dace.struct('vec3d', x=dace.float32, y=dace.float32, z=dace.float32)
sdfg = dace.SDFG('sred')
sdfg.add_array('A', [1], vec3d)
state = sdfg.add_state()
t = state.add_tasklet('sredtask', {}, {'a'},
'a = vec3d(x=float(1.0), y=float(2.0), z=float(3.0))')
a = state.add_write('A')
state.add_edge(
t, 'a', a, None,
dace.Memlet.simple(
'A',
'0',
wcr_str='lambda a, b: vec3d(x=a.x + b.x, y=a.y + b.y, z=a.z + b.z)'))
def test():
inout = np.ndarray([1], dtype=np.dtype(vec3d.as_ctypes()))
inout[0] = (4.0, 5.0, 6.0)
sdfg(A=inout)
expected = (5.0, 7.0, 9.0)
diff = tuple(abs(x - y) for x, y in zip(inout[0], expected))
print('Difference:', diff)
