DaCe Tensor operations not supported inside Map tasklets

user_8200440 0
Viewed 36

DaCe comes with a syntax which enables user to define a Map What’s under the Map is the tasklet users defined for computation. As the computation, we can do scalar operations such as integer addition.

However, if we do tensor(matrix) operation inside that tasklet, for example, A@B or A+B, where A and B are two matrices, DaCe will give us an error.

Here is the concrete code snippet which raises the error. You can try and run this example.

@dace.program
def fusion(A: dace.float32[10, 20], B: dace.float32[10, 20],
           out: dace.float32[1]):

    tmp = dace.define_local([10, 20], dtype=A.dtype)
    tmp_2 = dace.define_local([10, 20], dtype=A.dtype)
    for i, j in dace.map[0:10, 0:20]:
        with dace.tasklet:
            a << (A+A)[i, j]
            b >> tmp[i, j]
            b = a * a

Look at line #16, where there is a matrix addition A+A. I believe this is the cause of the error.

Therefore, I would like to know, is tensor operation allowed in such situation? If it is allowed, can I know the right syntax to write this? If not, I would like to know why this is not supported?

tensor python
1 Answers

The with dace.tasklet statement is the low-level SDFG API, in which the graph elements are written directly (i.e., memlets and tasklets). As tasklets only accept accesses to individual arrays or streams, the code indeed fails in line #16 (since there is no "A+A" array to make into an access node).

Since tasklets are also converted directly to generated code, writing a tensor operation in it like the @ operator would not work either, since everything has to go into/out of a tasklet through explicit memlets.

For the code in the question to be correct, all there is to do is not declare an explicit tasklet:

@dace.program
def option1(A: dace.float32[10, 20], B: dace.float32[10, 20]):
    for i, j in dace.map[0:10, 0:20]:
        B[i, j] = A[i, j] + A[i, j]

This can also be shortened to coarse-grained tensor operations:

@dace.program
def option2(A: dace.float32[10, 20], B: dace.float32[10, 20]):
    B[:] = A + A   # or "tmp = A + A", this will automatically create a new array for tmp
user_64891420
Related