From e745c48fd95856bf2049568f4fe5c8166a02700e Mon Sep 17 00:00:00 2001 From: Andrew Adams Date: Tue, 11 Aug 2026 11:29:53 -0700 Subject: [PATCH 1/3] Add peel_lets/rewrap_lets helpers and use them throughout lowering Peeling Lets or LetStmts into a vector of name/value pairs, doing something to the body, and then rewrapping is a common pattern in the compiler. Rewrapping conditionally was done by calling expr_uses_var on the partially-rebuilt body once per let, which is quadratic. Adds a Stmt overload of peel_lets, plus rewrap_used_lets and rewrap_all_lets for Expr and Stmt. rewrap_used_lets gathers the names the body mentions once and extends the set with the value of each let it keeps, so it is linearithmic. It conservatively treats every name mentioned as a possible reference to a peeled let, even where an inner let shadows it. Co-authored-by: Claude Opus 5 --- src/AddImageChecks.cpp | 16 ++---- src/AsyncProducers.cpp | 10 +--- src/BoundConstantExtentLoops.cpp | 4 +- src/CSE.cpp | 5 +- src/Closure.cpp | 26 +++++----- src/CodeGen_ARM.cpp | 12 +---- src/DerivativeUtils.cpp | 25 +++------- src/ExtractTileOperations.cpp | 12 ++--- src/IROperator.cpp | 86 ++++++++++++++++++++++++++++++++ src/IROperator.h | 36 ++++++++++--- src/LICM.cpp | 12 +---- src/LoopCarry.cpp | 15 ++---- src/LowerWarpShuffles.cpp | 11 ++-- src/PartitionLoops.cpp | 8 +-- src/ScheduleFunctions.cpp | 9 +--- src/Simplify_Stmts.cpp | 11 +--- src/SplitTuples.cpp | 37 +++----------- src/VectorizeLoops.cpp | 11 +--- 18 files changed, 172 insertions(+), 174 deletions(-) diff --git a/src/AddImageChecks.cpp b/src/AddImageChecks.cpp index 46dbaf9f2385..2878c7aacb74 100644 --- a/src/AddImageChecks.cpp +++ b/src/AddImageChecks.cpp @@ -681,14 +681,6 @@ Stmt add_image_checks_inner(Stmt s, } }; - auto prepend_lets = [&](vector> *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 @@ -724,13 +716,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); diff --git a/src/AsyncProducers.cpp b/src/AsyncProducers.cpp index 7038ddceff17..1751f1d88221 100644 --- a/src/AsyncProducers.cpp +++ b/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> lets; - Expr value = op->value; - while (const Let *l = value.as()) { - lets.emplace_back(l->name, l->value); - value = l->body; - } + Expr value = peel_lets(op->value, &lets); const Call *call = value.as(); 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); diff --git a/src/BoundConstantExtentLoops.cpp b/src/BoundConstantExtentLoops.cpp index 6ea10fbb17ec..ebc41007e6bf 100644 --- a/src/BoundConstantExtentLoops.cpp +++ b/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, diff --git a/src/CSE.cpp b/src/CSE.cpp index e46f94e5fa0e..ba94199da669 100644 --- a/src/CSE.cpp +++ b/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> lets; - while (const Let *let = dummy.as()) { - 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); diff --git a/src/Closure.cpp b/src/Closure.cpp index 00566127e322..5d72f031f66b 100644 --- a/src/Closure.cpp +++ b/src/Closure.cpp @@ -153,20 +153,20 @@ Stmt Closure::unpack_from_struct(const Expr &e, const Stmt &s) const { const Call *c = packed.as(); - Stmt result = s; - for (int idx = (int)c->args.size() - 1; idx >= 0; idx--) { - Expr arg = c->args[idx]; - const Variable *var = arg.as(); - 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> lets; + lets.reserve(c->args.size()); + for (int idx = 0; idx < (int)c->args.size(); idx++) { + const Variable *var = c->args[idx].as(); + 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; diff --git a/src/CodeGen_ARM.cpp b/src/CodeGen_ARM.cpp index 57501916f3d2..4fbf5c4a088d 100644 --- a/src/CodeGen_ARM.cpp +++ b/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> lets; - while (const Let *let = rhs.as()) { - rhs = let->body; - lets.emplace_back(let->name, let->value); - } + Expr rhs = peel_lets(op->value, &lets); const Shuffle *shuffle = rhs.as(); // 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()) { - 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; diff --git a/src/DerivativeUtils.cpp b/src/DerivativeUtils.cpp index 812eeeb7a517..f9dac4aeabb4 100644 --- a/src/DerivativeUtils.cpp +++ b/src/DerivativeUtils.cpp @@ -115,23 +115,14 @@ map gather_rvariables(const Expr &expr) { Expr add_let_expression(const Expr &expr, const map &let_var_mapping, const vector &let_variables) { - // TODO: find a faster way to do this - Expr ret = StripLets()(expr); - bool changed = true; - vector 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> 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 { diff --git a/src/ExtractTileOperations.cpp b/src/ExtractTileOperations.cpp index 5fee9c76c80b..c2211b086693 100644 --- a/src/ExtractTileOperations.cpp +++ b/src/ExtractTileOperations.cpp @@ -67,11 +67,7 @@ Matmul convert_to_matmul(const Store *op, const string &new_name) { // Peel lets std::vector> peeled_lets; - Expr value = op->value; - while (const Let *let = value.as()) { - 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(); @@ -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}; } diff --git a/src/IROperator.cpp b/src/IROperator.cpp index a676b7f9f933..927ba7e78914 100644 --- a/src/IROperator.cpp +++ b/src/IROperator.cpp @@ -2,6 +2,7 @@ #include #include #include +#include #include #include @@ -13,6 +14,7 @@ #include "IRMutator.h" #include "IROperator.h" #include "IRPrinter.h" +#include "IRVisitor.h" #include "Interval.h" #include "StrictifyFloat.h" #include "Util.h" @@ -1113,6 +1115,90 @@ Expr peel_lets(const Expr &e, std::vector> *lets) { return body; } +Stmt peel_lets(const Stmt &s, std::vector> *lets) { + Stmt body = s; + while (const LetStmt *let = body.as()) { + 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 names; +}; + +template +Body rewrap_used_lets_impl(const Body &body, + const std::vector> &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.erase(name)) { + value.accept(&used); + if constexpr (std::is_same_v) { + 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> &lets) { + return rewrap_used_lets_impl(body, lets); +} + +Stmt rewrap_used_lets(const Stmt &body, const std::vector> &lets) { + return rewrap_used_lets_impl(body, lets); +} + +Expr rewrap_all_lets(const Expr &body, const std::vector> &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> &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}); } diff --git a/src/IROperator.h b/src/IROperator.h index f377522d4b07..292efabbb748 100644 --- a/src/IROperator.h +++ b/src/IROperator.h @@ -320,14 +320,36 @@ inline double div_imp(double a, double b) { return a / b; } -/** Strip any Let nodes off the front of an Expr, appending the name and value - * of each to `lets` from outermost to innermost, and return what they wrapped. - * Analysing an Expr that CSE or LICM has lifted subexpressions out of means - * getting past the Lets first. Substituting them back in also does that, but it - * repeats each value at every use, which is what lifting them out avoided. - * Callers rewrap the Lets around whatever they build, or, in codegen, put them - * in scope while they build it. */ +/** Strip any Let nodes off the front of an Expr, or LetStmt nodes off the front + * of a Stmt, appending the name and value of each to `lets` from outermost to + * innermost, and return what they wrapped. Analysing an Expr or Stmt that CSE + * or LICM has lifted subexpressions out of means getting past the Lets first. + * Substituting them back in also does that, but it repeats each value at every + * use, which is what lifting them out avoided. Callers rewrap the Lets around + * whatever they build (see rewrap_used_lets and rewrap_all_lets), or, in + * codegen, put them in scope while they build it. */ +// @{ Expr peel_lets(const Expr &e, std::vector> *lets); +Stmt peel_lets(const Stmt &s, std::vector> *lets); +// @} + +/** Rewrap a list of lets produced by peel_lets around a new body, skipping any + * the body doesn't use. Conservatively treats every name the body mentions as a + * possible reference to a peeled let, even where an inner let shadows it. + * Gathers those names once and updates them as it goes, so it takes time + * linearithmic in the size of the result rather than the quadratic time taken + * by testing each let in turn with expr_uses_var. */ +// @{ +Expr rewrap_used_lets(const Expr &body, const std::vector> &lets); +Stmt rewrap_used_lets(const Stmt &body, const std::vector> &lets); +// @} + +/** Rewrap a list of lets produced by peel_lets around a new body, without + * checking whether the body uses them. */ +// @{ +Expr rewrap_all_lets(const Expr &body, const std::vector> &lets); +Stmt rewrap_all_lets(const Stmt &body, const std::vector> &lets); +// @} /** Return an Expr that is identical to the input Expr, but with * all calls to likely() and likely_if_innermost() removed. */ diff --git a/src/LICM.cpp b/src/LICM.cpp index 7a71f5f27cba..0a167768b76b 100644 --- a/src/LICM.cpp +++ b/src/LICM.cpp @@ -276,10 +276,7 @@ class LICM : public IRMutator { // Peel off containing lets. These will be lifted. vector> lets; - while (const Let *let = dummy_call.as()) { - lets.emplace_back(let->name, let->value); - dummy_call = let->body; - } + dummy_call = peel_lets(dummy_call, &lets); // Track the set of variables used by the inner loop set vars; @@ -327,12 +324,7 @@ class LICM : public IRMutator { } // Wrap the lets pulled out by CSE - while (!lets.empty()) { - new_stmt = LetStmt::make(lets.back().first, lets.back().second, new_stmt); - lets.pop_back(); - } - - return new_stmt; + return rewrap_all_lets(new_stmt, lets); } } }; diff --git a/src/LoopCarry.cpp b/src/LoopCarry.cpp index d2ce820002bd..fff1489ddf43 100644 --- a/src/LoopCarry.cpp +++ b/src/LoopCarry.cpp @@ -457,10 +457,7 @@ class LoopCarryOverLoop : public IRMutator { // Run CSE call = simplify(common_subexpression_elimination(call)); // Peel off lets - while (const Let *l = call.as()) { - initial_lets.emplace_back(l->name, l->value); - call = l->body; - } + call = peel_lets(call, &initial_lets); internal_assert(call.as()); initial_scratch_values = call.as()->args; @@ -476,16 +473,10 @@ class LoopCarryOverLoop : public IRMutator { Stmt initial_stores = Block::make(initial_scratch_stores); // Wrap them in the appropriate lets - for (const auto &[var, value] : reverse_view(initial_lets)) { - initial_stores = LetStmt::make(var, value, initial_stores); - } + initial_stores = rewrap_all_lets(initial_stores, initial_lets); // We may be lifting the initial stores out of let stmts, // so rewrap them in the necessary ones. - for (const auto &[var, value] : reverse_view(containing_lets)) { - if (stmt_uses_var(initial_stores, var)) { - initial_stores = LetStmt::make(var, value, initial_stores); - } - } + initial_stores = rewrap_used_lets(initial_stores, containing_lets); allocs.push_back({scratch, loads[c.front()][0]->type.element_of(), diff --git a/src/LowerWarpShuffles.cpp b/src/LowerWarpShuffles.cpp index e64e62cab46d..fb800c3f37e4 100644 --- a/src/LowerWarpShuffles.cpp +++ b/src/LowerWarpShuffles.cpp @@ -743,13 +743,10 @@ class HoistWarpShufflesFromSingleIfStmt : public IRMutator { public: bool success = true; - Stmt rewrap(Stmt s) { - while (!lifted_lets.empty()) { - const pair &p = lifted_lets.back(); - s = LetStmt::make(p.first, p.second, s); - lifted_lets.pop_back(); - } - return s; + Stmt rewrap(const Stmt &s) { + Stmt result = rewrap_all_lets(s, lifted_lets); + lifted_lets.clear(); + return result; } }; diff --git a/src/PartitionLoops.cpp b/src/PartitionLoops.cpp index 770978851912..1de800570eae 100644 --- a/src/PartitionLoops.cpp +++ b/src/PartitionLoops.cpp @@ -878,12 +878,8 @@ class RenormalizeGPULoops : public IRMutator { if (in_gpu_loop && !old_in_gpu_loop) { // This was the outermost GPU loop. Dump any lifted lets here. - while (!lifted_lets.empty()) { - stmt = LetStmt::make(lifted_lets.back().first, - lifted_lets.back().second, - stmt); - lifted_lets.pop_back(); - } + stmt = rewrap_all_lets(stmt, lifted_lets); + lifted_lets.clear(); } in_gpu_loop = old_in_gpu_loop; diff --git a/src/ScheduleFunctions.cpp b/src/ScheduleFunctions.cpp index 23b6e041106a..3daf50969881 100644 --- a/src/ScheduleFunctions.cpp +++ b/src/ScheduleFunctions.cpp @@ -1576,10 +1576,7 @@ class InjectFunctionRealization : public IRMutator { // Strip off the containing lets. The bounds of the parent fused loop // (i.e. the union bounds) might refer to them, so we need to move them // to the topmost position. - while (const auto *let = produce.as()) { - add_lets.emplace_back(let->name, let->value); - produce = let->body; - } + produce = peel_lets(produce, &add_lets); // Only one branch runs at a time, so a single trailing fence // after the whole (specialized) production is equivalent to, @@ -1853,9 +1850,7 @@ class InjectFunctionRealization : public IRMutator { internal_assert(producer.defined()); // Rewrap the loop in the containing lets. - for (const auto &[var, value] : reverse_view(add_lets)) { - producer = LetStmt::make(var, value, producer); - } + producer = rewrap_all_lets(producer, add_lets); // The original bounds of the loop nests (without any loop-fusion) auto bounds = CollectBounds::collect_bounds(producer); diff --git a/src/Simplify_Stmts.cpp b/src/Simplify_Stmts.cpp index 381ebb5a6177..c0e942a177ca 100644 --- a/src/Simplify_Stmts.cpp +++ b/src/Simplify_Stmts.cpp @@ -471,21 +471,14 @@ Stmt Simplify::visit(const Evaluate *op) { // Rewrite Lets inside an evaluate as LetStmts outside the Evaluate. vector> lets; - while (const Let *let = value.as()) { - lets.emplace_back(let->name, let->value); - value = let->body; - } + value = peel_lets(value, &lets); if (value.same_as(op->value)) { internal_assert(lets.empty()); return op; } else { // Rewrap the lets outside the evaluate node - Stmt stmt = Evaluate::make(value); - for (const auto &[var, value] : reverse_view(lets)) { - stmt = LetStmt::make(var, value, stmt); - } - return stmt; + return rewrap_all_lets(Evaluate::make(value), lets); } } diff --git a/src/SplitTuples.cpp b/src/SplitTuples.cpp index bdd9c001b593..77aee6b6dca8 100644 --- a/src/SplitTuples.cpp +++ b/src/SplitTuples.cpp @@ -181,11 +181,7 @@ class SplitTuples : public IRMutator { aliases = aliases && (a[i] == b[i]); } // Might need some of the containing lets - for (const auto &[var, value] : reverse_view(lets)) { - if (expr_uses_var(aliases, var)) { - aliases = Let::make(var, value, aliases); - } - } + aliases = rewrap_used_lets(aliases, lets); return !can_prove(!aliases); } @@ -281,13 +277,7 @@ class SplitTuples : public IRMutator { provides.push_back(Provide::make(name, {val}, args, op->predicate)); } - s = Block::make(provides); - - while (!lets.empty()) { - auto p = lets.back(); - lets.pop_back(); - s = LetStmt::make(p.first, p.second, s); - } + s = rewrap_all_lets(Block::make(provides), lets); } if (atomic && separate_atomic_nodes_per_store) { @@ -422,11 +412,7 @@ class SplitScatterGather : public IRMutator { body = substitute(op->name, gather_replacement, body); body = mutate(body); - for (const auto &[var, value] : reverse_view(lets)) { - body = LetStmt::make(var, value, body); - } - - return body; + return rewrap_all_lets(body, lets); } Stmt visit(const LetStmt *op) override { @@ -451,11 +437,7 @@ class SplitScatterGather : public IRMutator { body = mutate(body); } - for (const auto &[var, value] : reverse_view(lets)) { - body = LetStmt::make(var, value, body); - } - - return body; + return rewrap_all_lets(body, lets); } Stmt visit(const Provide *op) override { @@ -499,10 +481,7 @@ class SplitScatterGather : public IRMutator { bundle = common_subexpression_elimination(bundle); vector> lets; - while (const Let *let = bundle.as()) { - lets.emplace_back(let->name, let->value); - bundle = let->body; - } + bundle = peel_lets(bundle, &lets); const Call *c = bundle.as(); internal_assert(c && c->is_intrinsic(Call::bundle)); for (size_t i = 0; i < exprs.size(); i++) { @@ -515,11 +494,7 @@ class SplitScatterGather : public IRMutator { } } - for (const auto &[var, value] : reverse_view(lets)) { - s = LetStmt::make(var, value, s); - } - - return s; + return rewrap_all_lets(s, lets); } }; diff --git a/src/VectorizeLoops.cpp b/src/VectorizeLoops.cpp index 20e84b65a17e..82f77040a744 100644 --- a/src/VectorizeLoops.cpp +++ b/src/VectorizeLoops.cpp @@ -1359,9 +1359,7 @@ class VectorSubs : public IRMutator { s = SerializeLoops()(s); } // We'll need the original scalar versions of any containing lets. - for (const auto &[var, value] : reverse_view(containing_lets)) { - s = LetStmt::make(var, value, s); - } + s = rewrap_all_lets(s, containing_lets); for (int ix = vectorized_vars.size() - 1; ix >= 0; ix--) { s = For::make(vectorized_vars[ix].name, vectorized_vars[ix].min, @@ -1608,12 +1606,7 @@ class LiftVectorizableExprsOutOfAllAtomicNodes : public IRMutator { LiftVectorizableExprsOutOfSingleAtomicNode lifter(finder.liftable); Stmt new_body = lifter.mutate(op->body); new_body = op->with(new_body); - while (!lifter.lifted.empty()) { - auto p = lifter.lifted.back(); - new_body = LetStmt::make(p.first, p.second, new_body); - lifter.lifted.pop_back(); - } - return new_body; + return rewrap_all_lets(new_body, lifter.lifted); } const map &env; From 8f984733f2297ec7536936c1880e8fc8e51e9c23 Mon Sep 17 00:00:00 2001 From: Andrew Adams Date: Tue, 11 Aug 2026 11:40:00 -0700 Subject: [PATCH 2/3] Keep names in the used set when rewrapping lets IRGraphVisitor memoizes nodes with a refcount above one, so a name removed from the set can't be re-added by a later traversal that reaches it through a shared subexpression. Testing without removing keeps rewrap_used_lets as conservative as its documentation claims. Co-authored-by: Claude Opus 5 --- src/IROperator.cpp | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/src/IROperator.cpp b/src/IROperator.cpp index 927ba7e78914..1d7cb5510961 100644 --- a/src/IROperator.cpp +++ b/src/IROperator.cpp @@ -4,6 +4,7 @@ #include #include #include +#include #include #include "CSE.h" @@ -1161,7 +1162,7 @@ Body rewrap_used_lets_impl(const Body &body, body.accept(&used); Body result = body; for (const auto &[name, value] : reverse_view(lets)) { - if (used.names.erase(name)) { + if (used.names.count(name)) { value.accept(&used); if constexpr (std::is_same_v) { result = Let::make(name, value, result); From 90c1c4b8bdd72aeb90d9709cf69296bff00fe838 Mon Sep 17 00:00:00 2001 From: Andrew Adams Date: Tue, 11 Aug 2026 11:44:58 -0700 Subject: [PATCH 3/3] Peel lets as name/value pairs in Deinterleave Holding a vector of const LetStmt * only works while some handle keeps the chain alive, which is an easy thing to get wrong for no measurable gain here. Use the same name/value pairs as the rest of the compiler. Co-authored-by: Claude Opus 5 --- src/Deinterleave.cpp | 33 ++++++++++----------------------- 1 file changed, 10 insertions(+), 23 deletions(-) diff --git a/src/Deinterleave.cpp b/src/Deinterleave.cpp index 878f96fa493a..18a3d5f140cc 100644 --- a/src/Deinterleave.cpp +++ b/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 &let_stmts; + std::vector> &let_stmts; std::vector &stores; StoreCollector(const std::string &name, int stride, int ms, - std::vector &lets, std::vector &ss) + std::vector> &lets, + std::vector &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 potential_lets; + std::vector> 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())); + 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 lets, std::vector &stores) { + std::vector> lets, + std::vector &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(); - const Store *store = op->first.as(); - // Gather all the let stmts surrounding the first. - std::vector let_stmts; - while (let) { - let_stmts.emplace_back(let); - store = let->body.as(); - let = let->body.as(); - } + std::vector> let_stmts; + const Store *store = peel_lets(op->first, &let_stmts).as(); // 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(); - 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.