Add more let peeling and rewrapping helpers

src/AddImageChecks.cpp

@@ -681,14 +681,6 @@ Stmt add_image_checks_inner(Stmt s,
         }
     };
 
-    auto prepend_lets = [&](vector<pair<string, Expr>> *lets) {
-        while (!lets->empty()) {
-            auto &p = lets->back();
-            s = LetStmt::make(p.first, std::move(p.second), s);
-            lets->pop_back();
-        }
-    };
-
     // After all asserts, set host dirty on outputs if this is a CPU-only
     // pipeline
     prepend_stmts(&set_host_dirty);
@@ -698,7 +690,7 @@ Stmt add_image_checks_inner(Stmt s,
     prepend_stmts(&asserts_host_alignment);
     prepend_stmts(&asserts_device_not_dirty);
     prepend_stmts(&dims_no_overflow_asserts);
-    prepend_lets(&lets_overflow);
+    s = rewrap_all_lets(s, lets_overflow);
 
     // Replace uses of the var with the constrained versions in the
     // rest of the program. We also need to respect the existence of
@@ -726,13 +718,13 @@ Stmt add_image_checks_inner(Stmt s,
     prepend_stmts(&asserts_proposed);
 
     // Inject the code that defines the proposed sizes.
-    prepend_lets(&lets_proposed);
+    s = rewrap_all_lets(s, lets_proposed);
 
     // Inject the code that defines the constrained sizes.
-    prepend_lets(&lets_constrained);
+    s = rewrap_all_lets(s, lets_constrained);
 
     // Inject the code that defines the required sizes produced by bounds inference.
-    prepend_lets(&lets_required);
+    s = rewrap_all_lets(s, lets_required);
 
     // Inject the code that checks that does msan checks. (Note that this ignores no_asserts.)
     prepend_stmts(&msan_checks);

src/AsyncProducers.cpp

@@ -507,11 +507,7 @@ class InitializeSemaphores : public IRMutator {
             body = mutate(op->body);
             // Peel off any enclosing let expressions from the value
             vector<pair<string, Expr>> lets;
-            Expr value = op->value;
-            while (const Let *l = value.as<Let>()) {
-                lets.emplace_back(l->name, l->value);
-                value = l->body;
-            }
+            Expr value = peel_lets(op->value, &lets);
             const Call *call = value.as<Call>();
             if (call && call->name == "halide_make_semaphore") {
                 internal_assert(call->args.size() == 1);
@@ -525,9 +521,7 @@ class InitializeSemaphores : public IRMutator {
                 body = op->with(sema_allocate, body);
 
                 // Re-wrap any other lets
-                for (const auto &[var, value] : reverse_view(lets)) {
-                    body = LetStmt::make(var, value, std::move(body));
-                }
+                body = rewrap_all_lets(body, lets);
             }
         } else {
             body = mutate(frames.back()->body);

src/BoundConstantExtentLoops.cpp

@@ -62,9 +62,7 @@ class BoundLoops : public IRMutator {
             if (e == nullptr) {
                 // We're about to hard fail. Get really aggressive
                 // with the simplifier.
-                for (const auto &[var, value] : reverse_view(lets)) {
-                    extent = Let::make(var, value, extent);
-                }
+                extent = rewrap_used_lets(extent, lets);
                 extent = remove_likelies(extent);
                 extent = substitute_in_all_lets(extent);
                 extent = simplify(extent,

src/CSE.cpp

@@ -233,10 +233,7 @@ class CSEEveryExprInStmt : public IRMutator {
         Expr dummy = Call::make(Int(32), Call::bundle, {op->value, op->index}, Call::PureIntrinsic);
         dummy = common_subexpression_elimination(dummy, lift_all);
         vector<pair<string, Expr>> lets;
-        while (const Let *let = dummy.as<Let>()) {
-            lets.emplace_back(let->name, let->value);
-            dummy = let->body;
-        }
+        dummy = peel_lets(dummy, &lets);
         const Call *bundle = Call::as_intrinsic(dummy, {Call::bundle});
         internal_assert(bundle && bundle->args.size() == 2);
 

src/Closure.cpp

@@ -153,20 +153,20 @@ Stmt Closure::unpack_from_struct(const Expr &e, const Stmt &s) const {
 
     const Call *c = packed.as<Call>();
 
-    Stmt result = s;
-    for (int idx = (int)c->args.size() - 1; idx >= 0; idx--) {
-        Expr arg = c->args[idx];
-        const Variable *var = arg.as<Variable>();
-        Expr val = Call::make(var->type,
-                              Call::load_typed_struct_member,
-                              {e, prototype_var, idx},
-                              Call::Intrinsic);
-        if (stmt_uses_var(result, var->name)) {
-            // If a closure is generated for multiple consuming blocks of IR,
-            // then some of those blocks might only need some of the field.
-            result = LetStmt::make(var->name, val, result);
-        }
+    // If a closure is generated for multiple consuming blocks of IR, then some
+    // of those blocks might only need some of the fields, so only bind the ones
+    // that are used.
+    std::vector<std::pair<std::string, Expr>> lets;
+    lets.reserve(c->args.size());
+    for (int idx = 0; idx < (int)c->args.size(); idx++) {
+        const Variable *var = c->args[idx].as<Variable>();
+        lets.emplace_back(var->name,
+                          Call::make(var->type,
+                                     Call::load_typed_struct_member,
+                                     {e, prototype_var, idx},
+                                     Call::Intrinsic));
     }
+    Stmt result = rewrap_used_lets(s, lets);
     result = LetStmt::make(prototype_name, prototype, result);
 
     return result;

src/CodeGen_ARM.cpp

@@ -1577,12 +1577,8 @@ void CodeGen_ARM::visit(const Store *op) {
     }
 
     // First dig through let expressions
-    Expr rhs = op->value;
     vector<pair<string, Expr>> lets;
-    while (const Let *let = rhs.as<Let>()) {
-        rhs = let->body;
-        lets.emplace_back(let->name, let->value);
-    }
+    Expr rhs = peel_lets(op->value, &lets);
     const Shuffle *shuffle = rhs.as<Shuffle>();
 
     // Interleaving store instructions only exist for certain types.
@@ -1690,11 +1686,7 @@ void CodeGen_ARM::visit(const Store *op) {
             // And we make sure the deinterleaved predicates are all the same.
 
             // Dig through let expressions
-            Expr rhs = op->predicate;
-            while (const Let *let = rhs.as<Let>()) {
-                rhs = let->body;
-                lets_pred.emplace_back(let->name, let->value);
-            }
+            Expr rhs = peel_lets(op->predicate, &lets_pred);
 
             Expr vpred_predicated_store;
             bool predicates_are_same = true;

src/Deinterleave.cpp

@@ -21,11 +21,12 @@ class StoreCollector : public IRMutator {
 public:
     const std::string store_name;
     const int store_stride, max_stores;
-    std::vector<Stmt> &let_stmts;
+    std::vector<std::pair<std::string, Expr>> &let_stmts;
     std::vector<Stmt> &stores;
 
     StoreCollector(const std::string &name, int stride, int ms,
-                   std::vector<Stmt> &lets, std::vector<Stmt> &ss)
+                   std::vector<std::pair<std::string, Expr>> &lets,
+                   std::vector<Stmt> &ss)
         : store_name(name), store_stride(stride), max_stores(ms),
           let_stmts(lets), stores(ss) {
     }
@@ -59,7 +60,7 @@ class StoreCollector : public IRMutator {
     // These are lets that we've encountered since the last collected
     // store. If we collect another store, these "potential" lets
     // become lets used by the collected stores.
-    std::vector<Stmt> potential_lets;
+    std::vector<std::pair<std::string, Expr>> potential_lets;
 
     Expr visit(const Load *op) override {
         if (!collecting) {
@@ -143,11 +144,7 @@ class StoreCollector : public IRMutator {
 
         // If we're still collecting, we need to save the entire let chain as potential lets.
         if (collecting) {
-            Stmt body;
-            do {
-                potential_lets.emplace_back(op);
-                body = op->body;
-            } while ((op = body.as<LetStmt>()));
+            peel_lets(op, &potential_lets);
         }
         return stmt;
     }
@@ -168,7 +165,8 @@ class StoreCollector : public IRMutator {
 };
 
 Stmt collect_strided_stores(const Stmt &stmt, const std::string &name, int stride, int max_stores,
-                            std::vector<Stmt> lets, std::vector<Stmt> &stores) {
+                            std::vector<std::pair<std::string, Expr>> lets,
+                            std::vector<Stmt> &stores) {
 
     return StoreCollector(name, stride, max_stores, lets, stores)(stmt);
 }
@@ -645,16 +643,9 @@ class Interleaver : public IRMutator {
     }
 
     HALIDE_NEVER_INLINE Stmt gather_stores(const Block *op) {
-        const LetStmt *let = op->first.as<LetStmt>();
-        const Store *store = op->first.as<Store>();
-
         // Gather all the let stmts surrounding the first.
-        std::vector<Stmt> let_stmts;
-        while (let) {
-            let_stmts.emplace_back(let);
-            store = let->body.as<Store>();
-            let = let->body.as<LetStmt>();
-        }
+        std::vector<std::pair<std::string, Expr>> let_stmts;
+        const Store *store = peel_lets(op->first, &let_stmts).as<Store>();
 
         // There was no inner store.
         if (!store) {
@@ -771,11 +762,7 @@ class Interleaver : public IRMutator {
         Stmt new_store = store->with(value, index, predicate, ModulusRemainder());
 
         // Rewrap the let statements we pulled off.
-        while (!let_stmts.empty()) {
-            const LetStmt *let = let_stmts.back().as<LetStmt>();
-            new_store = let->with(let->value, new_store);
-            let_stmts.pop_back();
-        }
+        new_store = rewrap_all_lets(new_store, let_stmts);
 
         // Continue recursively into the stuff that
         // collect_strided_stores didn't collect.

src/DerivativeUtils.cpp

@@ -115,23 +115,14 @@ map<string, ReductionVariableInfo> gather_rvariables(const Expr &expr) {
 Expr add_let_expression(const Expr &expr,
                         const map<string, Expr> &let_var_mapping,
                         const vector<string> &let_variables) {
-    // TODO: find a faster way to do this
-    Expr ret = StripLets()(expr);
-    bool changed = true;
-    vector<bool> injected(let_variables.size(), false);
-    while (changed) {
-        changed = false;
-        for (size_t i = 0; i < let_variables.size(); i++) {
-            const auto &let_variable = let_variables[i];
-            if (!injected[i] && expr_uses_var(ret, let_variable)) {
-                auto value = let_var_mapping.find(let_variable)->second;
-                ret = Let::make(let_variable, value, ret);
-                injected[i] = true;
-                changed = true;
-            }
-        }
-    }
-    return ret;
+    // sort_expressions lists Lets innermost first, so reverse them to get the
+    // outermost-to-innermost order rewrap_used_lets expects.
+    vector<pair<string, Expr>> lets;
+    lets.reserve(let_variables.size());
+    for (const auto &let_variable : reverse_view(let_variables)) {
+        lets.emplace_back(let_variable, let_var_mapping.find(let_variable)->second);
+    }
+    return rewrap_used_lets(StripLets()(expr), lets);
 }
 
 namespace {

src/ExtractTileOperations.cpp

@@ -67,11 +67,7 @@ Matmul convert_to_matmul(const Store *op, const string &new_name) {
 
     // Peel lets
     std::vector<std::pair<std::string, Expr>> peeled_lets;
-    Expr value = op->value;
-    while (const Let *let = value.as<Let>()) {
-        peeled_lets.emplace_back(let->name, let->value);
-        value = let->body;
-    }
+    Expr value = peel_lets(op->value, &peeled_lets);
 
     // The RHS must be an add
     const auto *add = value.as<Add>();
@@ -317,10 +313,8 @@ Matmul convert_to_matmul(const Store *op, const string &new_name) {
     auto matmul = Call::make(res_type, "tile_matmul",
                              {I, col_bytes, K, out_load, lhs_call, rhs_call},
                              Call::Intrinsic);
-    auto store = Store::make(new_name, matmul, std::move(subtile_idx));
-    for (auto &[name, value] : reverse_view(peeled_lets)) {
-        store = LetStmt::make(name, std::move(value), store);
-    }
+    Stmt store = Store::make(new_name, matmul, std::move(subtile_idx));
+    store = rewrap_all_lets(store, peeled_lets);
     return {true, std::move(store), I, J, K};
 }
 

src/IROperator.cpp

@@ -2,7 +2,9 @@
 #include <atomic>
 #include <cmath>
 #include <iostream>
+#include <set>
 #include <sstream>
+#include <type_traits>
 #include <utility>
 
 #include "CSE.h"
@@ -13,6 +15,7 @@
 #include "IRMutator.h"
 #include "IROperator.h"
 #include "IRPrinter.h"
+#include "IRVisitor.h"
 #include "Interval.h"
 #include "StrictifyFloat.h"
 #include "Util.h"
@@ -1113,6 +1116,90 @@ Expr peel_lets(const Expr &e, std::vector<std::pair<std::string, Expr>> *lets) {
     return body;
 }
 
+Stmt peel_lets(const Stmt &s, std::vector<std::pair<std::string, Expr>> *lets) {
+    Stmt body = s;
+    while (const LetStmt *let = body.as<LetStmt>()) {
+        lets->emplace_back(let->name, let->value);
+        body = let->body;
+    }
+    return body;
+}
+
+namespace {
+
+/** Gather the names an Expr or Stmt might get from a wrapping let: the names of
+ * Variables, and the buffer names of Loads and Stores. Conservatively assumes
+ * every such name refers to one of the peeled lets, even where an inner let
+ * shadows it. One instance is reused across an entire rewrap so that shared
+ * subexpressions are only visited once. */
+class CollectUsedNames : public IRGraphVisitor {
+    using IRGraphVisitor::visit;
+
+    void visit(const Variable *op) override {
+        names.insert(op->name);
+    }
+
+    void visit(const Load *op) override {
+        names.insert(op->name);
+        IRGraphVisitor::visit(op);
+    }
+
+    void visit(const Store *op) override {
+        names.insert(op->name);
+        IRGraphVisitor::visit(op);
+    }
+
+public:
+    std::set<std::string> names;
+};
+
+template<typename Body>
+Body rewrap_used_lets_impl(const Body &body,
+                           const std::vector<std::pair<std::string, Expr>> &lets) {
+    // The set of names the growing body refers to. Maintaining it as we go
+    // avoids rescanning the body once per let.
+    CollectUsedNames used;
+    body.accept(&used);
+    Body result = body;
+    for (const auto &[name, value] : reverse_view(lets)) {
+        if (used.names.count(name)) {
+            value.accept(&used);
+            if constexpr (std::is_same_v<Body, Expr>) {
+                result = Let::make(name, value, result);
+            } else {
+                result = LetStmt::make(name, value, result);
+            }
+        }
+    }
+    return result;
+}
+
+}  // namespace
+
+Expr rewrap_used_lets(const Expr &body, const std::vector<std::pair<std::string, Expr>> &lets) {
+    return rewrap_used_lets_impl(body, lets);
+}
+
+Stmt rewrap_used_lets(const Stmt &body, const std::vector<std::pair<std::string, Expr>> &lets) {
+    return rewrap_used_lets_impl(body, lets);
+}
+
+Expr rewrap_all_lets(const Expr &body, const std::vector<std::pair<std::string, Expr>> &lets) {
+    Expr result = body;
+    for (const auto &[name, value] : reverse_view(lets)) {
+        result = Let::make(name, value, result);
+    }
+    return result;
+}
+
+Stmt rewrap_all_lets(const Stmt &body, const std::vector<std::pair<std::string, Expr>> &lets) {
+    Stmt result = body;
+    for (const auto &[name, value] : reverse_view(lets)) {
+        result = LetStmt::make(name, value, result);
+    }
+    return result;
+}
+
 Expr remove_likelies(const Expr &e) {
     return remove_intrinsics(e, {Call::likely, Call::likely_if_innermost});
 }