blob: 31fd05c683c8afd701ae0c10b45ec9df20971a4e [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 Lattnerf8cce872018-07-20 09:28:54 -070041template <typename Container, typename UnaryFunctor>
42inline void interleaveComma(raw_ostream &os, const Container &c,
43 UnaryFunctor each_fn) {
44 interleave(c.begin(), c.end(), each_fn, [&]() { os << ", "; });
45}
46
Chris Lattner4c95a502018-06-23 16:03:42 -070047//===----------------------------------------------------------------------===//
MLIR Team4718bc92018-07-17 16:56:54 -070048// Module printing
49//===----------------------------------------------------------------------===//
50
51namespace {
MLIR Team54b55a22018-07-18 10:16:05 -070052class ModuleState {
53public:
MLIR Team4718bc92018-07-17 16:56:54 -070054 ModuleState(raw_ostream &os);
55
56 void initialize(const Module *module);
57
58 void print(const Module *module);
MLIR Teamb61885d2018-07-18 16:29:21 -070059 void print(const Attribute *attr) const;
MLIR Team4718bc92018-07-17 16:56:54 -070060 void print(const Type *type) const;
MLIR Team54b55a22018-07-18 10:16:05 -070061 void print(const Function *fn);
62 void print(const ExtFunction *fn);
63 void print(const CFGFunction *fn);
64 void print(const MLFunction *fn);
MLIR Team4718bc92018-07-17 16:56:54 -070065
MLIR Team54b55a22018-07-18 10:16:05 -070066 void recordAffineMapReference(const AffineMap *affineMap) {
MLIR Team4718bc92018-07-17 16:56:54 -070067 if (affineMapIds.count(affineMap) == 0) {
68 affineMapIds[affineMap] = nextAffineMapId++;
69 }
70 }
71
MLIR Team54b55a22018-07-18 10:16:05 -070072 int getAffineMapId(const AffineMap *affineMap) const {
MLIR Team4718bc92018-07-17 16:56:54 -070073 auto it = affineMapIds.find(affineMap);
74 if (it == affineMapIds.end()) {
75 return -1;
76 }
77 return it->second;
78 }
79
MLIR Team54b55a22018-07-18 10:16:05 -070080private:
MLIR Team4718bc92018-07-17 16:56:54 -070081 // Visit functions.
82 void visitFunction(const Function *fn);
83 void visitExtFunction(const ExtFunction *fn);
84 void visitCFGFunction(const CFGFunction *fn);
85 void visitMLFunction(const MLFunction *fn);
86 void visitType(const Type *type);
MLIR Teamb61885d2018-07-18 16:29:21 -070087 void visitAttribute(const Attribute *attr);
88 void visitOperation(const Operation *op);
89
90 void printAffineMapId(int affineMapId) const;
91 void printAffineMapReference(const AffineMap* affineMap) const;
MLIR Team4718bc92018-07-17 16:56:54 -070092
93 raw_ostream &os;
MLIR Team54b55a22018-07-18 10:16:05 -070094 DenseMap<const AffineMap *, int> affineMapIds;
MLIR Team4718bc92018-07-17 16:56:54 -070095 int nextAffineMapId = 0;
96};
MLIR Team54b55a22018-07-18 10:16:05 -070097} // end anonymous namespace
MLIR Team4718bc92018-07-17 16:56:54 -070098
MLIR Team54b55a22018-07-18 10:16:05 -070099ModuleState::ModuleState(raw_ostream &os) : os(os) {}
MLIR Team4718bc92018-07-17 16:56:54 -0700100
101// Initializes module state, populating affine map state.
102void ModuleState::initialize(const Module *module) {
103 for (auto fn : module->functionList) {
104 visitFunction(fn);
105 }
106}
107
108// TODO Support visiting other types/instructions when implemented.
109void ModuleState::visitType(const Type *type) {
110 if (type->getKind() == Type::Kind::Function) {
111 // Visit input and result types for functions.
112 auto *funcType = cast<FunctionType>(type);
MLIR Team54b55a22018-07-18 10:16:05 -0700113 for (auto *input : funcType->getInputs()) {
MLIR Team4718bc92018-07-17 16:56:54 -0700114 visitType(input);
115 }
MLIR Team54b55a22018-07-18 10:16:05 -0700116 for (auto *result : funcType->getResults()) {
MLIR Team4718bc92018-07-17 16:56:54 -0700117 visitType(result);
118 }
119 } else if (type->getKind() == Type::Kind::MemRef) {
120 // Visit affine maps in memref type.
121 auto *memref = cast<MemRefType>(type);
MLIR Team54b55a22018-07-18 10:16:05 -0700122 for (AffineMap *map : memref->getAffineMaps()) {
MLIR Team4718bc92018-07-17 16:56:54 -0700123 recordAffineMapReference(map);
124 }
125 }
126}
127
MLIR Teamb61885d2018-07-18 16:29:21 -0700128void ModuleState::visitAttribute(const Attribute *attr) {
129 if (isa<AffineMapAttr>(attr)) {
130 recordAffineMapReference(cast<AffineMapAttr>(attr)->getValue());
131 } else if (isa<ArrayAttr>(attr)) {
132 for (auto elt : cast<ArrayAttr>(attr)->getValue()) {
133 visitAttribute(elt);
134 }
135 }
136}
137
138void ModuleState::visitOperation(const Operation *op) {
139 for (auto elt : op->getAttrs()) {
140 visitAttribute(elt.second);
141 }
142}
143
MLIR Team4718bc92018-07-17 16:56:54 -0700144void ModuleState::visitExtFunction(const ExtFunction *fn) {
145 visitType(fn->getType());
146}
147
148void ModuleState::visitCFGFunction(const CFGFunction *fn) {
149 visitType(fn->getType());
150 // TODO Visit function body instructions.
MLIR Teamb61885d2018-07-18 16:29:21 -0700151 for (auto &block : *fn) {
152 for (auto &op : block.getOperations()) {
153 visitOperation(&op);
154 }
155 }
MLIR Team4718bc92018-07-17 16:56:54 -0700156}
157
158void ModuleState::visitMLFunction(const MLFunction *fn) {
159 visitType(fn->getType());
MLIR Teamb61885d2018-07-18 16:29:21 -0700160 // TODO Visit function body statements (and attributes if required).
MLIR Team4718bc92018-07-17 16:56:54 -0700161}
162
163void ModuleState::visitFunction(const Function *fn) {
164 switch (fn->getKind()) {
MLIR Team54b55a22018-07-18 10:16:05 -0700165 case Function::Kind::ExtFunc:
166 return visitExtFunction(cast<ExtFunction>(fn));
167 case Function::Kind::CFGFunc:
168 return visitCFGFunction(cast<CFGFunction>(fn));
169 case Function::Kind::MLFunc:
170 return visitMLFunction(cast<MLFunction>(fn));
MLIR Team4718bc92018-07-17 16:56:54 -0700171 }
172}
173
MLIR Team4718bc92018-07-17 16:56:54 -0700174// Prints function with initialized module state.
MLIR Team54b55a22018-07-18 10:16:05 -0700175void ModuleState::print(const Function *fn) {
MLIR Team4718bc92018-07-17 16:56:54 -0700176 switch (fn->getKind()) {
MLIR Team54b55a22018-07-18 10:16:05 -0700177 case Function::Kind::ExtFunc:
178 return print(cast<ExtFunction>(fn));
179 case Function::Kind::CFGFunc:
180 return print(cast<CFGFunction>(fn));
181 case Function::Kind::MLFunc:
182 return print(cast<MLFunction>(fn));
MLIR Team4718bc92018-07-17 16:56:54 -0700183 }
184}
185
186// Prints affine map identifier.
MLIR Teamb61885d2018-07-18 16:29:21 -0700187void ModuleState::printAffineMapId(int affineMapId) const {
MLIR Team4718bc92018-07-17 16:56:54 -0700188 os << "#map" << affineMapId;
189}
190
MLIR Teamb61885d2018-07-18 16:29:21 -0700191void ModuleState::printAffineMapReference(const AffineMap* affineMap) const {
192 const int mapId = getAffineMapId(affineMap);
193 if (mapId >= 0) {
194 // Map will be printed at top of module so print reference to its id.
195 printAffineMapId(mapId);
196 } else {
197 // Map not in module state so print inline.
198 affineMap->print(os);
199 }
200}
201
MLIR Team4718bc92018-07-17 16:56:54 -0700202void ModuleState::print(const Module *module) {
MLIR Team54b55a22018-07-18 10:16:05 -0700203 for (const auto &mapAndId : affineMapIds) {
MLIR Teamb61885d2018-07-18 16:29:21 -0700204 printAffineMapId(mapAndId.second);
MLIR Team4718bc92018-07-17 16:56:54 -0700205 os << " = ";
206 mapAndId.first->print(os);
207 os << '\n';
208 }
MLIR Team54b55a22018-07-18 10:16:05 -0700209 for (auto *fn : module->functionList) print(fn);
MLIR Team4718bc92018-07-17 16:56:54 -0700210}
211
MLIR Teamb61885d2018-07-18 16:29:21 -0700212void ModuleState::print(const Attribute *attr) const {
213 switch (attr->getKind()) {
214 case Attribute::Kind::Bool:
215 os << (cast<BoolAttr>(attr)->getValue() ? "true" : "false");
216 break;
217 case Attribute::Kind::Integer:
218 os << cast<IntegerAttr>(attr)->getValue();
219 break;
220 case Attribute::Kind::Float:
221 // FIXME: this isn't precise, we should print with a hex format.
222 os << cast<FloatAttr>(attr)->getValue();
223 break;
224 case Attribute::Kind::String:
225 // FIXME: should escape the string.
226 os << '"' << cast<StringAttr>(attr)->getValue() << '"';
227 break;
228 case Attribute::Kind::Array: {
229 auto elts = cast<ArrayAttr>(attr)->getValue();
230 os << '[';
Chris Lattnerf8cce872018-07-20 09:28:54 -0700231 interleaveComma(os, elts, [&](Attribute *attr) { print(attr); });
MLIR Teamb61885d2018-07-18 16:29:21 -0700232 os << ']';
233 break;
234 }
235 case Attribute::Kind::AffineMap:
236 printAffineMapReference(cast<AffineMapAttr>(attr)->getValue());
237 break;
238 }
239}
240
MLIR Team4718bc92018-07-17 16:56:54 -0700241void ModuleState::print(const Type *type) const {
242 switch (type->getKind()) {
MLIR Team54b55a22018-07-18 10:16:05 -0700243 case Type::Kind::AffineInt:
244 os << "affineint";
245 return;
246 case Type::Kind::BF16:
247 os << "bf16";
248 return;
249 case Type::Kind::F16:
250 os << "f16";
251 return;
252 case Type::Kind::F32:
253 os << "f32";
254 return;
255 case Type::Kind::F64:
256 os << "f64";
257 return;
MLIR Team4718bc92018-07-17 16:56:54 -0700258
259 case Type::Kind::Integer: {
260 auto *integer = cast<IntegerType>(type);
261 os << 'i' << integer->getWidth();
262 return;
263 }
264 case Type::Kind::Function: {
265 auto *func = cast<FunctionType>(type);
266 os << '(';
Chris Lattnerf8cce872018-07-20 09:28:54 -0700267 interleaveComma(os, func->getInputs(), [&](Type *type) { os << *type; });
MLIR Team4718bc92018-07-17 16:56:54 -0700268 os << ") -> ";
269 auto results = func->getResults();
270 if (results.size() == 1)
271 os << *results[0];
272 else {
273 os << '(';
Chris Lattnerf8cce872018-07-20 09:28:54 -0700274 interleaveComma(os, results, [&](Type *type) { os << *type; });
MLIR Team4718bc92018-07-17 16:56:54 -0700275 os << ')';
276 }
277 return;
278 }
279 case Type::Kind::Vector: {
280 auto *v = cast<VectorType>(type);
281 os << "vector<";
MLIR Team54b55a22018-07-18 10:16:05 -0700282 for (auto dim : v->getShape()) os << dim << 'x';
MLIR Team4718bc92018-07-17 16:56:54 -0700283 os << *v->getElementType() << '>';
284 return;
285 }
286 case Type::Kind::RankedTensor: {
287 auto *v = cast<RankedTensorType>(type);
288 os << "tensor<";
289 for (auto dim : v->getShape()) {
290 if (dim < 0)
291 os << '?';
292 else
293 os << dim;
294 os << 'x';
295 }
296 os << *v->getElementType() << '>';
297 return;
298 }
299 case Type::Kind::UnrankedTensor: {
300 auto *v = cast<UnrankedTensorType>(type);
301 os << "tensor<??" << *v->getElementType() << '>';
302 return;
303 }
304 case Type::Kind::MemRef: {
305 auto *v = cast<MemRefType>(type);
306 os << "memref<";
307 for (auto dim : v->getShape()) {
308 if (dim < 0)
309 os << '?';
310 else
311 os << dim;
312 os << 'x';
313 }
314 os << *v->getElementType();
315 for (auto map : v->getAffineMaps()) {
316 os << ", ";
MLIR Teamb61885d2018-07-18 16:29:21 -0700317 printAffineMapReference(map);
MLIR Team4718bc92018-07-17 16:56:54 -0700318 }
319 os << ", " << v->getMemorySpace();
320 os << '>';
321 return;
322 }
323 }
324}
325
326//===----------------------------------------------------------------------===//
Chris Lattner4c95a502018-06-23 16:03:42 -0700327// Function printing
328//===----------------------------------------------------------------------===//
329
MLIR Team4718bc92018-07-17 16:56:54 -0700330static void printFunctionSignature(const Function *fn,
331 const ModuleState *moduleState,
332 raw_ostream &os) {
Chris Lattner4c95a502018-06-23 16:03:42 -0700333 auto type = fn->getType();
334
335 os << "@" << fn->getName() << '(';
Chris Lattnerf8cce872018-07-20 09:28:54 -0700336 interleaveComma(os, type->getInputs(),
337 [&](Type *eltType) { moduleState->print(eltType); });
Chris Lattner4c95a502018-06-23 16:03:42 -0700338 os << ')';
339
340 switch (type->getResults().size()) {
MLIR Team54b55a22018-07-18 10:16:05 -0700341 case 0:
342 break;
Chris Lattner4c95a502018-06-23 16:03:42 -0700343 case 1:
MLIR Team4718bc92018-07-17 16:56:54 -0700344 os << " -> ";
345 moduleState->print(type->getResults()[0]);
Chris Lattner4c95a502018-06-23 16:03:42 -0700346 break;
347 default:
348 os << " -> (";
Chris Lattnerf8cce872018-07-20 09:28:54 -0700349 interleaveComma(os, type->getResults(),
350 [&](Type *eltType) { moduleState->print(eltType); });
Chris Lattner4c95a502018-06-23 16:03:42 -0700351 os << ')';
352 break;
353 }
354}
355
MLIR Team54b55a22018-07-18 10:16:05 -0700356void ModuleState::print(const ExtFunction *fn) {
Chris Lattner4c95a502018-06-23 16:03:42 -0700357 os << "extfunc ";
MLIR Team54b55a22018-07-18 10:16:05 -0700358 printFunctionSignature(fn, this, os);
359 os << '\n';
Chris Lattner4c95a502018-06-23 16:03:42 -0700360}
361
Tatiana Shpeismanfa412f72018-07-09 17:42:46 -0700362namespace {
363
364// FunctionState contains common functionality for printing
365// CFG and ML functions.
366class FunctionState {
367public:
MLIR Team4718bc92018-07-17 16:56:54 -0700368 FunctionState(MLIRContext *context, const ModuleState *moduleState,
369 raw_ostream &os);
Tatiana Shpeismanfa412f72018-07-09 17:42:46 -0700370
371 void printOperation(const Operation *op);
372
373protected:
374 raw_ostream &os;
MLIR Team4718bc92018-07-17 16:56:54 -0700375 const ModuleState *moduleState;
Tatiana Shpeismanfa412f72018-07-09 17:42:46 -0700376 const OperationSet &operationSet;
Chris Lattnerf8cce872018-07-20 09:28:54 -0700377
378 void numberValueID(const SSAValue *value) {
379 assert(!valueIDs.count(value) && "Value numbered multiple times");
380 valueIDs[value] = nextValueID++;
381 }
382
383 void printValueID(const SSAValue *value) const {
384 // TODO: If this is the result of an operation with multiple results, look
385 // up the first result, and print the #32 syntax.
386 auto it = valueIDs.find(value);
387 if (it != valueIDs.end())
388 os << '%' << it->getSecond();
389 else
390 os << "<<INVALID SSA VALUE>>";
391 }
392
393private:
394 /// This is the value ID for each SSA value in the current function.
395 DenseMap<const SSAValue *, unsigned> valueIDs;
396 unsigned nextValueID = 0;
Tatiana Shpeismanfa412f72018-07-09 17:42:46 -0700397};
MLIR Team54b55a22018-07-18 10:16:05 -0700398} // end anonymous namespace
Tatiana Shpeismanfa412f72018-07-09 17:42:46 -0700399
MLIR Team4718bc92018-07-17 16:56:54 -0700400FunctionState::FunctionState(MLIRContext *context,
MLIR Team54b55a22018-07-18 10:16:05 -0700401 const ModuleState *moduleState, raw_ostream &os)
402 : os(os),
403 moduleState(moduleState),
MLIR Team4718bc92018-07-17 16:56:54 -0700404 operationSet(OperationSet::get(context)) {}
Tatiana Shpeismanfa412f72018-07-09 17:42:46 -0700405
406void FunctionState::printOperation(const Operation *op) {
Chris Lattnerf8cce872018-07-20 09:28:54 -0700407 os << " ";
408
409 // TODO: When we have SSAValue version of operands & results wired into
410 // Operation this check can go away.
411 if (auto *inst = dyn_cast<OperationInst>(op)) {
412 if (inst->getNumResults()) {
413 printValueID(inst->getResult(0));
414 os << " = ";
415 }
416 }
417
Tatiana Shpeismanfa412f72018-07-09 17:42:46 -0700418 // Check to see if this is a known operation. If so, use the registered
419 // custom printer hook.
420 if (auto opInfo = operationSet.lookup(op->getName().str())) {
Tatiana Shpeismanfa412f72018-07-09 17:42:46 -0700421 opInfo->printAssembly(op, os);
422 return;
423 }
424
Chris Lattnerf8cce872018-07-20 09:28:54 -0700425 // Otherwise use the standard verbose printing approach.
426
Tatiana Shpeismanfa412f72018-07-09 17:42:46 -0700427 // TODO: escape name if necessary.
Chris Lattnerf8cce872018-07-20 09:28:54 -0700428 os << "\"" << op->getName().str() << "\"(";
Tatiana Shpeismanfa412f72018-07-09 17:42:46 -0700429
Chris Lattnerf8cce872018-07-20 09:28:54 -0700430 // TODO: When we have SSAValue version of operands & results wired into
431 // Operation this check can go away.
432 if (auto *inst = dyn_cast<OperationInst>(op)) {
433 // TODO: Use getOperands() when we have it.
434 interleaveComma(
435 os, inst->getInstOperands(),
436 [&](const InstOperand &operand) { printValueID(operand.get()); });
437 }
Chris Lattner7f9cc272018-07-19 08:35:28 -0700438
Chris Lattnerf8cce872018-07-20 09:28:54 -0700439 os << ')';
Tatiana Shpeismanfa412f72018-07-09 17:42:46 -0700440 auto attrs = op->getAttrs();
441 if (!attrs.empty()) {
442 os << '{';
Chris Lattnerf8cce872018-07-20 09:28:54 -0700443 interleaveComma(os, attrs, [&](NamedAttribute attr) {
444 os << attr.first << ": ";
445 moduleState->print(attr.second);
446 });
Tatiana Shpeismanfa412f72018-07-09 17:42:46 -0700447 os << '}';
448 }
Chris Lattner3b2ef762018-07-18 15:31:25 -0700449
Chris Lattnerf8cce872018-07-20 09:28:54 -0700450 // TODO: When we have SSAValue version of operands & results wired into
451 // Operation this check can go away.
452 if (auto *inst = dyn_cast<OperationInst>(op)) {
453 // Print the type signature of the operation.
454 os << " : (";
455 // TODO: Switch to getOperands() when we have it.
456 interleaveComma(os, inst->getInstOperands(), [&](const InstOperand &op) {
457 moduleState->print(op.get()->getType());
458 });
459 os << ") -> ";
460
461 // TODO: Switch to getResults() when we have it.
462 if (inst->getNumResults() == 1) {
463 moduleState->print(inst->getInstResult(0).getType());
464 } else {
465 os << '(';
466 interleaveComma(os, inst->getInstResults(),
467 [&](const InstResult &result) {
468 moduleState->print(result.getType());
469 });
470 os << ')';
471 }
472 }
Tatiana Shpeismanfa412f72018-07-09 17:42:46 -0700473}
474
Chris Lattner4c95a502018-06-23 16:03:42 -0700475//===----------------------------------------------------------------------===//
476// CFG Function printing
477//===----------------------------------------------------------------------===//
478
479namespace {
Tatiana Shpeismanfa412f72018-07-09 17:42:46 -0700480class CFGFunctionState : public FunctionState {
Chris Lattner4c95a502018-06-23 16:03:42 -0700481public:
MLIR Team4718bc92018-07-17 16:56:54 -0700482 CFGFunctionState(const CFGFunction *function, const ModuleState *moduleState,
483 raw_ostream &os);
Chris Lattner4c95a502018-06-23 16:03:42 -0700484
485 const CFGFunction *getFunction() const { return function; }
486
487 void print();
488 void print(const BasicBlock *block);
Chris Lattnered65a732018-06-28 20:45:33 -0700489
490 void print(const Instruction *inst);
491 void print(const OperationInst *inst);
492 void print(const ReturnInst *inst);
493 void print(const BranchInst *inst);
Chris Lattner4c95a502018-06-23 16:03:42 -0700494
495 unsigned getBBID(const BasicBlock *block) {
496 auto it = basicBlockIDs.find(block);
497 assert(it != basicBlockIDs.end() && "Block not in this function?");
498 return it->second;
499 }
500
501private:
502 const CFGFunction *function;
MLIR Team54b55a22018-07-18 10:16:05 -0700503 DenseMap<const BasicBlock *, unsigned> basicBlockIDs;
Chris Lattnerf8cce872018-07-20 09:28:54 -0700504
505 void numberBlock(const BasicBlock *block);
Chris Lattner4c95a502018-06-23 16:03:42 -0700506};
MLIR Team54b55a22018-07-18 10:16:05 -0700507} // end anonymous namespace
Chris Lattner4c95a502018-06-23 16:03:42 -0700508
MLIR Team4718bc92018-07-17 16:56:54 -0700509CFGFunctionState::CFGFunctionState(const CFGFunction *function,
510 const ModuleState *moduleState,
511 raw_ostream &os)
512 : FunctionState(function->getContext(), moduleState, os),
513 function(function) {
Chris Lattner4c95a502018-06-23 16:03:42 -0700514 // Each basic block gets a unique ID per function.
515 unsigned blockID = 0;
Chris Lattnerf8cce872018-07-20 09:28:54 -0700516 for (auto &block : *function) {
517 basicBlockIDs[&block] = blockID++;
518 numberBlock(&block);
519 }
520}
521
522/// Number all of the SSA values in the specified basic block.
523void CFGFunctionState::numberBlock(const BasicBlock *block) {
524 // TODO: basic block arguments.
525 for (auto &op : *block) {
526 // We number instruction that have results, and we only number the first
527 // result.
528 if (op.getNumResults() != 0)
529 numberValueID(op.getResult(0));
530 }
531
532 // Terminators do not define values.
Chris Lattner4c95a502018-06-23 16:03:42 -0700533}
534
535void CFGFunctionState::print() {
536 os << "cfgfunc ";
MLIR Team4718bc92018-07-17 16:56:54 -0700537 printFunctionSignature(this->getFunction(), moduleState, os);
Chris Lattner4c95a502018-06-23 16:03:42 -0700538 os << " {\n";
539
MLIR Team54b55a22018-07-18 10:16:05 -0700540 for (auto &block : *function) print(&block);
Chris Lattner4c95a502018-06-23 16:03:42 -0700541 os << "}\n\n";
542}
543
544void CFGFunctionState::print(const BasicBlock *block) {
545 os << "bb" << getBBID(block) << ":\n";
546
Chris Lattnered65a732018-06-28 20:45:33 -0700547 // TODO Print arguments.
Jacques Pienaarb020c542018-07-15 00:06:54 -0700548 for (auto &inst : block->getOperations()) {
Chris Lattner3a467cc2018-07-01 20:28:00 -0700549 print(&inst);
Jacques Pienaarb020c542018-07-15 00:06:54 -0700550 os << "\n";
551 }
Chris Lattner4c95a502018-06-23 16:03:42 -0700552
553 print(block->getTerminator());
Jacques Pienaarb020c542018-07-15 00:06:54 -0700554 os << "\n";
Chris Lattner4c95a502018-06-23 16:03:42 -0700555}
556
Chris Lattnered65a732018-06-28 20:45:33 -0700557void CFGFunctionState::print(const Instruction *inst) {
Chris Lattner4c95a502018-06-23 16:03:42 -0700558 switch (inst->getKind()) {
Chris Lattnered65a732018-06-28 20:45:33 -0700559 case Instruction::Kind::Operation:
560 return print(cast<OperationInst>(inst));
Chris Lattnerf6d80a02018-06-24 11:18:29 -0700561 case TerminatorInst::Kind::Branch:
Chris Lattnered65a732018-06-28 20:45:33 -0700562 return print(cast<BranchInst>(inst));
Chris Lattner4c95a502018-06-23 16:03:42 -0700563 case TerminatorInst::Kind::Return:
Chris Lattnered65a732018-06-28 20:45:33 -0700564 return print(cast<ReturnInst>(inst));
Chris Lattner4c95a502018-06-23 16:03:42 -0700565 }
566}
567
Chris Lattnered65a732018-06-28 20:45:33 -0700568void CFGFunctionState::print(const OperationInst *inst) {
Tatiana Shpeismanfa412f72018-07-09 17:42:46 -0700569 printOperation(inst);
Chris Lattner3b2ef762018-07-18 15:31:25 -0700570}
Chris Lattnered65a732018-06-28 20:45:33 -0700571void CFGFunctionState::print(const BranchInst *inst) {
Jacques Pienaarb020c542018-07-15 00:06:54 -0700572 os << " br bb" << getBBID(inst->getDest());
Chris Lattnered65a732018-06-28 20:45:33 -0700573}
MLIR Team54b55a22018-07-18 10:16:05 -0700574void CFGFunctionState::print(const ReturnInst *inst) { os << " return"; }
575
576void ModuleState::print(const CFGFunction *fn) {
577 CFGFunctionState state(fn, this, os);
578 state.print();
Chris Lattnered65a732018-06-28 20:45:33 -0700579}
580
Chris Lattner4c95a502018-06-23 16:03:42 -0700581//===----------------------------------------------------------------------===//
Tatiana Shpeismanbf079c92018-07-03 17:51:28 -0700582// ML Function printing
Chris Lattner4c95a502018-06-23 16:03:42 -0700583//===----------------------------------------------------------------------===//
584
Tatiana Shpeismanbf079c92018-07-03 17:51:28 -0700585namespace {
Tatiana Shpeismanfa412f72018-07-09 17:42:46 -0700586class MLFunctionState : public FunctionState {
Tatiana Shpeismanbf079c92018-07-03 17:51:28 -0700587public:
MLIR Team4718bc92018-07-17 16:56:54 -0700588 MLFunctionState(const MLFunction *function, const ModuleState *moduleState,
589 raw_ostream &os);
Tatiana Shpeismanbf079c92018-07-03 17:51:28 -0700590
591 const MLFunction *getFunction() const { return function; }
592
Tatiana Shpeisman1bcfe982018-07-13 13:03:13 -0700593 // Prints ML function
Tatiana Shpeismanbf079c92018-07-03 17:51:28 -0700594 void print();
595
Tatiana Shpeisman1bcfe982018-07-13 13:03:13 -0700596 // Methods to print ML function statements
Tatiana Shpeismanbf079c92018-07-03 17:51:28 -0700597 void print(const Statement *stmt);
Tatiana Shpeismanfa412f72018-07-09 17:42:46 -0700598 void print(const OperationStmt *stmt);
Tatiana Shpeismanbf079c92018-07-03 17:51:28 -0700599 void print(const ForStmt *stmt);
600 void print(const IfStmt *stmt);
Tatiana Shpeisman1bcfe982018-07-13 13:03:13 -0700601 void print(const StmtBlock *block);
602
603 // Number of spaces used for indenting nested statements
604 const static unsigned indentWidth = 2;
Tatiana Shpeismanbf079c92018-07-03 17:51:28 -0700605
606private:
Tatiana Shpeismanbf079c92018-07-03 17:51:28 -0700607 const MLFunction *function;
Tatiana Shpeismanbf079c92018-07-03 17:51:28 -0700608 int numSpaces;
609};
MLIR Team54b55a22018-07-18 10:16:05 -0700610} // end anonymous namespace
Tatiana Shpeismanbf079c92018-07-03 17:51:28 -0700611
MLIR Team4718bc92018-07-17 16:56:54 -0700612MLFunctionState::MLFunctionState(const MLFunction *function,
613 const ModuleState *moduleState,
614 raw_ostream &os)
615 : FunctionState(function->getContext(), moduleState, os),
MLIR Team54b55a22018-07-18 10:16:05 -0700616 function(function),
617 numSpaces(0) {}
Tatiana Shpeismanbf079c92018-07-03 17:51:28 -0700618
619void MLFunctionState::print() {
620 os << "mlfunc ";
621 // FIXME: should print argument names rather than just signature
MLIR Team4718bc92018-07-17 16:56:54 -0700622 printFunctionSignature(function, moduleState, os);
Tatiana Shpeismanbf079c92018-07-03 17:51:28 -0700623 os << " {\n";
Tatiana Shpeisman1bcfe982018-07-13 13:03:13 -0700624 print(function);
Tatiana Shpeismanbf079c92018-07-03 17:51:28 -0700625 os << " return\n";
626 os << "}\n\n";
627}
628
Tatiana Shpeisman1bcfe982018-07-13 13:03:13 -0700629void MLFunctionState::print(const StmtBlock *block) {
630 numSpaces += indentWidth;
Jacques Pienaarb020c542018-07-15 00:06:54 -0700631 for (auto &stmt : block->getStatements()) {
Tatiana Shpeisman1bcfe982018-07-13 13:03:13 -0700632 print(&stmt);
Jacques Pienaarb020c542018-07-15 00:06:54 -0700633 os << "\n";
634 }
Tatiana Shpeisman1bcfe982018-07-13 13:03:13 -0700635 numSpaces -= indentWidth;
636}
637
Tatiana Shpeismanbf079c92018-07-03 17:51:28 -0700638void MLFunctionState::print(const Statement *stmt) {
Tatiana Shpeismanbf079c92018-07-03 17:51:28 -0700639 switch (stmt->getKind()) {
Tatiana Shpeisman565b9642018-07-16 11:47:09 -0700640 case Statement::Kind::Operation:
641 return print(cast<OperationStmt>(stmt));
Tatiana Shpeismanbf079c92018-07-03 17:51:28 -0700642 case Statement::Kind::For:
643 return print(cast<ForStmt>(stmt));
644 case Statement::Kind::If:
645 return print(cast<IfStmt>(stmt));
Tatiana Shpeismanbf079c92018-07-03 17:51:28 -0700646 }
647}
648
MLIR Team54b55a22018-07-18 10:16:05 -0700649void MLFunctionState::print(const OperationStmt *stmt) { printOperation(stmt); }
Tatiana Shpeismanbf079c92018-07-03 17:51:28 -0700650
651void MLFunctionState::print(const ForStmt *stmt) {
Tatiana Shpeisman1da50c42018-07-19 09:52:39 -0700652 os.indent(numSpaces) << "for x = " << *stmt->getLowerBound();
653 os << " to " << *stmt->getUpperBound();
654 if (stmt->getStep()->getValue() != 1)
655 os << " step " << *stmt->getStep();
656
657 os << " {\n";
Tatiana Shpeisman1bcfe982018-07-13 13:03:13 -0700658 print(static_cast<const StmtBlock *>(stmt));
Tatiana Shpeisman565b9642018-07-16 11:47:09 -0700659 os.indent(numSpaces) << "}";
Tatiana Shpeismanbf079c92018-07-03 17:51:28 -0700660}
661
662void MLFunctionState::print(const IfStmt *stmt) {
Tatiana Shpeisman1bcfe982018-07-13 13:03:13 -0700663 os.indent(numSpaces) << "if () {\n";
664 print(stmt->getThenClause());
665 os.indent(numSpaces) << "}";
666 if (stmt->hasElseClause()) {
667 os << " else {\n";
668 print(stmt->getElseClause());
669 os.indent(numSpaces) << "}";
670 }
Tatiana Shpeismanbf079c92018-07-03 17:51:28 -0700671}
672
MLIR Team54b55a22018-07-18 10:16:05 -0700673void ModuleState::print(const MLFunction *fn) {
674 MLFunctionState state(fn, this, os);
MLIR Team4718bc92018-07-17 16:56:54 -0700675 state.print();
676}
677
Tatiana Shpeismanbf079c92018-07-03 17:51:28 -0700678//===----------------------------------------------------------------------===//
679// print and dump methods
680//===----------------------------------------------------------------------===//
Chris Lattnered65a732018-06-28 20:45:33 -0700681
MLIR Teamb61885d2018-07-18 16:29:21 -0700682void Attribute::print(raw_ostream &os) const {
683 ModuleState moduleState(os);
684 moduleState.print(this);
685}
686
687void Attribute::dump() const {
688 print(llvm::errs());
689}
690
MLIR Team4718bc92018-07-17 16:56:54 -0700691void Type::print(raw_ostream &os) const {
692 ModuleState moduleState(os);
693 moduleState.print(this);
694}
695
MLIR Team54b55a22018-07-18 10:16:05 -0700696void Type::dump() const { print(llvm::errs()); }
MLIR Team4718bc92018-07-17 16:56:54 -0700697
Chris Lattnered65a732018-06-28 20:45:33 -0700698void Instruction::print(raw_ostream &os) const {
MLIR Team4718bc92018-07-17 16:56:54 -0700699 ModuleState moduleState(os);
700 CFGFunctionState state(getFunction(), &moduleState, os);
Chris Lattner4c95a502018-06-23 16:03:42 -0700701 state.print(this);
702}
703
Chris Lattnered65a732018-06-28 20:45:33 -0700704void Instruction::dump() const {
Chris Lattner4c95a502018-06-23 16:03:42 -0700705 print(llvm::errs());
Jacques Pienaarb020c542018-07-15 00:06:54 -0700706 llvm::errs() << "\n";
Chris Lattner4c95a502018-06-23 16:03:42 -0700707}
708
MLIR Team718c82f2018-07-16 09:45:22 -0700709void AffineMap::dump() const {
710 print(llvm::errs());
711 llvm::errs() << "\n";
712}
Uday Bondhugula3934d4d2018-07-09 09:00:25 -0700713
Uday Bondhugula015cbb12018-07-03 20:16:08 -0700714void AffineExpr::dump() const {
715 print(llvm::errs());
716 llvm::errs() << "\n";
717}
718
Uday Bondhugula015cbb12018-07-03 20:16:08 -0700719void AffineSymbolExpr::print(raw_ostream &os) const {
Uday Bondhugulac1faf662018-07-19 14:08:50 -0700720 os << 's' << getPosition();
Uday Bondhugula015cbb12018-07-03 20:16:08 -0700721}
722
Uday Bondhugulac1faf662018-07-19 14:08:50 -0700723void AffineDimExpr::print(raw_ostream &os) const { os << 'd' << getPosition(); }
Uday Bondhugula015cbb12018-07-03 20:16:08 -0700724
725void AffineConstantExpr::print(raw_ostream &os) const { os << getValue(); }
726
Uday Bondhugulac1faf662018-07-19 14:08:50 -0700727static void printAdd(const AffineBinaryOpExpr *addExpr, raw_ostream &os) {
728 os << '(' << *addExpr->getLHS();
729
730 // Pretty print addition to a product that has a negative operand as a
731 // subtraction.
732 if (auto *rhs = dyn_cast<AffineBinaryOpExpr>(addExpr->getRHS())) {
733 if (rhs->getKind() == AffineExpr::Kind::Mul) {
734 if (auto *rrhs = dyn_cast<AffineConstantExpr>(rhs->getRHS())) {
735 if (rrhs->getValue() < 0) {
736 os << " - (" << *rhs->getLHS() << " * " << -rrhs->getValue() << "))";
737 return;
738 }
739 }
740 }
741 }
742
743 // Pretty print addition to a negative number as a subtraction.
744 if (auto *rhs = dyn_cast<AffineConstantExpr>(addExpr->getRHS())) {
745 if (rhs->getValue() < 0) {
746 os << " - " << -rhs->getValue() << ")";
747 return;
748 }
749 }
750
751 os << " + " << *addExpr->getRHS() << ")";
752}
753
754void AffineBinaryOpExpr::print(raw_ostream &os) const {
755 switch (getKind()) {
756 case Kind::Add:
757 return printAdd(this, os);
758 case Kind::Mul:
759 os << "(" << *getLHS() << " * " << *getRHS() << ")";
760 return;
761 case Kind::FloorDiv:
762 os << "(" << *getLHS() << " floordiv " << *getRHS() << ")";
763 return;
764 case Kind::CeilDiv:
765 os << "(" << *getLHS() << " ceildiv " << *getRHS() << ")";
766 return;
767 case Kind::Mod:
768 os << "(" << *getLHS() << " mod " << *getRHS() << ")";
769 return;
770 default:
771 llvm_unreachable("unexpected affine binary op expression");
772 }
773}
774
Uday Bondhugulafaf37dd2018-06-29 18:09:29 -0700775void AffineExpr::print(raw_ostream &os) const {
Uday Bondhugula015cbb12018-07-03 20:16:08 -0700776 switch (getKind()) {
777 case Kind::SymbolId:
778 return cast<AffineSymbolExpr>(this)->print(os);
779 case Kind::DimId:
780 return cast<AffineDimExpr>(this)->print(os);
781 case Kind::Constant:
782 return cast<AffineConstantExpr>(this)->print(os);
783 case Kind::Add:
Uday Bondhugula015cbb12018-07-03 20:16:08 -0700784 case Kind::Mul:
Uday Bondhugula015cbb12018-07-03 20:16:08 -0700785 case Kind::FloorDiv:
Uday Bondhugula015cbb12018-07-03 20:16:08 -0700786 case Kind::CeilDiv:
Uday Bondhugula015cbb12018-07-03 20:16:08 -0700787 case Kind::Mod:
Uday Bondhugulac1faf662018-07-19 14:08:50 -0700788 return cast<AffineBinaryOpExpr>(this)->print(os);
Uday Bondhugula015cbb12018-07-03 20:16:08 -0700789 }
Uday Bondhugulafaf37dd2018-06-29 18:09:29 -0700790}
791
792void AffineMap::print(raw_ostream &os) const {
Uday Bondhugula015cbb12018-07-03 20:16:08 -0700793 // Dimension identifiers.
794 os << "(";
MLIR Team54b55a22018-07-18 10:16:05 -0700795 for (int i = 0; i < (int)getNumDims() - 1; i++) os << "d" << i << ", ";
796 if (getNumDims() >= 1) os << "d" << getNumDims() - 1;
Uday Bondhugula015cbb12018-07-03 20:16:08 -0700797 os << ")";
798
799 // Symbolic identifiers.
800 if (getNumSymbols() >= 1) {
801 os << " [";
MLIR Team54b55a22018-07-18 10:16:05 -0700802 for (int i = 0; i < (int)getNumSymbols() - 1; i++) os << "s" << i << ", ";
803 if (getNumSymbols() >= 1) os << "s" << getNumSymbols() - 1;
Uday Bondhugula015cbb12018-07-03 20:16:08 -0700804 os << "]";
805 }
806
807 // AffineMap should have at least one result.
808 assert(!getResults().empty());
809 // Result affine expressions.
810 os << " -> (";
Chris Lattnerf8cce872018-07-20 09:28:54 -0700811 interleaveComma(os, getResults(), [&](AffineExpr *expr) { os << *expr; });
Uday Bondhugula0115dbb2018-07-11 21:31:07 -0700812 os << ")";
813
814 if (!isBounded()) {
Uday Bondhugula0115dbb2018-07-11 21:31:07 -0700815 return;
816 }
817
818 // Print range sizes for bounded affine maps.
819 os << " size (";
Chris Lattnerf8cce872018-07-20 09:28:54 -0700820 interleaveComma(os, getRangeSizes(), [&](AffineExpr *expr) { os << *expr; });
MLIR Team718c82f2018-07-16 09:45:22 -0700821 os << ")";
Uday Bondhugulafaf37dd2018-06-29 18:09:29 -0700822}
823
Chris Lattner4c95a502018-06-23 16:03:42 -0700824void BasicBlock::print(raw_ostream &os) const {
MLIR Team4718bc92018-07-17 16:56:54 -0700825 ModuleState moduleState(os);
826 CFGFunctionState state(getFunction(), &moduleState, os);
Chris Lattner4c95a502018-06-23 16:03:42 -0700827 state.print();
828}
829
MLIR Team54b55a22018-07-18 10:16:05 -0700830void BasicBlock::dump() const { print(llvm::errs()); }
Chris Lattner4c95a502018-06-23 16:03:42 -0700831
Tatiana Shpeismanbf079c92018-07-03 17:51:28 -0700832void Statement::print(raw_ostream &os) const {
MLIR Team4718bc92018-07-17 16:56:54 -0700833 ModuleState moduleState(os);
834 MLFunctionState state(getFunction(), &moduleState, os);
Tatiana Shpeismanbf079c92018-07-03 17:51:28 -0700835 state.print(this);
Tatiana Shpeismanc96b5872018-06-28 17:02:32 -0700836}
837
MLIR Team54b55a22018-07-18 10:16:05 -0700838void Statement::dump() const { print(llvm::errs()); }
Jacques Pienaarb020c542018-07-15 00:06:54 -0700839
Chris Lattner4c95a502018-06-23 16:03:42 -0700840void Function::print(raw_ostream &os) const {
841 switch (getKind()) {
MLIR Team54b55a22018-07-18 10:16:05 -0700842 case Kind::ExtFunc:
843 return cast<ExtFunction>(this)->print(os);
844 case Kind::CFGFunc:
845 return cast<CFGFunction>(this)->print(os);
846 case Kind::MLFunc:
847 return cast<MLFunction>(this)->print(os);
Chris Lattner4c95a502018-06-23 16:03:42 -0700848 }
849}
850
MLIR Team54b55a22018-07-18 10:16:05 -0700851void Function::dump() const { print(llvm::errs()); }
852
853void ExtFunction::print(raw_ostream &os) const {
854 ModuleState moduleState(os);
855 os << "extfunc ";
856 printFunctionSignature(this, &moduleState, os);
857 os << "\n";
Chris Lattner4c95a502018-06-23 16:03:42 -0700858}
859
860void CFGFunction::print(raw_ostream &os) const {
MLIR Team4718bc92018-07-17 16:56:54 -0700861 ModuleState moduleState(os);
862 CFGFunctionState state(this, &moduleState, os);
Chris Lattner4c95a502018-06-23 16:03:42 -0700863 state.print();
864}
865
Tatiana Shpeismanc96b5872018-06-28 17:02:32 -0700866void MLFunction::print(raw_ostream &os) const {
MLIR Team4718bc92018-07-17 16:56:54 -0700867 ModuleState moduleState(os);
868 MLFunctionState state(this, &moduleState, os);
Tatiana Shpeismanbf079c92018-07-03 17:51:28 -0700869 state.print();
Tatiana Shpeismanc96b5872018-06-28 17:02:32 -0700870}
871
Chris Lattner4c95a502018-06-23 16:03:42 -0700872void Module::print(raw_ostream &os) const {
MLIR Team4718bc92018-07-17 16:56:54 -0700873 ModuleState moduleState(os);
874 moduleState.initialize(this);
875 moduleState.print(this);
Chris Lattner4c95a502018-06-23 16:03:42 -0700876}
877
MLIR Team54b55a22018-07-18 10:16:05 -0700878void Module::dump() const { print(llvm::errs()); }