blob: b56c775e11ec8b528fe3bd54c79541886b49a209 [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
522 void printValueID(const SSAValue *value) const {
523 // TODO: If this is the result of an operation with multiple results, look
524 // up the first result, and print the #32 syntax.
525 auto it = valueIDs.find(value);
526 if (it != valueIDs.end())
527 os << '%' << it->getSecond();
528 else
529 os << "<<INVALID SSA VALUE>>";
530 }
531
532private:
533 /// This is the value ID for each SSA value in the current function.
534 DenseMap<const SSAValue *, unsigned> valueIDs;
535 unsigned nextValueID = 0;
Tatiana Shpeismanfa412f72018-07-09 17:42:46 -0700536};
MLIR Team54b55a22018-07-18 10:16:05 -0700537} // end anonymous namespace
Tatiana Shpeismanfa412f72018-07-09 17:42:46 -0700538
Tatiana Shpeismanfa412f72018-07-09 17:42:46 -0700539void FunctionState::printOperation(const Operation *op) {
Chris Lattnerf8cce872018-07-20 09:28:54 -0700540 os << " ";
541
542 // TODO: When we have SSAValue version of operands & results wired into
543 // Operation this check can go away.
544 if (auto *inst = dyn_cast<OperationInst>(op)) {
545 if (inst->getNumResults()) {
546 printValueID(inst->getResult(0));
547 os << " = ";
548 }
549 }
550
Tatiana Shpeismanfa412f72018-07-09 17:42:46 -0700551 // Check to see if this is a known operation. If so, use the registered
552 // custom printer hook.
Chris Lattner4fd59b02018-07-20 09:35:47 -0700553 if (auto opInfo = state.operationSet->lookup(op->getName().str())) {
Tatiana Shpeismanfa412f72018-07-09 17:42:46 -0700554 opInfo->printAssembly(op, os);
555 return;
556 }
557
Chris Lattnerf8cce872018-07-20 09:28:54 -0700558 // Otherwise use the standard verbose printing approach.
559
Tatiana Shpeismanfa412f72018-07-09 17:42:46 -0700560 // TODO: escape name if necessary.
Chris Lattnerf8cce872018-07-20 09:28:54 -0700561 os << "\"" << op->getName().str() << "\"(";
Tatiana Shpeismanfa412f72018-07-09 17:42:46 -0700562
Chris Lattnerf8cce872018-07-20 09:28:54 -0700563 // TODO: When we have SSAValue version of operands & results wired into
564 // Operation this check can go away.
565 if (auto *inst = dyn_cast<OperationInst>(op)) {
566 // TODO: Use getOperands() when we have it.
Chris Lattner4fd59b02018-07-20 09:35:47 -0700567 interleaveComma(inst->getInstOperands(), [&](const InstOperand &operand) {
568 printValueID(operand.get());
569 });
Chris Lattnerf8cce872018-07-20 09:28:54 -0700570 }
Chris Lattner7f9cc272018-07-19 08:35:28 -0700571
Chris Lattnerf8cce872018-07-20 09:28:54 -0700572 os << ')';
Tatiana Shpeismanfa412f72018-07-09 17:42:46 -0700573 auto attrs = op->getAttrs();
574 if (!attrs.empty()) {
575 os << '{';
Chris Lattner4fd59b02018-07-20 09:35:47 -0700576 interleaveComma(attrs, [&](NamedAttribute attr) {
Chris Lattnerf8cce872018-07-20 09:28:54 -0700577 os << attr.first << ": ";
Chris Lattner4fd59b02018-07-20 09:35:47 -0700578 print(attr.second);
Chris Lattnerf8cce872018-07-20 09:28:54 -0700579 });
Tatiana Shpeismanfa412f72018-07-09 17:42:46 -0700580 os << '}';
581 }
Chris Lattner3b2ef762018-07-18 15:31:25 -0700582
Chris Lattnerf8cce872018-07-20 09:28:54 -0700583 // TODO: When we have SSAValue version of operands & results wired into
584 // Operation this check can go away.
585 if (auto *inst = dyn_cast<OperationInst>(op)) {
586 // Print the type signature of the operation.
587 os << " : (";
588 // TODO: Switch to getOperands() when we have it.
Chris Lattner4fd59b02018-07-20 09:35:47 -0700589 interleaveComma(inst->getInstOperands(),
590 [&](const InstOperand &op) { print(op.get()->getType()); });
Chris Lattnerf8cce872018-07-20 09:28:54 -0700591 os << ") -> ";
592
593 // TODO: Switch to getResults() when we have it.
594 if (inst->getNumResults() == 1) {
Chris Lattner4fd59b02018-07-20 09:35:47 -0700595 print(inst->getInstResult(0).getType());
Chris Lattnerf8cce872018-07-20 09:28:54 -0700596 } else {
597 os << '(';
Chris Lattner4fd59b02018-07-20 09:35:47 -0700598 interleaveComma(inst->getInstResults(), [&](const InstResult &result) {
599 print(result.getType());
600 });
Chris Lattnerf8cce872018-07-20 09:28:54 -0700601 os << ')';
602 }
603 }
Tatiana Shpeismanfa412f72018-07-09 17:42:46 -0700604}
605
Chris Lattner4c95a502018-06-23 16:03:42 -0700606//===----------------------------------------------------------------------===//
607// CFG Function printing
608//===----------------------------------------------------------------------===//
609
610namespace {
Chris Lattner4fd59b02018-07-20 09:35:47 -0700611class CFGFunctionPrinter : public FunctionState {
Chris Lattner4c95a502018-06-23 16:03:42 -0700612public:
Chris Lattner4fd59b02018-07-20 09:35:47 -0700613 CFGFunctionPrinter(const CFGFunction *function, const ModulePrinter &other);
Chris Lattner4c95a502018-06-23 16:03:42 -0700614
615 const CFGFunction *getFunction() const { return function; }
616
617 void print();
618 void print(const BasicBlock *block);
Chris Lattnered65a732018-06-28 20:45:33 -0700619
620 void print(const Instruction *inst);
621 void print(const OperationInst *inst);
622 void print(const ReturnInst *inst);
623 void print(const BranchInst *inst);
Chris Lattner4c95a502018-06-23 16:03:42 -0700624
625 unsigned getBBID(const BasicBlock *block) {
626 auto it = basicBlockIDs.find(block);
627 assert(it != basicBlockIDs.end() && "Block not in this function?");
628 return it->second;
629 }
630
631private:
632 const CFGFunction *function;
MLIR Team54b55a22018-07-18 10:16:05 -0700633 DenseMap<const BasicBlock *, unsigned> basicBlockIDs;
Chris Lattnerf8cce872018-07-20 09:28:54 -0700634
Chris Lattner4fd59b02018-07-20 09:35:47 -0700635 void numberValuesInBlock(const BasicBlock *block);
Chris Lattner4c95a502018-06-23 16:03:42 -0700636};
MLIR Team54b55a22018-07-18 10:16:05 -0700637} // end anonymous namespace
Chris Lattner4c95a502018-06-23 16:03:42 -0700638
Chris Lattner4fd59b02018-07-20 09:35:47 -0700639CFGFunctionPrinter::CFGFunctionPrinter(const CFGFunction *function,
640 const ModulePrinter &other)
641 : FunctionState(other), function(function) {
Chris Lattner4c95a502018-06-23 16:03:42 -0700642 // Each basic block gets a unique ID per function.
643 unsigned blockID = 0;
Chris Lattnerf8cce872018-07-20 09:28:54 -0700644 for (auto &block : *function) {
645 basicBlockIDs[&block] = blockID++;
Chris Lattner4fd59b02018-07-20 09:35:47 -0700646 numberValuesInBlock(&block);
Chris Lattnerf8cce872018-07-20 09:28:54 -0700647 }
648}
649
650/// Number all of the SSA values in the specified basic block.
Chris Lattner4fd59b02018-07-20 09:35:47 -0700651void CFGFunctionPrinter::numberValuesInBlock(const BasicBlock *block) {
Chris Lattnerf8cce872018-07-20 09:28:54 -0700652 // TODO: basic block arguments.
653 for (auto &op : *block) {
654 // We number instruction that have results, and we only number the first
655 // result.
656 if (op.getNumResults() != 0)
657 numberValueID(op.getResult(0));
658 }
659
660 // Terminators do not define values.
Chris Lattner4c95a502018-06-23 16:03:42 -0700661}
662
Chris Lattner4fd59b02018-07-20 09:35:47 -0700663void CFGFunctionPrinter::print() {
Chris Lattner4c95a502018-06-23 16:03:42 -0700664 os << "cfgfunc ";
Chris Lattner4fd59b02018-07-20 09:35:47 -0700665 printFunctionSignature(getFunction());
Chris Lattner4c95a502018-06-23 16:03:42 -0700666 os << " {\n";
667
MLIR Team54b55a22018-07-18 10:16:05 -0700668 for (auto &block : *function) print(&block);
Chris Lattner4c95a502018-06-23 16:03:42 -0700669 os << "}\n\n";
670}
671
Chris Lattner4fd59b02018-07-20 09:35:47 -0700672void CFGFunctionPrinter::print(const BasicBlock *block) {
Chris Lattner4c95a502018-06-23 16:03:42 -0700673 os << "bb" << getBBID(block) << ":\n";
674
Chris Lattnered65a732018-06-28 20:45:33 -0700675 // TODO Print arguments.
Jacques Pienaarb020c542018-07-15 00:06:54 -0700676 for (auto &inst : block->getOperations()) {
Chris Lattner3a467cc2018-07-01 20:28:00 -0700677 print(&inst);
Jacques Pienaarb020c542018-07-15 00:06:54 -0700678 os << "\n";
679 }
Chris Lattner4c95a502018-06-23 16:03:42 -0700680
681 print(block->getTerminator());
Jacques Pienaarb020c542018-07-15 00:06:54 -0700682 os << "\n";
Chris Lattner4c95a502018-06-23 16:03:42 -0700683}
684
Chris Lattner4fd59b02018-07-20 09:35:47 -0700685void CFGFunctionPrinter::print(const Instruction *inst) {
Chris Lattner4c95a502018-06-23 16:03:42 -0700686 switch (inst->getKind()) {
Chris Lattnered65a732018-06-28 20:45:33 -0700687 case Instruction::Kind::Operation:
688 return print(cast<OperationInst>(inst));
Chris Lattnerf6d80a02018-06-24 11:18:29 -0700689 case TerminatorInst::Kind::Branch:
Chris Lattnered65a732018-06-28 20:45:33 -0700690 return print(cast<BranchInst>(inst));
Chris Lattner4c95a502018-06-23 16:03:42 -0700691 case TerminatorInst::Kind::Return:
Chris Lattnered65a732018-06-28 20:45:33 -0700692 return print(cast<ReturnInst>(inst));
Chris Lattner4c95a502018-06-23 16:03:42 -0700693 }
694}
695
Chris Lattner4fd59b02018-07-20 09:35:47 -0700696void CFGFunctionPrinter::print(const OperationInst *inst) {
Tatiana Shpeismanfa412f72018-07-09 17:42:46 -0700697 printOperation(inst);
Chris Lattner3b2ef762018-07-18 15:31:25 -0700698}
Chris Lattner4fd59b02018-07-20 09:35:47 -0700699void CFGFunctionPrinter::print(const BranchInst *inst) {
Jacques Pienaarb020c542018-07-15 00:06:54 -0700700 os << " br bb" << getBBID(inst->getDest());
Chris Lattnered65a732018-06-28 20:45:33 -0700701}
Chris Lattner4fd59b02018-07-20 09:35:47 -0700702void CFGFunctionPrinter::print(const ReturnInst *inst) { os << " return"; }
MLIR Team54b55a22018-07-18 10:16:05 -0700703
Chris Lattner4fd59b02018-07-20 09:35:47 -0700704void ModulePrinter::print(const CFGFunction *fn) {
705 CFGFunctionPrinter(fn, *this).print();
Chris Lattnered65a732018-06-28 20:45:33 -0700706}
707
Chris Lattner4c95a502018-06-23 16:03:42 -0700708//===----------------------------------------------------------------------===//
Tatiana Shpeismanbf079c92018-07-03 17:51:28 -0700709// ML Function printing
Chris Lattner4c95a502018-06-23 16:03:42 -0700710//===----------------------------------------------------------------------===//
711
Tatiana Shpeismanbf079c92018-07-03 17:51:28 -0700712namespace {
Chris Lattner4fd59b02018-07-20 09:35:47 -0700713class MLFunctionPrinter : public FunctionState {
Tatiana Shpeismanbf079c92018-07-03 17:51:28 -0700714public:
Chris Lattner4fd59b02018-07-20 09:35:47 -0700715 MLFunctionPrinter(const MLFunction *function, const ModulePrinter &other);
Tatiana Shpeismanbf079c92018-07-03 17:51:28 -0700716
717 const MLFunction *getFunction() const { return function; }
718
Tatiana Shpeisman1bcfe982018-07-13 13:03:13 -0700719 // Prints ML function
Tatiana Shpeismanbf079c92018-07-03 17:51:28 -0700720 void print();
721
Tatiana Shpeisman1bcfe982018-07-13 13:03:13 -0700722 // Methods to print ML function statements
Tatiana Shpeismanbf079c92018-07-03 17:51:28 -0700723 void print(const Statement *stmt);
Tatiana Shpeismanfa412f72018-07-09 17:42:46 -0700724 void print(const OperationStmt *stmt);
Tatiana Shpeismanbf079c92018-07-03 17:51:28 -0700725 void print(const ForStmt *stmt);
726 void print(const IfStmt *stmt);
Tatiana Shpeisman1bcfe982018-07-13 13:03:13 -0700727 void print(const StmtBlock *block);
728
729 // Number of spaces used for indenting nested statements
730 const static unsigned indentWidth = 2;
Tatiana Shpeismanbf079c92018-07-03 17:51:28 -0700731
732private:
Tatiana Shpeismanbf079c92018-07-03 17:51:28 -0700733 const MLFunction *function;
Tatiana Shpeismanbf079c92018-07-03 17:51:28 -0700734 int numSpaces;
735};
MLIR Team54b55a22018-07-18 10:16:05 -0700736} // end anonymous namespace
Tatiana Shpeismanbf079c92018-07-03 17:51:28 -0700737
Chris Lattner4fd59b02018-07-20 09:35:47 -0700738MLFunctionPrinter::MLFunctionPrinter(const MLFunction *function,
739 const ModulePrinter &other)
740 : FunctionState(other), function(function), numSpaces(0) {}
Tatiana Shpeismanbf079c92018-07-03 17:51:28 -0700741
Chris Lattner4fd59b02018-07-20 09:35:47 -0700742void MLFunctionPrinter::print() {
Tatiana Shpeismanbf079c92018-07-03 17:51:28 -0700743 os << "mlfunc ";
744 // FIXME: should print argument names rather than just signature
Chris Lattner4fd59b02018-07-20 09:35:47 -0700745 printFunctionSignature(function);
Tatiana Shpeismanbf079c92018-07-03 17:51:28 -0700746 os << " {\n";
Tatiana Shpeisman1bcfe982018-07-13 13:03:13 -0700747 print(function);
Tatiana Shpeismanbf079c92018-07-03 17:51:28 -0700748 os << " return\n";
749 os << "}\n\n";
750}
751
Chris Lattner4fd59b02018-07-20 09:35:47 -0700752void MLFunctionPrinter::print(const StmtBlock *block) {
Tatiana Shpeisman1bcfe982018-07-13 13:03:13 -0700753 numSpaces += indentWidth;
Jacques Pienaarb020c542018-07-15 00:06:54 -0700754 for (auto &stmt : block->getStatements()) {
Tatiana Shpeisman1bcfe982018-07-13 13:03:13 -0700755 print(&stmt);
Jacques Pienaarb020c542018-07-15 00:06:54 -0700756 os << "\n";
757 }
Tatiana Shpeisman1bcfe982018-07-13 13:03:13 -0700758 numSpaces -= indentWidth;
759}
760
Chris Lattner4fd59b02018-07-20 09:35:47 -0700761void MLFunctionPrinter::print(const Statement *stmt) {
Tatiana Shpeismanbf079c92018-07-03 17:51:28 -0700762 switch (stmt->getKind()) {
Tatiana Shpeisman565b9642018-07-16 11:47:09 -0700763 case Statement::Kind::Operation:
764 return print(cast<OperationStmt>(stmt));
Tatiana Shpeismanbf079c92018-07-03 17:51:28 -0700765 case Statement::Kind::For:
766 return print(cast<ForStmt>(stmt));
767 case Statement::Kind::If:
768 return print(cast<IfStmt>(stmt));
Tatiana Shpeismanbf079c92018-07-03 17:51:28 -0700769 }
770}
771
Chris Lattner4fd59b02018-07-20 09:35:47 -0700772void MLFunctionPrinter::print(const OperationStmt *stmt) {
773 printOperation(stmt);
774}
Tatiana Shpeismanbf079c92018-07-03 17:51:28 -0700775
Chris Lattner4fd59b02018-07-20 09:35:47 -0700776void MLFunctionPrinter::print(const ForStmt *stmt) {
Tatiana Shpeisman1da50c42018-07-19 09:52:39 -0700777 os.indent(numSpaces) << "for x = " << *stmt->getLowerBound();
778 os << " to " << *stmt->getUpperBound();
779 if (stmt->getStep()->getValue() != 1)
780 os << " step " << *stmt->getStep();
781
782 os << " {\n";
Tatiana Shpeisman1bcfe982018-07-13 13:03:13 -0700783 print(static_cast<const StmtBlock *>(stmt));
Tatiana Shpeisman565b9642018-07-16 11:47:09 -0700784 os.indent(numSpaces) << "}";
Tatiana Shpeismanbf079c92018-07-03 17:51:28 -0700785}
786
Chris Lattner4fd59b02018-07-20 09:35:47 -0700787void MLFunctionPrinter::print(const IfStmt *stmt) {
Tatiana Shpeisman1bcfe982018-07-13 13:03:13 -0700788 os.indent(numSpaces) << "if () {\n";
789 print(stmt->getThenClause());
790 os.indent(numSpaces) << "}";
791 if (stmt->hasElseClause()) {
792 os << " else {\n";
793 print(stmt->getElseClause());
794 os.indent(numSpaces) << "}";
795 }
Tatiana Shpeismanbf079c92018-07-03 17:51:28 -0700796}
797
Chris Lattner4fd59b02018-07-20 09:35:47 -0700798void ModulePrinter::print(const MLFunction *fn) {
799 MLFunctionPrinter(fn, *this).print();
MLIR Team4718bc92018-07-17 16:56:54 -0700800}
801
Tatiana Shpeismanbf079c92018-07-03 17:51:28 -0700802//===----------------------------------------------------------------------===//
803// print and dump methods
804//===----------------------------------------------------------------------===//
Chris Lattnered65a732018-06-28 20:45:33 -0700805
MLIR Teamb61885d2018-07-18 16:29:21 -0700806void Attribute::print(raw_ostream &os) const {
Chris Lattner4fd59b02018-07-20 09:35:47 -0700807 ModuleState state(/*no context is known*/ nullptr);
808 ModulePrinter(os, state).print(this);
MLIR Teamb61885d2018-07-18 16:29:21 -0700809}
810
811void Attribute::dump() const {
812 print(llvm::errs());
813}
814
MLIR Team4718bc92018-07-17 16:56:54 -0700815void Type::print(raw_ostream &os) const {
Chris Lattner4fd59b02018-07-20 09:35:47 -0700816 ModuleState state(getContext());
817 ModulePrinter(os, state).print(this);
MLIR Team4718bc92018-07-17 16:56:54 -0700818}
819
MLIR Team54b55a22018-07-18 10:16:05 -0700820void Type::dump() const { print(llvm::errs()); }
MLIR Team4718bc92018-07-17 16:56:54 -0700821
MLIR Team718c82f2018-07-16 09:45:22 -0700822void AffineMap::dump() const {
823 print(llvm::errs());
824 llvm::errs() << "\n";
825}
Uday Bondhugula3934d4d2018-07-09 09:00:25 -0700826
Uday Bondhugula015cbb12018-07-03 20:16:08 -0700827void AffineExpr::dump() const {
828 print(llvm::errs());
829 llvm::errs() << "\n";
830}
831
Uday Bondhugulafaf37dd2018-06-29 18:09:29 -0700832void AffineExpr::print(raw_ostream &os) const {
Chris Lattner4fd59b02018-07-20 09:35:47 -0700833 ModuleState state(/*no context is known*/ nullptr);
834 ModulePrinter(os, state).print(this);
Uday Bondhugulafaf37dd2018-06-29 18:09:29 -0700835}
836
837void AffineMap::print(raw_ostream &os) const {
Chris Lattner4fd59b02018-07-20 09:35:47 -0700838 ModuleState state(/*no context is known*/ nullptr);
839 ModulePrinter(os, state).print(this);
840}
Uday Bondhugula015cbb12018-07-03 20:16:08 -0700841
Chris Lattner4fd59b02018-07-20 09:35:47 -0700842void Instruction::print(raw_ostream &os) const {
843 ModuleState state(getFunction()->getContext());
844 ModulePrinter modulePrinter(os, state);
845 CFGFunctionPrinter(getFunction(), modulePrinter).print(this);
846}
Uday Bondhugula015cbb12018-07-03 20:16:08 -0700847
Chris Lattner4fd59b02018-07-20 09:35:47 -0700848void Instruction::dump() const {
849 print(llvm::errs());
850 llvm::errs() << "\n";
Uday Bondhugulafaf37dd2018-06-29 18:09:29 -0700851}
852
Chris Lattner4c95a502018-06-23 16:03:42 -0700853void BasicBlock::print(raw_ostream &os) const {
Chris Lattner4fd59b02018-07-20 09:35:47 -0700854 ModuleState state(getFunction()->getContext());
855 ModulePrinter modulePrinter(os, state);
856 CFGFunctionPrinter(getFunction(), modulePrinter).print(this);
Chris Lattner4c95a502018-06-23 16:03:42 -0700857}
858
MLIR Team54b55a22018-07-18 10:16:05 -0700859void BasicBlock::dump() const { print(llvm::errs()); }
Chris Lattner4c95a502018-06-23 16:03:42 -0700860
Tatiana Shpeismanbf079c92018-07-03 17:51:28 -0700861void Statement::print(raw_ostream &os) const {
Chris Lattner4fd59b02018-07-20 09:35:47 -0700862 ModuleState state(getFunction()->getContext());
863 ModulePrinter modulePrinter(os, state);
864 MLFunctionPrinter(getFunction(), modulePrinter).print(this);
Tatiana Shpeismanc96b5872018-06-28 17:02:32 -0700865}
866
MLIR Team54b55a22018-07-18 10:16:05 -0700867void Statement::dump() const { print(llvm::errs()); }
Jacques Pienaarb020c542018-07-15 00:06:54 -0700868
Chris Lattner4c95a502018-06-23 16:03:42 -0700869void Function::print(raw_ostream &os) const {
Chris Lattner4fd59b02018-07-20 09:35:47 -0700870 ModuleState state(getContext());
871 ModulePrinter(os, state).print(this);
Chris Lattner4c95a502018-06-23 16:03:42 -0700872}
873
MLIR Team54b55a22018-07-18 10:16:05 -0700874void Function::dump() const { print(llvm::errs()); }
875
Chris Lattner4c95a502018-06-23 16:03:42 -0700876void Module::print(raw_ostream &os) const {
Chris Lattner4fd59b02018-07-20 09:35:47 -0700877 ModuleState state(getContext());
878 state.initialize(this);
879 ModulePrinter(os, state).print(this);
Chris Lattner4c95a502018-06-23 16:03:42 -0700880}
881
MLIR Team54b55a22018-07-18 10:16:05 -0700882void Module::dump() const { print(llvm::errs()); }