From c8d2c15261c3fb1551d2da5bea3192fc3a667f84 Mon Sep 17 00:00:00 2001
From: "A. Unique TensorFlower" <gardener@tensorflow.org>
Date: Fri, 30 Oct 2020 10:19:11 -0700
Subject: [PATCH] [MLIR][KernelGen] Use MLIR's `walk` to iterate over blocks.

PiperOrigin-RevId: 339893950
Change-Id: I6d3bf00d3296f5008ce2bd9f208dc0600a515665
---
 .../kernel_gen/transforms/buffer_reuse_pass.cc | 18 ++++--------------
 1 file changed, 4 insertions(+), 14 deletions(-)

diff --git a/tensorflow/compiler/mlir/tools/kernel_gen/transforms/buffer_reuse_pass.cc b/tensorflow/compiler/mlir/tools/kernel_gen/transforms/buffer_reuse_pass.cc
index d8ea4f52bd8..7e45d884ece 100644
--- a/tensorflow/compiler/mlir/tools/kernel_gen/transforms/buffer_reuse_pass.cc
+++ b/tensorflow/compiler/mlir/tools/kernel_gen/transforms/buffer_reuse_pass.cc
@@ -52,16 +52,6 @@ namespace kernel_gen {
 namespace transforms {
 namespace {
 
-// TODO(frgossen): Move this to MLIR.
-static void walkBlocks(Operation *op, std::function<void(Block &)> callback) {
-  for (Region &region : op->getRegions()) {
-    for (Block &block : region) {
-      callback(block);
-      for (Operation &nested_op : block) walkBlocks(&nested_op, callback);
-    }
-  }
-}
-
 /// A temporary buffer size analysis that is correct but may be incomplete.
 class BufferSizeAnalysis {
  public:
@@ -121,8 +111,8 @@ class BufferSizeAnalysis {
       for (Value val : op->getResults())
         if (val.getType().isa<BaseMemRefType>()) buffers.push_back(val);
     });
-    walkBlocks(f.getOperation(), [&](Block &block) {
-      for (Value val : block.getArguments()) {
+    f.walk([&](Block *block) {
+      for (Value val : block->getArguments()) {
         if (val.getType().isa<BaseMemRefType>()) buffers.push_back(val);
       }
     });
@@ -180,8 +170,8 @@ class BufferReuseAnalysis {
   void find_reuse_candiates(FuncOp &f, BufferAliasAnalysis &aliases) {
     Liveness liveness(f);
     BufferSizeAnalysis size_equivalences(f);
-    walkBlocks(f.getOperation(), [&](Block &block) {
-      find_reuse_candiates(block, aliases, liveness.getLiveness(&block),
+    f.walk([&](Block *block) {
+      find_reuse_candiates(*block, aliases, liveness.getLiveness(block),
                            size_equivalences, f.getArguments());
     });
   }
-- 
GitLab