1 | #include "taichi/ir/ir.h" |
2 | #include "taichi/ir/analysis.h" |
3 | #include "taichi/ir/statements.h" |
4 | |
5 | namespace taichi::lang { |
6 | |
7 | namespace irpass::analysis { |
8 | |
9 | AliasResult alias_analysis(Stmt *var1, Stmt *var2) { |
10 | // If both stmts are allocas, they have the same address iff var1 == var2. |
11 | // If only one of them is an alloca, they can never share the same address. |
12 | if (var1 == var2) |
13 | return AliasResult::same; |
14 | if (!var1 || !var2) |
15 | return AliasResult::different; |
16 | |
17 | // TODO: further optimize with offset inside MatrixPtrStmt |
18 | // If at least one of var1 and var2 is local, they will be treated here. |
19 | auto retrieve_local = [&](Stmt *var) { |
20 | if (var->is<AllocaStmt>()) { |
21 | return var; |
22 | } else if (var->is<MatrixPtrStmt>() && |
23 | var->cast<MatrixPtrStmt>()->offset_used_as_index()) { |
24 | return var->cast<MatrixPtrStmt>()->origin; |
25 | } else { |
26 | return (Stmt *)nullptr; |
27 | } |
28 | }; |
29 | Stmt *origin1 = retrieve_local(var1); |
30 | Stmt *origin2 = retrieve_local(var2); |
31 | if (origin1 != nullptr && origin2 != nullptr) { |
32 | if (origin1 == origin2) { |
33 | if (var1->is<MatrixPtrStmt>() && var2->is<MatrixPtrStmt>()) { |
34 | auto diff = value_diff_ptr_index(var1->cast<MatrixPtrStmt>()->offset, |
35 | var2->cast<MatrixPtrStmt>()->offset); |
36 | if (diff.is_diff_certain) { |
37 | return diff.diff_range == 0 ? AliasResult::same |
38 | : AliasResult::different; |
39 | } |
40 | } |
41 | return AliasResult::uncertain; |
42 | } |
43 | if (origin1->is<AllocaStmt>() || origin2->is<AllocaStmt>()) |
44 | return AliasResult::different; |
45 | TI_ASSERT(origin1->is<GlobalTemporaryStmt>() && |
46 | origin2->is<GlobalTemporaryStmt>()); |
47 | if (origin1->cast<GlobalTemporaryStmt>()->offset == |
48 | origin2->cast<GlobalTemporaryStmt>()->offset) { |
49 | return AliasResult::uncertain; |
50 | } else { |
51 | return AliasResult::different; |
52 | } |
53 | } |
54 | if (origin1 != nullptr || origin2 != nullptr) { |
55 | return AliasResult::different; |
56 | } |
57 | |
58 | if (var1->is<AdStackAllocaStmt>() || var2->is<AdStackAllocaStmt>()) |
59 | return AliasResult::different; |
60 | |
61 | // TODO(xumingkuan): Put GlobalTemporaryStmt, ThreadLocalPtrStmt and |
62 | // BlockLocalPtrStmt into GlobalPtrStmt. |
63 | // If both statements are global temps, they have the same address iff they |
64 | // have the same offset. If only one of them is a global temp, they can never |
65 | // share the same address. |
66 | if (var1->is<GlobalTemporaryStmt>() || var2->is<GlobalTemporaryStmt>()) { |
67 | if (!var1->is<GlobalTemporaryStmt>() || !var2->is<GlobalTemporaryStmt>()) |
68 | return AliasResult::different; |
69 | return var1->as<GlobalTemporaryStmt>()->offset == |
70 | var2->as<GlobalTemporaryStmt>()->offset |
71 | ? AliasResult::same |
72 | : AliasResult::different; |
73 | } |
74 | |
75 | if (var1->is<ThreadLocalPtrStmt>() || var2->is<ThreadLocalPtrStmt>()) { |
76 | if (!var1->is<ThreadLocalPtrStmt>() || !var2->is<ThreadLocalPtrStmt>()) |
77 | return AliasResult::different; |
78 | return var1->as<ThreadLocalPtrStmt>()->offset == |
79 | var2->as<ThreadLocalPtrStmt>()->offset |
80 | ? AliasResult::same |
81 | : AliasResult::different; |
82 | } |
83 | |
84 | if (var1->is<BlockLocalPtrStmt>() || var2->is<BlockLocalPtrStmt>()) { |
85 | if (!var1->is<BlockLocalPtrStmt>() || !var2->is<BlockLocalPtrStmt>()) |
86 | return AliasResult::different; |
87 | return irpass::analysis::same_statements( |
88 | var1->as<BlockLocalPtrStmt>()->offset, |
89 | var2->as<BlockLocalPtrStmt>()->offset) |
90 | ? AliasResult::same |
91 | : AliasResult::uncertain; |
92 | } |
93 | |
94 | if (var1->is<ExternalPtrStmt>() || var2->is<ExternalPtrStmt>()) { |
95 | if (!var1->is<ExternalPtrStmt>() || !var2->is<ExternalPtrStmt>()) |
96 | return AliasResult::different; |
97 | auto ptr1 = var1->as<ExternalPtrStmt>(); |
98 | auto ptr2 = var2->as<ExternalPtrStmt>(); |
99 | if (ptr1->base_ptr != ptr2->base_ptr) { |
100 | auto base1 = ptr1->base_ptr->as<ArgLoadStmt>(); |
101 | auto base2 = ptr2->base_ptr->as<ArgLoadStmt>(); |
102 | if (base1->is_grad != base2->is_grad || base1->arg_id != base2->arg_id) { |
103 | return AliasResult::different; |
104 | } |
105 | } |
106 | TI_ASSERT(ptr1->indices.size() == ptr2->indices.size()); |
107 | bool uncertain = false; |
108 | for (int i = 0; i < (int)ptr1->indices.size(); i++) { |
109 | auto diff = value_diff_ptr_index(ptr1->indices[i], ptr2->indices[i]); |
110 | if (!diff.is_diff_certain) { |
111 | uncertain = true; |
112 | } else if (diff.diff_range != 0) { |
113 | return AliasResult::different; |
114 | } |
115 | } |
116 | return uncertain ? AliasResult::uncertain : AliasResult::same; |
117 | } |
118 | |
119 | // If both statements are GlobalPtrStmts or GetChStmts, we can check by |
120 | // SNode::id. |
121 | auto get_snode_id = [](Stmt *s) { |
122 | if (auto ptr = s->cast<GlobalPtrStmt>()) { |
123 | return ptr->snode->id; |
124 | } else if (auto get_child = s->cast<GetChStmt>()) { |
125 | return get_child->output_snode->id; |
126 | } |
127 | return -1; |
128 | }; |
129 | int snode1 = get_snode_id(var1); |
130 | int snode2 = get_snode_id(var2); |
131 | if (snode1 != -1 && snode2 != -1 && snode1 != snode2) { |
132 | return AliasResult::different; |
133 | } |
134 | |
135 | // GlobalPtrStmts with guaranteed different indices cannot share the same |
136 | // address. |
137 | if (var1->is<GlobalPtrStmt>() && var2->is<GlobalPtrStmt>()) { |
138 | auto ptr1 = var1->as<GlobalPtrStmt>(); |
139 | auto ptr2 = var2->as<GlobalPtrStmt>(); |
140 | auto snode = ptr1->snode; |
141 | TI_ASSERT(snode == ptr2->snode); |
142 | TI_ASSERT(ptr1->indices.size() == ptr2->indices.size()); |
143 | bool uncertain = false; |
144 | for (int i = 0; i < (int)ptr1->indices.size(); i++) { |
145 | auto diff = value_diff_ptr_index(ptr1->indices[i], ptr2->indices[i]); |
146 | if (!diff.is_diff_certain || (diff.diff_range != 0)) { |
147 | uncertain = true; |
148 | } |
149 | if (std::abs(diff.diff_range) >= 1) { |
150 | return AliasResult::different; |
151 | } |
152 | } |
153 | return uncertain ? AliasResult::uncertain : AliasResult::same; |
154 | } |
155 | |
156 | // In other cases (probably after lower_access), we don't know if the two |
157 | // statements share the same address. |
158 | return AliasResult::uncertain; |
159 | } |
160 | |
161 | bool definitely_same_address(Stmt *var1, Stmt *var2) { |
162 | return alias_analysis(var1, var2) == AliasResult::same; |
163 | } |
164 | |
165 | bool maybe_same_address(Stmt *var1, Stmt *var2) { |
166 | return alias_analysis(var1, var2) != AliasResult::different; |
167 | } |
168 | |
169 | } // namespace irpass::analysis |
170 | |
171 | } // namespace taichi::lang |
172 | |