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
7namespace taichi::lang {
8
9class 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
45namespace irpass::analysis {
46std::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