1 | #include "taichi/ir/ir.h" |
---|---|
2 | #include "taichi/ir/statements.h" |
3 | #include "taichi/ir/transforms.h" |
4 | #include "taichi/ir/visitors.h" |
5 | #include "taichi/system/profiler.h" |
6 | |
7 | namespace taichi::lang { |
8 | |
9 | namespace { |
10 | |
11 | // Remove all the loop_unique statements. |
12 | |
13 | class RemoveLoopUnique : public BasicStmtVisitor { |
14 | public: |
15 | using BasicStmtVisitor::visit; |
16 | DelayedIRModifier modifier; |
17 | |
18 | void visit(LoopUniqueStmt *stmt) override { |
19 | stmt->replace_usages_with(stmt->input); |
20 | modifier.erase(stmt); |
21 | } |
22 | |
23 | static bool run(IRNode *node) { |
24 | RemoveLoopUnique pass; |
25 | node->accept(&pass); |
26 | return pass.modifier.modify_ir(); |
27 | } |
28 | }; |
29 | |
30 | } // namespace |
31 | |
32 | namespace irpass { |
33 | |
34 | bool remove_loop_unique(IRNode *root) { |
35 | TI_AUTO_PROF; |
36 | return RemoveLoopUnique::run(root); |
37 | } |
38 | |
39 | } // namespace irpass |
40 | |
41 | } // namespace taichi::lang |
42 |