blob: ff6ce3e4cc00e7d1d7a874eca54c046fa7b01d6a [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};
MLIR Team54b55a22018-07-18 10:16:05 -070089} // end anonymous namespace
MLIR Team4718bc92018-07-17 16:56:54 -070090
MLIR Team4718bc92018-07-17 16:56:54 -070091
92// TODO Support visiting other types/instructions when implemented.
93void ModuleState::visitType(const Type *type) {
94 if (type->getKind() == Type::Kind::Function) {
95 // Visit input and result types for functions.
96 auto *funcType = cast<FunctionType>(type);
MLIR Team54b55a22018-07-18 10:16:05 -070097 for (auto *input : funcType->getInputs()) {
MLIR Team4718bc92018-07-17 16:56:54 -070098 visitType(input);
99 }
MLIR Team54b55a22018-07-18 10:16:05 -0700100 for (auto *result : funcType->getResults()) {
MLIR Team4718bc92018-07-17 16:56:54 -0700101 visitType(result);
102 }
103 } else if (type->getKind() == Type::Kind::MemRef) {
104 // Visit affine maps in memref type.
105 auto *memref = cast<MemRefType>(type);
MLIR Team54b55a22018-07-18 10:16:05 -0700106 for (AffineMap *map : memref->getAffineMaps()) {
MLIR Team4718bc92018-07-17 16:56:54 -0700107 recordAffineMapReference(map);
108 }
109 }
110}
111
MLIR Teamb61885d2018-07-18 16:29:21 -0700112void ModuleState::visitAttribute(const Attribute *attr) {
113 if (isa<AffineMapAttr>(attr)) {
114 recordAffineMapReference(cast<AffineMapAttr>(attr)->getValue());
115 } else if (isa<ArrayAttr>(attr)) {
116 for (auto elt : cast<ArrayAttr>(attr)->getValue()) {
117 visitAttribute(elt);
118 }
119 }
120}
121
122void ModuleState::visitOperation(const Operation *op) {
123 for (auto elt : op->getAttrs()) {
124 visitAttribute(elt.second);
125 }
126}
127
MLIR Team4718bc92018-07-17 16:56:54 -0700128void ModuleState::visitExtFunction(const ExtFunction *fn) {
129 visitType(fn->getType());
130}
131
132void ModuleState::visitCFGFunction(const CFGFunction *fn) {
133 visitType(fn->getType());
134 // TODO Visit function body instructions.
MLIR Teamb61885d2018-07-18 16:29:21 -0700135 for (auto &block : *fn) {
136 for (auto &op : block.getOperations()) {
137 visitOperation(&op);
138 }
139 }
MLIR Team4718bc92018-07-17 16:56:54 -0700140}
141
142void ModuleState::visitMLFunction(const MLFunction *fn) {
143 visitType(fn->getType());
MLIR Teamb61885d2018-07-18 16:29:21 -0700144 // TODO Visit function body statements (and attributes if required).
MLIR Team4718bc92018-07-17 16:56:54 -0700145}
146
147void ModuleState::visitFunction(const Function *fn) {
148 switch (fn->getKind()) {
MLIR Team54b55a22018-07-18 10:16:05 -0700149 case Function::Kind::ExtFunc:
150 return visitExtFunction(cast<ExtFunction>(fn));
151 case Function::Kind::CFGFunc:
152 return visitCFGFunction(cast<CFGFunction>(fn));
153 case Function::Kind::MLFunc:
154 return visitMLFunction(cast<MLFunction>(fn));
MLIR Team4718bc92018-07-17 16:56:54 -0700155 }
156}
157
Chris Lattner4fd59b02018-07-20 09:35:47 -0700158// Initializes module state, populating affine map state.
159void ModuleState::initialize(const Module *module) {
160 for (auto fn : module->functionList) {
161 visitFunction(fn);
162 }
163}
164
165//===----------------------------------------------------------------------===//
166// ModulePrinter
167//===----------------------------------------------------------------------===//
168
169namespace {
170class ModulePrinter {
171public:
172 ModulePrinter(raw_ostream &os, ModuleState &state) : os(os), state(state) {}
173 explicit ModulePrinter(const ModulePrinter &printer)
174 : os(printer.os), state(printer.state) {}
175
176 template <typename Container, typename UnaryFunctor>
177 inline void interleaveComma(const Container &c, UnaryFunctor each_fn) const {
178 interleave(c.begin(), c.end(), each_fn, [&]() { os << ", "; });
179 }
180
181 void print(const Module *module);
182 void print(const Attribute *attr) const;
183 void print(const Type *type) const;
184 void print(const Function *fn);
185 void print(const ExtFunction *fn);
186 void print(const CFGFunction *fn);
187 void print(const MLFunction *fn);
188
189 void print(const AffineMap *map);
190 void print(const AffineExpr *expr) const;
191
192protected:
193 raw_ostream &os;
194 ModuleState &state;
195
196 void printFunctionSignature(const Function *fn);
197 void printAffineMapId(int affineMapId) const;
198 void printAffineMapReference(const AffineMap *affineMap) const;
199
200 void print(const AffineBinaryOpExpr *expr) const;
201};
202} // end anonymous namespace
203
MLIR Team4718bc92018-07-17 16:56:54 -0700204// Prints function with initialized module state.
Chris Lattner4fd59b02018-07-20 09:35:47 -0700205void ModulePrinter::print(const Function *fn) {
MLIR Team4718bc92018-07-17 16:56:54 -0700206 switch (fn->getKind()) {
MLIR Team54b55a22018-07-18 10:16:05 -0700207 case Function::Kind::ExtFunc:
208 return print(cast<ExtFunction>(fn));
209 case Function::Kind::CFGFunc:
210 return print(cast<CFGFunction>(fn));
211 case Function::Kind::MLFunc:
212 return print(cast<MLFunction>(fn));
MLIR Team4718bc92018-07-17 16:56:54 -0700213 }
214}
215
216// Prints affine map identifier.
Chris Lattner4fd59b02018-07-20 09:35:47 -0700217void ModulePrinter::printAffineMapId(int affineMapId) const {
MLIR Team4718bc92018-07-17 16:56:54 -0700218 os << "#map" << affineMapId;
219}
220
Chris Lattner4fd59b02018-07-20 09:35:47 -0700221void ModulePrinter::printAffineMapReference(const AffineMap *affineMap) const {
222 int mapId = state.getAffineMapId(affineMap);
MLIR Teamb61885d2018-07-18 16:29:21 -0700223 if (mapId >= 0) {
224 // Map will be printed at top of module so print reference to its id.
225 printAffineMapId(mapId);
226 } else {
227 // Map not in module state so print inline.
228 affineMap->print(os);
229 }
230}
231
Chris Lattner4fd59b02018-07-20 09:35:47 -0700232void ModulePrinter::print(const Module *module) {
233 for (const auto &mapAndId : state.getAffineMapIds()) {
MLIR Teamb61885d2018-07-18 16:29:21 -0700234 printAffineMapId(mapAndId.second);
MLIR Team4718bc92018-07-17 16:56:54 -0700235 os << " = ";
236 mapAndId.first->print(os);
237 os << '\n';
238 }
MLIR Team54b55a22018-07-18 10:16:05 -0700239 for (auto *fn : module->functionList) print(fn);
MLIR Team4718bc92018-07-17 16:56:54 -0700240}
241
Chris Lattner4fd59b02018-07-20 09:35:47 -0700242void ModulePrinter::print(const Attribute *attr) const {
MLIR Teamb61885d2018-07-18 16:29:21 -0700243 switch (attr->getKind()) {
244 case Attribute::Kind::Bool:
245 os << (cast<BoolAttr>(attr)->getValue() ? "true" : "false");
246 break;
247 case Attribute::Kind::Integer:
248 os << cast<IntegerAttr>(attr)->getValue();
249 break;
250 case Attribute::Kind::Float:
251 // FIXME: this isn't precise, we should print with a hex format.
252 os << cast<FloatAttr>(attr)->getValue();
253 break;
254 case Attribute::Kind::String:
255 // FIXME: should escape the string.
256 os << '"' << cast<StringAttr>(attr)->getValue() << '"';
257 break;
258 case Attribute::Kind::Array: {
259 auto elts = cast<ArrayAttr>(attr)->getValue();
260 os << '[';
Chris Lattner4fd59b02018-07-20 09:35:47 -0700261 interleaveComma(elts, [&](Attribute *attr) { print(attr); });
MLIR Teamb61885d2018-07-18 16:29:21 -0700262 os << ']';
263 break;
264 }
265 case Attribute::Kind::AffineMap:
266 printAffineMapReference(cast<AffineMapAttr>(attr)->getValue());
267 break;
268 }
269}
270
Chris Lattner4fd59b02018-07-20 09:35:47 -0700271void ModulePrinter::print(const Type *type) const {
MLIR Team4718bc92018-07-17 16:56:54 -0700272 switch (type->getKind()) {
MLIR Team54b55a22018-07-18 10:16:05 -0700273 case Type::Kind::AffineInt:
274 os << "affineint";
275 return;
276 case Type::Kind::BF16:
277 os << "bf16";
278 return;
279 case Type::Kind::F16:
280 os << "f16";
281 return;
282 case Type::Kind::F32:
283 os << "f32";
284 return;
285 case Type::Kind::F64:
286 os << "f64";
287 return;
MLIR Team4718bc92018-07-17 16:56:54 -0700288
289 case Type::Kind::Integer: {
290 auto *integer = cast<IntegerType>(type);
291 os << 'i' << integer->getWidth();
292 return;
293 }
294 case Type::Kind::Function: {
295 auto *func = cast<FunctionType>(type);
296 os << '(';
Chris Lattner4fd59b02018-07-20 09:35:47 -0700297 interleaveComma(func->getInputs(), [&](Type *type) { os << *type; });
MLIR Team4718bc92018-07-17 16:56:54 -0700298 os << ") -> ";
299 auto results = func->getResults();
300 if (results.size() == 1)
301 os << *results[0];
302 else {
303 os << '(';
Chris Lattner4fd59b02018-07-20 09:35:47 -0700304 interleaveComma(results, [&](Type *type) { os << *type; });
MLIR Team4718bc92018-07-17 16:56:54 -0700305 os << ')';
306 }
307 return;
308 }
309 case Type::Kind::Vector: {
310 auto *v = cast<VectorType>(type);
311 os << "vector<";
MLIR Team54b55a22018-07-18 10:16:05 -0700312 for (auto dim : v->getShape()) 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};
MLIR Team54b55a22018-07-18 10:16:05 -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
558 // TODO: When we have SSAValue version of operands & results wired into
559 // Operation this check can go away.
560 if (auto *inst = dyn_cast<OperationInst>(op)) {
561 if (inst->getNumResults()) {
Chris Lattner6119d382018-07-20 18:41:34 -0700562 printValueID(inst->getResult(0), /*dontPrintResultNo*/ true);
Chris Lattnerf8cce872018-07-20 09:28:54 -0700563 os << " = ";
564 }
565 }
566
Tatiana Shpeismanfa412f72018-07-09 17:42:46 -0700567 // Check to see if this is a known operation. If so, use the registered
568 // custom printer hook.
Chris Lattner4fd59b02018-07-20 09:35:47 -0700569 if (auto opInfo = state.operationSet->lookup(op->getName().str())) {
Tatiana Shpeismanfa412f72018-07-09 17:42:46 -0700570 opInfo->printAssembly(op, os);
571 return;
572 }
573
Chris Lattnerf8cce872018-07-20 09:28:54 -0700574 // Otherwise use the standard verbose printing approach.
575
Tatiana Shpeismanfa412f72018-07-09 17:42:46 -0700576 // TODO: escape name if necessary.
Chris Lattnerf8cce872018-07-20 09:28:54 -0700577 os << "\"" << op->getName().str() << "\"(";
Tatiana Shpeismanfa412f72018-07-09 17:42:46 -0700578
Chris Lattnerf8cce872018-07-20 09:28:54 -0700579 // TODO: When we have SSAValue version of operands & results wired into
580 // Operation this check can go away.
581 if (auto *inst = dyn_cast<OperationInst>(op)) {
582 // TODO: Use getOperands() when we have it.
Chris Lattner4fd59b02018-07-20 09:35:47 -0700583 interleaveComma(inst->getInstOperands(), [&](const InstOperand &operand) {
584 printValueID(operand.get());
585 });
Chris Lattnerf8cce872018-07-20 09:28:54 -0700586 }
Chris Lattner7f9cc272018-07-19 08:35:28 -0700587
Chris Lattnerf8cce872018-07-20 09:28:54 -0700588 os << ')';
Tatiana Shpeismanfa412f72018-07-09 17:42:46 -0700589 auto attrs = op->getAttrs();
590 if (!attrs.empty()) {
591 os << '{';
Chris Lattner4fd59b02018-07-20 09:35:47 -0700592 interleaveComma(attrs, [&](NamedAttribute attr) {
Chris Lattnerf8cce872018-07-20 09:28:54 -0700593 os << attr.first << ": ";
Chris Lattner4fd59b02018-07-20 09:35:47 -0700594 print(attr.second);
Chris Lattnerf8cce872018-07-20 09:28:54 -0700595 });
Tatiana Shpeismanfa412f72018-07-09 17:42:46 -0700596 os << '}';
597 }
Chris Lattner3b2ef762018-07-18 15:31:25 -0700598
Chris Lattnerf8cce872018-07-20 09:28:54 -0700599 // TODO: When we have SSAValue version of operands & results wired into
600 // Operation this check can go away.
601 if (auto *inst = dyn_cast<OperationInst>(op)) {
602 // Print the type signature of the operation.
603 os << " : (";
604 // TODO: Switch to getOperands() when we have it.
Chris Lattner4fd59b02018-07-20 09:35:47 -0700605 interleaveComma(inst->getInstOperands(),
606 [&](const InstOperand &op) { print(op.get()->getType()); });
Chris Lattnerf8cce872018-07-20 09:28:54 -0700607 os << ") -> ";
608
609 // TODO: Switch to getResults() when we have it.
610 if (inst->getNumResults() == 1) {
Chris Lattner4fd59b02018-07-20 09:35:47 -0700611 print(inst->getInstResult(0).getType());
Chris Lattnerf8cce872018-07-20 09:28:54 -0700612 } else {
613 os << '(';
Chris Lattner4fd59b02018-07-20 09:35:47 -0700614 interleaveComma(inst->getInstResults(), [&](const InstResult &result) {
615 print(result.getType());
616 });
Chris Lattnerf8cce872018-07-20 09:28:54 -0700617 os << ')';
618 }
619 }
Tatiana Shpeismanfa412f72018-07-09 17:42:46 -0700620}
621
Chris Lattner4c95a502018-06-23 16:03:42 -0700622//===----------------------------------------------------------------------===//
623// CFG Function printing
624//===----------------------------------------------------------------------===//
625
626namespace {
Chris Lattner4fd59b02018-07-20 09:35:47 -0700627class CFGFunctionPrinter : public FunctionState {
Chris Lattner4c95a502018-06-23 16:03:42 -0700628public:
Chris Lattner4fd59b02018-07-20 09:35:47 -0700629 CFGFunctionPrinter(const CFGFunction *function, const ModulePrinter &other);
Chris Lattner4c95a502018-06-23 16:03:42 -0700630
631 const CFGFunction *getFunction() const { return function; }
632
633 void print();
634 void print(const BasicBlock *block);
Chris Lattnered65a732018-06-28 20:45:33 -0700635
636 void print(const Instruction *inst);
637 void print(const OperationInst *inst);
638 void print(const ReturnInst *inst);
639 void print(const BranchInst *inst);
Chris Lattner4c95a502018-06-23 16:03:42 -0700640
641 unsigned getBBID(const BasicBlock *block) {
642 auto it = basicBlockIDs.find(block);
643 assert(it != basicBlockIDs.end() && "Block not in this function?");
644 return it->second;
645 }
646
647private:
648 const CFGFunction *function;
MLIR Team54b55a22018-07-18 10:16:05 -0700649 DenseMap<const BasicBlock *, unsigned> basicBlockIDs;
Chris Lattnerf8cce872018-07-20 09:28:54 -0700650
Chris Lattner4fd59b02018-07-20 09:35:47 -0700651 void numberValuesInBlock(const BasicBlock *block);
Chris Lattner4c95a502018-06-23 16:03:42 -0700652};
MLIR Team54b55a22018-07-18 10:16:05 -0700653} // end anonymous namespace
Chris Lattner4c95a502018-06-23 16:03:42 -0700654
Chris Lattner4fd59b02018-07-20 09:35:47 -0700655CFGFunctionPrinter::CFGFunctionPrinter(const CFGFunction *function,
656 const ModulePrinter &other)
657 : FunctionState(other), function(function) {
Chris Lattner4c95a502018-06-23 16:03:42 -0700658 // Each basic block gets a unique ID per function.
659 unsigned blockID = 0;
Chris Lattnerf8cce872018-07-20 09:28:54 -0700660 for (auto &block : *function) {
661 basicBlockIDs[&block] = blockID++;
Chris Lattner4fd59b02018-07-20 09:35:47 -0700662 numberValuesInBlock(&block);
Chris Lattnerf8cce872018-07-20 09:28:54 -0700663 }
664}
665
666/// Number all of the SSA values in the specified basic block.
Chris Lattner4fd59b02018-07-20 09:35:47 -0700667void CFGFunctionPrinter::numberValuesInBlock(const BasicBlock *block) {
Chris Lattnerf8cce872018-07-20 09:28:54 -0700668 // TODO: basic block arguments.
669 for (auto &op : *block) {
670 // We number instruction that have results, and we only number the first
671 // result.
672 if (op.getNumResults() != 0)
673 numberValueID(op.getResult(0));
674 }
675
676 // Terminators do not define values.
Chris Lattner4c95a502018-06-23 16:03:42 -0700677}
678
Chris Lattner4fd59b02018-07-20 09:35:47 -0700679void CFGFunctionPrinter::print() {
Chris Lattner4c95a502018-06-23 16:03:42 -0700680 os << "cfgfunc ";
Chris Lattner4fd59b02018-07-20 09:35:47 -0700681 printFunctionSignature(getFunction());
Chris Lattner4c95a502018-06-23 16:03:42 -0700682 os << " {\n";
683
MLIR Team54b55a22018-07-18 10:16:05 -0700684 for (auto &block : *function) print(&block);
Chris Lattner4c95a502018-06-23 16:03:42 -0700685 os << "}\n\n";
686}
687
Chris Lattner4fd59b02018-07-20 09:35:47 -0700688void CFGFunctionPrinter::print(const BasicBlock *block) {
Chris Lattner4c95a502018-06-23 16:03:42 -0700689 os << "bb" << getBBID(block) << ":\n";
690
Chris Lattnered65a732018-06-28 20:45:33 -0700691 // TODO Print arguments.
Jacques Pienaarb020c542018-07-15 00:06:54 -0700692 for (auto &inst : block->getOperations()) {
Chris Lattner3a467cc2018-07-01 20:28:00 -0700693 print(&inst);
Jacques Pienaarb020c542018-07-15 00:06:54 -0700694 os << "\n";
695 }
Chris Lattner4c95a502018-06-23 16:03:42 -0700696
697 print(block->getTerminator());
Jacques Pienaarb020c542018-07-15 00:06:54 -0700698 os << "\n";
Chris Lattner4c95a502018-06-23 16:03:42 -0700699}
700
Chris Lattner4fd59b02018-07-20 09:35:47 -0700701void CFGFunctionPrinter::print(const Instruction *inst) {
Chris Lattner4c95a502018-06-23 16:03:42 -0700702 switch (inst->getKind()) {
Chris Lattnered65a732018-06-28 20:45:33 -0700703 case Instruction::Kind::Operation:
704 return print(cast<OperationInst>(inst));
Chris Lattnerf6d80a02018-06-24 11:18:29 -0700705 case TerminatorInst::Kind::Branch:
Chris Lattnered65a732018-06-28 20:45:33 -0700706 return print(cast<BranchInst>(inst));
Chris Lattner4c95a502018-06-23 16:03:42 -0700707 case TerminatorInst::Kind::Return:
Chris Lattnered65a732018-06-28 20:45:33 -0700708 return print(cast<ReturnInst>(inst));
Chris Lattner4c95a502018-06-23 16:03:42 -0700709 }
710}
711
Chris Lattner4fd59b02018-07-20 09:35:47 -0700712void CFGFunctionPrinter::print(const OperationInst *inst) {
Tatiana Shpeismanfa412f72018-07-09 17:42:46 -0700713 printOperation(inst);
Chris Lattner3b2ef762018-07-18 15:31:25 -0700714}
Chris Lattner4fd59b02018-07-20 09:35:47 -0700715void CFGFunctionPrinter::print(const BranchInst *inst) {
Jacques Pienaarb020c542018-07-15 00:06:54 -0700716 os << " br bb" << getBBID(inst->getDest());
Chris Lattnered65a732018-06-28 20:45:33 -0700717}
Chris Lattner4fd59b02018-07-20 09:35:47 -0700718void CFGFunctionPrinter::print(const ReturnInst *inst) { os << " return"; }
MLIR Team54b55a22018-07-18 10:16:05 -0700719
Chris Lattner4fd59b02018-07-20 09:35:47 -0700720void ModulePrinter::print(const CFGFunction *fn) {
721 CFGFunctionPrinter(fn, *this).print();
Chris Lattnered65a732018-06-28 20:45:33 -0700722}
723
Chris Lattner4c95a502018-06-23 16:03:42 -0700724//===----------------------------------------------------------------------===//
Tatiana Shpeismanbf079c92018-07-03 17:51:28 -0700725// ML Function printing
Chris Lattner4c95a502018-06-23 16:03:42 -0700726//===----------------------------------------------------------------------===//
727
Tatiana Shpeismanbf079c92018-07-03 17:51:28 -0700728namespace {
Chris Lattner4fd59b02018-07-20 09:35:47 -0700729class MLFunctionPrinter : public FunctionState {
Tatiana Shpeismanbf079c92018-07-03 17:51:28 -0700730public:
Chris Lattner4fd59b02018-07-20 09:35:47 -0700731 MLFunctionPrinter(const MLFunction *function, const ModulePrinter &other);
Tatiana Shpeismanbf079c92018-07-03 17:51:28 -0700732
733 const MLFunction *getFunction() const { return function; }
734
Tatiana Shpeisman1bcfe982018-07-13 13:03:13 -0700735 // Prints ML function
Tatiana Shpeismanbf079c92018-07-03 17:51:28 -0700736 void print();
737
Tatiana Shpeisman1bcfe982018-07-13 13:03:13 -0700738 // Methods to print ML function statements
Tatiana Shpeismanbf079c92018-07-03 17:51:28 -0700739 void print(const Statement *stmt);
Tatiana Shpeismanfa412f72018-07-09 17:42:46 -0700740 void print(const OperationStmt *stmt);
Tatiana Shpeismanbf079c92018-07-03 17:51:28 -0700741 void print(const ForStmt *stmt);
742 void print(const IfStmt *stmt);
Tatiana Shpeisman1bcfe982018-07-13 13:03:13 -0700743 void print(const StmtBlock *block);
744
745 // Number of spaces used for indenting nested statements
746 const static unsigned indentWidth = 2;
Tatiana Shpeismanbf079c92018-07-03 17:51:28 -0700747
748private:
Tatiana Shpeismanbf079c92018-07-03 17:51:28 -0700749 const MLFunction *function;
Tatiana Shpeismanbf079c92018-07-03 17:51:28 -0700750 int numSpaces;
751};
MLIR Team54b55a22018-07-18 10:16:05 -0700752} // end anonymous namespace
Tatiana Shpeismanbf079c92018-07-03 17:51:28 -0700753
Chris Lattner4fd59b02018-07-20 09:35:47 -0700754MLFunctionPrinter::MLFunctionPrinter(const MLFunction *function,
755 const ModulePrinter &other)
756 : FunctionState(other), function(function), numSpaces(0) {}
Tatiana Shpeismanbf079c92018-07-03 17:51:28 -0700757
Chris Lattner4fd59b02018-07-20 09:35:47 -0700758void MLFunctionPrinter::print() {
Tatiana Shpeismanbf079c92018-07-03 17:51:28 -0700759 os << "mlfunc ";
760 // FIXME: should print argument names rather than just signature
Chris Lattner4fd59b02018-07-20 09:35:47 -0700761 printFunctionSignature(function);
Tatiana Shpeismanbf079c92018-07-03 17:51:28 -0700762 os << " {\n";
Tatiana Shpeisman1bcfe982018-07-13 13:03:13 -0700763 print(function);
Tatiana Shpeismanbf079c92018-07-03 17:51:28 -0700764 os << " return\n";
765 os << "}\n\n";
766}
767
Chris Lattner4fd59b02018-07-20 09:35:47 -0700768void MLFunctionPrinter::print(const StmtBlock *block) {
Tatiana Shpeisman1bcfe982018-07-13 13:03:13 -0700769 numSpaces += indentWidth;
Jacques Pienaarb020c542018-07-15 00:06:54 -0700770 for (auto &stmt : block->getStatements()) {
Tatiana Shpeisman1bcfe982018-07-13 13:03:13 -0700771 print(&stmt);
Jacques Pienaarb020c542018-07-15 00:06:54 -0700772 os << "\n";
773 }
Tatiana Shpeisman1bcfe982018-07-13 13:03:13 -0700774 numSpaces -= indentWidth;
775}
776
Chris Lattner4fd59b02018-07-20 09:35:47 -0700777void MLFunctionPrinter::print(const Statement *stmt) {
Tatiana Shpeismanbf079c92018-07-03 17:51:28 -0700778 switch (stmt->getKind()) {
Tatiana Shpeisman565b9642018-07-16 11:47:09 -0700779 case Statement::Kind::Operation:
780 return print(cast<OperationStmt>(stmt));
Tatiana Shpeismanbf079c92018-07-03 17:51:28 -0700781 case Statement::Kind::For:
782 return print(cast<ForStmt>(stmt));
783 case Statement::Kind::If:
784 return print(cast<IfStmt>(stmt));
Tatiana Shpeismanbf079c92018-07-03 17:51:28 -0700785 }
786}
787
Chris Lattner4fd59b02018-07-20 09:35:47 -0700788void MLFunctionPrinter::print(const OperationStmt *stmt) {
789 printOperation(stmt);
790}
Tatiana Shpeismanbf079c92018-07-03 17:51:28 -0700791
Chris Lattner4fd59b02018-07-20 09:35:47 -0700792void MLFunctionPrinter::print(const ForStmt *stmt) {
Tatiana Shpeisman1da50c42018-07-19 09:52:39 -0700793 os.indent(numSpaces) << "for x = " << *stmt->getLowerBound();
794 os << " to " << *stmt->getUpperBound();
795 if (stmt->getStep()->getValue() != 1)
796 os << " step " << *stmt->getStep();
797
798 os << " {\n";
Tatiana Shpeisman1bcfe982018-07-13 13:03:13 -0700799 print(static_cast<const StmtBlock *>(stmt));
Tatiana Shpeisman565b9642018-07-16 11:47:09 -0700800 os.indent(numSpaces) << "}";
Tatiana Shpeismanbf079c92018-07-03 17:51:28 -0700801}
802
Chris Lattner4fd59b02018-07-20 09:35:47 -0700803void MLFunctionPrinter::print(const IfStmt *stmt) {
Tatiana Shpeisman1bcfe982018-07-13 13:03:13 -0700804 os.indent(numSpaces) << "if () {\n";
805 print(stmt->getThenClause());
806 os.indent(numSpaces) << "}";
807 if (stmt->hasElseClause()) {
808 os << " else {\n";
809 print(stmt->getElseClause());
810 os.indent(numSpaces) << "}";
811 }
Tatiana Shpeismanbf079c92018-07-03 17:51:28 -0700812}
813
Chris Lattner4fd59b02018-07-20 09:35:47 -0700814void ModulePrinter::print(const MLFunction *fn) {
815 MLFunctionPrinter(fn, *this).print();
MLIR Team4718bc92018-07-17 16:56:54 -0700816}
817
Tatiana Shpeismanbf079c92018-07-03 17:51:28 -0700818//===----------------------------------------------------------------------===//
819// print and dump methods
820//===----------------------------------------------------------------------===//
Chris Lattnered65a732018-06-28 20:45:33 -0700821
MLIR Teamb61885d2018-07-18 16:29:21 -0700822void Attribute::print(raw_ostream &os) const {
Chris Lattner4fd59b02018-07-20 09:35:47 -0700823 ModuleState state(/*no context is known*/ nullptr);
824 ModulePrinter(os, state).print(this);
MLIR Teamb61885d2018-07-18 16:29:21 -0700825}
826
827void Attribute::dump() const {
828 print(llvm::errs());
829}
830
MLIR Team4718bc92018-07-17 16:56:54 -0700831void Type::print(raw_ostream &os) const {
Chris Lattner4fd59b02018-07-20 09:35:47 -0700832 ModuleState state(getContext());
833 ModulePrinter(os, state).print(this);
MLIR Team4718bc92018-07-17 16:56:54 -0700834}
835
MLIR Team54b55a22018-07-18 10:16:05 -0700836void Type::dump() const { print(llvm::errs()); }
MLIR Team4718bc92018-07-17 16:56:54 -0700837
MLIR Team718c82f2018-07-16 09:45:22 -0700838void AffineMap::dump() const {
839 print(llvm::errs());
840 llvm::errs() << "\n";
841}
Uday Bondhugula3934d4d2018-07-09 09:00:25 -0700842
Uday Bondhugula015cbb12018-07-03 20:16:08 -0700843void AffineExpr::dump() const {
844 print(llvm::errs());
845 llvm::errs() << "\n";
846}
847
Uday Bondhugulafaf37dd2018-06-29 18:09:29 -0700848void AffineExpr::print(raw_ostream &os) const {
Chris Lattner4fd59b02018-07-20 09:35:47 -0700849 ModuleState state(/*no context is known*/ nullptr);
850 ModulePrinter(os, state).print(this);
Uday Bondhugulafaf37dd2018-06-29 18:09:29 -0700851}
852
853void AffineMap::print(raw_ostream &os) const {
Chris Lattner4fd59b02018-07-20 09:35:47 -0700854 ModuleState state(/*no context is known*/ nullptr);
855 ModulePrinter(os, state).print(this);
856}
Uday Bondhugula015cbb12018-07-03 20:16:08 -0700857
Chris Lattner4fd59b02018-07-20 09:35:47 -0700858void Instruction::print(raw_ostream &os) const {
859 ModuleState state(getFunction()->getContext());
860 ModulePrinter modulePrinter(os, state);
861 CFGFunctionPrinter(getFunction(), modulePrinter).print(this);
862}
Uday Bondhugula015cbb12018-07-03 20:16:08 -0700863
Chris Lattner4fd59b02018-07-20 09:35:47 -0700864void Instruction::dump() const {
865 print(llvm::errs());
866 llvm::errs() << "\n";
Uday Bondhugulafaf37dd2018-06-29 18:09:29 -0700867}
868
Chris Lattner4c95a502018-06-23 16:03:42 -0700869void BasicBlock::print(raw_ostream &os) const {
Chris Lattner4fd59b02018-07-20 09:35:47 -0700870 ModuleState state(getFunction()->getContext());
871 ModulePrinter modulePrinter(os, state);
872 CFGFunctionPrinter(getFunction(), modulePrinter).print(this);
Chris Lattner4c95a502018-06-23 16:03:42 -0700873}
874
MLIR Team54b55a22018-07-18 10:16:05 -0700875void BasicBlock::dump() const { print(llvm::errs()); }
Chris Lattner4c95a502018-06-23 16:03:42 -0700876
Tatiana Shpeismanbf079c92018-07-03 17:51:28 -0700877void Statement::print(raw_ostream &os) const {
Chris Lattner4fd59b02018-07-20 09:35:47 -0700878 ModuleState state(getFunction()->getContext());
879 ModulePrinter modulePrinter(os, state);
880 MLFunctionPrinter(getFunction(), modulePrinter).print(this);
Tatiana Shpeismanc96b5872018-06-28 17:02:32 -0700881}
882
MLIR Team54b55a22018-07-18 10:16:05 -0700883void Statement::dump() const { print(llvm::errs()); }
Jacques Pienaarb020c542018-07-15 00:06:54 -0700884
Chris Lattner4c95a502018-06-23 16:03:42 -0700885void Function::print(raw_ostream &os) const {
Chris Lattner4fd59b02018-07-20 09:35:47 -0700886 ModuleState state(getContext());
887 ModulePrinter(os, state).print(this);
Chris Lattner4c95a502018-06-23 16:03:42 -0700888}
889
MLIR Team54b55a22018-07-18 10:16:05 -0700890void Function::dump() const { print(llvm::errs()); }
891
Chris Lattner4c95a502018-06-23 16:03:42 -0700892void Module::print(raw_ostream &os) const {
Chris Lattner4fd59b02018-07-20 09:35:47 -0700893 ModuleState state(getContext());
894 state.initialize(this);
895 ModulePrinter(os, state).print(this);
Chris Lattner4c95a502018-06-23 16:03:42 -0700896}
897
MLIR Team54b55a22018-07-18 10:16:05 -0700898void Module::dump() const { print(llvm::errs()); }