blob: 64b5a2c5026d6e377bc5afd4371178ef41fe419a [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());
MLIR Teamb61885d2018-07-18 16:29:21 -0700134 for (auto &block : *fn) {
135 for (auto &op : block.getOperations()) {
136 visitOperation(&op);
137 }
138 }
MLIR Team4718bc92018-07-17 16:56:54 -0700139}
140
141void ModuleState::visitMLFunction(const MLFunction *fn) {
142 visitType(fn->getType());
MLIR Teamb61885d2018-07-18 16:29:21 -0700143 // TODO Visit function body statements (and attributes if required).
MLIR Team4718bc92018-07-17 16:56:54 -0700144}
145
146void ModuleState::visitFunction(const Function *fn) {
147 switch (fn->getKind()) {
MLIR Team54b55a22018-07-18 10:16:05 -0700148 case Function::Kind::ExtFunc:
149 return visitExtFunction(cast<ExtFunction>(fn));
150 case Function::Kind::CFGFunc:
151 return visitCFGFunction(cast<CFGFunction>(fn));
152 case Function::Kind::MLFunc:
153 return visitMLFunction(cast<MLFunction>(fn));
MLIR Team4718bc92018-07-17 16:56:54 -0700154 }
155}
156
Chris Lattner4fd59b02018-07-20 09:35:47 -0700157// Initializes module state, populating affine map state.
158void ModuleState::initialize(const Module *module) {
159 for (auto fn : module->functionList) {
160 visitFunction(fn);
161 }
162}
163
164//===----------------------------------------------------------------------===//
165// ModulePrinter
166//===----------------------------------------------------------------------===//
167
168namespace {
169class ModulePrinter {
170public:
171 ModulePrinter(raw_ostream &os, ModuleState &state) : os(os), state(state) {}
172 explicit ModulePrinter(const ModulePrinter &printer)
173 : os(printer.os), state(printer.state) {}
174
175 template <typename Container, typename UnaryFunctor>
176 inline void interleaveComma(const Container &c, UnaryFunctor each_fn) const {
177 interleave(c.begin(), c.end(), each_fn, [&]() { os << ", "; });
178 }
179
180 void print(const Module *module);
181 void print(const Attribute *attr) const;
182 void print(const Type *type) const;
183 void print(const Function *fn);
184 void print(const ExtFunction *fn);
185 void print(const CFGFunction *fn);
186 void print(const MLFunction *fn);
187
188 void print(const AffineMap *map);
189 void print(const AffineExpr *expr) const;
190
191protected:
192 raw_ostream &os;
193 ModuleState &state;
194
195 void printFunctionSignature(const Function *fn);
196 void printAffineMapId(int affineMapId) const;
197 void printAffineMapReference(const AffineMap *affineMap) const;
198
199 void print(const AffineBinaryOpExpr *expr) const;
200};
201} // end anonymous namespace
202
MLIR Team4718bc92018-07-17 16:56:54 -0700203// Prints function with initialized module state.
Chris Lattner4fd59b02018-07-20 09:35:47 -0700204void ModulePrinter::print(const Function *fn) {
MLIR Team4718bc92018-07-17 16:56:54 -0700205 switch (fn->getKind()) {
MLIR Team54b55a22018-07-18 10:16:05 -0700206 case Function::Kind::ExtFunc:
207 return print(cast<ExtFunction>(fn));
208 case Function::Kind::CFGFunc:
209 return print(cast<CFGFunction>(fn));
210 case Function::Kind::MLFunc:
211 return print(cast<MLFunction>(fn));
MLIR Team4718bc92018-07-17 16:56:54 -0700212 }
213}
214
215// Prints affine map identifier.
Chris Lattner4fd59b02018-07-20 09:35:47 -0700216void ModulePrinter::printAffineMapId(int affineMapId) const {
MLIR Team4718bc92018-07-17 16:56:54 -0700217 os << "#map" << affineMapId;
218}
219
Chris Lattner4fd59b02018-07-20 09:35:47 -0700220void ModulePrinter::printAffineMapReference(const AffineMap *affineMap) const {
221 int mapId = state.getAffineMapId(affineMap);
MLIR Teamb61885d2018-07-18 16:29:21 -0700222 if (mapId >= 0) {
223 // Map will be printed at top of module so print reference to its id.
224 printAffineMapId(mapId);
225 } else {
226 // Map not in module state so print inline.
227 affineMap->print(os);
228 }
229}
230
Chris Lattner4fd59b02018-07-20 09:35:47 -0700231void ModulePrinter::print(const Module *module) {
232 for (const auto &mapAndId : state.getAffineMapIds()) {
MLIR Teamb61885d2018-07-18 16:29:21 -0700233 printAffineMapId(mapAndId.second);
MLIR Team4718bc92018-07-17 16:56:54 -0700234 os << " = ";
235 mapAndId.first->print(os);
236 os << '\n';
237 }
MLIR Team54b55a22018-07-18 10:16:05 -0700238 for (auto *fn : module->functionList) 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<";
MLIR Team54b55a22018-07-18 10:16:05 -0700311 for (auto dim : v->getShape()) os << dim << 'x';
MLIR Team4718bc92018-07-17 16:56:54 -0700312 os << *v->getElementType() << '>';
313 return;
314 }
315 case Type::Kind::RankedTensor: {
316 auto *v = cast<RankedTensorType>(type);
317 os << "tensor<";
318 for (auto dim : v->getShape()) {
319 if (dim < 0)
320 os << '?';
321 else
322 os << dim;
323 os << 'x';
324 }
325 os << *v->getElementType() << '>';
326 return;
327 }
328 case Type::Kind::UnrankedTensor: {
329 auto *v = cast<UnrankedTensorType>(type);
330 os << "tensor<??" << *v->getElementType() << '>';
331 return;
332 }
333 case Type::Kind::MemRef: {
334 auto *v = cast<MemRefType>(type);
335 os << "memref<";
336 for (auto dim : v->getShape()) {
337 if (dim < 0)
338 os << '?';
339 else
340 os << dim;
341 os << 'x';
342 }
343 os << *v->getElementType();
344 for (auto map : v->getAffineMaps()) {
345 os << ", ";
MLIR Teamb61885d2018-07-18 16:29:21 -0700346 printAffineMapReference(map);
MLIR Team4718bc92018-07-17 16:56:54 -0700347 }
348 os << ", " << v->getMemorySpace();
349 os << '>';
350 return;
351 }
352 }
353}
354
355//===----------------------------------------------------------------------===//
Chris Lattner4fd59b02018-07-20 09:35:47 -0700356// Affine expressions and maps
357//===----------------------------------------------------------------------===//
358
359void ModulePrinter::print(const AffineExpr *expr) const {
360 switch (expr->getKind()) {
361 case AffineExpr::Kind::SymbolId:
362 os << 's' << cast<AffineSymbolExpr>(expr)->getPosition();
363 return;
364 case AffineExpr::Kind::DimId:
365 os << 'd' << cast<AffineDimExpr>(expr)->getPosition();
366 return;
367 case AffineExpr::Kind::Constant:
368 os << cast<AffineConstantExpr>(expr)->getValue();
369 return;
370 case AffineExpr::Kind::Add:
371 case AffineExpr::Kind::Mul:
372 case AffineExpr::Kind::FloorDiv:
373 case AffineExpr::Kind::CeilDiv:
374 case AffineExpr::Kind::Mod:
375 return print(cast<AffineBinaryOpExpr>(expr));
376 }
377}
378
379void ModulePrinter::print(const AffineBinaryOpExpr *expr) const {
380 if (expr->getKind() != AffineExpr::Kind::Add) {
381 os << '(';
382 print(expr->getLHS());
383 switch (expr->getKind()) {
384 case AffineExpr::Kind::Mul:
385 os << " * ";
386 break;
387 case AffineExpr::Kind::FloorDiv:
388 os << " floordiv ";
389 break;
390 case AffineExpr::Kind::CeilDiv:
391 os << " ceildiv ";
392 break;
393 case AffineExpr::Kind::Mod:
394 os << " mod ";
395 break;
396 default:
397 llvm_unreachable("unexpected affine binary op expression");
398 }
399
400 print(expr->getRHS());
401 os << ')';
402 return;
403 }
404
405 // Print out special "pretty" forms for add.
406 os << '(';
407 print(expr->getLHS());
408
409 // Pretty print addition to a product that has a negative operand as a
410 // subtraction.
411 if (auto *rhs = dyn_cast<AffineBinaryOpExpr>(expr->getRHS())) {
412 if (rhs->getKind() == AffineExpr::Kind::Mul) {
413 if (auto *rrhs = dyn_cast<AffineConstantExpr>(rhs->getRHS())) {
414 if (rrhs->getValue() < 0) {
415 os << " - (";
416 print(rhs->getLHS());
417 os << " * " << -rrhs->getValue() << "))";
418 return;
419 }
420 }
421 }
422 }
423
424 // Pretty print addition to a negative number as a subtraction.
425 if (auto *rhs = dyn_cast<AffineConstantExpr>(expr->getRHS())) {
426 if (rhs->getValue() < 0) {
427 os << " - " << -rhs->getValue() << ")";
428 return;
429 }
430 }
431
432 os << " + ";
433 print(expr->getRHS());
434 os << ')';
435}
436
437void ModulePrinter::print(const AffineMap *map) {
438 // Dimension identifiers.
439 os << '(';
440 for (int i = 0; i < (int)map->getNumDims() - 1; i++)
441 os << "d" << i << ", ";
442 if (map->getNumDims() >= 1)
443 os << "d" << map->getNumDims() - 1;
444 os << ")";
445
446 // Symbolic identifiers.
447 if (map->getNumSymbols() >= 1) {
448 os << " [";
449 for (int i = 0; i < (int)map->getNumSymbols() - 1; i++)
450 os << "s" << i << ", ";
451 if (map->getNumSymbols() >= 1)
452 os << "s" << map->getNumSymbols() - 1;
453 os << "]";
454 }
455
456 // AffineMap should have at least one result.
457 assert(!map->getResults().empty());
458 // Result affine expressions.
459 os << " -> (";
460 interleaveComma(map->getResults(), [&](AffineExpr *expr) { print(expr); });
461 os << ")";
462
463 if (!map->isBounded()) {
464 return;
465 }
466
467 // Print range sizes for bounded affine maps.
468 os << " size (";
469 interleaveComma(map->getRangeSizes(), [&](AffineExpr *expr) { print(expr); });
470 os << ")";
471}
472
473//===----------------------------------------------------------------------===//
Chris Lattner4c95a502018-06-23 16:03:42 -0700474// Function printing
475//===----------------------------------------------------------------------===//
476
Chris Lattner4fd59b02018-07-20 09:35:47 -0700477void ModulePrinter::printFunctionSignature(const Function *fn) {
Chris Lattner4c95a502018-06-23 16:03:42 -0700478 auto type = fn->getType();
479
480 os << "@" << fn->getName() << '(';
Chris Lattner4fd59b02018-07-20 09:35:47 -0700481 interleaveComma(type->getInputs(), [&](Type *eltType) { print(eltType); });
Chris Lattner4c95a502018-06-23 16:03:42 -0700482 os << ')';
483
484 switch (type->getResults().size()) {
MLIR Team54b55a22018-07-18 10:16:05 -0700485 case 0:
486 break;
Chris Lattner4c95a502018-06-23 16:03:42 -0700487 case 1:
MLIR Team4718bc92018-07-17 16:56:54 -0700488 os << " -> ";
Chris Lattner4fd59b02018-07-20 09:35:47 -0700489 print(type->getResults()[0]);
Chris Lattner4c95a502018-06-23 16:03:42 -0700490 break;
491 default:
492 os << " -> (";
Chris Lattner4fd59b02018-07-20 09:35:47 -0700493 interleaveComma(type->getResults(), [&](Type *eltType) { print(eltType); });
Chris Lattner4c95a502018-06-23 16:03:42 -0700494 os << ')';
495 break;
496 }
497}
498
Chris Lattner4fd59b02018-07-20 09:35:47 -0700499void ModulePrinter::print(const ExtFunction *fn) {
Chris Lattner4c95a502018-06-23 16:03:42 -0700500 os << "extfunc ";
Chris Lattner4fd59b02018-07-20 09:35:47 -0700501 printFunctionSignature(fn);
MLIR Team54b55a22018-07-18 10:16:05 -0700502 os << '\n';
Chris Lattner4c95a502018-06-23 16:03:42 -0700503}
504
Tatiana Shpeismanfa412f72018-07-09 17:42:46 -0700505namespace {
506
507// FunctionState contains common functionality for printing
508// CFG and ML functions.
Chris Lattner4fd59b02018-07-20 09:35:47 -0700509class FunctionState : public ModulePrinter {
Tatiana Shpeismanfa412f72018-07-09 17:42:46 -0700510public:
Chris Lattner4fd59b02018-07-20 09:35:47 -0700511 FunctionState(const ModulePrinter &other) : ModulePrinter(other) {}
Tatiana Shpeismanfa412f72018-07-09 17:42:46 -0700512
513 void printOperation(const Operation *op);
514
515protected:
Chris Lattnerf8cce872018-07-20 09:28:54 -0700516 void numberValueID(const SSAValue *value) {
517 assert(!valueIDs.count(value) && "Value numbered multiple times");
518 valueIDs[value] = nextValueID++;
519 }
520
Chris Lattner6119d382018-07-20 18:41:34 -0700521 void printValueID(const SSAValue *value,
522 bool dontPrintResultNo = false) const {
523 int resultNo = -1;
524 auto lookupValue = value;
525
526 // If this is a reference to the result of a multi-result instruction, print
527 // out the # identifier and make sure to map our lookup to the first result
528 // of the instruction.
529 if (auto *result = dyn_cast<InstResult>(value)) {
530 if (result->getOwner()->getNumResults() != 1) {
531 resultNo = result->getResultNumber();
532 lookupValue = result->getOwner()->getResult(0);
533 }
534 }
535
536 auto it = valueIDs.find(lookupValue);
537 if (it == valueIDs.end()) {
Chris Lattnerf8cce872018-07-20 09:28:54 -0700538 os << "<<INVALID SSA VALUE>>";
Chris Lattner6119d382018-07-20 18:41:34 -0700539 return;
540 }
541
542 os << '%' << it->getSecond();
543 if (resultNo != -1 && !dontPrintResultNo)
544 os << '#' << resultNo;
Chris Lattnerf8cce872018-07-20 09:28:54 -0700545 }
546
547private:
548 /// This is the value ID for each SSA value in the current function.
549 DenseMap<const SSAValue *, unsigned> valueIDs;
550 unsigned nextValueID = 0;
Tatiana Shpeismanfa412f72018-07-09 17:42:46 -0700551};
MLIR Team54b55a22018-07-18 10:16:05 -0700552} // end anonymous namespace
Tatiana Shpeismanfa412f72018-07-09 17:42:46 -0700553
Tatiana Shpeismanfa412f72018-07-09 17:42:46 -0700554void FunctionState::printOperation(const Operation *op) {
Chris Lattnerf8cce872018-07-20 09:28:54 -0700555 os << " ";
556
Chris Lattnerac591f12018-07-22 21:02:26 -0700557 if (op->getNumResults()) {
558 printValueID(op->getResult(0), /*dontPrintResultNo*/ true);
559 os << " = ";
Chris Lattnerf8cce872018-07-20 09:28:54 -0700560 }
561
Tatiana Shpeismanfa412f72018-07-09 17:42:46 -0700562 // Check to see if this is a known operation. If so, use the registered
563 // custom printer hook.
Chris Lattner4fd59b02018-07-20 09:35:47 -0700564 if (auto opInfo = state.operationSet->lookup(op->getName().str())) {
Tatiana Shpeismanfa412f72018-07-09 17:42:46 -0700565 opInfo->printAssembly(op, os);
566 return;
567 }
568
Chris Lattnerf8cce872018-07-20 09:28:54 -0700569 // Otherwise use the standard verbose printing approach.
570
Tatiana Shpeismanfa412f72018-07-09 17:42:46 -0700571 // TODO: escape name if necessary.
Chris Lattnerf8cce872018-07-20 09:28:54 -0700572 os << "\"" << op->getName().str() << "\"(";
Tatiana Shpeismanfa412f72018-07-09 17:42:46 -0700573
Chris Lattnerac591f12018-07-22 21:02:26 -0700574 interleaveComma(op->getOperands(),
575 [&](const SSAValue *value) { printValueID(value); });
Chris Lattner7f9cc272018-07-19 08:35:28 -0700576
Chris Lattnerf8cce872018-07-20 09:28:54 -0700577 os << ')';
Tatiana Shpeismanfa412f72018-07-09 17:42:46 -0700578 auto attrs = op->getAttrs();
579 if (!attrs.empty()) {
580 os << '{';
Chris Lattner4fd59b02018-07-20 09:35:47 -0700581 interleaveComma(attrs, [&](NamedAttribute attr) {
Chris Lattnerf8cce872018-07-20 09:28:54 -0700582 os << attr.first << ": ";
Chris Lattner4fd59b02018-07-20 09:35:47 -0700583 print(attr.second);
Chris Lattnerf8cce872018-07-20 09:28:54 -0700584 });
Tatiana Shpeismanfa412f72018-07-09 17:42:46 -0700585 os << '}';
586 }
Chris Lattner3b2ef762018-07-18 15:31:25 -0700587
Chris Lattnerac591f12018-07-22 21:02:26 -0700588 // Print the type signature of the operation.
589 os << " : (";
590 interleaveComma(op->getOperands(),
591 [&](const SSAValue *value) { print(value->getType()); });
592 os << ") -> ";
Chris Lattnerf8cce872018-07-20 09:28:54 -0700593
Chris Lattnerac591f12018-07-22 21:02:26 -0700594 if (op->getNumResults() == 1) {
595 print(op->getResult(0)->getType());
596 } else {
597 os << '(';
598 interleaveComma(op->getResults(),
599 [&](const SSAValue *result) { print(result->getType()); });
600 os << ')';
Chris Lattnerf8cce872018-07-20 09:28:54 -0700601 }
Tatiana Shpeismanfa412f72018-07-09 17:42:46 -0700602}
603
Chris Lattner4c95a502018-06-23 16:03:42 -0700604//===----------------------------------------------------------------------===//
605// CFG Function printing
606//===----------------------------------------------------------------------===//
607
608namespace {
Chris Lattner4fd59b02018-07-20 09:35:47 -0700609class CFGFunctionPrinter : public FunctionState {
Chris Lattner4c95a502018-06-23 16:03:42 -0700610public:
Chris Lattner4fd59b02018-07-20 09:35:47 -0700611 CFGFunctionPrinter(const CFGFunction *function, const ModulePrinter &other);
Chris Lattner4c95a502018-06-23 16:03:42 -0700612
613 const CFGFunction *getFunction() const { return function; }
614
615 void print();
616 void print(const BasicBlock *block);
Chris Lattnered65a732018-06-28 20:45:33 -0700617
618 void print(const Instruction *inst);
619 void print(const OperationInst *inst);
620 void print(const ReturnInst *inst);
621 void print(const BranchInst *inst);
Chris Lattner4c95a502018-06-23 16:03:42 -0700622
623 unsigned getBBID(const BasicBlock *block) {
624 auto it = basicBlockIDs.find(block);
625 assert(it != basicBlockIDs.end() && "Block not in this function?");
626 return it->second;
627 }
628
629private:
630 const CFGFunction *function;
MLIR Team54b55a22018-07-18 10:16:05 -0700631 DenseMap<const BasicBlock *, unsigned> basicBlockIDs;
Chris Lattnerf8cce872018-07-20 09:28:54 -0700632
Chris Lattner4fd59b02018-07-20 09:35:47 -0700633 void numberValuesInBlock(const BasicBlock *block);
Chris Lattner4c95a502018-06-23 16:03:42 -0700634};
MLIR Team54b55a22018-07-18 10:16:05 -0700635} // end anonymous namespace
Chris Lattner4c95a502018-06-23 16:03:42 -0700636
Chris Lattner4fd59b02018-07-20 09:35:47 -0700637CFGFunctionPrinter::CFGFunctionPrinter(const CFGFunction *function,
638 const ModulePrinter &other)
639 : FunctionState(other), function(function) {
Chris Lattner4c95a502018-06-23 16:03:42 -0700640 // Each basic block gets a unique ID per function.
641 unsigned blockID = 0;
Chris Lattnerf8cce872018-07-20 09:28:54 -0700642 for (auto &block : *function) {
643 basicBlockIDs[&block] = blockID++;
Chris Lattner4fd59b02018-07-20 09:35:47 -0700644 numberValuesInBlock(&block);
Chris Lattnerf8cce872018-07-20 09:28:54 -0700645 }
646}
647
648/// Number all of the SSA values in the specified basic block.
Chris Lattner4fd59b02018-07-20 09:35:47 -0700649void CFGFunctionPrinter::numberValuesInBlock(const BasicBlock *block) {
James Molloy61a656c2018-07-22 15:45:24 -0700650 for (auto *arg : block->getArguments()) {
651 numberValueID(arg);
652 }
Chris Lattnerf8cce872018-07-20 09:28:54 -0700653 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) {
James Molloy61a656c2018-07-22 15:45:24 -0700673 os << "bb" << getBBID(block);
Chris Lattner4c95a502018-06-23 16:03:42 -0700674
James Molloy61a656c2018-07-22 15:45:24 -0700675 if (!block->args_empty()) {
676 os << '(';
677 interleaveComma(block->getArguments(), [&](const BBArgument *arg) {
678 printValueID(arg);
679 os << ": ";
680 ModulePrinter::print(arg->getType());
681 });
682 os << ')';
683 }
684 os << ":\n";
685
Jacques Pienaarb020c542018-07-15 00:06:54 -0700686 for (auto &inst : block->getOperations()) {
Chris Lattner3a467cc2018-07-01 20:28:00 -0700687 print(&inst);
James Molloy61a656c2018-07-22 15:45:24 -0700688 os << '\n';
Jacques Pienaarb020c542018-07-15 00:06:54 -0700689 }
Chris Lattner4c95a502018-06-23 16:03:42 -0700690
691 print(block->getTerminator());
James Molloy61a656c2018-07-22 15:45:24 -0700692 os << '\n';
Chris Lattner4c95a502018-06-23 16:03:42 -0700693}
694
Chris Lattner4fd59b02018-07-20 09:35:47 -0700695void CFGFunctionPrinter::print(const Instruction *inst) {
Chris Lattner4c95a502018-06-23 16:03:42 -0700696 switch (inst->getKind()) {
Chris Lattnered65a732018-06-28 20:45:33 -0700697 case Instruction::Kind::Operation:
698 return print(cast<OperationInst>(inst));
Chris Lattnerf6d80a02018-06-24 11:18:29 -0700699 case TerminatorInst::Kind::Branch:
Chris Lattnered65a732018-06-28 20:45:33 -0700700 return print(cast<BranchInst>(inst));
Chris Lattner4c95a502018-06-23 16:03:42 -0700701 case TerminatorInst::Kind::Return:
Chris Lattnered65a732018-06-28 20:45:33 -0700702 return print(cast<ReturnInst>(inst));
Chris Lattner4c95a502018-06-23 16:03:42 -0700703 }
704}
705
Chris Lattner4fd59b02018-07-20 09:35:47 -0700706void CFGFunctionPrinter::print(const OperationInst *inst) {
Tatiana Shpeismanfa412f72018-07-09 17:42:46 -0700707 printOperation(inst);
Chris Lattner3b2ef762018-07-18 15:31:25 -0700708}
Chris Lattner1604e472018-07-23 08:42:19 -0700709
Chris Lattner4fd59b02018-07-20 09:35:47 -0700710void CFGFunctionPrinter::print(const BranchInst *inst) {
Jacques Pienaarb020c542018-07-15 00:06:54 -0700711 os << " br bb" << getBBID(inst->getDest());
Chris Lattner1604e472018-07-23 08:42:19 -0700712
713 if (inst->getNumOperands() != 0) {
714 os << '(';
715 // TODO: Use getOperands() when we have it.
716 interleaveComma(inst->getInstOperands(), [&](const InstOperand &operand) {
717 printValueID(operand.get());
718 });
719 os << ") : ";
720 interleaveComma(inst->getInstOperands(), [&](const InstOperand &operand) {
721 ModulePrinter::print(operand.get()->getType());
722 });
723 }
Chris Lattnered65a732018-06-28 20:45:33 -0700724}
Chris Lattner1604e472018-07-23 08:42:19 -0700725
Chris Lattner40746442018-07-21 14:32:09 -0700726void CFGFunctionPrinter::print(const ReturnInst *inst) {
727 os << " return";
728
729 if (inst->getNumOperands() != 0)
730 os << ' ';
731
Chris Lattnerac591f12018-07-22 21:02:26 -0700732 interleaveComma(inst->getOperands(), [&](const CFGValue *operand) {
733 printValueID(operand);
Chris Lattner40746442018-07-21 14:32:09 -0700734 os << " : ";
Chris Lattnerac591f12018-07-22 21:02:26 -0700735 ModulePrinter::print(operand->getType());
Chris Lattner40746442018-07-21 14:32:09 -0700736 });
737}
MLIR Team54b55a22018-07-18 10:16:05 -0700738
Chris Lattner4fd59b02018-07-20 09:35:47 -0700739void ModulePrinter::print(const CFGFunction *fn) {
740 CFGFunctionPrinter(fn, *this).print();
Chris Lattnered65a732018-06-28 20:45:33 -0700741}
742
Chris Lattner4c95a502018-06-23 16:03:42 -0700743//===----------------------------------------------------------------------===//
Tatiana Shpeismanbf079c92018-07-03 17:51:28 -0700744// ML Function printing
Chris Lattner4c95a502018-06-23 16:03:42 -0700745//===----------------------------------------------------------------------===//
746
Tatiana Shpeismanbf079c92018-07-03 17:51:28 -0700747namespace {
Chris Lattner4fd59b02018-07-20 09:35:47 -0700748class MLFunctionPrinter : public FunctionState {
Tatiana Shpeismanbf079c92018-07-03 17:51:28 -0700749public:
Chris Lattner4fd59b02018-07-20 09:35:47 -0700750 MLFunctionPrinter(const MLFunction *function, const ModulePrinter &other);
Tatiana Shpeismanbf079c92018-07-03 17:51:28 -0700751
752 const MLFunction *getFunction() const { return function; }
753
Tatiana Shpeisman1bcfe982018-07-13 13:03:13 -0700754 // Prints ML function
Tatiana Shpeismanbf079c92018-07-03 17:51:28 -0700755 void print();
756
Tatiana Shpeisman1bcfe982018-07-13 13:03:13 -0700757 // Methods to print ML function statements
Tatiana Shpeismanbf079c92018-07-03 17:51:28 -0700758 void print(const Statement *stmt);
Tatiana Shpeismanfa412f72018-07-09 17:42:46 -0700759 void print(const OperationStmt *stmt);
Tatiana Shpeismanbf079c92018-07-03 17:51:28 -0700760 void print(const ForStmt *stmt);
761 void print(const IfStmt *stmt);
Tatiana Shpeisman1bcfe982018-07-13 13:03:13 -0700762 void print(const StmtBlock *block);
763
764 // Number of spaces used for indenting nested statements
765 const static unsigned indentWidth = 2;
Tatiana Shpeismanbf079c92018-07-03 17:51:28 -0700766
767private:
Tatiana Shpeismanbf079c92018-07-03 17:51:28 -0700768 const MLFunction *function;
Tatiana Shpeismanbf079c92018-07-03 17:51:28 -0700769 int numSpaces;
770};
MLIR Team54b55a22018-07-18 10:16:05 -0700771} // end anonymous namespace
Tatiana Shpeismanbf079c92018-07-03 17:51:28 -0700772
Chris Lattner4fd59b02018-07-20 09:35:47 -0700773MLFunctionPrinter::MLFunctionPrinter(const MLFunction *function,
774 const ModulePrinter &other)
775 : FunctionState(other), function(function), numSpaces(0) {}
Tatiana Shpeismanbf079c92018-07-03 17:51:28 -0700776
Chris Lattner4fd59b02018-07-20 09:35:47 -0700777void MLFunctionPrinter::print() {
Tatiana Shpeismanbf079c92018-07-03 17:51:28 -0700778 os << "mlfunc ";
779 // FIXME: should print argument names rather than just signature
Chris Lattner4fd59b02018-07-20 09:35:47 -0700780 printFunctionSignature(function);
Tatiana Shpeismanbf079c92018-07-03 17:51:28 -0700781 os << " {\n";
Tatiana Shpeisman1bcfe982018-07-13 13:03:13 -0700782 print(function);
Tatiana Shpeismanbf079c92018-07-03 17:51:28 -0700783 os << " return\n";
784 os << "}\n\n";
785}
786
Chris Lattner4fd59b02018-07-20 09:35:47 -0700787void MLFunctionPrinter::print(const StmtBlock *block) {
Tatiana Shpeisman1bcfe982018-07-13 13:03:13 -0700788 numSpaces += indentWidth;
Jacques Pienaarb020c542018-07-15 00:06:54 -0700789 for (auto &stmt : block->getStatements()) {
Tatiana Shpeisman1bcfe982018-07-13 13:03:13 -0700790 print(&stmt);
Jacques Pienaarb020c542018-07-15 00:06:54 -0700791 os << "\n";
792 }
Tatiana Shpeisman1bcfe982018-07-13 13:03:13 -0700793 numSpaces -= indentWidth;
794}
795
Chris Lattner4fd59b02018-07-20 09:35:47 -0700796void MLFunctionPrinter::print(const Statement *stmt) {
Tatiana Shpeismanbf079c92018-07-03 17:51:28 -0700797 switch (stmt->getKind()) {
Tatiana Shpeisman565b9642018-07-16 11:47:09 -0700798 case Statement::Kind::Operation:
799 return print(cast<OperationStmt>(stmt));
Tatiana Shpeismanbf079c92018-07-03 17:51:28 -0700800 case Statement::Kind::For:
801 return print(cast<ForStmt>(stmt));
802 case Statement::Kind::If:
803 return print(cast<IfStmt>(stmt));
Tatiana Shpeismanbf079c92018-07-03 17:51:28 -0700804 }
805}
806
Chris Lattner4fd59b02018-07-20 09:35:47 -0700807void MLFunctionPrinter::print(const OperationStmt *stmt) {
808 printOperation(stmt);
809}
Tatiana Shpeismanbf079c92018-07-03 17:51:28 -0700810
Chris Lattner4fd59b02018-07-20 09:35:47 -0700811void MLFunctionPrinter::print(const ForStmt *stmt) {
Tatiana Shpeisman1da50c42018-07-19 09:52:39 -0700812 os.indent(numSpaces) << "for x = " << *stmt->getLowerBound();
813 os << " to " << *stmt->getUpperBound();
814 if (stmt->getStep()->getValue() != 1)
815 os << " step " << *stmt->getStep();
816
817 os << " {\n";
Tatiana Shpeisman1bcfe982018-07-13 13:03:13 -0700818 print(static_cast<const StmtBlock *>(stmt));
Tatiana Shpeisman565b9642018-07-16 11:47:09 -0700819 os.indent(numSpaces) << "}";
Tatiana Shpeismanbf079c92018-07-03 17:51:28 -0700820}
821
Chris Lattner4fd59b02018-07-20 09:35:47 -0700822void MLFunctionPrinter::print(const IfStmt *stmt) {
Tatiana Shpeisman1bcfe982018-07-13 13:03:13 -0700823 os.indent(numSpaces) << "if () {\n";
824 print(stmt->getThenClause());
825 os.indent(numSpaces) << "}";
826 if (stmt->hasElseClause()) {
827 os << " else {\n";
828 print(stmt->getElseClause());
829 os.indent(numSpaces) << "}";
830 }
Tatiana Shpeismanbf079c92018-07-03 17:51:28 -0700831}
832
Chris Lattner4fd59b02018-07-20 09:35:47 -0700833void ModulePrinter::print(const MLFunction *fn) {
834 MLFunctionPrinter(fn, *this).print();
MLIR Team4718bc92018-07-17 16:56:54 -0700835}
836
Tatiana Shpeismanbf079c92018-07-03 17:51:28 -0700837//===----------------------------------------------------------------------===//
838// print and dump methods
839//===----------------------------------------------------------------------===//
Chris Lattnered65a732018-06-28 20:45:33 -0700840
MLIR Teamb61885d2018-07-18 16:29:21 -0700841void Attribute::print(raw_ostream &os) const {
Chris Lattner4fd59b02018-07-20 09:35:47 -0700842 ModuleState state(/*no context is known*/ nullptr);
843 ModulePrinter(os, state).print(this);
MLIR Teamb61885d2018-07-18 16:29:21 -0700844}
845
846void Attribute::dump() const {
847 print(llvm::errs());
848}
849
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()); }