2.2. Shapes
Open the notebook in Colab

The vector addition module defined in Section 1.2 only accepts vectors with 100-length. It’s too restrictive for real scenarios where inputs can have arbitrary shapes. In this section, we will show how to relax this constraint to deal with general cases.

2.2.1. Variable Shapes

Remember that we create symbolic placeholders for tensors A and B so we can feed with data later. We can do the same thing for the shape as well. In particular, the following code block uses te.var to create a symbolic variable for an int32 scalar, whose value can be specified later.

import d2ltvm
import numpy as np
import tvm
from tvm import te

n = te.var(name='n')
type(n), n.dtype
(tvm.tir.expr.Var, 'int32')

Now we can use (n,) to create a placeholder for an arbitrary length vector.

A = te.placeholder((n,), name='a')
B = te.placeholder((n,), name='b')
C = te.compute(A.shape, lambda i: A[i] + B[i], name='c')
s = te.create_schedule(C.op)
tvm.lower(s, [A, B, C], simple_mode=True)
produce c {
  for (i, 0, n) {
    c[(i*stride)] = (a[(i*stride)] + b[(i*stride)])
  }
}

Compared to the generated pseudo codes in Section 1.2, we can see the upper bound value of the for loop is changed from 100 to n.

Now we define a similar test function as before to verify that the compiled module is able to correctly execute on input vectors with different lengths.

def test_mod(mod, n):
    a, b, c = d2ltvm.get_abc(n, tvm.nd.array)
    mod(a, b, c)
    print('c.shape:', c.shape)
    np.testing.assert_equal(c.asnumpy(), a.asnumpy() + b.asnumpy())

mod = tvm.build(s, [A, B, C])
test_mod(mod, 5)
test_mod(mod, 1000)
c.shape: (5,)
c.shape: (1000,)

But note that we still place the constraint that A, B, and C must be in the same shape. So an error will occur if it is not satisfied.

2.2.2. Multi-dimensional Shapes

You may already notice that a shape is presented as a tuple. A single element tuple means a 1-D tensor, or a vector. We can extend it to multi-dimensional tensors by adding variables to the shape tuple.

The following method builds a module for multi-dimensional tensor addition, the number of dimensions is specified by ndim. For a 2-D tensor, we can access its element by A[i,j], similarly A[i,j,k] for 3-D tensors. Note that we use *i to handle the general multi-dimensional case in the following code.

def tvm_vector_add(ndim):
    A = te.placeholder([te.var() for _ in range(ndim)])
    B = te.placeholder(A.shape)
    C = te.compute(A.shape, lambda *i: A[i] + B[i])
    s = te.create_schedule(C.op)
    return tvm.build(s, [A, B, C])

Verify that it works beyond vectors.

mod = tvm_vector_add(2)
test_mod(mod, (2, 2))

mod = tvm_vector_add(4)
test_mod(mod, (2, 3, 4, 5))
c.shape: (2, 2)
c.shape: (2, 3, 4, 5)

2.2.3. Summary

  • We can use te.var() to specify the dimension(s) of a shape when we don’t know the concrete data shape before execution.

  • The shape of an \(n\)-dimensional tensor is presented as an \(n\)-length tuple.