Hybrid Frontend Language Reference

Hybrid Frontend Language Reference#

Overview#

This hybrid frontend allows users to write preliminary versions of some idioms that yet have been supported by TVM officially.

Features#

Software Emulation#

Both software emulation and compilation are supported. To define a function, you need to use tvm.te.hybrid.script decorator to indicate this is a hybrid function:

@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)

This decorator will import Keywords required spontaneously when software emulation. After software emulation is done, the imported keywords will be cleaned up. Users do not need worry about keyword conflict and pollution.

Every element passed for software emulation in the argument list is either a python variable or numpy numeric type.

Backend Compilation#

This function is not encouraged to use, users are encouraged to use the second interface. The current parse interface looks like:

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

If we pass these tvm data structures, like Tensor, Var, Expr.*Imm, or tvm.container.Array, to this function, it returns a op node:

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

You can use any methods that can be applied on a TVM OpNode, like create_schedule, although so far, the functionality of schedule is as limited as ExternOpNode. At least, it can be built to LLVM module.

Tuning#

Follow up the example above, you can use some tvm like interfaces to tune the code:

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

For now, you can use loop annotations (unroll, parallel, vectorize, and bind), loop manipulation (split and fuse), and reorder.

备注

This is a preliminary function, so users should be in charge of the correctness of the functionality after tuning. Specifically, users should be careful when fusing and reorderding imperfect loops.

Loops#

In HalideIR, loops have in total 4 types: serial, unrolled, parallel, and 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.

Besides all the loop types supported in Halide, const_range is supported for some specific conditions. Sometimes, tvm.container.Array is desired to pass as an argument, but in TVM-HalideIR, there is no such support that converts tvm.container.Array to an Expr. Thus, a limited feature is supported. Users can access containers by either constants or 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

Variables#

All the mutable variables will be lowered to an array with size 1. It regards the first store of a variable as its declaration.

备注

Unlike conventional Python, in hybrid script, the declared variable can only be used in the scope level it is declared.

备注

Currently, you can ONLY use basic-typed variables, i.e. the type of the variable should be either float32, or 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

Attributes#

So far, ONLY tensors' shape and dtype attribute are supported! The shape attribute is essentially a tuple, so you MUST access it as an array. Currently, only constant-indexed access is supported.

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

Conditional Statement and Expression#

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

However, NO True and False keyword supported yet.

Math Intrinsics#

So far, these math intrinsics, log, exp, sigmoid, tanh, power, and popcount, are supported. No import is required, just as it is mentioned in Software Emulation, just use it!

Array Allocation#

Under construction, this function will be supported later!

Use a function call allocation(shape, type, share/local) to declare an array buffer. The basic usage is roughly the same as a normal numpy.array, and you should access high-dim array in a[i, j, k] fashion instead of a[i][j][k], even for tvm.container.Array for compilation.

Thread Bind#

You can also do loop-thread bind by writing code like this:

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

Assert Statement#

Assert statement is supported, you can simply use it as it is in standard Python.

assert cond, mesg

备注

Assert is NOT a function call. Users are encouraged to use assert in the way presented above --- condition followed by message. It fits both Python AST and HalideIR.

Keywords#

  • 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