Hybrid 前端语言参考#

概述#

这种混合前端允许用户编写一些 TVM 官方支持的习语的初步版本。

特性#

软件仿真#

同时支持软件仿真和编译。要定义函数，你需要使用 tvm.te.hybrid.script 装饰器来表明这是 hybrid 函数：

@tvm.te.hybrid.script
def outer_product(a, b):
    c = output_tensor((100, 99), 'float32')
    for i in range(a.shape[0]):
        for j in range(b.shape[0]):
            c[i, j] = a[i] * b[j]
    return c
a = numpy.random.randn(100)
b = numpy.random.randn(99)
c = outer_product(a, b)

此装饰器将导入软件仿真时自发需要的 Keywords。软件仿真完成后，导入的关键字将被清除。用户不需要担心关键字冲突和污染。

参数列表中传递给软件仿真的每个元素要么是 python 变量，要么是 numpy 数值类型。

后端编译#

这个函数不鼓励使用，鼓励用户使用第二个 interface。当前的 parse 接口如下：

a = tvm.te.placeholder((100, ), name='a')
b = tvm.te.placeholder((99, ), name='b')
parser = tvm.hybrid.parse(outer_product, [a, b]) # return the parser of this function

如果传递这些 tvm 数据结构，比如 Tensor、 Var、 Expr.*Imm 或 tvm.container.Array，此函数返回 op 节点：

a = tvm.te.placeholder((100, ), name='a')
b = tvm.te.placeholder((99, ), name='b')
c = outer_product(a, b) # return the output tensor(s) of the operator

你可以使用任何可以应用在 TVM OpNode 上的方法，比如 create_schedule，尽管到目前为止，schedule 的功能和 ExternOpNode 一样有限。至少，它可以构建到 LLVM 模块。

调优#

按照上面的例子，你可以使用一些类似 tvm 的接口来调优代码：

i, j = c.op.axis
sch = te.create_schedule(op)
jo, ji = sch.split(j, 4)
sch.vectorize(ji)

现在，您可以使用 loop annotations(unroll、parallel、vectorize 和 bind)、loop manipulation(split 和 fuse) 和 reorder。

备注

这是初步的功能，所以调优后功能的正确性应该由用户负责。具体来说，用户在融合和重新排序不完美的循环时应该小心。

循环#

在 HalideIR 中，循环总共有 4 种类型：serial、unrolled、parallel 和 vectorized。

Here we use range aka serial, unroll, parallel, and vectorize, these 4 keywords to annotate the corresponding types of for loops. The usage is roughly the same as Python standard range.

除了 Halide 中支持的所有循环类型，const_range 在某些特定条件下也被支持。有时候，tvm.container.Array 希望作为参数传递，但在 TVM-HalideIR 中，没有转换 tvm.container.Array 到 Expr 这样的支持。因此，只支持有限的特性。用户可以通过 constants 或 constants loops annotated 来访问容器。

@tvm.te.hybrid.script
def foo(a, b): # b is a tvm.container.Array
    c = output_tensor(a.shape, a.dtype)
    for i in const_range(len(a)): # because you have b access, i should be explicitly annotated as const_range
        c[i] = a[i] + b[i]
    return c

变量#

所有可变变量将 lower 为大小为 1 的数组。它将变量的第一个存储区视为它的声明。

备注

与传统的 Python 不同，在 hybrid script 中，声明的变量只能在它声明的作用域级别中使用。

备注

目前，只能使用基本类型的变量，即变量的类型应该是 float32 或 int32。

for i in range(5):
    s = 0 # declaration, this s will be a 1-array in lowered IR
    for j in range(5):
      s += a[i, j] # do something with s
    b[i] = s # you can still use s in this level
a[0] = s # you CANNOT use s here, even though it is allowed in conventional Python

属性#

到目前为止，只支持张量 shape 和 dtype 属性！shape 属性本质上是元组，所以你必须以数组的形式访问它。目前，只支持常量索引访问。

x = a.shape[2] # OK!
for i in range(3):
   for j in a.shape[i]: # BAD! i is not a constant!
       # do something

条件语句和表达式#

if condition1 and condition2 and condition3:
    # do something
else:
    # do something else
# Select
a = b if condition else c

但是，至今还不支持 True 和 False 关键字。

数学指令#

到目前为止，这些数学指令(log、exp、sigmoid、tanh、power 和 popcount) 都得到了支持。不需要导入，就像在 Software Emulation 中提到的那样，只需使用它！

数组分配#

正在建设中，后续将支持此功能！

使用函数调用 allocation(shape, type, share/local) 来声明 array buffer。基本用法与普通的 numpy.array 大致相同，你应该以 a[i, j, k] 方式而不是 a[i][j][k] 方式访问高维数组，即使是用于 tvm.container.Array 编译。

线程绑定#

你也可以写这样的代码来做 loop-thread 绑定：

for tx in bind("threadIdx.x", 100):
    a[tx] = b[tx]

断言声明#

支持断言声明，你可以像在标准 Python 中那样简单地使用它。

assert cond, mesg

备注

Assert 不是函数调用。鼓励用户以上述方式使用 assert ——条件后跟消息。它适合 Python AST 和 HalideIR。

关键字#

  • For keywords: serial, range, unroll, parallel, vectorize, bind, const_range

  • Math keywords: log, exp, sqrt, rsqrt, sigmoid, tanh, power, popcount, round, ceil_div

  • Allocate keywords: allocate, output_tensor

  • Data type keywords: uint8, uint16, uint32, uint64, int8, int16, int32, int64, float16, float32, float64

  • Others: max_num_threads