move

taichi-dev · Feb 18, 2024 · 231e395 · 231e395
1 parent f448b35
commit 231e395
Show file tree

Hide file tree

Showing 2 changed files with 10 additions and 7 deletions.
diff --git a/taichi/transforms/alg_simp.cpp b/taichi/transforms/alg_simp.cpp
@@ -351,13 +351,6 @@ class AlgSimp : public BasicStmtVisitor {
           modifier.mark_as_modified();
         }
       }
-    } else if (stmt->op_type == UnaryOpType::abs) {
-      auto operand_type = stmt->operand->ret_type;
-      if (is_integral(operand_type) && is_unsigned(operand_type)) {
-        // abs(u) -> u
-        stmt->replace_usages_with(stmt->operand);
-        modifier.erase(stmt);
-      }
     }
   }
 

diff --git a/taichi/transforms/simplify.cpp b/taichi/transforms/simplify.cpp
@@ -133,6 +133,16 @@ class BasicBlockSimplify : public IRVisitor {
     }
   }
 
+  void visit(UnaryOpStmt *stmt) override {
+    if (stmt->op_type == UnaryOpType::abs) {
+      auto operand_type = stmt->operand->ret_type;
+      if (is_integral(operand_type) && is_unsigned(operand_type)) {
+        // abs(u) -> u
+        stmt->replace_usages_with(stmt->operand);
+        modifier.erase(stmt);
+      }
+    }
+  }
   template <typename T>
   static bool identical_vectors(const std::vector<T> &a,
                                 const std::vector<T> &b) {