stmt.py 1.27 KB
Newer Older
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15
"""Statement AST Node in TVM.

User do not need to deal with AST node directly.
But they can be helpful for developer to do quick proptyping.
While not displayed in the document and python file.
Each statement node have subfields that can be visited from python side.

.. code-block:: python

    x = tvm.var("n")
    a = tvm.var("array", tvm.handle)
    st = tvm.make.Store(a, x + 1, 1)
    assert isinstance(st, tvm.stmt.Store)
    assert(st.buffer_var == a)
"""
tqchen committed
16
from __future__ import absolute_import as _abs
17
from ._ffi.node import NodeBase, register_node
tqchen committed
18 19

class Stmt(NodeBase):
20
    pass
tqchen committed
21 22 23 24 25 26 27 28 29 30 31 32 33 34 35

@register_node
class LetStmt(Stmt):
    pass

@register_node
class AssertStmt(Stmt):
    pass

@register_node
class ProducerConsumer(Stmt):
    pass

@register_node
class For(Stmt):
tqchen committed
36 37 38 39
    Serial = 0
    Parallel = 1
    Vectorized = 2
    Unrolled = 3
tqchen committed
40 41 42 43 44 45 46 47 48 49 50 51 52 53

@register_node
class Store(Stmt):
    pass

@register_node
class Provide(Stmt):
    pass

@register_node
class Allocate(Stmt):
    pass

@register_node
54 55 56 57
class AttrStmt(Stmt):
    pass

@register_node
tqchen committed
58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75
class Free(Stmt):
    pass

@register_node
class Realize(Stmt):
    pass

@register_node
class Block(Stmt):
    pass

@register_node
class IfThenElse(Stmt):
    pass

@register_node
class Evaluate(Stmt):
    pass
76 77 78 79

@register_node
class Prefetch(Stmt):
    pass