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)