1 | #include "taichi/ir/ir.h" |
---|---|
2 | #include "taichi/ir/analysis.h" |
3 | #include "taichi/ir/visitors.h" |
4 | |
5 | namespace taichi::lang { |
6 | |
7 | // Count all statements (including containers) |
8 | class StmtCounter : public BasicStmtVisitor { |
9 | private: |
10 | StmtCounter() { |
11 | counter_ = 0; |
12 | allow_undefined_visitor = true; |
13 | invoke_default_visitor = true; |
14 | } |
15 | |
16 | using BasicStmtVisitor::visit; |
17 | |
18 | public: |
19 | void preprocess_container_stmt(Stmt *stmt) override { |
20 | counter_++; |
21 | } |
22 | |
23 | void visit(Stmt *stmt) override { |
24 | counter_++; |
25 | } |
26 | |
27 | static int run(IRNode *root) { |
28 | StmtCounter stmt_counter; |
29 | root->accept(&stmt_counter); |
30 | return stmt_counter.counter_; |
31 | } |
32 | |
33 | private: |
34 | int counter_; |
35 | }; |
36 | |
37 | namespace irpass::analysis { |
38 | int count_statements(IRNode *root) { |
39 | TI_ASSERT(root); |
40 | return StmtCounter::run(root); |
41 | } |
42 | } // namespace irpass::analysis |
43 | |
44 | } // namespace taichi::lang |
45 |