Map tasklet 中不支持 DaCe Tensor 操作答案

【问题标题】：DaCe Tensor operations not supported inside Map taskletsMap tasklet 中不支持 DaCe Tensor 操作
【发布时间】：2021-05-24 01:35:55
【问题描述】：

DaCe 带有一种语法，可以让用户定义一个 Map Map 下方是为计算定义的 tasklet 用户。作为计算，我们可以进行标量运算，例如整数加法。

但是，如果我们在该 tasklet 中进行张量（矩阵）运算，例如 A@B 或 A+B，其中 A 和 B 是两个矩阵，DaCe 会报错。

这是引发错误的具体代码 sn-p。您可以尝试运行此示例。

@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

查看第 16 行，其中有一个矩阵加法 A+A。我相信这是错误的原因。

因此，我想知道，在这种情况下是否允许张量操作？如果允许，我能知道写这个的正确语法吗？如果没有，我想知道为什么不支持？

【问题讨论】：

标签： python tensor

【解决方案1】：

with dace.tasklet 语句是低级 SDFG API，其中直接编写图形元素（即 memlet 和 tasklet）。由于 tasklet 只接受对单个数组或流的访问，因此代码在第 16 行确实失败了（因为没有“A+A”数组可以构成访问节点）。

由于 tasklet 也直接转换为生成的代码，因此在其中编写张量操作（如 @ 运算符）也行不通，因为所有内容都必须通过显式 memlet 进出 tasklet。

为了使问题中的代码正确，只需不声明一个显式的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]

这也可以简化为粗粒度张量操作：

@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

【讨论】：