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
14namespace triton{
15namespace ir{
16
17namespace {
18class 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
38public:
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
59private:
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
71class AssemblyWriter {
72 std::ostream &os;
73 SlotTracker &slot_tracker;
74
75public:
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//-------------------------
93void SlotTracker::process_module() {
94 // Nothing to do at the moment.
95 // Create slots for global variable & unnamed functions & ...
96 module_processed = true;
97}
98
99void 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
121void 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
128int 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
138void 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//-------------------------------
150void 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
175void 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
187void 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
211void 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
230void 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
265void 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
298void AssemblyWriter::print_value(const value *v) {
299 // Not implemented
300}
301
302
303//-------------------------------
304// External interface
305//-------------------------------
306void module::print(std::ostream &os) {
307 SlotTracker slot_tracker(this);
308 AssemblyWriter writer(os, slot_tracker);
309 writer.print_module(this);
310}
311
312void function::print(std::ostream &os) {
313 SlotTracker slot_tracker(this);
314 AssemblyWriter writer(os, slot_tracker);
315 writer.print_function(this);
316}
317
318void 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
324void 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//-------------------------------
333std::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
342void 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
400void print(function &fn, std::ostream &os) {
401 //
402}
403
404void 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
420void 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