1 | #pragma once |
2 | |
3 | #include <optional> |
4 | #include <unordered_set> |
5 | #include <unordered_map> |
6 | |
7 | #include "taichi/ir/statements.h" |
8 | #include "taichi/ir/type.h" |
9 | |
10 | namespace taichi::lang { |
11 | |
12 | /** |
13 | * Interprets a sequence of CHI IR statements within a block (acts like a |
14 | * VM based on CHI). |
15 | */ |
16 | class ArithmeticInterpretor { |
17 | public: |
18 | /** |
19 | * Evaluation context that maps from a Stmt to a constant value. |
20 | */ |
21 | class EvalContext { |
22 | public: |
23 | /** |
24 | * Pre-defines a value for statement @param s. |
25 | * |
26 | * @param s: Statement to be evaluated |
27 | * @param c: Predefined value |
28 | */ |
29 | EvalContext &insert(const Stmt *s, TypedConstant c) { |
30 | map_[s] = c; |
31 | return *this; |
32 | } |
33 | |
34 | /** |
35 | * Tries to get the evaluated value for statement @param s. |
36 | * |
37 | * @param s: Statement to get |
38 | * @return: The evaluated value, empty if not found. |
39 | */ |
40 | std::optional<TypedConstant> maybe_get(const Stmt *s) const { |
41 | auto itr = map_.find(s); |
42 | if (itr == map_.end()) { |
43 | return std::nullopt; |
44 | } |
45 | return itr->second; |
46 | } |
47 | |
48 | /** |
49 | * Tells the interpreter to ignore statement @param s. |
50 | * |
51 | * This is effective only for statements that are not supported by |
52 | * ArithmeticInterpretor. |
53 | * |
54 | * @param s: Statement to ignore |
55 | */ |
56 | void ignore(const Stmt *s) { |
57 | ignored_.insert(s); |
58 | } |
59 | |
60 | /** |
61 | * Checks if statement @param s is ignored. |
62 | * |
63 | * @return: True if ignored |
64 | */ |
65 | bool should_ignore(const Stmt *s) { |
66 | return ignored_.count(s) > 0; |
67 | } |
68 | |
69 | private: |
70 | std::unordered_map<const Stmt *, TypedConstant> map_; |
71 | std::unordered_set<const Stmt *> ignored_; |
72 | }; |
73 | |
74 | /** |
75 | * Defines the region of CHI statements to be evaluated. |
76 | */ |
77 | struct CodeRegion { |
78 | // Defines the sequence of CHI statements. |
79 | Block *block{nullptr}; |
80 | // The beginning statement within |block| to be evaluated. If nullptr, |
81 | // evaluates from the beginning of |block|. |
82 | Stmt *begin{nullptr}; |
83 | // The ending statement (exclusive) within |block| to be evaluated. If |
84 | // nullptr, evaluates to the end of |block|. |
85 | Stmt *end{nullptr}; |
86 | }; |
87 | |
88 | /** |
89 | * Evaluates the sequence of CHI as defined in |region|. |
90 | * @param region: A sequence of CHI statements to be evaluated |
91 | * @param init_ctx: This context can mock the result for certain types of |
92 | * statements that are not supported, or cannot be evaluated statically. |
93 | */ |
94 | std::optional<TypedConstant> evaluate(const CodeRegion ®ion, |
95 | const EvalContext &init_ctx) const; |
96 | }; |
97 | |
98 | } // namespace taichi::lang |
99 | |