Faster vars used tracking in simplify let visitor

src/IntrusivePtr.h

@@ -32,6 +32,9 @@ class RefCount {
     bool is_const_zero() const {
         return count == 0;
     }
+    int atomic_get() const {
+        return count;
+    }
 };
 
 /**
@@ -173,6 +176,11 @@ struct IntrusivePtr {
     bool operator<(const IntrusivePtr<T> &other) const {
         return ptr < other.ptr;
     }
+
+    HALIDE_ALWAYS_INLINE
+    bool is_sole_reference() const {
+        return ptr && ref_count(ptr).atomic_get() == 1;
+    }
 };
 
 }  // namespace Internal

src/Simplify_Let.cpp

@@ -1,6 +1,8 @@
 #include "Simplify_Internal.h"
 #include "Substitute.h"
 
+#include <unordered_set>
+
 namespace Halide {
 namespace Internal {
 
@@ -9,34 +11,50 @@ using std::vector;
 
 namespace {
 
-class CountVarUses : public IRVisitor {
-    std::map<std::string, int> &var_uses;
+class FindVarUses : public IRVisitor {
+    std::unordered_set<std::string> &unused_vars;
 
     void visit(const Variable *var) override {
-        var_uses[var->name]++;
+        unused_vars.erase(var->name);
     }
 
     void visit(const Load *op) override {
-        var_uses[op->name]++;
-        IRVisitor::visit(op);
+        if (!unused_vars.empty()) {
+            unused_vars.erase(op->name);
+            IRVisitor::visit(op);
+        }
     }
 
     void visit(const Store *op) override {
-        var_uses[op->name]++;
-        IRVisitor::visit(op);
+        if (!unused_vars.empty()) {
+            unused_vars.erase(op->name);
+            IRVisitor::visit(op);
+        }
+    }
+
+    void visit(const Block *op) override {
+        // Early out at Block nodes if we've already seen every name we're
+        // interested in. In principal we could early-out at every node, but
+        // blocks, loads, and stores seem to be enough.
+        if (!unused_vars.empty()) {
+            op->first.accept(this);
+            if (!unused_vars.empty()) {
+                op->rest.accept(this);
+            }
+        }
     }
 
     using IRVisitor::visit;
 
 public:
-    CountVarUses(std::map<std::string, int> &var_uses)
-        : var_uses(var_uses) {
+    FindVarUses(std::unordered_set<std::string> &unused_vars)
+        : unused_vars(unused_vars) {
     }
 };
 
 template<typename StmtOrExpr>
-void count_var_uses(StmtOrExpr x, std::map<std::string, int> &var_uses) {
-    CountVarUses counter(var_uses);
+void find_var_uses(StmtOrExpr x, std::unordered_set<std::string> &unused_vars) {
+    FindVarUses counter(unused_vars);
     x.accept(&counter);
 }
 
@@ -49,10 +67,11 @@ Body Simplify::simplify_let(const LetOrLetStmt *op, ExprInfo *bounds) {
     // the call stack where it could overflow onto an explicit stack.
     struct Frame {
         const LetOrLetStmt *op;
-        Expr value, new_value;
+        Expr value, new_value, new_var;
         string new_name;
         bool new_value_alignment_tracked = false, new_value_bounds_tracked = false;
         bool value_alignment_tracked = false, value_bounds_tracked = false;
+        VarInfo info;
         Frame(const LetOrLetStmt *op)
             : op(op) {
         }
@@ -189,6 +208,7 @@ Body Simplify::simplify_let(const LetOrLetStmt *op, ExprInfo *bounds) {
             // Nothing to substitute
             f.new_value = Expr();
             replacement = Expr();
+            new_var = Expr();
         } else {
             debug(4) << "new let " << f.new_name << " = " << f.new_value << " in ... " << replacement << " ...\n";
         }
@@ -197,6 +217,7 @@ Body Simplify::simplify_let(const LetOrLetStmt *op, ExprInfo *bounds) {
         info.old_uses = 0;
         info.new_uses = 0;
         info.replacement = replacement;
+        f.new_var = new_var;
 
         var_info.push(op->name, info);
 
@@ -226,14 +247,35 @@ Body Simplify::simplify_let(const LetOrLetStmt *op, ExprInfo *bounds) {
 
     result = mutate_let_body(result, bounds);
 
-    // TODO: var_info and vars_used are pretty redundant; however, at the time
+    // TODO: var_info and unused_vars are pretty redundant; however, at the time
     // of writing, both cover cases that the other does not:
     // - var_info prevents duplicate lets from being generated, even
     //   from different Frame objects.
-    // - vars_used avoids dead lets being generated in cases where vars are
+    // - unused_vars avoids dead lets being generated in cases where vars are
     //   seen as used by var_info, and then later removed.
-    std::map<std::string, int> vars_used;
-    count_var_uses(result, vars_used);
+
+    std::unordered_set<std::string> unused_vars(frames.size());
+    // Insert everything we think *might* be used, and then visit the body,
+    // removing things from the set as we find uses of them.
+    for (auto &f : frames) {
+        f.info = var_info.get(f.op->name);
+        // Drop any reference to new_var held by the replacement expression so
+        // that the only references are either f.new_var, or ones in the body or
+        // new_values of other lets.
+        f.info.replacement = Expr();
+        if (f.new_var.is_sole_reference()) {
+            // Any new_uses must have been eliminated by later mutations.
+            f.info.new_uses = 0;
+        }
+        var_info.pop(f.op->name);
+        if (f.info.old_uses) {
+            internal_assert(f.info.new_uses == 0);
+            unused_vars.insert(f.op->name);
+        } else if (f.info.new_uses && f.new_value.defined()) {
+            unused_vars.insert(f.new_name);
+        }
+    }
+    find_var_uses(result, unused_vars);
 
     for (auto it = frames.rbegin(); it != frames.rend(); it++) {
         if (it->value_bounds_tracked) {
@@ -243,20 +285,17 @@ Body Simplify::simplify_let(const LetOrLetStmt *op, ExprInfo *bounds) {
             bounds_and_alignment_info.pop(it->new_name);
         }
 
-        VarInfo info = var_info.get(it->op->name);
-        var_info.pop(it->op->name);
-
-        if (it->new_value.defined() && (info.new_uses > 0 && vars_used.count(it->new_name) > 0)) {
+        if (it->new_value.defined() && (it->info.new_uses > 0 && !unused_vars.count(it->new_name))) {
             // The new name/value may be used
             result = LetOrLetStmt::make(it->new_name, it->new_value, result);
-            count_var_uses(it->new_value, vars_used);
+            find_var_uses(it->new_value, unused_vars);
         }
 
         if ((!remove_dead_code && std::is_same<LetOrLetStmt, LetStmt>::value) ||
-            (info.old_uses > 0 && vars_used.count(it->op->name) > 0)) {
+            (it->info.old_uses > 0 && !unused_vars.count(it->op->name))) {
             // The old name is still in use. We'd better keep it as well.
             result = LetOrLetStmt::make(it->op->name, it->value, result);
-            count_var_uses(it->value, vars_used);
+            find_var_uses(it->value, unused_vars);
         }
 
         const LetOrLetStmt *new_op = result.template as<LetOrLetStmt>();