[MLIR] Support walks over regions and blocks

[frgossen]

Summary

Relands

• [MLIR] Support walks over regions and blocks (dbae3d50f114a8ec0a7c3211e3b1b9fb6ef22dbd)
• [MLIR] Use llvm::is_one_of in walk templates (56299b1e58bf3720dff2fe60163739ee1554a371)

[bkramer]

Can you add a note to the commit description that this is a reland of D90379?

mlir/include/mlir/Analysis/Liveness.h

public:
   /// Dumps the liveness information in a human readable format.
   void dump() const;
 
   /// Dumps the liveness information to the given stream.
   void print(raw_ostream &os) const;
 
 private:
   /// Initializes the internal mappings.
-  void build(MutableArrayRef<Region> regions);
+  void build();
 
 private:
   /// The operation this analysis was constructed from.
   Operation *operation;
 
   /// Maps blocks to internal liveness information.
   BlockMapT blockMapping;
 };

mlir/include/mlir/IR/Block.h

public:
   /// Walk the operations in the specified [begin, end) range of this block in
   /// postorder, calling the callback for each operation. This method is invoked
   /// for void return callbacks.
   /// See Operation::walk for more details.
   template <typename FnT, typename RetT = detail::walkResultType<FnT>>
   typename std::enable_if<std::is_same<RetT, void>::value, RetT>::type
   walk(Block::iterator begin, Block::iterator end, FnT &&callback) {
     for (auto &op : llvm::make_early_inc_range(llvm::make_range(begin, end)))
-      detail::walkOperations(&op, callback);
+      detail::walk(&op, callback);
   }
 
   /// Walk the operations in the specified [begin, end) range of this block in
   /// postorder, calling the callback for each operation. This method is invoked
   /// for interruptible callbacks.
   /// See Operation::walk for more details.
   template <typename FnT, typename RetT = detail::walkResultType<FnT>>
   typename std::enable_if<std::is_same<RetT, WalkResult>::value, RetT>::type
   walk(Block::iterator begin, Block::iterator end, FnT &&callback) {
     for (auto &op : llvm::make_early_inc_range(llvm::make_range(begin, end)))
-      if (detail::walkOperations(&op, callback).wasInterrupted())
+      if (detail::walk(&op, callback).wasInterrupted())
         return WalkResult::interrupt();
     return WalkResult::advance();
   }
 
   //===--------------------------------------------------------------------===//
   // Other
   //===--------------------------------------------------------------------===//
 

mlir/include/mlir/IR/Operation.h

