@@ -239,10 +239,39 @@ class CSEEveryExprInStmt : public IRMutator {
239239 }
240240 const Call *bundle = Call::as_intrinsic (dummy, {Call::bundle});
241241 internal_assert (bundle && bundle->args .size () == 2 );
242- Stmt s = Store::make (op->name , bundle->args [0 ], bundle->args [1 ],
242+
243+ Expr value = bundle->args [0 ], index = bundle->args [1 ];
244+
245+ // Figure out which ones are actually needed by the index
246+
247+ auto add_all_vars_to_set = [&](const Expr &e, std::set<std::string> &s) {
248+ visit_with (e, [&](auto *, const Variable *var) {
249+ s.insert (var->name );
250+ });
251+ };
252+
253+ std::set<string> index_lets;
254+ add_all_vars_to_set (index, index_lets);
255+ for (const auto &[var, val] : reverse_view (lets)) {
256+ if (index_lets.count (var)) {
257+ add_all_vars_to_set (val, index_lets);
258+ }
259+ }
260+
261+ vector<pair<string, Expr>> deferred;
262+ for (const auto &[var, val] : reverse_view (lets)) {
263+ if (index_lets.count (var)) {
264+ deferred.emplace_back (var, val);
265+ } else {
266+ value = Let::make (var, val, value);
267+ }
268+ }
269+
270+ Stmt s = Store::make (op->name , value, index,
243271 op->param , mutate (op->predicate ), op->alignment );
244- for (const auto &[var, value] : reverse_view (lets)) {
245- s = LetStmt::make (var, value, s);
272+
273+ for (const auto &[var, val] : deferred) {
274+ s = LetStmt::make (var, val, s);
246275 }
247276 return s;
248277 }
0 commit comments