[RFC] Generalize compute to tensor region

[ZihengJiang]

Target:

• Avoid verification in the original Tensorize, which we found it is not easy while detecting some patterns, like specifying abosolute/relative indexing.
• Achieve better encapsulation. Vendors can build external package for their hardware, and user do not need to know the detail while writing their TVM declaration.
• Propose a way to represent a block in tensor, which can be scheduled as an unit.

Motivation

The current tensorize method use TensorIntrin as a contract. While lowering, it will compare the body of the intrinsic‘s declaration and the body of the original declaration. After verifying the both are absolutely identical, the body of the original declaration will be replaced by intrinsic.

However, we found that it is not easy to guarantee/express the identity in some situation, especially involving the value of an index. This problem is related to the implementation of tensorize: normally, tensorize will replace the index in intrinsic with the index in the original declaration, so that the lowered code can be perfectly matched. For example, if we want to tensorize an addition intrinsic:

# intrinsic declaration
for k in 0...16
    A[k] = B[k] + C[k]

# original declaration
for i in 0...32
    for j in 0...16
        A[i*16+j] = B[i*16+j] + C[i*16+j]

We will replace k in intrinsic declaration with i*16+j, then to compare the ir whether is identical. However, if we require the value of index, things will change. The verification will fail unless we express absolute indexing and relative indexing clearly.

Design

Instead of proposing a way to express absolute indexing, we would like to come up with generalizing tvm.compute to the region of tensor. That means we can view a tensor intrinsic as an operation between regions, just like + as an operation between elements.

• region: following numpy, we use slice to express a region in tensor: A[i, 0:16]
• tensor_op: different with extern operation, we hope that we can keep the ability to schedule the part out of tensorized region. It means there are two kinds of axes: 1. axes in tensorized region, which cannot be scheduled; 2. axes out of the tensorized region, which can be scheduled through split, reorder, etc.

With this generalization, we can use tensor intrinsic in declaration directly, so verification is no longer needed. Also, this way provides better encapsulation, user do not need to know the detail of tensor intrinsic.

Interface

def tensor_op(out_dims,
              in_dims,
              finputs,
              tensor_intrin,
              reduce_axis=[],
              name='tensor_op',
              tag=""):
    """Construct new tensors with intrinsic.

    Parameters
    ----------
    out_dims: tuple
        The dimensions out of the tensorized region, which can be scheduled through `reorder`, `split`.

    in_dims: tuple
        The dimensions inside of the tensorized region, which cannot be manipulated.

    finputs: lambda function of out_dims -> list of TensorSlice
        Specifies involved regions of input tensors.

    tensor_intrin : TensorIntrin
        The tensor intrinsic used for computation.

    reduce_axis : IterVar
        An iteration variable representing the value.

    name: str, optional
        The name hint of the tensor

    tag: str, optional
        Additonal tag information about the compute.

Demo:

# tensor intrinsic
def intrin_vadd(n):
    x = tvm.placeholder((n,), dtype=dtype, name='vx')
    y = tvm.placeholder((n,), dtype=dtype, name='vy')
    z = tvm.compute(x.shape, lambda i: x[i] + y[i], name='z')

    def intrin_func(ins, outs):
        ib = tvm.ir_builder.create()
        ib.emit(tvm.call_extern(outs[0].dtype, 'vadd', ins[0].access_ptr("r"), ins[1].access_ptr('r'), outs[0].access_ptr('wr'), 8, 1, 1, 1, 8, 8, 8))
        return ib.get()

    return tvm.decl_tensor_intrin(z.op, intrin_func)


# tensorize way
def tensorize_vadd():
    A = tvm.placeholder((m/factor, factor), name='A')
    B = tvm.placeholder((m/factor, factor), name='B')
    C = tvm.compute(A.shape, lambda *i: A(*i) + B(*i), name='C')

    s = tvm.create_schedule(C.op)
    xo, xi = C.op.axis
    vadd = intrin_vadd(factor)
    s[C].tensorize(xi, vadd)
    print(tvm.lower(s, [A, B, C], simple_mode=True))


# tensor_op way
def tensor_op_vadd():
    A = tvm.placeholder((m/factor, factor), name="A", dtype=dtype)
    B = tvm.placeholder((m/factor, factor), name="B", dtype=dtype)

    intrin = intrin_vadd(16)
    C = tvm.tensor_op([m/factor,], [factor,],
                      lambda i: [A[i, 0:factor], B[i, 0:factor]], 
                      intrin, name='C')

    s = tvm.create_schedule(C.op)
    print(tvm.lower(s, [A, B, C], simple_mode=True))

link to the PR: #1476

@tqchen @xqdan

[tqchen]

cc @dmlc/tvm-team

[xqdan]

any update on this? @tqchen

[tqchen]

@ZihengJiang is going to send more updates on the proposal and the API behavior, we will start discussing from there.

[ZihengJiang]

@tqchen @yzhliu @tmoreau89 @merrymercy @cowanmeg welcome to comment

[tmoreau89]

I like this proposal as it provides better control to the programmer over how tensorization is applied.

[tmoreau89]

Another thing that is tricky with current tenorization, is error reporting (when the intrinsic does not match the code we apply tensorization on). Do you know if this new approach will lead to clearer error reporting?

[b-y-e]

the result's shape from tensor_op is not obvious and inconsistent with compute.

[ZihengJiang]

@b-y-e Because we need to distinguish the axis in tensorized region and others, so out_dims and in_dims are required. The final shape would be the concatenation of them. Welcome if you have other suggestion for the interface.

[ZihengJiang]

@tmoreau89 The only thing need to check in the new interface is the shape of tensorized region. It will report an error if they do not match. And the body part would be replaced directly.

[tqchen]

Thanks for all the discussions so far, this is a very important feature, as @xqdan and @ZihengJiang suggested, there are certain use cases for it. Here is a summary of my takes:

Significance

So far, there is a two-level separation of "expressions", the scalar index expression as in the compute body, and the tensor relation dataflow as in the chains of operators defined by the compute (topi).

Ideally, we want to ask: is it possible to also use tensor data type expressions within the compute themselves. Conceptually this offers a sense of "recursion". That is why this proposal is quite significant as it takes a step toward the direction.

Pros

It offers an alternative to the current tensorize:

• The current tensorization requires everything being declared in scalar form, and mix match the tensor intrinsics
• The proposed "tensor_op", allows us to directly put tensor_intrin(which implies a tensor level relation) as part of the body to specify the compute rule for a "region" instead of a single scalar index.
  • could be useful when the region is hard to pattern-match, because essentially we are doing the pattern matching directly by invoking the intrinsic and pass in the variables.
• It offers the first step toward "recursive, region level", can be useful for cases for example

Cons

There are certain drawbacks of abusing this notion vs tensorize

• It is somewhat a violation of compute/schedule separation because the pattern matching step is done at the compute stage.
  • This is, however, a practical consideration that is good for short-term but nevertheless need to keep in mind.
  • We do need to keep this in mind
• Divergence with the current tensorization API, and making tensor intrinsic black box
  • The current proposal offers a nice solution to this problem and reuses tensor intrinsic for both tensorize and "tensor_op", which I really like
  • This also means "tensor_op" is no longer blackbox, because if we want, we can replace the compute declaration and recover the original form