[API] Does the tvm has use-def chain?

I take a look at the function StmtSimplifier::VisitStmt_ in simplify.cc, it don’t use an api similar to getUser in llvm to get the user op of store op. so my question: Is there similar api in tvm to get the def-use chain ?

// eliminate useless stores
Stmt StmtSimplifier::VisitStmt_(const StoreNode* op) {
  Stmt stmt = Parent::VisitStmt_(op);
  op = stmt.as<StoreNode>();
  if (const LoadNode* load = op->value.as<LoadNode>()) {
    if (load->buffer_var.same_as(op->buffer_var) && ExprDeepEqual()(load->index, op->index)) {
      return Evaluate(0);
    }
  }
  return GetRef<Stmt>(op);
}

tvm/src/relax/analysis/udchain.cc at main · apache/tvm may be helpful.

2 Likes
  1. on schedule module(python), can get the op chain base on self._origin_op in vector_schedule.py:def _do_emit_insn
(Pdb) p self._origin_op 记录了数据流,可以看到操作之间的依赖关系
[
{'op': 'elewise_single_exp', 'dst_buffer': Tensor(shape=[280000, 280000], op.name=exp_2), 'src_buffer': [Tensor(shape=[280000, 280000], op.name=mul_1)], 'args': [], 'effective_op': True}, 
{'op': 'elewise_single_VS_mul', 'dst_buffer': Tensor(shape=[280000, 280000], op.name=mul_1), 'src_buffer': [Tensor(shape=[280000, 280000], op.name=log_0)], 'args': [log_0[i0, i1]], 'effective_op': True}, 
{'op': 'elewise_single_log', 'dst_buffer': Tensor(shape=[280000, 280000], op.name=log_0), 'src_buffer': [Tensor(shape=[280000, 280000], op.name=data)], 'args': [], 'effective_op': True}
]
  1. on the tir pass module (c++), the use-def chain can be get base on the StoreNode op itself (membase architecture), .ie: the dst_buffer depends all the src_buffer in op->value and op->index, so can create a edge between them. Note: attr::pragma_buffer_index ,attr::reinterpret_cast_begin等需要特别考虑,解析const AttrStmtNode* op*