TVM Compiler中文教程：TVM中Compute和Reduction如何使用元组输入

一时失言乱红尘 2022-01-21 01:07 270阅读 0赞

# Compute和Reduction使用元组输入 #

我们通常希望在单个循环内计算具有相同维度的多个输出，或者，执行涉及argmax等多个值的缩减。

在这篇教程，我们将介绍在TVM中元组输入。

from __future__ import absolute_import, print_function
    
    import tvm
    import numpy as np

## batch计算 ##

对于具有相同维度的运算符，如果我们希望它们在下一个调度程序中一起调度，我们可以将它们放在一起作为tvm.compute的输入。

n = tvm.var("n")
    m = tvm.var("m")
    A0 = tvm.placeholder((m,n),name='A0')
    A1 = tvm.placeholder((m,n),name='A1')
    B0, B1 = tvm.compute((m,n),lambda i,j:(A0[i,j]+2,A1[i,j]*3),name='B')
    #生成IR中间表示代码
    s = tvm.create_schedule(B0.op)
    print(tvm.lower(s, [A0,A1,B0,B1],simple_mode=True))

## Reduction使用协同输入 ##

有时，我们需要多个输入来表示一些Reduction算子，输入将协同工作，例如argmax。在Reduction过程中，argmax需要比较操作数的值，也需要去保存操作数的索引。这能使用comm\_reducer来表示：

# xy是Reduction的操作数，他们是索引和值的元组
    def fcombine(x,y):
        lhs = tvm.expr.Select((x[1]>=y[1]),x[0],y[0])
        rhs = tvm.expr.Select((x[1]>=y[1]),x[1],y[1])
        return lhs,rhs
    # 标识元素也需要是一个元组，所以‘fidentity’接受两种类型数据作为输入
    def fidentity(t0,t1):#t0,t1为类型dtype
        return tvm.const(-1,t0), tvm.min_value(t1)
    argmax = tvm.comm_reducer(fcombine, fidentity, name='argmax')
    
    #描述Reduction计算
    m = tvm.var('m')
    n = tvm.var('n')
    idx = tvm.placeholder((m,n),name='idx',dtype='int32')
    val = tvm.placeholder((m,n),name='val',dtype='int32')
    k = tvm.reduce_axis((0,n), 'k')
    T0, T1 = tvm.compute((m, ), lambda i: argmax(idx[i,k], val(i,k), axis=k),name='T')
    
    #生成IR代码
    s = tvm.create_schedule(T0.op)
    print(tvm.lower(s, [idx, val, T0, T1], simple_mode=True))

## 总结 ##

教程介绍元组输入算子的使用：

*  描述标准batch计算
 *  描述Reduction算子使用元组输入
 *  注意，只能使用操作而不是张量来调度计算