Imported Upstream version 0.4upstream/0.4

1 files changed, 23 insertions, 4 deletions

diff --git a/compiler/moco-tf/src/Canonicalization/StopGradientCanonicalizer.cpp b/compiler/moco-tf/src/Canonicalization/StopGradientCanonicalizer.cpp
index 574fa3993..a52af05a5 100644
--- a/compiler/moco-tf/src/Canonicalization/StopGradientCanonicalizer.cpp
+++ b/compiler/moco-tf/src/Canonicalization/StopGradientCanonicalizer.cpp
@@ -16,14 +16,17 @@
 
 #include "StopGradientCanonicalizer.h"
 
-#include <moco/IR/TFDialect.h>
+#include "Dialect/TFDialect.h"
+#include "Dialect/TFNodes.h"
+#include "Dialect/TFNodeVisitor.h"
+#include "Dialect/TFNodeImpl.h"
 
 #include <moco/Log.h>
 
 namespace
 {
 
-bool canonicalize_stopgradient(loco::Graph *graph, moco::TFStopGradient *node)
+bool canonicalize_stopgradient(loco::Graph *graph, moco::tf::TFStopGradient *node)
 {
   LOGGER(l);
 
@@ -62,9 +65,25 @@ namespace moco
 namespace tf
 {
 
-bool StopGradientCanonicalizer::transform(TFStopGradient *node) const
+bool StopGradientCanonicalizer::run(loco::Graph *graph)
 {
-  return canonicalize_stopgradient(node->graph(), node);
+  auto active_nodes = loco::active_nodes(loco::output_nodes(graph));
+  bool changed = false;
+
+  for (auto node : active_nodes)
+  {
+    if (node->dialect() == TFDialect::get())
+    {
+      auto tf_stopgradient = dynamic_cast<moco::tf::TFStopGradient *>(node);
+      if (tf_stopgradient != nullptr)
+      {
+        if (canonicalize_stopgradient(graph, tf_stopgradient))
+          changed = true;
+      }
+    }
+  }
+
+  return changed;
 }
 
 } // namespace tf