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
15namespace triton{
16
17namespace lang{
18
19class iteration_statement;
20class compound_statement;
21
22}
23
24namespace ir{
25
26class basic_block;
27class phi_node;
28class value;
29class context;
30class function;
31class attribute;
32class function_type;
33class constant;
34class global_value;
35class alloc_const;
36
37class value_constructor {
38 typedef std::pair<std::string, basic_block*> val_key_t;
39
40private:
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
46public:
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
60private:
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
71class 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
76public:
77 typedef std::map<std::string, global_value*> symbols_map_t;
78 typedef std::vector<function*> functions_list_t;
79
80private:
81 void push_function(function *fn) { functions_.push_back(fn); }
82
83public:
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
116private:
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