mlir/docs/Tutorials/UnderstandingTheIRStructure.md

   1 # Understanding the IR Structure
   2
   3 The MLIR Language Reference describes the
   4 [High Level Structure](../LangRef.md/#high-level-structure), this document
   5 illustrates this structure through examples, and introduces at the same time the
   6 C++ APIs involved in manipulating it.
   7
   8 We will implement a [pass](../PassManagement.md/#operation-pass) that traverses any
   9 MLIR input and prints the entity inside the IR. A pass (or in general almost any
  10 piece of IR) is always rooted with an operation. Most of the time the top-level
  11 operation is a `ModuleOp`, the MLIR `PassManager` is actually limited to
  12 operation on a top-level `ModuleOp`. As such a pass starts with an operation,
  13 and so will our traversal:
  14
  15 ```
  16   void runOnOperation() override {
  17     Operation *op = getOperation();
  18     resetIndent();
  19     printOperation(op);
  20   }
  21 ```
  22
  23 ## Traversing the IR Nesting
  24
  25 The IR is recursively nested, an `Operation` can have one or multiple nested
  26 `Region`s, each of which is actually a list of `Blocks`, each of which itself
  27 wraps a list of `Operation`s. Our traversal will follow this structure with
  28 three methods: `printOperation()`, `printRegion()`, and `printBlock()`.
  29
  30 The first method inspects the properties of an operation, before iterating on
  31 the nested regions and print them individually:
  32
  33 ```c++
  34   void printOperation(Operation *op) {
  35     // Print the operation itself and some of its properties
  36     printIndent() << "visiting op: '" << op->getName() << "' with "
  37                   << op->getNumOperands() << " operands and "
  38                   << op->getNumResults() << " results\n";
  39     // Print the operation attributes
  40     if (!op->getAttrs().empty()) {
  41       printIndent() << op->getAttrs().size() << " attributes:\n";
  42       for (NamedAttribute attr : op->getAttrs())
  43         printIndent() << " - '" << attr.getName() << "' : '"
  44                       << attr.getValue() << "'\n";
  45     }
  46
  47     // Recurse into each of the regions attached to the operation.
  48     printIndent() << " " << op->getNumRegions() << " nested regions:\n";
  49     auto indent = pushIndent();
  50     for (Region &region : op->getRegions())
  51       printRegion(region);
  52   }
  53 ```
  54
  55 A `Region` does not hold anything other than a list of `Block`s:
  56
  57 ```c++
  58   void printRegion(Region &region) {
  59     // A region does not hold anything by itself other than a list of blocks.
  60     printIndent() << "Region with " << region.getBlocks().size()
  61                   << " blocks:\n";
  62     auto indent = pushIndent();
  63     for (Block &block : region.getBlocks())
  64       printBlock(block);
  65   }
  66 ```
  67
  68 Finally, a `Block` has a list of arguments, and holds a list of `Operation`s:
  69
  70 ```c++
  71   void printBlock(Block &block) {
  72     // Print the block intrinsics properties (basically: argument list)
  73     printIndent()
  74         << "Block with " << block.getNumArguments() << " arguments, "
  75         << block.getNumSuccessors()
  76         << " successors, and "
  77         // Note, this `.size()` is traversing a linked-list and is O(n).
  78         << block.getOperations().size() << " operations\n";
  79
  80     // A block main role is to hold a list of Operations: let's recurse into
  81     // printing each operation.
  82     auto indent = pushIndent();
  83     for (Operation &op : block.getOperations())
  84       printOperation(&op);
  85   }
  86 ```
  87
  88 The code for the pass is available
  89 [here in the repo](https://github.com/llvm/llvm-project/blob/main/mlir/test/lib/IR/TestPrintNesting.cpp)
  90 and can be exercised with `mlir-opt -test-print-nesting`.
  91
  92 ### Example
  93
  94 The Pass introduced in the previous section can be applied on the following IR
  95 with `mlir-opt -test-print-nesting -allow-unregistered-dialect
  96 llvm-project/mlir/test/IR/print-ir-nesting.mlir`:
  97
  98 ```mlir
  99 "builtin.module"() ( {
 100   %results:4 = "dialect.op1"() {"attribute name" = 42 : i32} : () -> (i1, i16, i32, i64)
 101   "dialect.op2"() ( {
 102     "dialect.innerop1"(%results#0, %results#1) : (i1, i16) -> ()
 103   },  {
 104     "dialect.innerop2"() : () -> ()
 105     "dialect.innerop3"(%results#0, %results#2, %results#3)[^bb1, ^bb2] : (i1, i32, i64) -> ()
 106   ^bb1(%1: i32):  // pred: ^bb0
 107     "dialect.innerop4"() : () -> ()
 108     "dialect.innerop5"() : () -> ()
 109   ^bb2(%2: i64):  // pred: ^bb0
 110     "dialect.innerop6"() : () -> ()
 111     "dialect.innerop7"() : () -> ()
 112   }) {"other attribute" = 42 : i64} : () -> ()
 113 }) : () -> ()
 114 ```
 115
 116 And will yield the following output:
 117
 118 ```
 119 visiting op: 'builtin.module' with 0 operands and 0 results
 120  1 nested regions:
 121   Region with 1 blocks:
 122     Block with 0 arguments, 0 successors, and 2 operations
 123       visiting op: 'dialect.op1' with 0 operands and 4 results
 124       1 attributes:
 125        - 'attribute name' : '42 : i32'
 126        0 nested regions:
 127       visiting op: 'dialect.op2' with 0 operands and 0 results
 128       1 attributes:
 129        - 'other attribute' : '42 : i64'
 130        2 nested regions:
 131         Region with 1 blocks:
 132           Block with 0 arguments, 0 successors, and 1 operations
 133             visiting op: 'dialect.innerop1' with 2 operands and 0 results
 134              0 nested regions:
 135         Region with 3 blocks:
 136           Block with 0 arguments, 2 successors, and 2 operations
 137             visiting op: 'dialect.innerop2' with 0 operands and 0 results
 138              0 nested regions:
 139             visiting op: 'dialect.innerop3' with 3 operands and 0 results
 140              0 nested regions:
 141           Block with 1 arguments, 0 successors, and 2 operations
 142             visiting op: 'dialect.innerop4' with 0 operands and 0 results
 143              0 nested regions:
 144             visiting op: 'dialect.innerop5' with 0 operands and 0 results
 145              0 nested regions:
 146           Block with 1 arguments, 0 successors, and 2 operations
 147             visiting op: 'dialect.innerop6' with 0 operands and 0 results
 148              0 nested regions:
 149             visiting op: 'dialect.innerop7' with 0 operands and 0 results
 150              0 nested regions:
 151 ```
 152
 153 ## Other IR Traversal Methods
 154
 155 In many cases, unwrapping the recursive structure of the IR is cumbersome and
 156 you may be interested in using other helpers.
 157
 158 ### Filtered iterator: `getOps<OpTy>()`
 159
 160 For example the `Block` class exposes a convenient templated method
 161 `getOps<OpTy>()` that provided a filtered iterator. Here is an example:
 162
 163 ```c++
 164   auto varOps = entryBlock.getOps<spirv::GlobalVariableOp>();
 165   for (spirv::GlobalVariableOp gvOp : varOps) {
 166      // process each GlobalVariable Operation in the block.
 167      ...
 168   }
 169 ```
 170
 171 Similarly, the `Region` class exposes the same `getOps` method that will iterate
 172 on all the blocks in the region.
 173
 174 ### Walkers
 175
 176 The `getOps<OpTy>()` is useful to iterate on some Operations immediately listed
 177 inside a single block (or a single region), however it is frequently interesting
 178 to traverse the IR in a nested fashion. To this end MLIR exposes the `walk()`
 179 helper on `Operation`, `Block`, and `Region`. This helper takes a single
 180 argument: a callback method that will be invoked for every operation recursively
 181 nested under the provided entity.
 182
 183 ```c++
 184   // Recursively traverse all the regions and blocks nested inside the function
 185   // and apply the callback on every single operation in post-order.
 186   getFunction().walk([&](mlir::Operation *op) {
 187     // process Operation `op`.
 188   });
 189 ```
 190
 191 The provided callback can be specialized to filter on a particular type of
 192 Operation, for example the following will apply the callback only on `LinalgOp`
 193 operations nested inside the function:
 194
 195 ```c++
 196   getFunction().walk([](LinalgOp linalgOp) {
 197     // process LinalgOp `linalgOp`.
 198   });
 199 ```
 200
 201 Finally, the callback can optionally stop the walk by returning a
 202 `WalkResult::interrupt()` value. For example the following walk will find all
 203 `AllocOp` nested inside the function and interrupt the traversal if one of them
 204 does not satisfy a criteria:
 205
 206 ```c++
 207   WalkResult result = getFunction().walk([&](AllocOp allocOp) {
 208     if (!isValid(allocOp))
 209       return WalkResult::interrupt();
 210     return WalkResult::advance();
 211   });
 212   if (result.wasInterrupted())
 213     // One alloc wasn't matching.
 214     ...
 215 ```
 216
 217 ## Traversing the def-use chains
 218
 219 Another relationship in the IR is the one that links a `Value` with its users.
 220 As defined in the
 221 [language reference](../LangRef.md/#high-level-structure),
 222 each Value is either a `BlockArgument` or the result of exactly one `Operation`
 223 (an `Operation` can have multiple results, each of them is a separate `Value`).
 224 The users of a `Value` are `Operation`s, through their arguments: each
 225 `Operation` argument references a single `Value`.
 226
 227 Here is a code sample that inspects the operands of an `Operation` and prints
 228 some information about them:
 229
 230 ```c++
 231   // Print information about the producer of each of the operands.
 232   for (Value operand : op->getOperands()) {
 233     if (Operation *producer = operand.getDefiningOp()) {
 234       llvm::outs() << "  - Operand produced by operation '"
 235                    << producer->getName() << "'\n";
 236     } else {
 237       // If there is no defining op, the Value is necessarily a Block
 238       // argument.
 239       auto blockArg = operand.cast<BlockArgument>();
 240       llvm::outs() << "  - Operand produced by Block argument, number "
 241                    << blockArg.getArgNumber() << "\n";
 242     }
 243   }
 244 ```
 245
 246 Similarly, the following code sample iterates through the result `Value`s
 247 produced by an `Operation` and for each result will iterate the users of these
 248 results and print informations about them:
 249
 250 ```c++
 251   // Print information about the user of each of the result.
 252   llvm::outs() << "Has " << op->getNumResults() << " results:\n";
 253   for (auto indexedResult : llvm::enumerate(op->getResults())) {
 254     Value result = indexedResult.value();
 255     llvm::outs() << "  - Result " << indexedResult.index();
 256     if (result.use_empty()) {
 257       llvm::outs() << " has no uses\n";
 258       continue;
 259     }
 260     if (result.hasOneUse()) {
 261       llvm::outs() << " has a single use: ";
 262     } else {
 263       llvm::outs() << " has "
 264                    << std::distance(result.getUses().begin(),
 265                                     result.getUses().end())
 266                    << " uses:\n";
 267     }
 268     for (Operation *userOp : result.getUsers()) {
 269       llvm::outs() << "    - " << userOp->getName() << "\n";
 270     }
 271   }
 272 ```
 273
 274 The illustrating code for this pass is available
 275 [here in the repo](https://github.com/llvm/llvm-project/blob/main/mlir/test/lib/IR/TestPrintDefUse.cpp)
 276 and can be exercised with `mlir-opt -test-print-defuse`.
 277
 278 The chaining of `Value`s and their uses can be viewed as following:
 279
 280 ![Index Map Example](/includes/img/DefUseChains.svg)
 281
 282 The uses of a `Value` (`OpOperand` or `BlockOperand`) are also chained in a
 283 doubly linked-list, which is particularly useful when replacing all uses of a
 284 `Value` with a new one ("RAUW"):
 285
 286 ![Index Map Example](/includes/img/Use-list.svg)