1 | #pragma once |
2 | |
3 | #ifndef _TRITON_IR_MODULE_H_ |
4 | #define _TRITON_IR_MODULE_H_ |
5 | |
6 | #include <map> |
7 | #include <set> |
8 | #include <stack> |
9 | #include <string> |
10 | #include <functional> |
11 | #include "triton/ir/builder.h" |
12 | #include "triton/ir/metadata.h" |
13 | #include "triton/ir/context.h" |
14 | |
15 | namespace triton{ |
16 | |
17 | namespace lang{ |
18 | |
19 | class iteration_statement; |
20 | class compound_statement; |
21 | |
22 | } |
23 | |
24 | namespace ir{ |
25 | |
26 | class basic_block; |
27 | class phi_node; |
28 | class value; |
29 | class context; |
30 | class function; |
31 | class attribute; |
32 | class function_type; |
33 | class constant; |
34 | class global_value; |
35 | class alloc_const; |
36 | |
37 | class value_constructor { |
38 | typedef std::pair<std::string, basic_block*> val_key_t; |
39 | |
40 | private: |
41 | phi_node *make_phi(type *ty, unsigned num_values, basic_block *block); |
42 | value *try_remove_trivial_phis(ir::phi_node *&phi); |
43 | value *add_phi_operands(const std::string& name, phi_node *&phi); |
44 | value *get_value_recursive(const std::string& name, basic_block *block); |
45 | |
46 | public: |
47 | value_constructor(builder &builder); |
48 | |
49 | void set_value(const std::string& name, basic_block* block, value *x); |
50 | void set_value(const std::string& name, value* x); |
51 | const std::map<val_key_t, value*>& get_values() { return values_; } |
52 | void set_values(const std::map<val_key_t, value*>& values) { values_ = values; } |
53 | value *get_value(const std::string& name, basic_block* block); |
54 | value *get_value(const std::string& name); |
55 | void set_type(const std::string& name, ir::type* ty) { types_[name] = ty; } |
56 | // Seal block -- no more predecessors will be added |
57 | void seal_block(basic_block *block); |
58 | // Metadata |
59 | |
60 | private: |
61 | ir::builder& builder_; |
62 | std::map<val_key_t, value*> values_; |
63 | std::map<std::string, type*> types_; |
64 | std::set<basic_block*> sealed_blocks_; |
65 | std::map<basic_block*, std::map<std::string, phi_node*>> incomplete_phis_; |
66 | std::map<value*, value**> current_phi_; |
67 | }; |
68 | |
69 | /* Module */ |
70 | |
71 | class module { |
72 | typedef std::pair<std::string, basic_block*> val_key_t; |
73 | typedef std::pair<ir::metadata::kind_t, std::vector<unsigned>> md_pair_t; |
74 | friend class function; |
75 | |
76 | public: |
77 | typedef std::map<std::string, global_value*> symbols_map_t; |
78 | typedef std::vector<function*> functions_list_t; |
79 | |
80 | private: |
81 | void push_function(function *fn) { functions_.push_back(fn); } |
82 | |
83 | public: |
84 | module(const std::string &name, builder &builder): name_(name), builder_(builder) {} |
85 | builder &get_builder() { return builder_; }; |
86 | const std::string& get_name() { return name_; }; |
87 | |
88 | // Functions |
89 | const functions_list_t &get_function_list() const { return functions_; } |
90 | function *get_function(const std::string& name) { |
91 | if(symbols_.find(name) == symbols_.end()) |
92 | throw std::runtime_error("function " + name + " is not declared" ); |
93 | return (function*)symbols_.at(name); |
94 | } |
95 | function *get_or_insert_function(const std::string &name, function_type *ty); |
96 | bool has_function(const std::string& name){ |
97 | return symbols_.find(name) != symbols_.end(); |
98 | } |
99 | void remove_function(ir::function* fn){ |
100 | functions_.erase(std::remove(functions_.begin(), functions_.end(), fn), functions_.end()); |
101 | } |
102 | |
103 | void reset_ret_ty(const std::string& name, type* ty); |
104 | |
105 | // Const allocation |
106 | void add_alloc(ir::alloc_const* x) { allocs_.push_back(x); } |
107 | const std::vector<ir::alloc_const*>& allocs() { return allocs_; } |
108 | // Register global |
109 | void register_global(const std::string& name, ir::value *x) { globals_[name] = x; } |
110 | const std::map<std::string, ir::value*>& globals() const { return globals_; } |
111 | // Metadata |
112 | void print(std::ostream &os); |
113 | void add_metadata(const std::string &name, md_pair_t x) { metadatas_[name] = x; } |
114 | const std::map<std::string, md_pair_t> &get_metadatas() const { return metadatas_; } |
115 | |
116 | private: |
117 | std::string name_; |
118 | builder &builder_; |
119 | functions_list_t functions_; |
120 | symbols_map_t symbols_; |
121 | std::vector<ir::alloc_const*> allocs_; |
122 | std::map<std::string, ir::value*> globals_; |
123 | std::map<std::string, md_pair_t> metadatas_; |
124 | }; |
125 | |
126 | } |
127 | } |
128 | |
129 | #endif |
130 | |