TVM InferType

1. TVM InferType
- 1.1. Example

1. TVM InferType

TVM 进行优化和 schedule 时, 各个 op 的 output_shape 是至关重要的. InferType 是根据各个 Relay OP 注册的 TypeRelation 来推测它们的 output_type 和 output_shape.

另外, 在 Relay 优化过程中, 每个 pass 都可以修改/添加/删除 op, 所以每个 pass 之后都需要重新 InferType

1.1. Example

import tvm
from tvm import relay
import numpy as np


def get_demo_mod():
    a = relay.var("a", shape=(2, 3, 10), dtype="float32")
    b = relay.var("b", shape=(1, 10), dtype="float32")
    c = relay.add(a, b)
    func = relay.Function([a, b], c)
    mod = tvm.IRModule.from_expr(func)
    return mod


mod = get_demo_mod()

print("------before InferType------")
try:
    print(mod["main"].body.checked_type)
except Exception:
    print("can't get checked_type")

print("------after InferType------")
mod = relay.transform.InferType()(mod)
print(mod["main"].body.checked_type)

-–—before InferType-–— can't get checked_type -–—after InferType-–— Tensor[(2, 3, 10), float32]

TVM InferType

Table of Contents

1. TVM InferType

1.1. Example