1 | #include <iostream> |
2 | #include "triton/ir/basic_block.h" |
3 | #include "triton/ir/module.h" |
4 | #include "triton/ir/type.h" |
5 | #include "triton/ir/value.h" |
6 | #include "triton/ir/constant.h" |
7 | #include "triton/ir/function.h" |
8 | #include "triton/ir/instructions.h" |
9 | #include "triton/ir/print.h" |
10 | |
11 | #include <map> |
12 | #include <iomanip> |
13 | |
14 | namespace triton{ |
15 | namespace ir{ |
16 | |
17 | namespace { |
18 | class SlotTracker { |
19 | // A mapping of values to slot numbers. |
20 | using value_map = std::map<const value*, unsigned>; |
21 | |
22 | // The module for which we are holding slot numbers. |
23 | const module *mod_; |
24 | bool module_processed = false; |
25 | |
26 | // The function for which we are holding slot numbers. |
27 | const function *func_ = nullptr; |
28 | bool function_processed = false; |
29 | |
30 | // m_map - The slot map for the module level data. |
31 | value_map m_map; |
32 | unsigned m_next = 0; |
33 | |
34 | // f_map - The slot map for the function level data. |
35 | value_map f_map; |
36 | unsigned f_next = 0; |
37 | |
38 | public: |
39 | // Construct from a module |
40 | explicit SlotTracker(const module *mod) : mod_(mod) {} |
41 | |
42 | // Construct from a function |
43 | explicit SlotTracker(const function *f) |
44 | : mod_(f? f->get_parent() : nullptr), func_(f) {} |
45 | |
46 | // Return the slot number of the specified value. If something is not in |
47 | // the SlotTracker, return -1 |
48 | int get_local_slot(const value *v); |
49 | |
50 | void initialize_if_needed(); |
51 | |
52 | // If you'd like to deal with a function instead of just a module, use |
53 | // this method to get its data into the SlotTracker |
54 | void incorporate_function(const function *f) { |
55 | func_ = f; |
56 | function_processed = false; |
57 | } |
58 | |
59 | private: |
60 | // Add all of the module level global variables (and their initializers) |
61 | // and function declarations, but not contents of those functions. |
62 | void process_module(); |
63 | |
64 | // Add all of the functions arguments, basic blocks, and instructions. |
65 | void process_function(); |
66 | |
67 | // Insert specified value* into the slot table |
68 | void create_function_slot(const value *v); |
69 | }; |
70 | |
71 | class AssemblyWriter { |
72 | std::ostream &os; |
73 | SlotTracker &slot_tracker; |
74 | |
75 | public: |
76 | AssemblyWriter(std::ostream &os, SlotTracker &slot_tracker) |
77 | : os(os), slot_tracker(slot_tracker) {} |
78 | |
79 | void print_module(const module *mod); |
80 | void print_function(const function *f); |
81 | void print_argument(const argument *arg); |
82 | void print_basic_block(const basic_block *bb); |
83 | void print_instruction(const instruction *instr); |
84 | void print_value(const value *v); |
85 | |
86 | void write_operand(const value *op, bool print_type = false); |
87 | }; |
88 | } // anonymous namespace |
89 | |
90 | //------------------------- |
91 | // SlotTracker |
92 | //------------------------- |
93 | void SlotTracker::process_module() { |
94 | // Nothing to do at the moment. |
95 | // Create slots for global variable & unnamed functions & ... |
96 | module_processed = true; |
97 | } |
98 | |
99 | void SlotTracker::process_function() { |
100 | f_next = 0; |
101 | |
102 | // Add all the function arguments with no names. |
103 | for (const argument *arg : func_->args()) |
104 | if (!arg->has_name()) |
105 | create_function_slot(arg); |
106 | |
107 | // Add all of the basic blocks and instructions with no names. |
108 | for (const basic_block *bb : func_->blocks()) { |
109 | if (!bb->has_name()) |
110 | create_function_slot(bb); |
111 | |
112 | for (const instruction *instr : bb->get_inst_list()) { |
113 | if (!instr->get_type()->is_void_ty() && !instr->has_name()) |
114 | create_function_slot(instr); |
115 | } |
116 | } |
117 | |
118 | function_processed = true; |
119 | } |
120 | |
121 | void SlotTracker::create_function_slot(const value *v) { |
122 | assert(!v->get_type()->is_void_ty() && !v->has_name() && "Doesn't need a slot" ); |
123 | |
124 | unsigned dst_slot = f_next++; |
125 | f_map[v] = dst_slot; |
126 | } |
127 | |
128 | int SlotTracker::get_local_slot(const value *v) { |
129 | assert(dynamic_cast<const constant*>(v) == nullptr && "Can't get a constant slot" ); |
130 | |
131 | // Check for uninitialized state and do lazy initialization. |
132 | initialize_if_needed(); |
133 | |
134 | value_map::iterator f_iter = f_map.find(v); |
135 | return f_iter == f_map.end() ? -1 : (int)f_iter->second; |
136 | } |
137 | |
138 | void SlotTracker::initialize_if_needed() { |
139 | if (mod_ && !module_processed) |
140 | process_module(); |
141 | |
142 | if (func_ && !function_processed) |
143 | process_function(); |
144 | } |
145 | |
146 | |
147 | //------------------------------- |
148 | // AssemblyWriter |
149 | //------------------------------- |
150 | void AssemblyWriter::write_operand(const value *operand, bool print_type) { |
151 | if (!operand) { |
152 | os << "<null operand!>" ; |
153 | return; |
154 | } |
155 | |
156 | if (auto *c = dynamic_cast<const ir::constant*>(operand)) { |
157 | os << c->repr(); |
158 | return; |
159 | } |
160 | |
161 | if (operand->has_name()) { |
162 | os << operand->get_name(); |
163 | return; |
164 | } |
165 | |
166 | // Print the normal way |
167 | int slot_num = slot_tracker.get_local_slot(operand); |
168 | |
169 | if (slot_num != -1) |
170 | os << "%" << slot_num; |
171 | else |
172 | os << "<badref>" ; |
173 | } |
174 | |
175 | void AssemblyWriter::print_module(const module *mod) { |
176 | slot_tracker.initialize_if_needed(); |
177 | // ;ModuleID = ... |
178 | // source_filename = ... |
179 | |
180 | // Print all of the functions. |
181 | for (function *f : mod->get_function_list()) { |
182 | os << "\n" ; |
183 | print_function(f); |
184 | } |
185 | } |
186 | |
187 | void AssemblyWriter::print_function(const function *f) { |
188 | // Annotation & Attributes |
189 | |
190 | slot_tracker.incorporate_function(f); |
191 | |
192 | os << "def " ; |
193 | ir::type *rt_type = f->get_fn_type()->get_return_ty(); |
194 | // Functions must have names. |
195 | os << rt_type->repr() << " " << f->get_name() << "(" ; |
196 | // Print arguments |
197 | for (ir::argument *arg : f->args()) { |
198 | if (arg->get_arg_no() > 0) |
199 | os << ", " ; |
200 | print_argument(arg); |
201 | } |
202 | os << ")" ; |
203 | |
204 | // Print function body |
205 | os << "{" ; |
206 | for (const basic_block *bb : f->blocks()) |
207 | print_basic_block(bb); |
208 | os << "}\n" ; |
209 | } |
210 | |
211 | void AssemblyWriter::print_argument(const argument *arg) { |
212 | // Print type |
213 | os << arg->get_type()->repr(); |
214 | |
215 | // Print name, if available. |
216 | if (arg->has_name()) |
217 | os << " " << arg->get_name(); |
218 | else { |
219 | int slot_num = slot_tracker.get_local_slot(arg); |
220 | assert(slot_num != -1 && "expect argument in function here" ); |
221 | os << " %" << slot_num; |
222 | } |
223 | |
224 | // Print attributes |
225 | std::set<attribute> attrs = arg->get_parent()->get_attributes(arg); |
226 | for (attribute attr : attrs) |
227 | os << " " << attr.repr(); |
228 | } |
229 | |
230 | void AssemblyWriter::print_basic_block(const basic_block *bb) { |
231 | // bb label |
232 | if (bb->has_name()) { |
233 | os << "\n" ; |
234 | os << bb->get_name() << ":" ; |
235 | } else { |
236 | os << "\n" ; |
237 | int slot_num = slot_tracker.get_local_slot(bb); |
238 | if (slot_num != -1) |
239 | os << slot_num << ":" ; |
240 | else |
241 | os << "<badref>:" ; |
242 | } |
243 | |
244 | // Print predecessors for the block |
245 | auto const &predecessors = bb->get_predecessors(); |
246 | if (!predecessors.empty()) { |
247 | os << std::setw(50) << std::setfill(' ') |
248 | << "; preds = " ; |
249 | for (size_t i=0; i<predecessors.size(); ++i) { |
250 | if (i) |
251 | os << ", " ; |
252 | write_operand(predecessors[i]); |
253 | } |
254 | } |
255 | |
256 | os << "\n" ; |
257 | |
258 | // Annotation? |
259 | |
260 | // Print all of the instructions in the basic block |
261 | for (const ir::instruction *instr : bb->get_inst_list()) |
262 | print_instruction(instr); |
263 | } |
264 | |
265 | void AssemblyWriter::print_instruction(const instruction *instr) { |
266 | // Print out indentation for an instruction. |
267 | os << " " ; |
268 | |
269 | ir::type *type = instr->get_type(); |
270 | if (instr->has_name()) { |
271 | os << instr->get_name(); |
272 | os << " = " ; |
273 | } else if (!type->is_void_ty()) { |
274 | // Print out the def slot taken. |
275 | int slot_num = slot_tracker.get_local_slot(instr); |
276 | if (slot_num == -1) |
277 | os << "<badref> = " ; |
278 | else |
279 | os << "%" << slot_num << " = " ; |
280 | } |
281 | |
282 | // Print out opcode |
283 | os << instr->repr() << " " << type->repr(); |
284 | |
285 | size_t num_ops = instr->get_num_operands(); |
286 | if (num_ops > 0) |
287 | os << " " ; |
288 | ir::instruction::ops_t ops = instr->ops(); |
289 | for (unsigned i = 0; i < num_ops; ++i) { |
290 | if (i) |
291 | os << ", " ; |
292 | write_operand(ops[i]); |
293 | } |
294 | |
295 | os << ";\n" ; |
296 | } |
297 | |
298 | void AssemblyWriter::print_value(const value *v) { |
299 | // Not implemented |
300 | } |
301 | |
302 | |
303 | //------------------------------- |
304 | // External interface |
305 | //------------------------------- |
306 | void module::print(std::ostream &os) { |
307 | SlotTracker slot_tracker(this); |
308 | AssemblyWriter writer(os, slot_tracker); |
309 | writer.print_module(this); |
310 | } |
311 | |
312 | void function::print(std::ostream &os) { |
313 | SlotTracker slot_tracker(this); |
314 | AssemblyWriter writer(os, slot_tracker); |
315 | writer.print_function(this); |
316 | } |
317 | |
318 | void basic_block::print(std::ostream &os) { |
319 | SlotTracker slot_tracker(this->get_parent()); |
320 | AssemblyWriter writer(os, slot_tracker); |
321 | writer.print_basic_block(this); |
322 | } |
323 | |
324 | void instruction::print(std::ostream &os) { |
325 | SlotTracker slot_tracker(this->get_parent()->get_parent()); |
326 | AssemblyWriter writer(os, slot_tracker); |
327 | writer.print_instruction(this); |
328 | } |
329 | |
330 | //------------------------------- |
331 | // legacy print interface |
332 | //------------------------------- |
333 | std::string get_name(ir::value *v, unsigned i) { |
334 | if(v->get_name().empty()){ |
335 | std::string name = "%" + std::to_string(i); |
336 | v->set_name(name); |
337 | } |
338 | return v->get_name(); |
339 | } |
340 | |
341 | |
342 | void print(module &mod, std::ostream& os) { |
343 | unsigned cnt = 0; |
344 | for(ir::function *fn: mod.get_function_list()){ |
345 | os << "def " << fn->get_fn_type()->get_return_ty()->repr() << " " << fn->get_name() << "(" ; |
346 | for(ir::argument* arg: fn->args()) { |
347 | if(arg->get_arg_no() > 0) |
348 | os << ", " ; |
349 | os << arg->get_type()->repr() << " " << arg->get_name(); |
350 | auto attrs = fn->get_attributes(arg); |
351 | if(attrs.size() > 0) |
352 | os << " " ; |
353 | for(ir::attribute attr: attrs) |
354 | os << attr.repr() << " " ; |
355 | } |
356 | os << ")" << std::endl; |
357 | os << "{" << std::endl; |
358 | for(ir::basic_block *block: fn->blocks()){ |
359 | auto const &predecessors = block->get_predecessors(); |
360 | os << block->get_name() << ":" ; |
361 | if(!predecessors.empty()){ |
362 | os << " " ; |
363 | os << "; preds = " ; |
364 | auto const &predecessors = block->get_predecessors(); |
365 | for(ir::basic_block *pred: predecessors) |
366 | os << pred->get_name() << (pred!=predecessors.back()?", " :"" ); |
367 | } |
368 | os << std::endl; |
369 | for(ir::instruction *inst: block->get_inst_list()){ |
370 | os << " " ; |
371 | if(!inst->get_type()->is_void_ty()){ |
372 | os << get_name(inst, cnt++); |
373 | os << " = " ; |
374 | } |
375 | ir::type* type = inst->get_type(); |
376 | os << inst->repr() << " " << type->repr(); |
377 | ir::instruction::ops_t ops = inst->ops(); |
378 | size_t num_ops = inst->get_num_operands(); |
379 | if(num_ops > 0) |
380 | os << " " ;; |
381 | for(unsigned i = 0; i < num_ops; i++){ |
382 | if(auto *x = dynamic_cast<ir::constant*>(ops[i])) |
383 | os << x->repr(); |
384 | else |
385 | os << get_name(ops[i], cnt++); |
386 | os << (i < num_ops - 1?", " :"" ); |
387 | } |
388 | os << ";" ; |
389 | // os << " ("; |
390 | // for(ir::user* usr: inst->get_users()) |
391 | // os << get_name(usr, cnt++) << ", " ; |
392 | // os << " )"; |
393 | os << std::endl; |
394 | } |
395 | } |
396 | os << "}" << std::endl; |
397 | } |
398 | } |
399 | |
400 | void print(function &fn, std::ostream &os) { |
401 | // |
402 | } |
403 | |
404 | void print(basic_block &bb, std::ostream &os) { |
405 | auto const &predecessors = bb.get_predecessors(); |
406 | os << bb.get_name() << ":" ; |
407 | if(!predecessors.empty()){ |
408 | os << " " ; |
409 | os << "; preds = " ; |
410 | auto const &predecessors = bb.get_predecessors(); |
411 | for(ir::basic_block *pred: predecessors) |
412 | os << pred->get_name() << (pred!=predecessors.back()?", " :"" ); |
413 | } |
414 | os << std::endl; |
415 | for(ir::instruction *inst: bb.get_inst_list()){ |
416 | print(*inst, os); |
417 | } |
418 | } |
419 | |
420 | void print(instruction &instr, std::ostream &os) { |
421 | instruction *inst = &instr; |
422 | os << " " ; |
423 | if(!inst->get_type()->is_void_ty()){ |
424 | os << instr.get_name(); |
425 | os << " = " ; |
426 | } |
427 | ir::type* type = inst->get_type(); |
428 | os << inst->repr() << " " << type->repr(); |
429 | ir::instruction::ops_t ops = inst->ops(); |
430 | size_t num_ops = inst->get_num_operands(); |
431 | if(num_ops > 0) |
432 | os << " " ;; |
433 | for(unsigned i = 0; i < num_ops; i++){ |
434 | if(auto *x = dynamic_cast<ir::constant*>(ops[i])) |
435 | os << x->repr(); |
436 | else |
437 | os << ops[i]->get_name(); |
438 | os << (i < num_ops - 1?", " :"" ); |
439 | } |
440 | os << ";" ; |
441 | // os << " ("; |
442 | // for(ir::user* usr: inst->get_users()) |
443 | // os << get_name(usr, cnt++) << ", " ; |
444 | // os << " )"; |
445 | os << std::endl; |
446 | } |
447 | |
448 | |
449 | } |
450 | } |
451 | |