2.2. Shapes¶
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.