1#include "taichi/ir/ir.h"
2#include "taichi/ir/analysis.h"
3#include "taichi/ir/statements.h"
4
5namespace taichi::lang {
6
7namespace irpass::analysis {
8
9AliasResult 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
161bool definitely_same_address(Stmt *var1, Stmt *var2) {
162 return alias_analysis(var1, var2) == AliasResult::same;
163}
164
165bool 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