public:
   ///     * op->walk([](... op) {
   ///         // Interrupt, i.e cancel, the walk based on some invariant.
   ///         if (some_invariant)
   ///           return WalkResult::interrupt();
   ///         return WalkResult::advance();
   ///       });
   template <typename FnT, typename RetT = detail::walkResultType<FnT>>
   RetT walk(FnT &&callback) {
-    return detail::walkOperations(this, std::forward<FnT>(callback));
+    return detail::walk(this, std::forward<FnT>(callback));
   }
 
   //===--------------------------------------------------------------------===//
   // Uses
   //===--------------------------------------------------------------------===//
 
   /// Drop all uses of results of this operation.
   void dropAllUses() {

mlir/include/mlir/IR/Visitors.h

 #include "mlir/Support/LLVM.h"
 #include "mlir/Support/LogicalResult.h"
 #include "llvm/ADT/STLExtras.h"
 
 namespace mlir {
 class Diagnostic;
 class InFlightDiagnostic;
 class Operation;
+class Block;
+class Region;
 
 /// A utility result that is used to signal if a walk method should be
 /// interrupted or advance.
 class WalkResult {
   enum ResultEnum { Interrupt, Advance } result;
 
 public:
   WalkResult(ResultEnum result) : result(result) {}

 Arg first_argument_type(Ret (F::*)(Arg) const);
 template <typename F>
 decltype(first_argument_type(&F::operator())) first_argument_type(F);
 
 /// Type definition of the first argument to the given callable 'T'.
 template <typename T>
 using first_argument = decltype(first_argument_type(std::declval<T>()));
 
-/// Walk all of the operations nested under and including the given operation.
-void walkOperations(Operation *op, function_ref<void(Operation *op)> callback);
-
-/// Walk all of the operations nested under and including the given operation.
-/// This methods walks operations until an interrupt result is returned by the
-/// callback.
-WalkResult walkOperations(Operation *op,
-                          function_ref<WalkResult(Operation *op)> callback);
+/// Walk all of the regions, blocks, or operations nested under (and including)
+/// the given operation.
+void walk(Operation *op, function_ref<void(Region *)> callback);
+void walk(Operation *op, function_ref<void(Block *)> callback);
+void walk(Operation *op, function_ref<void(Operation *)> callback);
+
+/// Walk all of the regions, blocks, or operations nested under (and including)
+/// the given operation. These functions walk until an interrupt result is
+/// returned by the callback.
+WalkResult walk(Operation *op, function_ref<WalkResult(Region *)> callback);
+WalkResult walk(Operation *op, function_ref<WalkResult(Block *)> callback);
+WalkResult walk(Operation *op, function_ref<WalkResult(Operation *)> callback);
 
 // Below are a set of functions to walk nested operations. Users should favor
 // the direct `walk` methods on the IR classes(Operation/Block/etc) over these
 // methods. They are also templated to allow for statically dispatching based
 // upon the type of the callback function.
 
-/// Walk all of the operations nested under and including the given operation.
-/// This method is selected for callbacks that operate on Operation*.
+/// Walk all of the regions, blocks, or operations nested under (and including)
+/// the given operation. This method is selected for callbacks that operate on
+/// Region*, Block*, and Operation*.
 ///
 /// Example:
+///   op->walk([](Region *r) { ... });
+///   op->walk([](Block *b) { ... });
 ///   op->walk([](Operation *op) { ... });
 template <
     typename FuncTy, typename ArgT = detail::first_argument<FuncTy>,
     typename RetT = decltype(std::declval<FuncTy>()(std::declval<ArgT>()))>
-typename std::enable_if<std::is_same<ArgT, Operation *>::value, RetT>::type
-walkOperations(Operation *op, FuncTy &&callback) {
-  return detail::walkOperations(op, function_ref<RetT(ArgT)>(callback));
+typename std::enable_if<
+    llvm::is_one_of<ArgT, Operation *, Region *, Block *>::value, RetT>::type
+walk(Operation *op, FuncTy &&callback) {
+  return walk(op, function_ref<RetT(ArgT)>(callback));
 }
 
 /// Walk all of the operations of type 'ArgT' nested under and including the
 /// given operation. This method is selected for void returning callbacks that
 /// operate on a specific derived operation type.
 ///
 /// Example:
 ///   op->walk([](ReturnOp op) { ... });
 template <
     typename FuncTy, typename ArgT = detail::first_argument<FuncTy>,
     typename RetT = decltype(std::declval<FuncTy>()(std::declval<ArgT>()))>
-typename std::enable_if<!std::is_same<ArgT, Operation *>::value &&
-                            std::is_same<RetT, void>::value,
-                        RetT>::type
-walkOperations(Operation *op, FuncTy &&callback) {
+typename std::enable_if<
+    !llvm::is_one_of<ArgT, Operation *, Region *, Block *>::value &&
+walk(Operation *op, FuncTy &&callback) {
   auto wrapperFn = [&](Operation *op) {
     if (auto derivedOp = dyn_cast<ArgT>(op))
       callback(derivedOp);
   };
-  return detail::walkOperations(op, function_ref<RetT(Operation *)>(wrapperFn));
+  return detail::walk(op, function_ref<RetT(Operation *)>(wrapperFn));
 }
 
 /// Walk all of the operations of type 'ArgT' nested under and including the
 /// given operation. This method is selected for WalkReturn returning
 /// interruptible callbacks that operate on a specific derived operation type.
 ///
 /// Example:
 ///   op->walk([](ReturnOp op) {
 ///     if (some_invariant)
 ///       return WalkResult::interrupt();
 ///     return WalkResult::advance();
 ///   });
 template <
     typename FuncTy, typename ArgT = detail::first_argument<FuncTy>,
     typename RetT = decltype(std::declval<FuncTy>()(std::declval<ArgT>()))>
-typename std::enable_if<!std::is_same<ArgT, Operation *>::value &&
-                            std::is_same<RetT, WalkResult>::value,
-                        RetT>::type
-walkOperations(Operation *op, FuncTy &&callback) {
+typename std::enable_if<
+    !llvm::is_one_of<ArgT, Operation *, Region *, Block *>::value &&
+walk(Operation *op, FuncTy &&callback) {
   auto wrapperFn = [&](Operation *op) {
     if (auto derivedOp = dyn_cast<ArgT>(op))
       return callback(derivedOp);
     return WalkResult::advance();
   };
-  return detail::walkOperations(op, function_ref<RetT(Operation *)>(wrapperFn));
+  return detail::walk(op, function_ref<RetT(Operation *)>(wrapperFn));
 }
 
 /// Utility to provide the return type of a templated walk method.
 template <typename FnT>
-using walkResultType = decltype(walkOperations(nullptr, std::declval<FnT>()));
+using walkResultType = decltype(walk(nullptr, std::declval<FnT>()));
 } // end namespace detail
 
 } // namespace mlir
 
 #endif

mlir/lib/Analysis/Liveness.cpp

struct BlockInfoBuilder {
   /// The set of all defined values.
   ValueSetT defValues;
 
   /// The set of all used values.
   ValueSetT useValues;
 };
 } // namespace
 
-/// Walks all regions (including nested regions recursively) and invokes the
-/// given function for every block.
-template <typename FuncT>
-static void walkRegions(MutableArrayRef<Region> regions, const FuncT &func) {
-  for (Region &region : regions)
-    for (Block &block : region) {
-      func(block);
-
-      // Traverse all nested regions.
-      for (Operation &operation : block)
-        walkRegions(operation.getRegions(), func);
-    }
-}
-
 /// Builds the internal liveness block mapping.
-static void buildBlockMapping(MutableArrayRef<Region> regions,
+static void buildBlockMapping(Operation *operation,
                               DenseMap<Block *, BlockInfoBuilder> &builders) {
   llvm::SetVector<Block *> toProcess;
 
-  walkRegions(regions, [&](Block &block) {
+  operation->walk([&](Block *block) {
     BlockInfoBuilder &builder =
-        builders.try_emplace(&block, &block).first->second;
+        builders.try_emplace(block, block).first->second;
 
     if (builder.updateLiveIn())
-      toProcess.insert(block.pred_begin(), block.pred_end());
+      toProcess.insert(block->pred_begin(), block->pred_end());
   });
 
   // Propagate the in and out-value sets (fixpoint iteration)
   while (!toProcess.empty()) {
     Block *current = toProcess.pop_back_val();
     BlockInfoBuilder &builder = builders[current];
 
     // Update the current out values.
     builder.updateLiveOut(builders);
 
     // Compute (potentially) updated live in values.
     if (builder.updateLiveIn())
       toProcess.insert(current->pred_begin(), current->pred_end());
   }
 }
 
 //===----------------------------------------------------------------------===//
 // Liveness
 //===----------------------------------------------------------------------===//
 
 /// Creates a new Liveness analysis that computes liveness information for all
 /// associated regions.
-Liveness::Liveness(Operation *op) : operation(op) { build(op->getRegions()); }
+Liveness::Liveness(Operation *op) : operation(op) { build(); }
 
 /// Initializes the internal mappings.
-void Liveness::build(MutableArrayRef<Region> regions) {
+void Liveness::build() {
 
   // Build internal block mapping.
   DenseMap<Block *, BlockInfoBuilder> builders;
-  buildBlockMapping(regions, builders);
+  buildBlockMapping(operation, builders);
 
   // Store internal block data.
   for (auto &entry : builders) {
     BlockInfoBuilder &builder = entry.second;
     LivenessBlockInfo &info = blockMapping[entry.first];
 
     info.block = builder.block;
     info.inValues = std::move(builder.inValues);

 /// Dumps the liveness information to the given stream.
 void Liveness::print(raw_ostream &os) const {
   os << "// ---- Liveness -----\n";
 
   // Builds unique block/value mappings for testing purposes.
   DenseMap<Block *, size_t> blockIds;
   DenseMap<Operation *, size_t> operationIds;
   DenseMap<Value, size_t> valueIds;
-  walkRegions(operation->getRegions(), [&](Block &block) {
-    blockIds.insert({&block, blockIds.size()});
-    for (BlockArgument argument : block.getArguments())
+  operation->walk([&](Block *block) {
+    blockIds.insert({block, blockIds.size()});
+    for (BlockArgument argument : block->getArguments())
       valueIds.insert({argument, valueIds.size()});
-    for (Operation &operation : block) {
+    for (Operation &operation : *block) {
       operationIds.insert({&operation, operationIds.size()});
       for (Value result : operation.getResults())
         valueIds.insert({result, valueIds.size()});
     }
   });
 
   // Local printing helpers
   auto printValueRef = [&](Value value) {
std::sort(orderedValues.begin(), orderedValues.end(),
               [&](Value left, Value right) {
                 return valueIds[left] < valueIds[right];
               });
     for (Value value : orderedValues)
       printValueRef(value);
   };
 
   // Dump information about in and out values.
-  walkRegions(operation->getRegions(), [&](Block &block) {
-    os << "// - Block: " << blockIds[&block] << "\n";
-    auto liveness = getLiveness(&block);
+  operation->walk([&](Block *block) {
+    os << "// - Block: " << blockIds[block] << "\n";
+    const auto *liveness = getLiveness(block);
     os << "// --- LiveIn: ";
     printValueRefs(liveness->inValues);
     os << "\n// --- LiveOut: ";
     printValueRefs(liveness->outValues);
     os << "\n";
 
     // Print liveness intervals.
     os << "// --- BeginLiveness";
-    for (Operation &op : block) {
+    for (Operation &op : *block) {
       if (op.getNumResults() < 1)
         continue;
       os << "\n";
       for (Value result : op.getResults()) {
         os << "// ";
         printValueRef(result);
         os << ":";
         auto liveOperations = resolveLiveness(result);

mlir/lib/IR/Visitors.cpp

 //===- Visitors.cpp - MLIR Visitor Utilities ------------------------------===//
 //
 // Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
 // See https://llvm.org/LICENSE.txt for license information.
 // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
 //
 //===----------------------------------------------------------------------===//
 
 #include "mlir/IR/Visitors.h"
 #include "mlir/IR/Operation.h"
 
 using namespace mlir;
 
-/// Walk all of the operations nested under and including the given operations.
-void detail::walkOperations(Operation *op,
-                            function_ref<void(Operation *op)> callback) {
+/// Walk all of the regions/blocks/operations nested under and including the
+/// given operation.
+void detail::walk(Operation *op, function_ref<void(Region *)> callback) {
+  for (auto &region : op->getRegions()) {
+    callback(&region);
+    for (auto &block : region) {
+      for (auto &nestedOp : block)
+        walk(&nestedOp, callback);
+    }
+  }
+}
+
+void detail::walk(Operation *op, function_ref<void(Block *)> callback) {
+  for (auto &region : op->getRegions()) {
+    for (auto &block : region) {
+      callback(&block);
+      for (auto &nestedOp : block)
+        walk(&nestedOp, callback);
+    }
+  }
+}
+
+void detail::walk(Operation *op, function_ref<void(Operation *op)> callback) {
   // TODO: This walk should be iterative over the operations.
-  for (auto &region : op->getRegions())
-    for (auto &block : region)
+  for (auto &region : op->getRegions()) {
+    for (auto &block : region) {
       // Early increment here in the case where the operation is erased.
       for (auto &nestedOp : llvm::make_early_inc_range(block))
-        walkOperations(&nestedOp, callback);
-
+        walk(&nestedOp, callback);
+    }
+  }
   callback(op);
 }
 
-/// Walk all of the operations nested under and including the given operations.
-/// This methods walks operations until an interrupt signal is received.
-WalkResult
-detail::walkOperations(Operation *op,
+/// Walk all of the regions/blocks/operations nested under and including the
+/// given operation. These functions walk operations until an interrupt result
+/// is returned by the callback.
+WalkResult detail::walk(Operation *op,
+                        function_ref<WalkResult(Region *op)> callback) {
+  for (auto &region : op->getRegions()) {
+    if (callback(&region).wasInterrupted())
+      return WalkResult::interrupt();
+    for (auto &block : region) {
+      for (auto &nestedOp : block)
+        walk(&nestedOp, callback);
+    }
+  }
+  return WalkResult::advance();
+}
+
+WalkResult detail::walk(Operation *op,
+                        function_ref<WalkResult(Block *op)> callback) {
+  for (auto &region : op->getRegions()) {
+    for (auto &block : region) {
+      if (callback(&block).wasInterrupted())
+        return WalkResult::interrupt();
+      for (auto &nestedOp : block)
+        walk(&nestedOp, callback);
+    }
+  }
+  return WalkResult::advance();
+}
+
+WalkResult detail::walk(Operation *op,
                         function_ref<WalkResult(Operation *op)> callback) {
   // TODO: This walk should be iterative over the operations.
   for (auto &region : op->getRegions()) {
     for (auto &block : region) {
       // Early increment here in the case where the operation is erased.
-      for (auto &nestedOp : llvm::make_early_inc_range(block))
-        if (walkOperations(&nestedOp, callback).wasInterrupted())
+      for (auto &nestedOp : llvm::make_early_inc_range(block)) {
+        if (walk(&nestedOp, callback).wasInterrupted())
           return WalkResult::interrupt();
       }
     }
+  }
   return callback(op);
 }