[MLIR] Support walks over regions and blocks

[frgossen]

Summary

Add specializations for walk to allow traversal of regions and blocks.

[jpienaar]

mlir/include/mlir/IR/Visitors.h

68

Shouldn't this be walkRegions or walkRegionsOfOperation? (This also feels like it may deserve a different comment as the operations one does invoke the callback on the operation, but with regions and blocks it can't)

[frgossen]

Fix clang-tidy

[frgossen]

Address comment

mlir/include/mlir/IR/Visitors.h

68

Different names for regions, blocks, and operations would mean to also split the template below into 3 variants.
I renamed it to walk (which is also the name used when using this through the IR classes) and tried to clarify in the comments.

[rriddle]

mlir/include/mlir/IR/Visitors.h

100

Use llvm::is_one_of here.

[frgossen]

mlir/include/mlir/IR/Visitors.h

100

See https://reviews.llvm.org/D90449

[stella.stamenova]

We are seeing a failure in the mlir build on Ubuntu with gcc/g++ since these two changes went in:

FAILED: tools/mlir/lib/Dialect/GPU/CMakeFiles/obj.MLIRGPU.dir/Transforms/AsyncRegionRewriter.cpp.o 
/usr/bin/c++ -DBUILD_EXAMPLES -DMLIR_CUDA_CONVERSIONS_ENABLED=0 -DMLIR_ROCM_CONVERSIONS_ENABLED=0 -D_DEBUG -D_GNU_SOURCE -D__STDC_CONSTANT_MACROS -D__STDC_FORMAT_MACROS -D__STDC_LIMIT_MACROS -Itools/mlir/lib/Dialect/GPU -I/vstsdrive/_work/20/s/llvm-project/mlir/lib/Dialect/GPU -Iinclude -I/vstsdrive/_work/20/s/llvm-project/llvm/include -I/vstsdrive/_work/20/s/llvm-project/mlir/include -Itools/mlir/include -fPIC -fvisibility-inlines-hidden -Werror=date-time -Wall -Wextra -Wno-unused-parameter -Wwrite-strings -Wcast-qual -Wno-missing-field-initializers -pedantic -Wno-long-long -Wimplicit-fallthrough -Wno-maybe-uninitialized -Wno-noexcept-type -Wdelete-non-virtual-dtor -Wno-comment -fdiagnostics-color -g  -fno-exceptions -std=c++14 -MD -MT tools/mlir/lib/Dialect/GPU/CMakeFiles/obj.MLIRGPU.dir/Transforms/AsyncRegionRewriter.cpp.o -MF tools/mlir/lib/Dialect/GPU/CMakeFiles/obj.MLIRGPU.dir/Transforms/AsyncRegionRewriter.cpp.o.d -o tools/mlir/lib/Dialect/GPU/CMakeFiles/obj.MLIRGPU.dir/Transforms/AsyncRegionRewriter.cpp.o -c /vstsdrive/_work/20/s/llvm-project/mlir/lib/Dialect/GPU/Transforms/AsyncRegionRewriter.cpp
/vstsdrive/_work/20/s/llvm-project/mlir/lib/Dialect/GPU/Transforms/AsyncRegionRewriter.cpp: In member function ‘virtual void {anonymous}::GpuAsyncRegionPass::runOnFunction()’:
/vstsdrive/_work/20/s/llvm-project/mlir/lib/Dialect/GPU/Transforms/AsyncRegionRewriter.cpp:113:16: internal compiler error: in replace_placeholders_r, at cp/tree.c:2804
   if (getFunction()
       ~~~~~~~~~~~~~
           .getRegion()
           ~~~~~~~~~~~~
           .walk(Callback{OpBuilder(&getContext())})
           ~~~~~^~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
Please submit a full bug report,
with preprocessed source if appropriate.
See <file:///usr/share/doc/gcc-7/README.Bugs> for instructions.

[frgossen]

Unfortunately, I could not reproduce this failure,
I tried with gcc/g++ (Debian 10.2.0-9) 10.2.0.

[stella.stamenova]

In D90379#2367830, @frgossen wrote:

Unfortunately, I could not reproduce this failure,
I tried with gcc/g++ (Debian 10.2.0-9) 10.2.0.

This reproduces reliably on Ubuntu 18.04.5 LTS with gcc/g++ 7.5.0

[stella.stamenova]

@frgossen , Thanks! Let me know if you need any help investigating.

[bkramer]

Building with GCC7 is fixed by 5ece2115d85326e3b24336a7ccb4c191932ccb4a, but looks like this was reverted already.

[frgossen]

Yes, I reverted it yesterday as it caused problems with the integrate. I will re-land this when I figure out how to install an old GCC version.

[bkramer]

In D90379#2370654, @frgossen wrote:

Yes, I reverted it yesterday as it caused problems with the integrate. I will re-land this when I figure out how to install an old GCC version.

Sorry, what I was trying to say was that 5ece2115d85326e3b24336a7ccb4c191932ccb4a fixed the GCC7 crash. Your change seems mostly unrelated and should be safe to reland.

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/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.
+/// Walk all of the regions/blocks/operations nested under and including the
+/// given operation.
+void walkOperations(Operation *op, function_ref<void(Region *op)> callback);
+void walkOperations(Operation *op, function_ref<void(Block *op)> callback);
 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.
+/// 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 walkOperations(Operation *op,
+                          function_ref<WalkResult(Region *op)> callback);
+WalkResult walkOperations(Operation *op,
+                          function_ref<WalkResult(Block *op)> callback);
 WalkResult walkOperations(Operation *op,
                           function_ref<WalkResult(Operation *op)> 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*.
 ///
 /// Example:
 ///   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
+typename std::enable_if<std::is_same<ArgT, Operation *>::value ||
+                            std::is_same<ArgT, Region *>::value ||
+                            std::is_same<ArgT, Block *>::value,
+                        RetT>::type
 walkOperations(Operation *op, FuncTy &&callback) {
   return detail::walkOperations(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<ArgT, Region *>::value &&
+                            !std::is_same<ArgT, Block *>::value &&
                             std::is_same<RetT, void>::value,
                         RetT>::type
 walkOperations(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));
 }
 
 /// 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<ArgT, Region *>::value &&
+                            !std::is_same<ArgT, Block *>::value &&
                             std::is_same<RetT, WalkResult>::value,
                         RetT>::type
 walkOperations(Operation *op, FuncTy &&callback) {
   auto wrapperFn = [&](Operation *op) {
     if (auto derivedOp = dyn_cast<ArgT>(op))
       return callback(derivedOp);
     return WalkResult::advance();
   };

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";
+    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.
+/// Walk all of the regions/blocks/operations nested under and including the
+/// given operation.
+void detail::walkOperations(Operation *op,
+                            function_ref<void(Region *)> callback) {
+  for (auto &region : op->getRegions()) {
+    callback(&region);
+    for (auto &block : region) {
+      for (auto &nestedOp : block)
+        walkOperations(&nestedOp, callback);
+    }
+  }
+}
+
+void detail::walkOperations(Operation *op,
+                            function_ref<void(Block *)> callback) {
+  for (auto &region : op->getRegions()) {
+    for (auto &block : region) {
+      callback(&block);
+      for (auto &nestedOp : block)
+        walkOperations(&nestedOp, callback);
+    }
+  }
+}
+
 void detail::walkOperations(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);
-
+    }
+  }
   callback(op);
 }
 
-/// Walk all of the operations nested under and including the given operations.
-/// This methods walks operations until an interrupt signal is received.
+/// 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::walkOperations(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)
+        walkOperations(&nestedOp, callback);
+    }
+  }
+  return WalkResult::advance();
+}
+
+WalkResult
+detail::walkOperations(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)
+        walkOperations(&nestedOp, callback);
+    }
+  }
+  return WalkResult::advance();
+}
+
 WalkResult
 detail::walkOperations(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))
+      for (auto &nestedOp : llvm::make_early_inc_range(block)) {
         if (walkOperations(&nestedOp, callback).wasInterrupted())
           return WalkResult::interrupt();
       }
     }
+  }
   return callback(op);
 }