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
10namespace taichi::lang {
11
12/**
13 * Interprets a sequence of CHI IR statements within a block (acts like a
14 * VM based on CHI).
15 */
16class 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 &region,
95 const EvalContext &init_ctx) const;
96};
97
98} // namespace taichi::lang
99