blob: 55cc98ce34cd4289c28718c234ddae1cafb27300 [file] [log] [blame]
Chris Lattner4c95a502018-06-23 16:03:42 -07001//===- AsmPrinter.cpp - MLIR Assembly Printer Implementation --------------===//
2//
3// Copyright 2019 The MLIR Authors.
4//
5// Licensed under the Apache License, Version 2.0 (the "License");
6// you may not use this file except in compliance with the License.
7// You may obtain a copy of the License at
8//
9// http://www.apache.org/licenses/LICENSE-2.0
10//
11// Unless required by applicable law or agreed to in writing, software
12// distributed under the License is distributed on an "AS IS" BASIS,
13// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
14// See the License for the specific language governing permissions and
15// limitations under the License.
16// =============================================================================
17//
18// This file implements the MLIR AsmPrinter class, which is used to implement
19// the various print() methods on the core IR objects.
20//
21//===----------------------------------------------------------------------===//
22
Uday Bondhugulafaf37dd2018-06-29 18:09:29 -070023#include "mlir/IR/AffineExpr.h"
24#include "mlir/IR/AffineMap.h"
Chris Lattner7121b802018-07-04 20:45:39 -070025#include "mlir/IR/Attributes.h"
Chris Lattner4c95a502018-06-23 16:03:42 -070026#include "mlir/IR/CFGFunction.h"
Tatiana Shpeismanc96b5872018-06-28 17:02:32 -070027#include "mlir/IR/MLFunction.h"
Chris Lattner4c95a502018-06-23 16:03:42 -070028#include "mlir/IR/Module.h"
Chris Lattnerff0d5902018-07-05 09:12:11 -070029#include "mlir/IR/OperationSet.h"
Tatiana Shpeisman1bcfe982018-07-13 13:03:13 -070030#include "mlir/IR/Statements.h"
Chris Lattner4c95a502018-06-23 16:03:42 -070031#include "mlir/IR/Types.h"
32#include "mlir/Support/STLExtras.h"
Chris Lattner4c95a502018-06-23 16:03:42 -070033#include "llvm/ADT/DenseMap.h"
Uday Bondhugulafaf37dd2018-06-29 18:09:29 -070034#include "llvm/Support/raw_ostream.h"
Chris Lattner4c95a502018-06-23 16:03:42 -070035using namespace mlir;
36
MLIR Team54b55a22018-07-18 10:16:05 -070037void Identifier::print(raw_ostream &os) const { os << str(); }
Chris Lattner4c95a502018-06-23 16:03:42 -070038
MLIR Team54b55a22018-07-18 10:16:05 -070039void Identifier::dump() const { print(llvm::errs()); }
Chris Lattner7121b802018-07-04 20:45:39 -070040
Chris Lattner4c95a502018-06-23 16:03:42 -070041//===----------------------------------------------------------------------===//
Chris Lattner4fd59b02018-07-20 09:35:47 -070042// ModuleState
MLIR Team4718bc92018-07-17 16:56:54 -070043//===----------------------------------------------------------------------===//
44
45namespace {
MLIR Team54b55a22018-07-18 10:16:05 -070046class ModuleState {
47public:
Chris Lattner4fd59b02018-07-20 09:35:47 -070048 /// This is the operation set for the current context if it is knowable (a
49 /// context could be determined), otherwise this is null.
50 OperationSet *const operationSet;
MLIR Team4718bc92018-07-17 16:56:54 -070051
Chris Lattner4fd59b02018-07-20 09:35:47 -070052 explicit ModuleState(MLIRContext *context)
53 : operationSet(context ? &OperationSet::get(context) : nullptr) {}
54
55 // Initializes module state, populating affine map state.
MLIR Team4718bc92018-07-17 16:56:54 -070056 void initialize(const Module *module);
57
MLIR Team54b55a22018-07-18 10:16:05 -070058 int getAffineMapId(const AffineMap *affineMap) const {
MLIR Team4718bc92018-07-17 16:56:54 -070059 auto it = affineMapIds.find(affineMap);
60 if (it == affineMapIds.end()) {
61 return -1;
62 }
63 return it->second;
64 }
65
Chris Lattner4fd59b02018-07-20 09:35:47 -070066 const DenseMap<const AffineMap *, int> &getAffineMapIds() const {
67 return affineMapIds;
68 }
69
MLIR Team54b55a22018-07-18 10:16:05 -070070private:
Chris Lattner4fd59b02018-07-20 09:35:47 -070071 void recordAffineMapReference(const AffineMap *affineMap) {
72 if (affineMapIds.count(affineMap) == 0) {
73 affineMapIds[affineMap] = nextAffineMapId++;
74 }
75 }
76
MLIR Team4718bc92018-07-17 16:56:54 -070077 // Visit functions.
78 void visitFunction(const Function *fn);
79 void visitExtFunction(const ExtFunction *fn);
80 void visitCFGFunction(const CFGFunction *fn);
81 void visitMLFunction(const MLFunction *fn);
82 void visitType(const Type *type);
MLIR Teamb61885d2018-07-18 16:29:21 -070083 void visitAttribute(const Attribute *attr);
84 void visitOperation(const Operation *op);
85
MLIR Team54b55a22018-07-18 10:16:05 -070086 DenseMap<const AffineMap *, int> affineMapIds;
MLIR Team4718bc92018-07-17 16:56:54 -070087 int nextAffineMapId = 0;
88};
James Molloy87d81022018-07-23 11:44:40 -070089} // end anonymous namespace
MLIR Team4718bc92018-07-17 16:56:54 -070090
91// TODO Support visiting other types/instructions when implemented.
92void ModuleState::visitType(const Type *type) {
93 if (type->getKind() == Type::Kind::Function) {
94 // Visit input and result types for functions.
95 auto *funcType = cast<FunctionType>(type);
MLIR Team54b55a22018-07-18 10:16:05 -070096 for (auto *input : funcType->getInputs()) {
MLIR Team4718bc92018-07-17 16:56:54 -070097 visitType(input);
98 }
MLIR Team54b55a22018-07-18 10:16:05 -070099 for (auto *result : funcType->getResults()) {
MLIR Team4718bc92018-07-17 16:56:54 -0700100 visitType(result);
101 }
102 } else if (type->getKind() == Type::Kind::MemRef) {
103 // Visit affine maps in memref type.
104 auto *memref = cast<MemRefType>(type);
MLIR Team54b55a22018-07-18 10:16:05 -0700105 for (AffineMap *map : memref->getAffineMaps()) {
MLIR Team4718bc92018-07-17 16:56:54 -0700106 recordAffineMapReference(map);
107 }
108 }
109}
110
MLIR Teamb61885d2018-07-18 16:29:21 -0700111void ModuleState::visitAttribute(const Attribute *attr) {
112 if (isa<AffineMapAttr>(attr)) {
113 recordAffineMapReference(cast<AffineMapAttr>(attr)->getValue());
114 } else if (isa<ArrayAttr>(attr)) {
115 for (auto elt : cast<ArrayAttr>(attr)->getValue()) {
116 visitAttribute(elt);
117 }
118 }
119}
120
121void ModuleState::visitOperation(const Operation *op) {
122 for (auto elt : op->getAttrs()) {
123 visitAttribute(elt.second);
124 }
125}
126
MLIR Team4718bc92018-07-17 16:56:54 -0700127void ModuleState::visitExtFunction(const ExtFunction *fn) {
128 visitType(fn->getType());
129}
130
131void ModuleState::visitCFGFunction(const CFGFunction *fn) {
132 visitType(fn->getType());
MLIR Teamb61885d2018-07-18 16:29:21 -0700133 for (auto &block : *fn) {
134 for (auto &op : block.getOperations()) {
135 visitOperation(&op);
136 }
137 }
MLIR Team4718bc92018-07-17 16:56:54 -0700138}
139
140void ModuleState::visitMLFunction(const MLFunction *fn) {
141 visitType(fn->getType());
MLIR Teamb61885d2018-07-18 16:29:21 -0700142 // TODO Visit function body statements (and attributes if required).
MLIR Team4718bc92018-07-17 16:56:54 -0700143}
144
145void ModuleState::visitFunction(const Function *fn) {
146 switch (fn->getKind()) {
MLIR Team54b55a22018-07-18 10:16:05 -0700147 case Function::Kind::ExtFunc:
148 return visitExtFunction(cast<ExtFunction>(fn));
149 case Function::Kind::CFGFunc:
150 return visitCFGFunction(cast<CFGFunction>(fn));
151 case Function::Kind::MLFunc:
152 return visitMLFunction(cast<MLFunction>(fn));
MLIR Team4718bc92018-07-17 16:56:54 -0700153 }
154}
155
Chris Lattner4fd59b02018-07-20 09:35:47 -0700156// Initializes module state, populating affine map state.
157void ModuleState::initialize(const Module *module) {
158 for (auto fn : module->functionList) {
159 visitFunction(fn);
160 }
161}
162
163//===----------------------------------------------------------------------===//
164// ModulePrinter
165//===----------------------------------------------------------------------===//
166
167namespace {
168class ModulePrinter {
169public:
170 ModulePrinter(raw_ostream &os, ModuleState &state) : os(os), state(state) {}
171 explicit ModulePrinter(const ModulePrinter &printer)
172 : os(printer.os), state(printer.state) {}
173
174 template <typename Container, typename UnaryFunctor>
175 inline void interleaveComma(const Container &c, UnaryFunctor each_fn) const {
176 interleave(c.begin(), c.end(), each_fn, [&]() { os << ", "; });
177 }
178
179 void print(const Module *module);
180 void print(const Attribute *attr) const;
181 void print(const Type *type) const;
182 void print(const Function *fn);
183 void print(const ExtFunction *fn);
184 void print(const CFGFunction *fn);
185 void print(const MLFunction *fn);
186
187 void print(const AffineMap *map);
188 void print(const AffineExpr *expr) const;
189
190protected:
191 raw_ostream &os;
192 ModuleState &state;
193
194 void printFunctionSignature(const Function *fn);
195 void printAffineMapId(int affineMapId) const;
196 void printAffineMapReference(const AffineMap *affineMap) const;
197
198 void print(const AffineBinaryOpExpr *expr) const;
199};
200} // end anonymous namespace
201
MLIR Team4718bc92018-07-17 16:56:54 -0700202// Prints function with initialized module state.
Chris Lattner4fd59b02018-07-20 09:35:47 -0700203void ModulePrinter::print(const Function *fn) {
MLIR Team4718bc92018-07-17 16:56:54 -0700204 switch (fn->getKind()) {
MLIR Team54b55a22018-07-18 10:16:05 -0700205 case Function::Kind::ExtFunc:
206 return print(cast<ExtFunction>(fn));
207 case Function::Kind::CFGFunc:
208 return print(cast<CFGFunction>(fn));
209 case Function::Kind::MLFunc:
210 return print(cast<MLFunction>(fn));
MLIR Team4718bc92018-07-17 16:56:54 -0700211 }
212}
213
214// Prints affine map identifier.
Chris Lattner4fd59b02018-07-20 09:35:47 -0700215void ModulePrinter::printAffineMapId(int affineMapId) const {
MLIR Team4718bc92018-07-17 16:56:54 -0700216 os << "#map" << affineMapId;
217}
218
Chris Lattner4fd59b02018-07-20 09:35:47 -0700219void ModulePrinter::printAffineMapReference(const AffineMap *affineMap) const {
220 int mapId = state.getAffineMapId(affineMap);
MLIR Teamb61885d2018-07-18 16:29:21 -0700221 if (mapId >= 0) {
222 // Map will be printed at top of module so print reference to its id.
223 printAffineMapId(mapId);
224 } else {
225 // Map not in module state so print inline.
226 affineMap->print(os);
227 }
228}
229
Chris Lattner4fd59b02018-07-20 09:35:47 -0700230void ModulePrinter::print(const Module *module) {
231 for (const auto &mapAndId : state.getAffineMapIds()) {
MLIR Teamb61885d2018-07-18 16:29:21 -0700232 printAffineMapId(mapAndId.second);
MLIR Team4718bc92018-07-17 16:56:54 -0700233 os << " = ";
234 mapAndId.first->print(os);
235 os << '\n';
236 }
James Molloy87d81022018-07-23 11:44:40 -0700237 for (auto *fn : module->functionList)
238 print(fn);
MLIR Team4718bc92018-07-17 16:56:54 -0700239}
240
Chris Lattner4fd59b02018-07-20 09:35:47 -0700241void ModulePrinter::print(const Attribute *attr) const {
MLIR Teamb61885d2018-07-18 16:29:21 -0700242 switch (attr->getKind()) {
243 case Attribute::Kind::Bool:
244 os << (cast<BoolAttr>(attr)->getValue() ? "true" : "false");
245 break;
246 case Attribute::Kind::Integer:
247 os << cast<IntegerAttr>(attr)->getValue();
248 break;
249 case Attribute::Kind::Float:
250 // FIXME: this isn't precise, we should print with a hex format.
251 os << cast<FloatAttr>(attr)->getValue();
252 break;
253 case Attribute::Kind::String:
254 // FIXME: should escape the string.
255 os << '"' << cast<StringAttr>(attr)->getValue() << '"';
256 break;
257 case Attribute::Kind::Array: {
258 auto elts = cast<ArrayAttr>(attr)->getValue();
259 os << '[';
Chris Lattner4fd59b02018-07-20 09:35:47 -0700260 interleaveComma(elts, [&](Attribute *attr) { print(attr); });
MLIR Teamb61885d2018-07-18 16:29:21 -0700261 os << ']';
262 break;
263 }
264 case Attribute::Kind::AffineMap:
265 printAffineMapReference(cast<AffineMapAttr>(attr)->getValue());
266 break;
267 }
268}
269
Chris Lattner4fd59b02018-07-20 09:35:47 -0700270void ModulePrinter::print(const Type *type) const {
MLIR Team4718bc92018-07-17 16:56:54 -0700271 switch (type->getKind()) {
MLIR Team54b55a22018-07-18 10:16:05 -0700272 case Type::Kind::AffineInt:
273 os << "affineint";
274 return;
275 case Type::Kind::BF16:
276 os << "bf16";
277 return;
278 case Type::Kind::F16:
279 os << "f16";
280 return;
281 case Type::Kind::F32:
282 os << "f32";
283 return;
284 case Type::Kind::F64:
285 os << "f64";
286 return;
MLIR Team4718bc92018-07-17 16:56:54 -0700287
288 case Type::Kind::Integer: {
289 auto *integer = cast<IntegerType>(type);
290 os << 'i' << integer->getWidth();
291 return;
292 }
293 case Type::Kind::Function: {
294 auto *func = cast<FunctionType>(type);
295 os << '(';
Chris Lattner4fd59b02018-07-20 09:35:47 -0700296 interleaveComma(func->getInputs(), [&](Type *type) { os << *type; });
MLIR Team4718bc92018-07-17 16:56:54 -0700297 os << ") -> ";
298 auto results = func->getResults();
299 if (results.size() == 1)
300 os << *results[0];
301 else {
302 os << '(';
Chris Lattner4fd59b02018-07-20 09:35:47 -0700303 interleaveComma(results, [&](Type *type) { os << *type; });
MLIR Team4718bc92018-07-17 16:56:54 -0700304 os << ')';
305 }
306 return;
307 }
308 case Type::Kind::Vector: {
309 auto *v = cast<VectorType>(type);
310 os << "vector<";
James Molloy87d81022018-07-23 11:44:40 -0700311 for (auto dim : v->getShape())
312 os << dim << 'x';
MLIR Team4718bc92018-07-17 16:56:54 -0700313 os << *v->getElementType() << '>';
314 return;
315 }
316 case Type::Kind::RankedTensor: {
317 auto *v = cast<RankedTensorType>(type);
318 os << "tensor<";
319 for (auto dim : v->getShape()) {
320 if (dim < 0)
321 os << '?';
322 else
323 os << dim;
324 os << 'x';
325 }
326 os << *v->getElementType() << '>';
327 return;
328 }
329 case Type::Kind::UnrankedTensor: {
330 auto *v = cast<UnrankedTensorType>(type);
331 os << "tensor<??" << *v->getElementType() << '>';
332 return;
333 }
334 case Type::Kind::MemRef: {
335 auto *v = cast<MemRefType>(type);
336 os << "memref<";
337 for (auto dim : v->getShape()) {
338 if (dim < 0)
339 os << '?';
340 else
341 os << dim;
342 os << 'x';
343 }
344 os << *v->getElementType();
345 for (auto map : v->getAffineMaps()) {
346 os << ", ";
MLIR Teamb61885d2018-07-18 16:29:21 -0700347 printAffineMapReference(map);
MLIR Team4718bc92018-07-17 16:56:54 -0700348 }
349 os << ", " << v->getMemorySpace();
350 os << '>';
351 return;
352 }
353 }
354}
355
356//===----------------------------------------------------------------------===//
Chris Lattner4fd59b02018-07-20 09:35:47 -0700357// Affine expressions and maps
358//===----------------------------------------------------------------------===//
359
360void ModulePrinter::print(const AffineExpr *expr) const {
361 switch (expr->getKind()) {
362 case AffineExpr::Kind::SymbolId:
363 os << 's' << cast<AffineSymbolExpr>(expr)->getPosition();
364 return;
365 case AffineExpr::Kind::DimId:
366 os << 'd' << cast<AffineDimExpr>(expr)->getPosition();
367 return;
368 case AffineExpr::Kind::Constant:
369 os << cast<AffineConstantExpr>(expr)->getValue();
370 return;
371 case AffineExpr::Kind::Add:
372 case AffineExpr::Kind::Mul:
373 case AffineExpr::Kind::FloorDiv:
374 case AffineExpr::Kind::CeilDiv:
375 case AffineExpr::Kind::Mod:
376 return print(cast<AffineBinaryOpExpr>(expr));
377 }
378}
379
380void ModulePrinter::print(const AffineBinaryOpExpr *expr) const {
381 if (expr->getKind() != AffineExpr::Kind::Add) {
382 os << '(';
383 print(expr->getLHS());
384 switch (expr->getKind()) {
385 case AffineExpr::Kind::Mul:
386 os << " * ";
387 break;
388 case AffineExpr::Kind::FloorDiv:
389 os << " floordiv ";
390 break;
391 case AffineExpr::Kind::CeilDiv:
392 os << " ceildiv ";
393 break;
394 case AffineExpr::Kind::Mod:
395 os << " mod ";
396 break;
397 default:
398 llvm_unreachable("unexpected affine binary op expression");
399 }
400
401 print(expr->getRHS());
402 os << ')';
403 return;
404 }
405
406 // Print out special "pretty" forms for add.
407 os << '(';
408 print(expr->getLHS());
409
410 // Pretty print addition to a product that has a negative operand as a
411 // subtraction.
412 if (auto *rhs = dyn_cast<AffineBinaryOpExpr>(expr->getRHS())) {
413 if (rhs->getKind() == AffineExpr::Kind::Mul) {
414 if (auto *rrhs = dyn_cast<AffineConstantExpr>(rhs->getRHS())) {
415 if (rrhs->getValue() < 0) {
416 os << " - (";
417 print(rhs->getLHS());
418 os << " * " << -rrhs->getValue() << "))";
419 return;
420 }
421 }
422 }
423 }
424
425 // Pretty print addition to a negative number as a subtraction.
426 if (auto *rhs = dyn_cast<AffineConstantExpr>(expr->getRHS())) {
427 if (rhs->getValue() < 0) {
428 os << " - " << -rhs->getValue() << ")";
429 return;
430 }
431 }
432
433 os << " + ";
434 print(expr->getRHS());
435 os << ')';
436}
437
438void ModulePrinter::print(const AffineMap *map) {
439 // Dimension identifiers.
440 os << '(';
441 for (int i = 0; i < (int)map->getNumDims() - 1; i++)
442 os << "d" << i << ", ";
443 if (map->getNumDims() >= 1)
444 os << "d" << map->getNumDims() - 1;
445 os << ")";
446
447 // Symbolic identifiers.
448 if (map->getNumSymbols() >= 1) {
449 os << " [";
450 for (int i = 0; i < (int)map->getNumSymbols() - 1; i++)
451 os << "s" << i << ", ";
452 if (map->getNumSymbols() >= 1)
453 os << "s" << map->getNumSymbols() - 1;
454 os << "]";
455 }
456
457 // AffineMap should have at least one result.
458 assert(!map->getResults().empty());
459 // Result affine expressions.
460 os << " -> (";
461 interleaveComma(map->getResults(), [&](AffineExpr *expr) { print(expr); });
462 os << ")";
463
464 if (!map->isBounded()) {
465 return;
466 }
467
468 // Print range sizes for bounded affine maps.
469 os << " size (";
470 interleaveComma(map->getRangeSizes(), [&](AffineExpr *expr) { print(expr); });
471 os << ")";
472}
473
474//===----------------------------------------------------------------------===//
Chris Lattner4c95a502018-06-23 16:03:42 -0700475// Function printing
476//===----------------------------------------------------------------------===//
477
Chris Lattner4fd59b02018-07-20 09:35:47 -0700478void ModulePrinter::printFunctionSignature(const Function *fn) {
Chris Lattner4c95a502018-06-23 16:03:42 -0700479 auto type = fn->getType();
480
481 os << "@" << fn->getName() << '(';
Chris Lattner4fd59b02018-07-20 09:35:47 -0700482 interleaveComma(type->getInputs(), [&](Type *eltType) { print(eltType); });
Chris Lattner4c95a502018-06-23 16:03:42 -0700483 os << ')';
484
485 switch (type->getResults().size()) {
MLIR Team54b55a22018-07-18 10:16:05 -0700486 case 0:
487 break;
Chris Lattner4c95a502018-06-23 16:03:42 -0700488 case 1:
MLIR Team4718bc92018-07-17 16:56:54 -0700489 os << " -> ";
Chris Lattner4fd59b02018-07-20 09:35:47 -0700490 print(type->getResults()[0]);
Chris Lattner4c95a502018-06-23 16:03:42 -0700491 break;
492 default:
493 os << " -> (";
Chris Lattner4fd59b02018-07-20 09:35:47 -0700494 interleaveComma(type->getResults(), [&](Type *eltType) { print(eltType); });
Chris Lattner4c95a502018-06-23 16:03:42 -0700495 os << ')';
496 break;
497 }
498}
499
Chris Lattner4fd59b02018-07-20 09:35:47 -0700500void ModulePrinter::print(const ExtFunction *fn) {
Chris Lattner4c95a502018-06-23 16:03:42 -0700501 os << "extfunc ";
Chris Lattner4fd59b02018-07-20 09:35:47 -0700502 printFunctionSignature(fn);
MLIR Team54b55a22018-07-18 10:16:05 -0700503 os << '\n';
Chris Lattner4c95a502018-06-23 16:03:42 -0700504}
505
Tatiana Shpeismanfa412f72018-07-09 17:42:46 -0700506namespace {
507
508// FunctionState contains common functionality for printing
509// CFG and ML functions.
Chris Lattner4fd59b02018-07-20 09:35:47 -0700510class FunctionState : public ModulePrinter {
Tatiana Shpeismanfa412f72018-07-09 17:42:46 -0700511public:
Chris Lattner4fd59b02018-07-20 09:35:47 -0700512 FunctionState(const ModulePrinter &other) : ModulePrinter(other) {}
Tatiana Shpeismanfa412f72018-07-09 17:42:46 -0700513
514 void printOperation(const Operation *op);
515
516protected:
Chris Lattnerf8cce872018-07-20 09:28:54 -0700517 void numberValueID(const SSAValue *value) {
518 assert(!valueIDs.count(value) && "Value numbered multiple times");
519 valueIDs[value] = nextValueID++;
520 }
521
Chris Lattner6119d382018-07-20 18:41:34 -0700522 void printValueID(const SSAValue *value,
523 bool dontPrintResultNo = false) const {
524 int resultNo = -1;
525 auto lookupValue = value;
526
527 // If this is a reference to the result of a multi-result instruction, print
528 // out the # identifier and make sure to map our lookup to the first result
529 // of the instruction.
530 if (auto *result = dyn_cast<InstResult>(value)) {
531 if (result->getOwner()->getNumResults() != 1) {
532 resultNo = result->getResultNumber();
533 lookupValue = result->getOwner()->getResult(0);
534 }
535 }
536
537 auto it = valueIDs.find(lookupValue);
538 if (it == valueIDs.end()) {
Chris Lattnerf8cce872018-07-20 09:28:54 -0700539 os << "<<INVALID SSA VALUE>>";
Chris Lattner6119d382018-07-20 18:41:34 -0700540 return;
541 }
542
543 os << '%' << it->getSecond();
544 if (resultNo != -1 && !dontPrintResultNo)
545 os << '#' << resultNo;
Chris Lattnerf8cce872018-07-20 09:28:54 -0700546 }
547
548private:
549 /// This is the value ID for each SSA value in the current function.
550 DenseMap<const SSAValue *, unsigned> valueIDs;
551 unsigned nextValueID = 0;
Tatiana Shpeismanfa412f72018-07-09 17:42:46 -0700552};
James Molloy87d81022018-07-23 11:44:40 -0700553} // end anonymous namespace
Tatiana Shpeismanfa412f72018-07-09 17:42:46 -0700554
Tatiana Shpeismanfa412f72018-07-09 17:42:46 -0700555void FunctionState::printOperation(const Operation *op) {
Chris Lattnerf8cce872018-07-20 09:28:54 -0700556 os << " ";
557
Chris Lattnerac591f12018-07-22 21:02:26 -0700558 if (op->getNumResults()) {
559 printValueID(op->getResult(0), /*dontPrintResultNo*/ true);
560 os << " = ";
Chris Lattnerf8cce872018-07-20 09:28:54 -0700561 }
562
Tatiana Shpeismanfa412f72018-07-09 17:42:46 -0700563 // Check to see if this is a known operation. If so, use the registered
564 // custom printer hook.
Chris Lattner4fd59b02018-07-20 09:35:47 -0700565 if (auto opInfo = state.operationSet->lookup(op->getName().str())) {
Tatiana Shpeismanfa412f72018-07-09 17:42:46 -0700566 opInfo->printAssembly(op, os);
567 return;
568 }
569
Chris Lattnerf8cce872018-07-20 09:28:54 -0700570 // Otherwise use the standard verbose printing approach.
571
Tatiana Shpeismanfa412f72018-07-09 17:42:46 -0700572 // TODO: escape name if necessary.
Chris Lattnerf8cce872018-07-20 09:28:54 -0700573 os << "\"" << op->getName().str() << "\"(";
Tatiana Shpeismanfa412f72018-07-09 17:42:46 -0700574
Chris Lattnerac591f12018-07-22 21:02:26 -0700575 interleaveComma(op->getOperands(),
576 [&](const SSAValue *value) { printValueID(value); });
Chris Lattner7f9cc272018-07-19 08:35:28 -0700577
Chris Lattnerf8cce872018-07-20 09:28:54 -0700578 os << ')';
Tatiana Shpeismanfa412f72018-07-09 17:42:46 -0700579 auto attrs = op->getAttrs();
580 if (!attrs.empty()) {
581 os << '{';
Chris Lattner4fd59b02018-07-20 09:35:47 -0700582 interleaveComma(attrs, [&](NamedAttribute attr) {
Chris Lattnerf8cce872018-07-20 09:28:54 -0700583 os << attr.first << ": ";
Chris Lattner4fd59b02018-07-20 09:35:47 -0700584 print(attr.second);
Chris Lattnerf8cce872018-07-20 09:28:54 -0700585 });
Tatiana Shpeismanfa412f72018-07-09 17:42:46 -0700586 os << '}';
587 }
Chris Lattner3b2ef762018-07-18 15:31:25 -0700588
Chris Lattnerac591f12018-07-22 21:02:26 -0700589 // Print the type signature of the operation.
590 os << " : (";
591 interleaveComma(op->getOperands(),
592 [&](const SSAValue *value) { print(value->getType()); });
593 os << ") -> ";
Chris Lattnerf8cce872018-07-20 09:28:54 -0700594
Chris Lattnerac591f12018-07-22 21:02:26 -0700595 if (op->getNumResults() == 1) {
596 print(op->getResult(0)->getType());
597 } else {
598 os << '(';
599 interleaveComma(op->getResults(),
600 [&](const SSAValue *result) { print(result->getType()); });
601 os << ')';
Chris Lattnerf8cce872018-07-20 09:28:54 -0700602 }
Tatiana Shpeismanfa412f72018-07-09 17:42:46 -0700603}
604
Chris Lattner4c95a502018-06-23 16:03:42 -0700605//===----------------------------------------------------------------------===//
606// CFG Function printing
607//===----------------------------------------------------------------------===//
608
609namespace {
Chris Lattner4fd59b02018-07-20 09:35:47 -0700610class CFGFunctionPrinter : public FunctionState {
Chris Lattner4c95a502018-06-23 16:03:42 -0700611public:
Chris Lattner4fd59b02018-07-20 09:35:47 -0700612 CFGFunctionPrinter(const CFGFunction *function, const ModulePrinter &other);
Chris Lattner4c95a502018-06-23 16:03:42 -0700613
614 const CFGFunction *getFunction() const { return function; }
615
616 void print();
617 void print(const BasicBlock *block);
Chris Lattnered65a732018-06-28 20:45:33 -0700618
619 void print(const Instruction *inst);
620 void print(const OperationInst *inst);
621 void print(const ReturnInst *inst);
622 void print(const BranchInst *inst);
Chris Lattner4c95a502018-06-23 16:03:42 -0700623
624 unsigned getBBID(const BasicBlock *block) {
625 auto it = basicBlockIDs.find(block);
626 assert(it != basicBlockIDs.end() && "Block not in this function?");
627 return it->second;
628 }
629
630private:
631 const CFGFunction *function;
MLIR Team54b55a22018-07-18 10:16:05 -0700632 DenseMap<const BasicBlock *, unsigned> basicBlockIDs;
Chris Lattnerf8cce872018-07-20 09:28:54 -0700633
Chris Lattner4fd59b02018-07-20 09:35:47 -0700634 void numberValuesInBlock(const BasicBlock *block);
Chris Lattner4c95a502018-06-23 16:03:42 -0700635};
James Molloy87d81022018-07-23 11:44:40 -0700636} // end anonymous namespace
Chris Lattner4c95a502018-06-23 16:03:42 -0700637
Chris Lattner4fd59b02018-07-20 09:35:47 -0700638CFGFunctionPrinter::CFGFunctionPrinter(const CFGFunction *function,
639 const ModulePrinter &other)
640 : FunctionState(other), function(function) {
Chris Lattner4c95a502018-06-23 16:03:42 -0700641 // Each basic block gets a unique ID per function.
642 unsigned blockID = 0;
Chris Lattnerf8cce872018-07-20 09:28:54 -0700643 for (auto &block : *function) {
644 basicBlockIDs[&block] = blockID++;
Chris Lattner4fd59b02018-07-20 09:35:47 -0700645 numberValuesInBlock(&block);
Chris Lattnerf8cce872018-07-20 09:28:54 -0700646 }
647}
648
649/// Number all of the SSA values in the specified basic block.
Chris Lattner4fd59b02018-07-20 09:35:47 -0700650void CFGFunctionPrinter::numberValuesInBlock(const BasicBlock *block) {
James Molloy61a656c2018-07-22 15:45:24 -0700651 for (auto *arg : block->getArguments()) {
652 numberValueID(arg);
653 }
Chris Lattnerf8cce872018-07-20 09:28:54 -0700654 for (auto &op : *block) {
655 // We number instruction that have results, and we only number the first
656 // result.
657 if (op.getNumResults() != 0)
658 numberValueID(op.getResult(0));
659 }
660
661 // Terminators do not define values.
Chris Lattner4c95a502018-06-23 16:03:42 -0700662}
663
Chris Lattner4fd59b02018-07-20 09:35:47 -0700664void CFGFunctionPrinter::print() {
Chris Lattner4c95a502018-06-23 16:03:42 -0700665 os << "cfgfunc ";
Chris Lattner4fd59b02018-07-20 09:35:47 -0700666 printFunctionSignature(getFunction());
Chris Lattner4c95a502018-06-23 16:03:42 -0700667 os << " {\n";
668
James Molloy87d81022018-07-23 11:44:40 -0700669 for (auto &block : *function)
670 print(&block);
Chris Lattner4c95a502018-06-23 16:03:42 -0700671 os << "}\n\n";
672}
673
Chris Lattner4fd59b02018-07-20 09:35:47 -0700674void CFGFunctionPrinter::print(const BasicBlock *block) {
James Molloy61a656c2018-07-22 15:45:24 -0700675 os << "bb" << getBBID(block);
Chris Lattner4c95a502018-06-23 16:03:42 -0700676
James Molloy61a656c2018-07-22 15:45:24 -0700677 if (!block->args_empty()) {
678 os << '(';
679 interleaveComma(block->getArguments(), [&](const BBArgument *arg) {
680 printValueID(arg);
681 os << ": ";
682 ModulePrinter::print(arg->getType());
683 });
684 os << ')';
685 }
686 os << ":\n";
687
Jacques Pienaarb020c542018-07-15 00:06:54 -0700688 for (auto &inst : block->getOperations()) {
Chris Lattner3a467cc2018-07-01 20:28:00 -0700689 print(&inst);
James Molloy61a656c2018-07-22 15:45:24 -0700690 os << '\n';
Jacques Pienaarb020c542018-07-15 00:06:54 -0700691 }
Chris Lattner4c95a502018-06-23 16:03:42 -0700692
693 print(block->getTerminator());
James Molloy61a656c2018-07-22 15:45:24 -0700694 os << '\n';
Chris Lattner4c95a502018-06-23 16:03:42 -0700695}
696
Chris Lattner4fd59b02018-07-20 09:35:47 -0700697void CFGFunctionPrinter::print(const Instruction *inst) {
Chris Lattner4c95a502018-06-23 16:03:42 -0700698 switch (inst->getKind()) {
Chris Lattnered65a732018-06-28 20:45:33 -0700699 case Instruction::Kind::Operation:
700 return print(cast<OperationInst>(inst));
Chris Lattnerf6d80a02018-06-24 11:18:29 -0700701 case TerminatorInst::Kind::Branch:
Chris Lattnered65a732018-06-28 20:45:33 -0700702 return print(cast<BranchInst>(inst));
Chris Lattner4c95a502018-06-23 16:03:42 -0700703 case TerminatorInst::Kind::Return:
Chris Lattnered65a732018-06-28 20:45:33 -0700704 return print(cast<ReturnInst>(inst));
Chris Lattner4c95a502018-06-23 16:03:42 -0700705 }
706}
707
Chris Lattner4fd59b02018-07-20 09:35:47 -0700708void CFGFunctionPrinter::print(const OperationInst *inst) {
Tatiana Shpeismanfa412f72018-07-09 17:42:46 -0700709 printOperation(inst);
Chris Lattner3b2ef762018-07-18 15:31:25 -0700710}
Chris Lattner1604e472018-07-23 08:42:19 -0700711
Chris Lattner4fd59b02018-07-20 09:35:47 -0700712void CFGFunctionPrinter::print(const BranchInst *inst) {
Jacques Pienaarb020c542018-07-15 00:06:54 -0700713 os << " br bb" << getBBID(inst->getDest());
Chris Lattner1604e472018-07-23 08:42:19 -0700714
715 if (inst->getNumOperands() != 0) {
716 os << '(';
717 // TODO: Use getOperands() when we have it.
718 interleaveComma(inst->getInstOperands(), [&](const InstOperand &operand) {
719 printValueID(operand.get());
720 });
721 os << ") : ";
722 interleaveComma(inst->getInstOperands(), [&](const InstOperand &operand) {
723 ModulePrinter::print(operand.get()->getType());
724 });
725 }
Chris Lattnered65a732018-06-28 20:45:33 -0700726}
Chris Lattner1604e472018-07-23 08:42:19 -0700727
Chris Lattner40746442018-07-21 14:32:09 -0700728void CFGFunctionPrinter::print(const ReturnInst *inst) {
729 os << " return";
730
731 if (inst->getNumOperands() != 0)
732 os << ' ';
733
Chris Lattnerac591f12018-07-22 21:02:26 -0700734 interleaveComma(inst->getOperands(), [&](const CFGValue *operand) {
735 printValueID(operand);
Chris Lattner40746442018-07-21 14:32:09 -0700736 os << " : ";
Chris Lattnerac591f12018-07-22 21:02:26 -0700737 ModulePrinter::print(operand->getType());
Chris Lattner40746442018-07-21 14:32:09 -0700738 });
739}
MLIR Team54b55a22018-07-18 10:16:05 -0700740
Chris Lattner4fd59b02018-07-20 09:35:47 -0700741void ModulePrinter::print(const CFGFunction *fn) {
742 CFGFunctionPrinter(fn, *this).print();
Chris Lattnered65a732018-06-28 20:45:33 -0700743}
744
Chris Lattner4c95a502018-06-23 16:03:42 -0700745//===----------------------------------------------------------------------===//
Tatiana Shpeismanbf079c92018-07-03 17:51:28 -0700746// ML Function printing
Chris Lattner4c95a502018-06-23 16:03:42 -0700747//===----------------------------------------------------------------------===//
748
Tatiana Shpeismanbf079c92018-07-03 17:51:28 -0700749namespace {
Chris Lattner4fd59b02018-07-20 09:35:47 -0700750class MLFunctionPrinter : public FunctionState {
Tatiana Shpeismanbf079c92018-07-03 17:51:28 -0700751public:
Chris Lattner4fd59b02018-07-20 09:35:47 -0700752 MLFunctionPrinter(const MLFunction *function, const ModulePrinter &other);
Tatiana Shpeismanbf079c92018-07-03 17:51:28 -0700753
754 const MLFunction *getFunction() const { return function; }
755
Tatiana Shpeisman1bcfe982018-07-13 13:03:13 -0700756 // Prints ML function
Tatiana Shpeismanbf079c92018-07-03 17:51:28 -0700757 void print();
758
Tatiana Shpeisman1bcfe982018-07-13 13:03:13 -0700759 // Methods to print ML function statements
Tatiana Shpeismanbf079c92018-07-03 17:51:28 -0700760 void print(const Statement *stmt);
Tatiana Shpeismanfa412f72018-07-09 17:42:46 -0700761 void print(const OperationStmt *stmt);
Tatiana Shpeismanbf079c92018-07-03 17:51:28 -0700762 void print(const ForStmt *stmt);
763 void print(const IfStmt *stmt);
Tatiana Shpeisman1bcfe982018-07-13 13:03:13 -0700764 void print(const StmtBlock *block);
765
766 // Number of spaces used for indenting nested statements
767 const static unsigned indentWidth = 2;
Tatiana Shpeismanbf079c92018-07-03 17:51:28 -0700768
769private:
Tatiana Shpeismanbf079c92018-07-03 17:51:28 -0700770 const MLFunction *function;
Tatiana Shpeismanbf079c92018-07-03 17:51:28 -0700771 int numSpaces;
772};
James Molloy87d81022018-07-23 11:44:40 -0700773} // end anonymous namespace
Tatiana Shpeismanbf079c92018-07-03 17:51:28 -0700774
Chris Lattner4fd59b02018-07-20 09:35:47 -0700775MLFunctionPrinter::MLFunctionPrinter(const MLFunction *function,
776 const ModulePrinter &other)
777 : FunctionState(other), function(function), numSpaces(0) {}
Tatiana Shpeismanbf079c92018-07-03 17:51:28 -0700778
Chris Lattner4fd59b02018-07-20 09:35:47 -0700779void MLFunctionPrinter::print() {
Tatiana Shpeismanbf079c92018-07-03 17:51:28 -0700780 os << "mlfunc ";
781 // FIXME: should print argument names rather than just signature
Chris Lattner4fd59b02018-07-20 09:35:47 -0700782 printFunctionSignature(function);
Tatiana Shpeismanbf079c92018-07-03 17:51:28 -0700783 os << " {\n";
Tatiana Shpeisman1bcfe982018-07-13 13:03:13 -0700784 print(function);
Tatiana Shpeismanbf079c92018-07-03 17:51:28 -0700785 os << " return\n";
786 os << "}\n\n";
787}
788
Chris Lattner4fd59b02018-07-20 09:35:47 -0700789void MLFunctionPrinter::print(const StmtBlock *block) {
Tatiana Shpeisman1bcfe982018-07-13 13:03:13 -0700790 numSpaces += indentWidth;
Jacques Pienaarb020c542018-07-15 00:06:54 -0700791 for (auto &stmt : block->getStatements()) {
Tatiana Shpeisman1bcfe982018-07-13 13:03:13 -0700792 print(&stmt);
Jacques Pienaarb020c542018-07-15 00:06:54 -0700793 os << "\n";
794 }
Tatiana Shpeisman1bcfe982018-07-13 13:03:13 -0700795 numSpaces -= indentWidth;
796}
797
Chris Lattner4fd59b02018-07-20 09:35:47 -0700798void MLFunctionPrinter::print(const Statement *stmt) {
Tatiana Shpeismanbf079c92018-07-03 17:51:28 -0700799 switch (stmt->getKind()) {
Tatiana Shpeisman565b9642018-07-16 11:47:09 -0700800 case Statement::Kind::Operation:
801 return print(cast<OperationStmt>(stmt));
Tatiana Shpeismanbf079c92018-07-03 17:51:28 -0700802 case Statement::Kind::For:
803 return print(cast<ForStmt>(stmt));
804 case Statement::Kind::If:
805 return print(cast<IfStmt>(stmt));
Tatiana Shpeismanbf079c92018-07-03 17:51:28 -0700806 }
807}
808
Chris Lattner4fd59b02018-07-20 09:35:47 -0700809void MLFunctionPrinter::print(const OperationStmt *stmt) {
810 printOperation(stmt);
811}
Tatiana Shpeismanbf079c92018-07-03 17:51:28 -0700812
Chris Lattner4fd59b02018-07-20 09:35:47 -0700813void MLFunctionPrinter::print(const ForStmt *stmt) {
Tatiana Shpeisman1da50c42018-07-19 09:52:39 -0700814 os.indent(numSpaces) << "for x = " << *stmt->getLowerBound();
815 os << " to " << *stmt->getUpperBound();
816 if (stmt->getStep()->getValue() != 1)
817 os << " step " << *stmt->getStep();
818
819 os << " {\n";
Tatiana Shpeisman1bcfe982018-07-13 13:03:13 -0700820 print(static_cast<const StmtBlock *>(stmt));
Tatiana Shpeisman565b9642018-07-16 11:47:09 -0700821 os.indent(numSpaces) << "}";
Tatiana Shpeismanbf079c92018-07-03 17:51:28 -0700822}
823
Chris Lattner4fd59b02018-07-20 09:35:47 -0700824void MLFunctionPrinter::print(const IfStmt *stmt) {
Tatiana Shpeisman1bcfe982018-07-13 13:03:13 -0700825 os.indent(numSpaces) << "if () {\n";
826 print(stmt->getThenClause());
827 os.indent(numSpaces) << "}";
828 if (stmt->hasElseClause()) {
829 os << " else {\n";
830 print(stmt->getElseClause());
831 os.indent(numSpaces) << "}";
832 }
Tatiana Shpeismanbf079c92018-07-03 17:51:28 -0700833}
834
Chris Lattner4fd59b02018-07-20 09:35:47 -0700835void ModulePrinter::print(const MLFunction *fn) {
836 MLFunctionPrinter(fn, *this).print();
MLIR Team4718bc92018-07-17 16:56:54 -0700837}
838
Tatiana Shpeismanbf079c92018-07-03 17:51:28 -0700839//===----------------------------------------------------------------------===//
840// print and dump methods
841//===----------------------------------------------------------------------===//
Chris Lattnered65a732018-06-28 20:45:33 -0700842
MLIR Teamb61885d2018-07-18 16:29:21 -0700843void Attribute::print(raw_ostream &os) const {
Chris Lattner4fd59b02018-07-20 09:35:47 -0700844 ModuleState state(/*no context is known*/ nullptr);
845 ModulePrinter(os, state).print(this);
MLIR Teamb61885d2018-07-18 16:29:21 -0700846}
847
James Molloy87d81022018-07-23 11:44:40 -0700848void Attribute::dump() const { print(llvm::errs()); }
MLIR Teamb61885d2018-07-18 16:29:21 -0700849
MLIR Team4718bc92018-07-17 16:56:54 -0700850void Type::print(raw_ostream &os) const {
Chris Lattner4fd59b02018-07-20 09:35:47 -0700851 ModuleState state(getContext());
852 ModulePrinter(os, state).print(this);
MLIR Team4718bc92018-07-17 16:56:54 -0700853}
854
MLIR Team54b55a22018-07-18 10:16:05 -0700855void Type::dump() const { print(llvm::errs()); }
MLIR Team4718bc92018-07-17 16:56:54 -0700856
MLIR Team718c82f2018-07-16 09:45:22 -0700857void AffineMap::dump() const {
858 print(llvm::errs());
859 llvm::errs() << "\n";
860}
Uday Bondhugula3934d4d2018-07-09 09:00:25 -0700861
Uday Bondhugula015cbb12018-07-03 20:16:08 -0700862void AffineExpr::dump() const {
863 print(llvm::errs());
864 llvm::errs() << "\n";
865}
866
Uday Bondhugulafaf37dd2018-06-29 18:09:29 -0700867void AffineExpr::print(raw_ostream &os) const {
Chris Lattner4fd59b02018-07-20 09:35:47 -0700868 ModuleState state(/*no context is known*/ nullptr);
869 ModulePrinter(os, state).print(this);
Uday Bondhugulafaf37dd2018-06-29 18:09:29 -0700870}
871
872void AffineMap::print(raw_ostream &os) const {
Chris Lattner4fd59b02018-07-20 09:35:47 -0700873 ModuleState state(/*no context is known*/ nullptr);
874 ModulePrinter(os, state).print(this);
875}
Uday Bondhugula015cbb12018-07-03 20:16:08 -0700876
Chris Lattner4fd59b02018-07-20 09:35:47 -0700877void Instruction::print(raw_ostream &os) const {
878 ModuleState state(getFunction()->getContext());
879 ModulePrinter modulePrinter(os, state);
880 CFGFunctionPrinter(getFunction(), modulePrinter).print(this);
881}
Uday Bondhugula015cbb12018-07-03 20:16:08 -0700882
Chris Lattner4fd59b02018-07-20 09:35:47 -0700883void Instruction::dump() const {
884 print(llvm::errs());
885 llvm::errs() << "\n";
Uday Bondhugulafaf37dd2018-06-29 18:09:29 -0700886}
887
Chris Lattner4c95a502018-06-23 16:03:42 -0700888void BasicBlock::print(raw_ostream &os) const {
Chris Lattner4fd59b02018-07-20 09:35:47 -0700889 ModuleState state(getFunction()->getContext());
890 ModulePrinter modulePrinter(os, state);
891 CFGFunctionPrinter(getFunction(), modulePrinter).print(this);
Chris Lattner4c95a502018-06-23 16:03:42 -0700892}
893
MLIR Team54b55a22018-07-18 10:16:05 -0700894void BasicBlock::dump() const { print(llvm::errs()); }
Chris Lattner4c95a502018-06-23 16:03:42 -0700895
Tatiana Shpeismanbf079c92018-07-03 17:51:28 -0700896void Statement::print(raw_ostream &os) const {
Chris Lattner4fd59b02018-07-20 09:35:47 -0700897 ModuleState state(getFunction()->getContext());
898 ModulePrinter modulePrinter(os, state);
899 MLFunctionPrinter(getFunction(), modulePrinter).print(this);
Tatiana Shpeismanc96b5872018-06-28 17:02:32 -0700900}
901
MLIR Team54b55a22018-07-18 10:16:05 -0700902void Statement::dump() const { print(llvm::errs()); }
Jacques Pienaarb020c542018-07-15 00:06:54 -0700903
Chris Lattner4c95a502018-06-23 16:03:42 -0700904void Function::print(raw_ostream &os) const {
Chris Lattner4fd59b02018-07-20 09:35:47 -0700905 ModuleState state(getContext());
906 ModulePrinter(os, state).print(this);
Chris Lattner4c95a502018-06-23 16:03:42 -0700907}
908
MLIR Team54b55a22018-07-18 10:16:05 -0700909void Function::dump() const { print(llvm::errs()); }
910
Chris Lattner4c95a502018-06-23 16:03:42 -0700911void Module::print(raw_ostream &os) const {
Chris Lattner4fd59b02018-07-20 09:35:47 -0700912 ModuleState state(getContext());
913 state.initialize(this);
914 ModulePrinter(os, state).print(this);
Chris Lattner4c95a502018-06-23 16:03:42 -0700915}
916
MLIR Team54b55a22018-07-18 10:16:05 -0700917void Module::dump() const { print(llvm::errs()); }