The lifecycle of opt_gemm in tvm

发表于 | 阅读次数

0x0 简介

TVM是目前比较热门的深度学习编译框架，本文对tvm的函数注册机制和执行机制进行介绍，这些函数作为在python和c++代码之间交互的接口，可以理解成python和c++之间有统一的command交互接口。理解了这个交互接口，对理解tvm从python到c++的完整执行流程有很大的帮助。
另外本文后面以opt_gemm.py为例子，对tvm的从python到c++的执行流程进行了分析。

0x1 函数注册

TVM中所有上下层交互函数都封装到PackedFunc中，并且所有函数都保存在下面所示的Manager中，每一个函数都封装在Registry中，通过string作为key可以找到注册的函数并调用。tvm中使用的函数为什么要这样设计呢？好处是可以从python调用到这些函数的时候可以用统一的接口，简化接口层代码的编写，只需要把函数名作为key从Manager中找到对应的函数并调用。

1
2
3
4
5
6
7
8
// fmap以函数名称为key，保存函数列表
struct Registry::Manager {
  std::unordered_map<std::string, Registry*> fmap;
  // mutex
  std::mutex mutex;
  Manager() {
}

Registry的定义如下，Registry提供接口把函数都封装到PackedFunc类型的变量中，后面会详细介绍这些接口。

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
class Registry {
 public:
  Registry& set_body(PackedFunc::FType f) {  // NOLINT(*)
    return set_body(PackedFunc(f));
  }
  template<typename FLambda>
  Registry& set_body_typed(FLambda f) {
    using FType = typename detail::function_signature<FLambda>::FType;
    return set_body(TypedPackedFunc<FType>(std::move(f)).packed());
  }
  template<typename T, typename R, typename ...Args>
  Registry& set_body_method(R (T::*f)(Args...) const) {
    auto fwrap = [f](const T target, Args... params) -> R {
      // call method pointer
      return (target.*f)(params...);
    };
    return set_body(TypedPackedFunc<R(const T, Args...)>(fwrap));
  }
  template<typename TObjectRef, typename TNode, typename R, typename ...Args,
    typename = typename std::enable_if<std::is_base_of<ObjectRef, TObjectRef>::value>::type>
  Registry& set_body_method(R (TNode::*f)(Args...)) {
    auto fwrap = [f](TObjectRef ref, Args... params) {
      TNode* target = ref.operator->();
      // call method pointer
      return (target->*f)(params...);
    };
    return set_body(TypedPackedFunc<R(TObjectRef, Args...)>(fwrap));
  }
protected:
  /*! \brief name of the function */
  std::string name_;
  /*! \brief internal packed function */
  PackedFunc func_;
  friend struct Manager;
};

PackedFunc的定义如下。
PackedFunc内部使用std::function来保存函数对象。
TVMArgs提供了对函数参数的封装，可以包括多个参数。
TVMRetValue提供了对函数返回值的封装。

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
/*!
 * \brief Packed function is a type-erased function.
 *  The arguments are passed by packed format.
 *
 *  This is an useful unified interface to call generated functions,
 *  It is the unified function function type of TVM.
 *  It corresponds to TVMFunctionHandle in C runtime API.
 */
class PackedFunc {
 public:
   /*!
   * \brief The internal std::function
   * \param args The arguments to the function.
   * \param rv The return value.
   */
  using FType = std::function<void (TVMArgs args, TVMRetValue* rv)>;
  /*! \brief default constructor */
  PackedFunc() {}
  /*! \brief constructor from null */
  PackedFunc(std::nullptr_t null) {}  // NOLINT(*)
  /*!
   * \brief constructing a packed function from a std::function.
   * \param body the internal container of packed function.
   */
  explicit PackedFunc(FType body) : body_(body) {}
 private:
  /*! \brief internal container of packed function */
  FType body_;
};

PackedFunc中的函数参数TVMArgs定义如下，可以支持可变长度的函数参数。

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
/*!
 * \brief Union type of values
 *  being passed through API and function calls.
 */
typedef union {
  int64_t v_int64;
  double v_float64;
  void* v_handle;
  const char* v_str;
  TVMType v_type;
  TVMContext v_ctx;
} TVMValue;
/*! \brief Arguments into TVM functions. */
class TVMArgs {
 public:
  const TVMValue* values;
  const int* type_codes;
  int num_args;

下面来分析在tvm中注册执行函数的具体执行流程，包括三种方式。

0x11 通过set_body_typed的注册

下图说明了通过set_body_typed来注册函数对象到Registry::Manager的详细过程。

下面是set_body_typed注册函数对象的举例说明。
relay模块中通过set_body_typed的注册方式来注册函数_make.relu。
注意这个函数内部还通过Op::Get()来得到op operator，这个op operator的管理在本文0x14小节中介绍。

1
2
3
4
5
6
// src\relay\op\nn\nn.cc
TVM_REGISTER_GLOBAL("relay.op.nn._make.relu")
.set_body_typed([](Expr data) {
    static const Op& op = Op::Get("nn.relu");
    return CallNode::make(op, {data}, Attrs(), {});
  });

relay模块依赖于topi层的实现，topi层中对relu的定义如下。

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
// topi\include\topi\nn.h
template <typename T>
inline tvm::Tensor relu(const tvm::Tensor& t,
                        T threshold = static_cast<T>(0),
                        std::string name = "T_relu",
                        std::string tag = kElementWise) {
  return tvm::compute(
      t->shape,
      [&](const tvm::Array<tvm::Var>& i) {
        auto threshold_const = tvm::make_const(t->dtype, threshold);
        return tvm::max(t(i), threshold_const);
      },
      name,
      tag);
}

最后调用到lang模块中实现的tvm::max，构造出相应的TVM IR。

1
2
3
4
5
6
7
8
9
10
11
12
13
14
// src\lang\expr_operator.cc
Expr max(Expr a, Expr b) {
  // inf-aware simplificaiton
  using arith::is_pos_inf;
  using arith::is_neg_inf;
  if (is_pos_inf(a)) return a;
  if (is_neg_inf(a)) return b;
  if (is_pos_inf(b)) return b;
  if (is_neg_inf(b)) return a;
  BinaryOpMatchTypes(a, b);
  Expr ret = arith::TryConstFold<ir::Max>(a, b);
  if (ret.defined()) return ret;
  return ir::Max::make(a, b);
}

最后python模块中的relay层对relu的实现提供了python封装。这样python就可以调用到前面介绍的C++模块中的对应实现。

1
2
3
4
// python\tvm\relay\op\nn\nn.py
def relu(data):
    """Rectified linear unit.
    return _make.relu(data)

0x12 通过set_body_method的注册

下图说明了通过set_body_method来注册函数对象到Registry::Manager的详细过程。

1
2
3
4
5
下面是一个set_body_method的调用示例。
// src\api\api_lang.cc
TVM_REGISTER_GLOBAL("_BijectiveLayoutBackwardShape")
.set_body_method(&BijectiveLayout::BackwardShape);

0x13 通过set_body的注册

下图说明了通过set_body来注册函数对象到Registry::Manager的详细过程。

下面是一个set_body的调用示例。

1
2
3
4
5
TVM_REGISTER_GLOBAL("device_api.opencl")
.set_body([](TVMArgs args, TVMRetValue* rv) {
    DeviceAPI* ptr = OpenCLWorkspace::Global().get();
    *rv = static_cast<void*>(ptr);
  });

0x14 Relay OP注册

Relay OP保存在另外一个OpManager中，和前面的函数注册不是一个地方，这是因为这个OpManager管理的是relay operator函数，这些函数不会直接从python中调用过来。

OpManager的相关代码如下。

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
// src\relay\ir\op.cc
::dmlc::Registry<OpRegistry>* OpRegistry::Registry() {
  return ::dmlc::Registry<OpRegistry>::Get();
}
// single manager of operator information.
struct OpManager {
  // mutex to avoid registration from multiple threads.
  std::mutex mutex;
  // global operator counter
  std::atomic<int> op_counter{0};
  // storage of additional attribute table.
  std::unordered_map<std::string, std::unique_ptr<GenericOpMap>> attr;
  // frontend functions
  std::vector<PackedFunc*> frontend_funcs;
  // get singleton of the op manager
  static OpManager* Global() {
    static OpManager* inst = new OpManager();
    return inst;
  }
};
// find operator by name
const Op& Op::Get(const std::string& name) {
  const OpRegistry* reg = dmlc::Registry<OpRegistry>::Find(name);
  CHECK(reg != nullptr) << "Operator " << name << " is not registered";
  return reg->op();
}
OpRegistry::OpRegistry() {
  OpManager* mgr = OpManager::Global();
  ObjectPtr<OpNode> n = make_object<OpNode>();
  n->index_ = mgr->op_counter++;
  op_ = Op(n);
}
template<typename EntryType>
class Registry {
}

下面是调用OpManager来注册op函数的代码示例。

1
2
3
4
5
6
7
8
9
RELAY_REGISTER_OP("argsort")
.describe(R"doc(Returns the indices that would sort an
input array along the given axis.
)doc" TVM_ADD_FILELINE)
.set_num_inputs(1)
.set_attrs_type<ArgsortAttrs>()
.add_argument("data", "Tensor", "Input data.")
.set_support_level(6)
.add_type_rel("Argsort", ArgsortRel);

如何取得OpManager中注册的op函数呢？答案是像下面这段代码所示通过调用Op::Get()来得到。

1
2
3
4
5
6
// Cache the operators that are checked 
// recursively to reduce lookup overhead.
static const auto& expand_dims_op = Op::Get("expand_dims");
static const auto& reshape_op = Op::Get("reshape");
static const auto& transpose_op = Op::Get("transpose");
static const auto& squeeze_op = Op::Get("squeeze");

0x2 执行流程介绍

前面介绍了执行函数是如何注册的，那这些函数是怎样被python调用到的呢？下面来介绍一下。
我们知道从python调用过来的接口定义在c_runtime_api.h头文件中，这其中比较重要的是这两个API，TVMFuncCall()和TVMArrayAlloc，从API的名称上可知，TVMFuncCall对应于从python来的函数调用，TVMArrayAlloc对应于数组的分配。
这两个API接口定义如下。

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
// tvm\include\tvm\runtime\c_runtime_api.h
TVM_DLL int TVMFuncCall(TVMFunctionHandle func,
                        TVMValue* arg_values,
                        int* type_codes,
                        int num_args,
                        TVMValue* ret_val,
                        int* ret_type_code);
TVM_DLL int TVMArrayAlloc(const tvm_index_t* shape,
                          int ndim,
                          int dtype_code,
                          int dtype_bits,
                          int dtype_lanes,
                          int device_type,
                          int device_id,
                          TVMArrayHandle* out);

然而python是如何知道前面的函数Manager中包括了哪些函数呢？这个时候Manager提供了一个函数供上层来调用得到所有注册函数名称列表，这个函数的定义如下。

1
2
3
4
5
6
7
8
9
10
11
12
13
int TVMFuncListGlobalNames(int *out_size,
                           const char*** out_array) {
  API_BEGIN();
  TVMFuncThreadLocalEntry *ret = TVMFuncThreadLocalStore::Get();
  ret->ret_vec_str = tvm::runtime::Registry::ListNames();
  ret->ret_vec_charp.clear();
  for (size_t i = 0; i < ret->ret_vec_str.size(); ++i) {
    ret->ret_vec_charp.push_back(ret->ret_vec_str[i].c_str());
  }
  *out_array = dmlc::BeginPtr(ret->ret_vec_charp);
  *out_size = static_cast<int>(ret->ret_vec_str.size());
  API_END();
}

然后python层根据前面的函数名称列表，通过下面的函数来得到每一个函数名称所对应的函数对象，python层拿到了这些函数对象以后，会在python层也创建相应的函数对象。这样python和c++层的函数操作就可以对应起来了。

1
2
3
4
5
6
7
8
9
10
11
int TVMFuncGetGlobal(const char* name, TVMFunctionHandle* out) {
  API_BEGIN();
  const tvm::runtime::PackedFunc* fp =
      tvm::runtime::Registry::Get(name);
  if (fp != nullptr) {
    *out = new tvm::runtime::PackedFunc(*fp);  // NOLINT(*)
  } else {
    *out = nullptr;
  }
  API_END();
}

在python中创建函数对象的代码如下，先根据函数TVMFuncGetGlobal()来查找底层的函数对象，找到以后根据返回的handle创建上层函数对象Function。

1
2
3
4
5
6
7
8
9
10
def get_global_func(name, allow_missing=False):
    handle = FunctionHandle()
    check_call(_LIB.TVMFuncGetGlobal(c_str(name), ctypes.byref(handle)))
    if handle.value:
        return Function(handle, False)
    if allow_missing:
        return None
    raise ValueError("Cannot find global function %s" % name)

下面来介绍一下函数调用是如何从python层调用到C++层的，其中核心是要理解TVMFuncCall的调用过程。

0x21 TVMFuncCall的调用过程

Python中调用TVMFuncCall的代码如下所示，其中包括了函数参数的封装。
self.handle是python中持有的C++ PackedFunc对象。

1
2
3
4
5
6
7
8
9
10
11
12
13
14
class FunctionBase(object):
    def __call__(self, *args):
        temp_args = []
        values, tcodes, num_args = _make_tvm_args(args, temp_args)
        ret_val = TVMValue()
        ret_tcode = ctypes.c_int()
        if _LIB.TVMFuncCall(
                self.handle, values, tcodes, ctypes.c_int(num_args),
                ctypes.byref(ret_val), ctypes.byref(ret_tcode)) != 0:
            raise get_last_ffi_error()
        _ = temp_args
        _ = args
        return RETURN_SWITCH[ret_tcode.value](ret_val)

从Python代码调用到C++代码的入口函数如下。
函数参数func是封装好的PackedFunc对象。

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
// src\runtime\c_runtime_api.cc
int TVMFuncCall(TVMFunctionHandle func,
                TVMValue* args,
                int* arg_type_codes,
                int num_args,
                TVMValue* ret_val,
                int* ret_type_code) {
  API_BEGIN();
  TVMRetValue rv;
  (*static_cast<const PackedFunc*>(func)).CallPacked(
      TVMArgs(args, arg_type_codes, num_args), &rv);
  // handle return string.
  if (rv.type_code() == kStr ||
      rv.type_code() == kTVMType ||
      rv.type_code() == kBytes) {
    TVMRuntimeEntry* e = TVMAPIRuntimeStore::Get();
    if (rv.type_code() != kTVMType) {
      e->ret_str = *rv.ptr<std::string>();
    } else {
      e->ret_str = rv.operator std::string();
    }
    if (rv.type_code() == kBytes) {
      e->ret_bytes.data = e->ret_str.c_str();
      e->ret_bytes.size = e->ret_str.length();
      *ret_type_code = kBytes;
      ret_val->v_handle = &(e->ret_bytes);
    } else {
      *ret_type_code = kStr;
      ret_val->v_str = e->ret_str.c_str();
    }
  } else {
    rv.MoveToCHost(ret_val, ret_type_code);
  }
  API_END();
}
// tvm\include\tvm\runtime\packed_func.h
inline void PackedFunc::CallPacked(TVMArgs args, TVMRetValue* rv) const {
  body_(args, rv);
}

下面执行对body_的调用，利用可变参数模板的递归展开来实现，这样就可以调用到真正的注册函数了。

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
// include\tvm\runtime\packed_func.h
template<typename R, typename ...Args>
template<typename FType>
inline void TypedPackedFunc<R(Args...)>::AssignTypedLambda(FType flambda) {
  packed_ = PackedFunc([flambda](const TVMArgs& args, TVMRetValue* rv) {
      detail::unpack_call<R, sizeof...(Args)>(flambda, args, rv);
    });
}
template<typename R, int nargs, typename F>
inline void unpack_call(const F& f, const TVMArgs& args, TVMRetValue* rv) {
  unpack_call_dispatcher<R, nargs, 0, F>::run(f, args, rv);
}
template<typename R, int nleft, int index, typename F>
struct unpack_call_dispatcher {
  template<typename ...Args>
  static void run(const F& f,
                  const TVMArgs& args_pack,
                  TVMRetValue* rv,
                  Args&&... unpacked_args) {
    unpack_call_dispatcher<R, nleft - 1, index + 1, F>
        ::run(f, args_pack, rv,
              std::forward<Args>(unpacked_args)...,
              args_pack[index]);
  }
};
template<typename R, int index, typename F>
struct unpack_call_dispatcher<R, 0, index, F> {
  template<typename ...Args>
  static void run(const F& f,
                  const TVMArgs& args_pack,
                  TVMRetValue* rv,
                  Args&&... unpacked_args) {
    *rv = R(f(std::forward<Args>(unpacked_args)...));
  }
};

最后可以看到调用的是Variable::make来生成tvm IR。

1
2
3
4
5
// src\api\api_ir.cc
TVM_REGISTER_GLOBAL("_Var")
.set_body_typed([](std::string s, DataType t) {
    return Variable::make(t, s);
  });

0x22 TVMArrayAlloc的调用过程

TVMArrayAlloc的调用过程比较简单，直接调用NDArray的接口来分配Array。

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
  int TVMArrayAlloc(const tvm_index_t* shape,
                  int ndim,
                  int dtype_code,
                  int dtype_bits,
                  int dtype_lanes,
                  int device_type,
                  int device_id,
                  TVMArrayHandle* out) {
  API_BEGIN();
  DLDataType dtype;
  dtype.code = static_cast<uint8_t>(dtype_code);
  dtype.bits = static_cast<uint8_t>(dtype_bits);
  dtype.lanes = static_cast<uint16_t>(dtype_lanes);
  DLContext ctx;
  ctx.device_type = static_cast<DLDeviceType>(device_type);
  ctx.device_id = device_id;
  *out = NDArray::Internal::MoveToFFIHandle(
      NDArray::Empty(std::vector<int64_t>(shape, shape + ndim), dtype, ctx));
  API_END();
}

0x03 opt_gemm.py执行流程分析

下面来分析tvm自带的矩阵优化测试程序opt_gemm.py的执行流程。

0x31 根据算法创建数据流图

python代码如下，这部分描述了算法是两个矩阵相乘，根据矩阵的大小分配了相应的占位符placeholder,placeholder和tensorflow中的概念类似。然后调用compute()函数创建tvm IR。

1
2
3
4
5
6
7
8
9
10
11
M = 1024
K = 1024
N = 1024
k = tvm.reduce_axis((0, K), 'k')
A = tvm.placeholder((M, K), name='A')
B = tvm.placeholder((K, N), name='B')
C = tvm.compute(
           (M, N),
           lambda x, y: tvm.sum(A[x, k] * B[k, y], axis=k),
           name='C')

上面算法描述对应到C++代码中，会创建相应的tvm node，这部分可以理解成是tvm的IR的生成。
下面的代码描述了上述python流程执行的最后创建ComputeOpNode的过程。

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
Operation ComputeOpNode::make(std::string name,
                              std::string tag,
                              Map<std::string, ObjectRef> attrs,
                              Array<IterVar> axis,
                              Array<Expr> body) {
  if (!attrs.defined()) {
    attrs = Map<std::string, ObjectRef>();
  }
  auto n = make_object<ComputeOpNode>();
  n->name = std::move(name);
  n->tag = std::move(tag);
  n->attrs = std::move(attrs);
  n->axis = std::move(axis);
  n->body = std::move(body);
  if (n->body[0]->IsInstance<ir::Reduce>()) {
    const ir::Reduce* reduce = n->body[0].as<ir::Reduce>();
    n->reduce_axis = reduce->axis;
  }
  VerifyComputeOp(n.get());
  return Operation(n);
}

0x32 创建Schedule

python代码如下。

1
s = tvm.create_schedule(C.op)

对应的C++代码如下，这个时候会根据前面创建的tvm IR来生成reader graph，reader graph中描述了node之间的数据依赖关系。

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
Schedule ScheduleNode::make(Array<Operation> ops) {
  auto n = make_object<ScheduleNode>();
  Schedule sch(n);
  n->outputs = ops;
  auto g = schedule::CreateReadGraph(n->outputs);
  Array<Operation> post_order = schedule::PostDFSOrder(n->outputs, g);
  // output set.
  std::unordered_set<Operation> output_set;
  for (Operation x : ops) {
    output_set.insert(x);
  }
  for (Operation op : post_order) {
    Stage stage(op);
    stage->is_output = output_set.count(op) != 0;
    n->stages.push_back(stage);
    n->stage_map.Set(op, stage);
    // mark scan updates.
    if (const ScanOpNode* scan = op.as<ScanOpNode>()) {
      Array<Tensor> inputs;
      for (Tensor t : scan->state_placeholder) {
        inputs.push_back(t);
      }
      for (Tensor t : scan->inputs) {
        inputs.push_back(t);
      }
      // Create the scan group.
      Stage scan_group = sch.create_group(scan->update, inputs, false);
      scan_group->attach_type = kScanUpdate;
      scan_group->attach_stage = stage;
      for (size_t i = 0; i < scan->update.size(); ++i) {
        Stage s = n->stage_map[scan->update[i]->op];
        CHECK(scan_group.same_as(s->group));
      }
    }
  }
  return sch;
}

0x32 用TVM Pass来处理schedule生成的graph

执行下面的测试代码以后会调用下面的语句来创建stmt。
func = tvm.build(s, [A, B, C], target=target, name=’mmult’)

其Python代码如下。

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
def lower(sch,
          args,
          name="default_function",
          binds=None,
          simple_mode=False):
 
    cfg = current_build_config()
    add_lower_pass = cfg.add_lower_pass if cfg.add_lower_pass else []
    if cfg.dump_pass_ir:
        add_lower_pass = BuildConfig._dump_ir.decorate_custompass(add_lower_pass)
    lower_phase0 = [x[1] for x in add_lower_pass if x[0] == 0]
    lower_phase1 = [x[1] for x in add_lower_pass if x[0] == 1]
    lower_phase2 = [x[1] for x in add_lower_pass if x[0] == 2]
    lower_phase3 = [x[1] for x in add_lower_pass if x[0] > 2]
    # Phase 0
    if isinstance(sch, schedule.Schedule):
        stmt = form_body(sch)
    for f in lower_phase0:
        stmt = f(stmt)
    compact = ir_pass.VerifyCompactBuffer(stmt)
    binds, arg_list = get_binds(args, compact, binds)
    # Phase 1
    stmt = ir_pass.RewriteForTensorCore(stmt, sch, binds)
    stmt = ir_pass.StorageFlatten(stmt, binds, 64, cfg.instrument_bound_checkers)
    stmt = ir_pass.CanonicalSimplify(stmt)
    for f in lower_phase1:
        stmt = f(stmt)
    # Phase 2
    if not simple_mode:
        stmt = ir_pass.LoopPartition(stmt, cfg.partition_const_loop)
    if cfg.disable_vectorize:
        stmt = ir_pass.SkipVectorize(stmt)
    else:
        stmt = ir_pass.VectorizeLoop(stmt)
    stmt = ir_pass.InjectVirtualThread(stmt)
    stmt = ir_pass.InjectDoubleBuffer(stmt, cfg.double_buffer_split_loop)
    stmt = ir_pass.StorageRewrite(stmt)
    stmt = ir_pass.UnrollLoop(
        stmt,
        cfg.auto_unroll_max_step,
        cfg.auto_unroll_max_depth,
        cfg.auto_unroll_max_extent,
        cfg.unroll_explicit)
    for f in lower_phase2:
        stmt = f(stmt)
    # Phase 3
    stmt = ir_pass.Simplify(stmt)
    stmt = ir_pass.RemoveNoOp(stmt)
    if not cfg.disable_select_rewriting:
        stmt = ir_pass.RewriteUnsafeSelect(stmt)
    for f in lower_phase3:
        stmt = f(stmt)
    # Instrument BoundCheckers
    if cfg.instrument_bound_checkers:
        stmt = ir_pass.InstrumentBoundCheckers(stmt)
    if simple_mode:
        return stmt
    return ir_pass.MakeAPI(stmt, name, arg_list, 0, cfg.restricted_func)

前面的ir_pass.xxx函数调用都会对应到C++的实现，这些pass是tvm中实现的中间流程处理操作，
例如前面的函数中执行的下列python代码，
stmt = ir_pass.RemoveNoOp(stmt)
其对应的C++代码如下所示。

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
// src\pass\remove_no_op.cc
// Mark the statment of each stage.
class NoOpRemover : public StmtMutator {
 public:
  Stmt VisitStmt_(const LetStmt* op) final {
    Stmt stmt = StmtMutator::VisitStmt_(op);
    op = stmt.as<LetStmt>();
    return is_no_op(op->body) ? MakeEvaluate(op->value) : stmt;
  }
  Stmt VisitStmt_(const AttrStmt* op) final {
    if (op->attr_key == "pragma_debug_skip_region") {
      return MakeEvaluate(0);
    }
    Stmt stmt = StmtMutator::VisitStmt_(op);
    op = stmt.as<AttrStmt>();
    return is_no_op(op->body) ? MakeEvaluate(op->value) : stmt;
  }
  ......
};
Stmt RemoveNoOp(Stmt stmt) {
  return NoOpRemover()(std::move(stmt));
}

0x33 把前面优化过的Pass调用LLVM codegen来生成LLVM IR

python代码如下。

1
2
def build_module(lowered_func, target):
    return _Build(lowered_func, target)

C++代码如下，该段代码把tvm IR翻译成LLVM IR。
在其调用的Finish()函数中还会采用LLVM PassManager对已经生成的LLVM IR进行进一步优化。

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
void Init(const Array<LoweredFunc>& funcs, std::string target) {
    InitializeLLVM();
    tm_ = GetLLVMTargetMachine(target);
    bool system_lib = (target.find("-system-lib") != std::string::npos);
    CHECK_NE(funcs.size(), 0U);
    ctx_ = std::make_shared<llvm::LLVMContext>();
    std::unique_ptr<CodeGenLLVM> cg = CodeGenLLVM::Create(tm_.get());
    entry_func_ = funcs[0]->name;
    cg->Init(funcs[0]->name, tm_.get(), ctx_.get(), system_lib, system_lib);
    for (LoweredFunc f :  funcs) {
      cg->AddFunction(f);
    }
    cg->AddMainFunction(funcs[0]->name);
    module_ = cg->Finish();
    module_->addModuleFlag(llvm::Module::Warning, "tvm_target", llvm::MDString::get(*ctx_, target));
    module_->addModuleFlag(llvm::Module::Override, "Debug Info Version",
                           llvm::DEBUG_METADATA_VERSION);
    if (tm_->getTargetTriple().isOSDarwin()) {
      module_->addModuleFlag(llvm::Module::Override, "Dwarf Version", 2);
    }
    std::string verify_errors_storage;
    llvm::raw_string_ostream verify_errors(verify_errors_storage);
    LOG_IF(FATAL, llvm::verifyModule(*module_, &verify_errors))
        << "LLVM module verification failed with the following errors: \n"
        << verify_errors.str();
    target_ = target;
    mptr_ = module_.get();
  }

0x34 根据LLVM IR生成对应的机器指令

调用了如下python代码就触发了机器指令的生成。
func(a, b, c)

func为前面返回的LLVM IR module对应的地址，a, b, c为对应的执行参数，也就是矩阵运算的输入。
对应到C++中的下述实现，调用LLVM模块来生成对应target的机器代码。

1
2
3
4
5
6
7
8
9
10
PackedFunc WrapPackedFunc(BackendPackedCFunc faddr,
                          const ObjectPtr<Object>& sptr_to_self) {
  return PackedFunc([faddr, sptr_to_self](TVMArgs args, TVMRetValue* rv) {
      int ret = (*faddr)(
          const_cast<TVMValue*>(args.values),
          const_cast<int*>(args.type_codes),
          args.num_args);
      CHECK_EQ(ret, 0) << TVMGetLastError();
    });
}

0x35 评估执行时间

python代码。

1
2
evaluator = func.time_evaluator(func.entry_name, ctx, number=1)
print('Baseline: %f' % evaluator(a, b, c).mean)

c++代码如下，调用LLVM生成的机器指令来执行具体的运算。这部分还包括了把运算调用到其他机器的rpc操作。

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
PackedFunc WrapTimeEvaluator(PackedFunc pf,
                             TVMContext ctx,
                             int number,
                             int repeat,
                             int min_repeat_ms) {
    ......
    for (int i = 0; i < repeat; ++i) {
        ......
        // start timing
        for (int i = 0; i < number; ++i) {
          pf.CallPacked(args, &temp);
        }
        ......
    }
}