TVM TIR 各种Node类型举例

yuanfz1998

已于 2022-03-27 02:19:00 修改

阅读量845

点赞数

CC 4.0 BY-SA版权

分类专栏： tvm 文章标签： compiler ai ml

于 2022-03-27 02:05:46 首次发布

本文链接：https://blog.csdn.net/Eurypterid/article/details/123767428

tvm 专栏收录该内容

10 篇文章

订阅专栏

本文介绍了TVMTIR中的多种Node类型及其在代码中的具体表现形式，包括StmtNode、AnyNode等，并通过实例详细展示了每种Node的应用场景。

摘要生成于 C知道，由 DeepSeek-R1 满血版支持，前往体验 >

TVM TIR中有很多类型的Node。这里收录了全部可打印的类型：

StmtNode
AnyNode
PrimExprNode
TypeNode
PrimFuncNode
IRModuleNode
ArrayNode
IterVarNode
RangeNode
BufferNode
DataProducerNode
StringObj
BufferRegionNode
TargetNode

为了更好地理解这些Node。一方面可以去看它的数据结构，另一方面，也希望可以把它们和TIR代码对应起来。
例1:

@main = primfn(n: int32) -> () {
  producer_realize(compute: DataProducer("compute", float32, [n]), [0:n], True,  {
    for (i.outer: int32, 0, floordiv((n + 3), 4)) {
      for (i.inner: int32, 0, 4) {
        if @tir.likely(((i.inner + (i.outer*4)) < n), dtype=bool) {
          compute_1: DataProducer("compute", float32, [n])[(i.inner + (i.outer*4))] = (A[(i.inner + (i.outer*4))] + B[(i.inner + (i.outer*4))])
        }
      }
    }
  })
}

分解出每一个元素：

PrimExprNode : 
(i.outer*4)
A[(i.inner + (i.outer*4))]
({n|n>=0} + 3)
i.outer
B[(i.inner + (i.outer*4))]
(i.inner + (i.outer*4))
tir.likely(((i.inner + (i.outer*4)) < {n|n>=0}))
((i.inner + (i.outer*4)) < {n|n>=0})
0
(bool)1
{n|n>=0}
i.inner
3
4
(A[(i.inner + (i.outer*4))] + B[(i.inner + (i.outer*4))])
floordiv(({n|n>=0} + 3), 4)

TypeNode : 
TupleTypeNode([])
int32

StmtNode : 
for (i.outer, 0, floordiv(({n|n>=0} + 3), 4)) {
  for (i.inner, 0, 4) {
    if (tir.likely(((i.inner + (i.outer*4)) < {n|n>=0}))) {
      compute[(i.inner + (i.outer*4))] =(A[(i.inner + (i.outer*4))] + B[(i.inner + (i.outer*4))])
    }
  }
}

for (i.inner, 0, 4) {
  if (tir.likely(((i.inner + (i.outer*4)) < {n|n>=0}))) {
    compute[(i.inner + (i.outer*4))] =(A[(i.inner + (i.outer*4))] + B[(i.inner + (i.outer*4))])
  }
}

compute[(i.inner + (i.outer*4))] =(A[(i.inner + (i.outer*4))] + B[(i.inner + (i.outer*4))])

if (tir.likely(((i.inner + (i.outer*4)) < {n|n>=0}))) {
  compute[(i.inner + (i.outer*4))] =(A[(i.inner + (i.outer*4))] + B[(i.inner + (i.outer*4))])
}

producer_realize compute([0, {n|n>=0}]) {
  for (i.outer, 0, floordiv(({n|n>=0} + 3), 4)) {
    for (i.inner, 0, 4) {
      if (tir.likely(((i.inner + (i.outer*4)) < {n|n>=0}))) {
        compute[(i.inner + (i.outer*4))] =(A[(i.inner + (i.outer*4))] + B[(i.inner + (i.outer*4))])
      }
    }
  }
}


DataProducerNode : 
Tensor(shape=[{n|n>=0}], op.name=compute)

StringObj : 
"compute"

ArrayNode : 
[{n|n>=0}]
[(i.inner + (i.outer*4))]
[range(min=0, ext={n|n>=0})]

RangeNode : 
range(min=0, ext={n|n>=0})

注意，子类型的Node统一以其父类的名字展示。比如ForNode继承自StmtNode，上面的StmtNode中就包含了ForNode。