1 //===- Visitors.cpp - MLIR Visitor Utilities ------------------------------===//
2 //
3 // Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
4 // See https://llvm.org/LICENSE.txt for license information.
5 // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
6 //
7 //===----------------------------------------------------------------------===//
8 
9 #include "mlir/IR/Visitors.h"
10 #include "mlir/IR/Operation.h"
11 
12 using namespace mlir;
13 
14 /// Walk all of the regions/blocks/operations nested under and including the
15 /// given operation. Regions, blocks and operations at the same nesting level
16 /// are visited in lexicographical order. The walk order for enclosing regions,
17 /// blocks and operations with respect to their nested ones is specified by
18 /// 'order'. These methods are invoked for void-returning callbacks. A callback
19 /// on a block or operation is allowed to erase that block or operation only if
20 /// the walk is in post-order. See non-void method for pre-order erasure.
21 void detail::walk(Operation *op, function_ref<void(Region *)> callback,
22                   WalkOrder order) {
23   // We don't use early increment for regions because they can't be erased from
24   // a callback.
25   for (auto &region : op->getRegions()) {
26     if (order == WalkOrder::PreOrder)
27       callback(&region);
28     for (auto &block : region) {
29       for (auto &nestedOp : block)
30         walk(&nestedOp, callback, order);
31     }
32     if (order == WalkOrder::PostOrder)
33       callback(&region);
34   }
35 }
36 
37 void detail::walk(Operation *op, function_ref<void(Block *)> callback,
38                   WalkOrder order) {
39   for (auto &region : op->getRegions()) {
40     // Early increment here in the case where the block is erased.
41     for (auto &block : llvm::make_early_inc_range(region)) {
42       if (order == WalkOrder::PreOrder)
43         callback(&block);
44       for (auto &nestedOp : block)
45         walk(&nestedOp, callback, order);
46       if (order == WalkOrder::PostOrder)
47         callback(&block);
48     }
49   }
50 }
51 
52 void detail::walk(Operation *op, function_ref<void(Operation *)> callback,
53                   WalkOrder order) {
54   if (order == WalkOrder::PreOrder)
55     callback(op);
56 
57   // TODO: This walk should be iterative over the operations.
58   for (auto &region : op->getRegions()) {
59     for (auto &block : region) {
60       // Early increment here in the case where the operation is erased.
61       for (auto &nestedOp : llvm::make_early_inc_range(block))
62         walk(&nestedOp, callback, order);
63     }
64   }
65 
66   if (order == WalkOrder::PostOrder)
67     callback(op);
68 }
69 
70 /// Walk all of the regions/blocks/operations nested under and including the
71 /// given operation. These functions walk operations until an interrupt result
72 /// is returned by the callback. Walks on regions, blocks and operations may
73 /// also be skipped if the callback returns a skip result. Regions, blocks and
74 /// operations at the same nesting level are visited in lexicographical order.
75 /// The walk order for enclosing regions, blocks and operations with respect to
76 /// their nested ones is specified by 'order'. A callback on a block or
77 /// operation is allowed to erase that block or operation if either:
78 ///   * the walk is in post-order, or
79 ///   * the walk is in pre-order and the walk is skipped after the erasure.
80 WalkResult detail::walk(Operation *op,
81                         function_ref<WalkResult(Region *)> callback,
82                         WalkOrder order) {
83   // We don't use early increment for regions because they can't be erased from
84   // a callback.
85   for (auto &region : op->getRegions()) {
86     if (order == WalkOrder::PreOrder) {
87       WalkResult result = callback(&region);
88       if (result.wasSkipped())
89         continue;
90       if (result.wasInterrupted())
91         return WalkResult::interrupt();
92     }
93     for (auto &block : region) {
94       for (auto &nestedOp : block)
95         walk(&nestedOp, callback, order);
96     }
97     if (order == WalkOrder::PostOrder) {
98       if (callback(&region).wasInterrupted())
99         return WalkResult::interrupt();
100       // We don't check if this region was skipped because its walk already
101       // finished and the walk will continue with the next region.
102     }
103   }
104   return WalkResult::advance();
105 }
106 
107 WalkResult detail::walk(Operation *op,
108                         function_ref<WalkResult(Block *)> callback,
109                         WalkOrder order) {
110   for (auto &region : op->getRegions()) {
111     // Early increment here in the case where the block is erased.
112     for (auto &block : llvm::make_early_inc_range(region)) {
113       if (order == WalkOrder::PreOrder) {
114         WalkResult result = callback(&block);
115         if (result.wasSkipped())
116           continue;
117         if (result.wasInterrupted())
118           return WalkResult::interrupt();
119       }
120       for (auto &nestedOp : block)
121         walk(&nestedOp, callback, order);
122       if (order == WalkOrder::PostOrder) {
123         if (callback(&block).wasInterrupted())
124           return WalkResult::interrupt();
125         // We don't check if this block was skipped because its walk already
126         // finished and the walk will continue with the next block.
127       }
128     }
129   }
130   return WalkResult::advance();
131 }
132 
133 WalkResult detail::walk(Operation *op,
134                         function_ref<WalkResult(Operation *)> callback,
135                         WalkOrder order) {
136   if (order == WalkOrder::PreOrder) {
137     WalkResult result = callback(op);
138     // If skipped, caller will continue the walk on the next operation.
139     if (result.wasSkipped())
140       return WalkResult::advance();
141     if (result.wasInterrupted())
142       return WalkResult::interrupt();
143   }
144 
145   // TODO: This walk should be iterative over the operations.
146   for (auto &region : op->getRegions()) {
147     for (auto &block : region) {
148       // Early increment here in the case where the operation is erased.
149       for (auto &nestedOp : llvm::make_early_inc_range(block)) {
150         if (walk(&nestedOp, callback, order).wasInterrupted())
151           return WalkResult::interrupt();
152       }
153     }
154   }
155 
156   if (order == WalkOrder::PostOrder)
157     return callback(op);
158   return WalkResult::advance();
159 }
160