2.1. Data Types
Open the notebook in Colab

Every tensor has a data type, which is typically float32 in deep learning, but also could be int8 (e.g. for model quantization) and others. The tvm_vector_add module we created in Section 1.2 only accepts float32 tensors. Let’s extend it to other data types in this section.

2.1.1. Specifying a Data Type

To use a data type different to the default float32, we can specify it explicitly when creating the placeholders. In the following code block, we generalize the vector addition expression defined in Section 1.2 to accept an argument dtype to specify the data type. In particular, we pass dtype to te.placeholder when creating A and B. The result C then obtains the same data type as A and B.

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

n = 100

def tvm_vector_add(dtype):
    A = te.placeholder((n,), dtype=dtype)
    B = te.placeholder((n,), dtype=dtype)
    C = te.compute(A.shape, lambda i: A[i] + B[i])
    print('expression dtype:', A.dtype, B.dtype, C.dtype)
    s = te.create_schedule(C.op)
    return tvm.build(s, [A, B, C])

Let’s compile a module that accepts int32 tensors.

mod = tvm_vector_add('int32')

expression dtype: int32 int32 int32

Then we define a method to verify the results with a particular data type. Note that we pass a constructor that modifies the tensor data type by astype.

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

test_mod(mod, 'int32')

tensor dtype: int32 int32 int32

Let’s try other data types as well

for dtype in ['float16', 'float64', 'int8','int16', 'int64']:
    mod = tvm_vector_add(dtype)
    test_mod(mod, dtype)

expression dtype: float16 float16 float16
tensor dtype: float16 float16 float16
expression dtype: float64 float64 float64
tensor dtype: float64 float64 float64
expression dtype: int8 int8 int8
tensor dtype: int8 int8 int8
expression dtype: int16 int16 int16
tensor dtype: int16 int16 int16
expression dtype: int64 int64 int64
tensor dtype: int64 int64 int64

2.1.2. Converting Elements Data Types

Besides constructing a tensor with a particular data type, we can also cast the data type of a tensor element during the computation. The following method is the same as tvm_vector_add except that it casts the data type of A and B in te.compute, leaving the data type defined in te.placeholder as default (float32). Because of the casting done by astype, the result C will have the data type specified by dtype.

def tvm_vector_add_2(dtype):
    A = te.placeholder((n,))
    B = te.placeholder((n,))
    C = te.compute(A.shape,
                    lambda i: A[i].astype(dtype) + B[i].astype(dtype))
    print('expression dtype:', A.dtype, B.dtype, C.dtype)
    s = te.create_schedule(C.op)
    return tvm.build(s, [A, B, C])

Then we define a similar test function to verify the results.

def test_mod_2(mod, dtype):
    a, b, c = d2ltvm.get_abc(n)
    # by default `get_abc` returns NumPy ndarray in float32
    a_tvm, b_tvm = tvm.nd.array(a), tvm.nd.array(b)
    c_tvm = tvm.nd.array(c.astype(dtype))
    print('tensor dtype:', a_tvm.dtype, b_tvm.dtype, c_tvm.dtype)
    mod(a_tvm, b_tvm, c_tvm)
    np.testing.assert_equal(c_tvm.asnumpy(), a.astype(dtype) + b.astype(dtype))

mod = tvm_vector_add_2('int32')
test_mod_2(mod, 'int32')

expression dtype: float32 float32 int32
tensor dtype: float32 float32 int32

2.1.3. Summary

  • We can specify the data type by dtype when creating TVM placeholders.

  • The data type of a tensor element can be cast by astype in TVM compute.