1 | #include "taichi/ir/ir.h" |
---|---|
2 | #include "taichi/ir/analysis.h" |
3 | #include "taichi/ir/statements.h" |
4 | #include "taichi/ir/visitors.h" |
5 | #include <unordered_set> |
6 | |
7 | namespace taichi::lang { |
8 | |
9 | class UsedAtomicsSearcher : public BasicStmtVisitor { |
10 | private: |
11 | std::unique_ptr<std::unordered_set<AtomicOpStmt *>> used_atomics_; |
12 | |
13 | public: |
14 | using BasicStmtVisitor::visit; |
15 | |
16 | UsedAtomicsSearcher() { |
17 | allow_undefined_visitor = true; |
18 | invoke_default_visitor = true; |
19 | used_atomics_ = std::make_unique<std::unordered_set<AtomicOpStmt *>>(); |
20 | } |
21 | |
22 | void search_operands(Stmt *stmt) { |
23 | for (auto &op : stmt->get_operands()) { |
24 | if (op != nullptr && op->is<AtomicOpStmt>()) { |
25 | used_atomics_->insert(op->as<AtomicOpStmt>()); |
26 | } |
27 | } |
28 | } |
29 | |
30 | void preprocess_container_stmt(Stmt *stmt) override { |
31 | search_operands(stmt); |
32 | } |
33 | |
34 | void visit(Stmt *stmt) override { |
35 | search_operands(stmt); |
36 | } |
37 | |
38 | static std::unique_ptr<std::unordered_set<AtomicOpStmt *>> run(IRNode *root) { |
39 | UsedAtomicsSearcher searcher; |
40 | root->accept(&searcher); |
41 | return std::move(searcher.used_atomics_); |
42 | } |
43 | }; |
44 | |
45 | namespace irpass::analysis { |
46 | std::unique_ptr<std::unordered_set<AtomicOpStmt *>> gather_used_atomics( |
47 | IRNode *root) { |
48 | return UsedAtomicsSearcher::run(root); |
49 | } |
50 | } // namespace irpass::analysis |
51 | |
52 | } // namespace taichi::lang |
53 |