From 4e4612d751ccdd8ca53f0ae6a79190239ca08532 Mon Sep 17 00:00:00 2001 From: Alex Reinking Date: Wed, 19 Aug 2026 12:31:30 -0400 Subject: [PATCH 1/2] Add rename_ir() primitive and consolidate duplicate IR-renaming code substitute() only rewrites Variable/Let/LetStmt/For, so several call sites that needed to rename Allocate/Free/Load/Store identifiers (or apply a deterministic prefix/rename policy to a whole subtree) each hand-rolled their own small IRMutator. Add Renamer/rename_ir() in Rename.h, a template (not std::function-boxed, matching the LambdaMutator/mutate_with idiom in IRMutator.h) that applies a memoized name -> name policy across every name-bearing IR node. Migrate the six call sites whose renaming is a flat, injective function of the name (no shadowing elimination needed): Qualify.cpp, ParallelRVar.cpp's RenameFreeVars, Simplify.cpp's can_prove debug canonicalizer, FuseGPUThreadLoops.cpp's register allocation renamer, and both of CodeGen_D3D12Compute_Dev.cpp's renamers. UniquifyVariableNames.cpp is intentionally left alone: its job is to eliminate shadowing (force simultaneously-live same-named bindings to diverge), which needs per-binding-occurrence scope tracking that a pure name->name policy can't express. The D3D12 shared-allocation renamer keeps its original per-Allocate- occurrence unique_name() call (via a small mutate_with lambda) rather than a single whole-Stmt rename_ir call, since two distinct shared allocations can legitimately share an original name (e.g. GuardWithIf-duplicated boundary branches), and HLSL requires globally unique top-level declaration names even across mutually-exclusive branches. Co-Authored-By: Claude Sonnet 5 --- src/CMakeLists.txt | 1 + src/CodeGen_D3D12Compute_Dev.cpp | 173 +++++++---------------------- src/FuseGPUThreadLoops.cpp | 28 +---- src/ParallelRVar.cpp | 40 ++----- src/Qualify.cpp | 33 +----- src/Rename.h | 185 +++++++++++++++++++++++++++++++ src/Simplify.cpp | 57 +++++----- 7 files changed, 271 insertions(+), 246 deletions(-) create mode 100644 src/Rename.h diff --git a/src/CMakeLists.txt b/src/CMakeLists.txt index 7727588fe764..34d9ebbe271c 100644 --- a/src/CMakeLists.txt +++ b/src/CMakeLists.txt @@ -190,6 +190,7 @@ target_sources( RemoveDeadAllocations.h RemoveExternLoops.h RemoveUndef.h + Rename.h runtime/HalideBuffer.h runtime/HalideRuntime.h Schedule.h diff --git a/src/CodeGen_D3D12Compute_Dev.cpp b/src/CodeGen_D3D12Compute_Dev.cpp index f58b43452271..394a33c84077 100644 --- a/src/CodeGen_D3D12Compute_Dev.cpp +++ b/src/CodeGen_D3D12Compute_Dev.cpp @@ -1,5 +1,6 @@ #include #include +#include #include #include @@ -15,6 +16,8 @@ #include "IREquality.h" #include "IRMutator.h" #include "IROperator.h" +#include "IRVisitor.h" +#include "Rename.h" #include "Simplify.h" #include "StrictifyFloat.h" @@ -1758,79 +1761,41 @@ void CodeGen_D3D12Compute_Dev::CodeGen_D3D12Compute_C::add_kernel(Stmt s, constants[i].size += constants[i - 1].size; } - // Find all the shared allocations, uniquify their names, - // and declare them at global scope. - class FindSharedAllocationsAndUniquify : public IRMutator { - using IRMutator::visit; - Stmt visit(const Allocate *op) override { - if (is_shared_allocation(op)) { - // Because these will go in global scope, - // we need to ensure they have unique names. - std::string new_name = unique_name(op->name); - replacements[op->name] = new_name; - - auto new_extents = mutate(op->extents); - Stmt new_body = mutate(op->body); - Expr new_condition = mutate(op->condition); - Expr new_new_expr; - if (op->new_expr.defined()) { - new_new_expr = mutate(op->new_expr); - } - - Stmt new_alloc = Allocate::make(new_name, op->type, op->memory_type, new_extents, - std::move(new_condition), std::move(new_body), - std::move(new_new_expr), op->free_function, op->padding); - - allocs.push_back(new_alloc); - replacements.erase(op->name); - return new_alloc; - } else { - return IRMutator::visit(op); + // Find all the shared allocations, uniquify their names, and declare + // them at global scope. Two shared allocations can legitimately share + // an original name (e.g. mutually-exclusive GuardWithIf boundary + // branches realizing the same Func), so each occurrence needs its own + // fresh unique_name() call, not one shared across the whole Stmt -- + // otherwise two such allocations would collide on one new name and + // emit duplicate groupshared declarations. + vector allocs; + s = mutate_with( + s, + [&](auto *self, const Allocate *op) -> Stmt { + if (!is_shared_allocation(op)) { + return self->visit_base(op); } - } - - Stmt visit(const Free *op) override { - auto it = replacements.find(op->name); - if (it != replacements.end()) { - return Free::make(it->second); - } else { - return IRMutator::visit(op); - } - } - - Expr visit(const Load *op) override { - auto it = replacements.find(op->name); - if (it != replacements.end()) { - return Load::make(op->type, it->second, - mutate(op->index), op->image, op->param, - mutate(op->predicate), op->alignment, op->is_streaming); - } else { - return IRMutator::visit(op); - } - } - - Stmt visit(const Store *op) override { - auto it = replacements.find(op->name); - if (it != replacements.end()) { - return Store::make(it->second, mutate(op->value), - mutate(op->index), op->param, - mutate(op->predicate), op->alignment, op->is_streaming); - } else { - return IRMutator::visit(op); - } - } - - std::map replacements; - friend class CodeGen_D3D12Compute_Dev::CodeGen_D3D12Compute_C; - vector allocs; - }; - - FindSharedAllocationsAndUniquify fsa; - s = fsa(s); + auto new_extents = self->mutate(op->extents); + Stmt new_body = self->mutate(op->body); + Expr new_condition = self->mutate(op->condition); + Expr new_new_expr = op->new_expr.defined() ? self->mutate(op->new_expr) : Expr(); + + // Because this will go in global scope, we need to ensure it + // has a unique name. + string new_name = unique_name(op->name); + new_body = rename_ir(new_body, [&](const string &name) { + return name == op->name ? new_name : name; + }); + + Stmt new_alloc = Allocate::make(new_name, op->type, op->memory_type, new_extents, + new_condition, new_body, new_new_expr, op->free_function, op->padding); + allocs.push_back(new_alloc); + return new_alloc; + }); const int sm = target.get_d3d12compute_capability_lower_bound(); uint32_t total_shared_bytes = 0; - for (const Stmt &sop : fsa.allocs) { + for (const Stmt &sop : allocs) { const Allocate *op = sop.as(); internal_assert(op->extents.size() == 1); internal_assert(op->type.lanes() == 1); @@ -1922,70 +1887,12 @@ void CodeGen_D3D12Compute_Dev::CodeGen_D3D12Compute_C::add_kernel(Stmt s, dxc_renames[arg.name] = name + "_" + arg.name; } - // Mutate Load/Store/Variable nodes in the body to use the prefixed names. - class RenameKernelArgs : public IRMutator { - using IRMutator::visit; - const std::map &renames; - Expr visit(const Load *op) override { - auto it = renames.find(op->name); - if (it != renames.end()) { - return Load::make(op->type, it->second, - mutate(op->index), op->image, op->param, - mutate(op->predicate), op->alignment, op->is_streaming); - } - return IRMutator::visit(op); - } - Stmt visit(const Store *op) override { - auto it = renames.find(op->name); - if (it != renames.end()) { - return Store::make(it->second, mutate(op->value), - mutate(op->index), op->param, - mutate(op->predicate), op->alignment, op->is_streaming); - } - return IRMutator::visit(op); - } - Expr visit(const Variable *op) override { - auto it = renames.find(op->name); - if (it != renames.end()) { - return Variable::make(op->type, it->second, - op->image, op->param, op->reduction_domain); - } - return IRMutator::visit(op); - } - Expr visit(const Call *op) override { - // image_load/image_store carry the buffer name as args[0] StringImm. - if ((op->is_intrinsic(Call::image_load) || op->is_intrinsic(Call::image_store)) && - !op->args.empty()) { - // args[0] is the buffer name as StringImm, possibly - // wrapped in Broadcast for vectorized texture access. - const StringImm *name_imm = op->args[0].as(); - if (!name_imm) { - if (const Broadcast *b = op->args[0].as()) { - name_imm = b->value.as(); - } - } - if (name_imm) { - auto it = renames.find(name_imm->value); - if (it != renames.end()) { - vector new_args = op->args; - Expr renamed = StringImm::make(it->second); - new_args[0] = op->args[0].as() ? Broadcast::make(renamed, op->args[0].as()->lanes) : renamed; - for (size_t i = 1; i < new_args.size(); ++i) { - new_args[i] = mutate(new_args[i]); - } - return op->with(new_args); - } - } - } - return IRMutator::visit(op); - } - - public: - RenameKernelArgs(const std::map &r) - : renames(r) { - } - }; - s = RenameKernelArgs(dxc_renames)(s); + // Rename Load/Store/Variable/image_load/image_store references in + // the body to use the prefixed names. + s = rename_ir(s, [&](const string &name) { + auto it = dxc_renames.find(name); + return it != dxc_renames.end() ? it->second : name; + }); // Declare all resources globally with explicit register bindings and // put scalar uniforms in a per-kernel constant buffer. The runtime binds: diff --git a/src/FuseGPUThreadLoops.cpp b/src/FuseGPUThreadLoops.cpp index a444273c9a46..3556d9d3a5c3 100644 --- a/src/FuseGPUThreadLoops.cpp +++ b/src/FuseGPUThreadLoops.cpp @@ -14,6 +14,7 @@ #include "IROperator.h" #include "IRPrinter.h" #include "Monotonic.h" +#include "Rename.h" #include "Simplify.h" #include "Solve.h" #include "Substitute.h" @@ -1302,30 +1303,9 @@ class ExtractRegisterAllocations : public IRMutator { Stmt body = op->body; if (block_alloc_names.count(op->name)) { name = unique_name(op->name); - const string &from = op->name; - const string &to = name; - body = mutate_with( - body, - [&](auto *self, const Load *load) -> Expr { - if (load->name == from) { - return Load::make(load->type, to, self->mutate(load->index), load->image, load->param, - self->mutate(load->predicate), load->alignment, load->is_streaming); - } - return self->visit_base(load); - }, - [&](auto *self, const Store *store) -> Stmt { - if (store->name == from) { - return Store::make(to, self->mutate(store->value), self->mutate(store->index), store->param, - self->mutate(store->predicate), store->alignment, store->is_streaming); - } - return self->visit_base(store); - }, - [&](auto *, const Free *free) -> Stmt { - if (free->name == from) { - return Free::make(to); - } - return free; - }); + body = rename_ir(body, [&](const string &n) { + return n == op->name ? name : n; + }); } RegisterAllocation alloc; diff --git a/src/ParallelRVar.cpp b/src/ParallelRVar.cpp index 538cd144f449..25edef8a90e0 100644 --- a/src/ParallelRVar.cpp +++ b/src/ParallelRVar.cpp @@ -7,13 +7,13 @@ #include "IRMutator.h" #include "IROperator.h" #include "IRVisitor.h" +#include "Rename.h" #include "Simplify.h" #include "Substitute.h" namespace Halide { namespace Internal { -using std::map; using std::string; using std::vector; @@ -49,32 +49,12 @@ class FindLoads : public IRVisitor { vector> loads; }; -/** Rename all free variables to unique new names. */ -class RenameFreeVars : public IRMutator { - using IRMutator::visit; - - map new_names; - - Expr visit(const Variable *op) override { - if (!op->param.defined() && !op->image.defined()) { - return Variable::make(op->type, get_new_name(op->name)); - } else { - return op; - } - } - -public: - string get_new_name(const string &s) { - map::iterator iter = new_names.find(s); - if (iter != new_names.end()) { - return iter->second; - } else { - string new_name = s + "$_"; - new_names[s] = new_name; - return new_name; - } - } -}; +/** A renaming policy that appends a fixed suffix to every free variable, so + * that it can stand in for a hypothetical distinct thread's copy of the + * same variable. */ +string rename_free_var(const string &name) { + return name + "$_"; +} /** Substitute in boolean expressions. */ class SubstituteInBooleanLets : public IRMutator { @@ -109,14 +89,14 @@ bool can_parallelize_rvar(const string &v, } // Make an expr representing the store done by a different thread. - RenameFreeVars renamer; + Renamer renamer(rename_free_var); auto other_store = renamer(args); // Construct an expression which is true when the two threads are // in fact two different threads. We'll use this liberally in the // following conditions to give the simplifier the best chance. Expr distinct_v = (Variable::make(Int(32), v) != - Variable::make(Int(32), renamer.get_new_name(v))); + Variable::make(Int(32), renamer.new_name_for(v))); // Construct an expression which is true if there's a collision // between this thread's store and the other thread's store. @@ -141,7 +121,7 @@ bool can_parallelize_rvar(const string &v, for (const auto &rv : rvars) { Interval in = Interval(rv.min, simplify(rv.min + rv.extent - 1)); bounds.push(rv.var, in); - bounds.push(renamer.get_new_name(rv.var), in); + bounds.push(renamer.new_name_for(rv.var), in); } // Add the definition's predicate if there is any diff --git a/src/Qualify.cpp b/src/Qualify.cpp index ef4a32abb355..94d5d065e265 100644 --- a/src/Qualify.cpp +++ b/src/Qualify.cpp @@ -1,42 +1,13 @@ #include "Qualify.h" -#include "IRMutator.h" +#include "Rename.h" namespace Halide { namespace Internal { using std::string; -namespace { - -// Prefix all names in an expression with some string. -class QualifyExpr : public IRMutator { - using IRMutator::visit; - - const string &prefix; - - Expr visit(const Variable *v) override { - if (v->param.defined()) { - return v; - } else { - return Variable::make(v->type, prefix + v->name, v->reduction_domain); - } - } - Expr visit(const Let *op) override { - Expr value = mutate(op->value); - Expr body = mutate(op->body); - return Let::make(prefix + op->name, value, body); - } - -public: - QualifyExpr(const string &p) - : prefix(p) { - } -}; - -} // namespace - Expr qualify(const string &prefix, const Expr &value) { - return QualifyExpr(prefix)(value); + return rename_ir(value, [&](const string &name) { return prefix + name; }); } } // namespace Internal diff --git a/src/Rename.h b/src/Rename.h new file mode 100644 index 000000000000..c40479aef3e6 --- /dev/null +++ b/src/Rename.h @@ -0,0 +1,185 @@ +#ifndef HALIDE_RENAME_H +#define HALIDE_RENAME_H + +/** \file + * Defines a generic primitive for renaming the identifiers that IR nodes + * carry directly, as opposed to substitute() (Substitute.h), which only + * follows references to a name via Variable nodes. */ + +#include +#include +#include +#include + +#include "IR.h" +#include "IRMutator.h" + +namespace Halide { +namespace Internal { + +/** Applies a renaming policy to every IR node that carries a name directly: + * Variable, Let, LetStmt, For, Allocate, Free, Load, Store, and the buffer + * name argument of image_load/image_store Call nodes. Never renames a + * Variable that carries a Parameter or Buffer, since those are referenced + * by identity, not by name. + * + * The policy is called at most once per distinct name; its result is + * memoized and reused for every later occurrence of that name, so it is + * safe (and typically the point) for the policy to have side effects, e.g. + * calling unique_name(). + * + * This is a flat, non-scope-aware rename: every occurrence of a given + * original name anywhere in the subtree -- whether it's a binding + * occurrence (the name on a Let/For/Allocate) or a use -- is replaced with + * the same new name, regardless of which (if any) enclosing binder + * introduced it. For this to be safe, the policy must be injective: it + * must never map two different input names to the same output name. Given + * that, renaming is safe under arbitrary nesting/shadowing, because it + * relabels symbols without changing which binder each use resolves to. + * + * This does not, and cannot, eliminate shadowing: if two bindings that are + * simultaneously in scope must end up with provably distinct names (e.g. to + * satisfy some later pass that assumes distinct names never collide), use + * uniquify_variable_names instead. */ +template +class Renamer : public IRMutator { + Policy policy; + std::unordered_map renamed; + + using IRMutator::visit; + + Expr visit(const Variable *op) override { + if (op->param.defined() || op->image.defined()) { + return op; + } + const std::string &new_name = new_name_for(op->name); + if (new_name == op->name) { + return op; + } + return Variable::make(op->type, new_name, op->image, op->param, op->reduction_domain); + } + + template + auto visit_let(const LetOrLetStmt *op) -> decltype(op->body) { + const std::string &new_name = new_name_for(op->name); + if (new_name == op->name) { + return IRMutator::visit(op); + } + auto new_value = mutate(op->value); + auto new_body = mutate(op->body); + return LetOrLetStmt::make(new_name, new_value, new_body); + } + + Expr visit(const Let *op) override { + return visit_let(op); + } + + Stmt visit(const LetStmt *op) override { + return visit_let(op); + } + + Stmt visit(const For *op) override { + const std::string &new_name = new_name_for(op->name); + if (new_name == op->name) { + return IRMutator::visit(op); + } + Expr new_min = mutate(op->min); + Expr new_max = mutate(op->max); + Stmt new_body = mutate(op->body); + return For::make(new_name, new_min, new_max, op->for_type, op->partition_policy, op->device_api, new_body); + } + + Stmt visit(const Allocate *op) override { + const std::string &new_name = new_name_for(op->name); + if (new_name == op->name) { + return IRMutator::visit(op); + } + std::vector new_extents = mutate(op->extents); + Stmt new_body = mutate(op->body); + Expr new_condition = mutate(op->condition); + Expr new_new_expr = op->new_expr.defined() ? mutate(op->new_expr) : Expr(); + return Allocate::make(new_name, op->type, op->memory_type, new_extents, + new_condition, new_body, new_new_expr, op->free_function, op->padding); + } + + Stmt visit(const Free *op) override { + const std::string &new_name = new_name_for(op->name); + if (new_name == op->name) { + return IRMutator::visit(op); + } + return Free::make(new_name); + } + + Expr visit(const Load *op) override { + const std::string &new_name = new_name_for(op->name); + if (new_name == op->name) { + return IRMutator::visit(op); + } + return Load::make(op->type, new_name, mutate(op->index), op->image, op->param, + mutate(op->predicate), op->alignment, op->is_streaming); + } + + Stmt visit(const Store *op) override { + const std::string &new_name = new_name_for(op->name); + if (new_name == op->name) { + return IRMutator::visit(op); + } + return Store::make(new_name, mutate(op->value), mutate(op->index), op->param, + mutate(op->predicate), op->alignment, op->is_streaming); + } + + Expr visit(const Call *op) override { + // image_load/image_store carry the buffer name as args[0], a + // StringImm possibly wrapped in a Broadcast for vectorized access. + if (op->args.empty() || + !(op->is_intrinsic(Call::image_load) || op->is_intrinsic(Call::image_store))) { + return IRMutator::visit(op); + } + const Broadcast *broadcast = op->args[0].as(); + const StringImm *name_imm = broadcast ? broadcast->value.as() : op->args[0].as(); + if (!name_imm) { + return IRMutator::visit(op); + } + const std::string &new_name = new_name_for(name_imm->value); + if (new_name == name_imm->value) { + return IRMutator::visit(op); + } + std::vector new_args = op->args; + Expr new_imm = StringImm::make(new_name); + new_args[0] = broadcast ? Broadcast::make(new_imm, broadcast->lanes) : new_imm; + for (size_t i = 1; i < new_args.size(); i++) { + new_args[i] = mutate(new_args[i]); + } + return op->with(new_args); + } + +public: + explicit Renamer(Policy policy) + : policy(std::move(policy)) { + } + + /** Get the new name for a given old name, computing and memoizing it via + * the policy if this is the first time this name has been seen. */ + const std::string &new_name_for(const std::string &name) { + auto it = renamed.find(name); + if (it != renamed.end()) { + return it->second; + } + return renamed.emplace(name, policy(name)).first->second; + } +}; + +template +Expr rename_ir(const Expr &e, Policy policy) { + return Renamer(std::move(policy))(e); +} + +template +Stmt rename_ir(const Stmt &s, Policy policy) { + return Renamer(std::move(policy))(s); +} + +} // namespace Internal +} // namespace Halide + +#endif diff --git a/src/Simplify.cpp b/src/Simplify.cpp index 75b908ce8202..97e1ca2b2629 100644 --- a/src/Simplify.cpp +++ b/src/Simplify.cpp @@ -3,6 +3,7 @@ #include "CSE.h" #include "IRMutator.h" +#include "Rename.h" #include "Substitute.h" namespace Halide { @@ -446,47 +447,47 @@ bool can_prove(Expr e, const Scope &bounds) { // simplifier weaknesses if (!is_const(e)) { debug(1) << [&]() -> std::string { - struct RenameVariables : public IRMutator { - using IRMutator::visit; - - Expr visit(const Variable *op) override { - auto it = vars.find(op->name); - if (const std::string *n = lets.find(op->name)) { - return Variable::make(op->type, *n); - } else if (it == vars.end()) { - std::string name = "v" + std::to_string(count++); - vars[op->name] = name; - out_vars.emplace_back(op->type, name); - return Variable::make(op->type, name); - } else { - return Variable::make(op->type, it->second); + // Canonicalize every variable name to v0, v1, ... in encounter + // order, so unrelated failed proofs render identically. + int count = 0; + Expr renamed = rename_ir(e, [&](const string &) { + return "v" + std::to_string(count++); + }); + + // Collect the free variables of the renamed expression, for + // the random probing below. Needs its own scope-tracking, since + // renaming can (harmlessly) reuse a name across unrelated + // bindings that were never simultaneously live. + struct FindFreeVars : public IRVisitor { + using IRVisitor::visit; + + void visit(const Variable *op) override { + if (!lets.contains(op->name)) { + free_vars[op->name] = op->type; } } - Expr visit(const Let *op) override { - std::string name = "v" + std::to_string(count++); - ScopedBinding bind(lets, op->name, name); - return Let::make(name, mutate(op->value), mutate(op->body)); + void visit(const Let *op) override { + op->value.accept(this); + ScopedBinding<> bind(lets, op->name); + op->body.accept(this); } - int count = 0; - map vars; - Scope lets; - std::vector> out_vars; - } renamer; - - Expr renamed = renamer(e); + Scope<> lets; + map free_vars; + } finder; + renamed.accept(&finder); // Look for a concrete counter-example with random probing static std::mt19937 rng(0); for (int i = 0; i < 100; i++) { map s; - for (const auto &p : renamer.out_vars) { - if (p.first.is_handle()) { + for (const auto &p : finder.free_vars) { + if (p.second.is_handle()) { // This aint gonna work return ""; } - s[p.second] = make_const(p.first, (int)(rng() & 0xffff) - 0x7fff); + s[p.first] = make_const(p.second, (int)(rng() & 0xffff) - 0x7fff); } Expr probe = unwrap_tags(simplify(substitute(s, renamed))); if (!is_const_one(probe)) { From 087f1d48a8664617a345ccf74f54e0b1da746e76 Mon Sep 17 00:00:00 2001 From: Alex Reinking Date: Wed, 19 Aug 2026 13:08:12 -0400 Subject: [PATCH 2/2] Consolidate free-variable-collector duplicates; fix Let-shadowing bugs Add FreeVariables (src/FreeVariables.h/.cpp), a Scope-aware IRVisitor that collects the free variables of an Expr/Stmt subtree, mirroring the accumulator style of Closure::include(). Replace three near- identical private reimplementations of the same algorithm with it: UniquifyVariableNames.cpp's FindFreeVars, Simplify.cpp's can_prove debug helper, and SimplifyCorrelatedDifferences.cpp's TrackFreeVars. Also add UnboundVarChecker, a graph-aware sibling (IRGraphVisitor, since these validate raw user-authored condition Exprs that may share subexpressions) used to check whether an Expr references a free Var/RVar and, optionally, calls a Halide Func. RDom.cpp's CheckRDomBounds already did this correctly with its own Scope tracking; Func.cpp's CheckForFreeVars (guarding Stage::specialize) and Pipeline.cpp's Checker (guarding Pipeline::add_requirement) did not track Let-shadowing at all, so they rejected any condition containing so much as a self-contained Let, since they flag every non-Param/Image Variable node unconditionally. Both are migrated to UnboundVarChecker, fixing the bug. A check_func_calls flag preserves each site's prior func-call-checking behavior exactly (RDom.cpp and Pipeline.cpp checked for it; Func.cpp did not and still doesn't). Pipeline.cpp's Variable check previously exempted only Param-backed variables, not Image-backed ones; standardized to exempt both, matching the other two sites. Add test/correctness/unbound_var_checker_let_shadowing.cpp, verified to fail against the pre-fix logic before being added. Co-Authored-By: Claude Sonnet 5 --- src/CMakeLists.txt | 2 + src/FreeVariables.cpp | 86 ++++++++++++++++++ src/FreeVariables.h | 88 +++++++++++++++++++ src/Func.cpp | 18 +--- src/Pipeline.cpp | 31 ++----- src/RDom.cpp | 39 +------- src/Simplify.cpp | 27 +----- src/SimplifyCorrelatedDifferences.cpp | 34 ++----- src/UniquifyVariableNames.cpp | 56 ++---------- test/correctness/CMakeLists.txt | 1 + .../unbound_var_checker_let_shadowing.cpp | 42 +++++++++ 11 files changed, 248 insertions(+), 176 deletions(-) create mode 100644 src/FreeVariables.cpp create mode 100644 src/FreeVariables.h create mode 100644 test/correctness/unbound_var_checker_let_shadowing.cpp diff --git a/src/CMakeLists.txt b/src/CMakeLists.txt index 34d9ebbe271c..fea885aa1a69 100644 --- a/src/CMakeLists.txt +++ b/src/CMakeLists.txt @@ -123,6 +123,7 @@ target_sources( FindIntrinsics.h FlattenNestedRamps.h Float16.h + FreeVariables.h Func.h Function.h FunctionPtr.h @@ -307,6 +308,7 @@ target_sources( FindIntrinsics.cpp FlattenNestedRamps.cpp Float16.cpp + FreeVariables.cpp Func.cpp Function.cpp FuseGPUThreadLoops.cpp diff --git a/src/FreeVariables.cpp b/src/FreeVariables.cpp new file mode 100644 index 000000000000..296f0f268ec5 --- /dev/null +++ b/src/FreeVariables.cpp @@ -0,0 +1,86 @@ +#include "FreeVariables.h" + +namespace Halide { +namespace Internal { + +using std::string; +using std::vector; + +void FreeVariables::include(const Expr &e) { + e.accept(this); +} + +void FreeVariables::include(const Stmt &s) { + s.accept(this); +} + +void FreeVariables::visit(const Variable *op) { + if (!scope.contains(op->name)) { + vars[op->name] = op->type; + } +} + +template +void FreeVariables::visit_let(const LetOrLetStmt *op) { + vector> frame; + decltype(op->body) body; + do { + op->value.accept(this); + frame.emplace_back(scope, op->name); + body = op->body; + op = body.template as(); + } while (op); + body.accept(this); +} + +void FreeVariables::visit(const Let *op) { + visit_let(op); +} + +void FreeVariables::visit(const LetStmt *op) { + visit_let(op); +} + +void FreeVariables::visit(const For *op) { + op->min.accept(this); + op->max.accept(this); + ScopedBinding<> bind(scope, op->name); + op->body.accept(this); +} + +std::map find_free_vars(const Expr &e) { + FreeVariables finder; + finder.include(e); + return finder.vars; +} + +std::map find_free_vars(const Stmt &s) { + FreeVariables finder; + finder.include(s); + return finder.vars; +} + +UnboundVarChecker::UnboundVarChecker(bool check_func_calls) + : check_func_calls(check_func_calls) { +} + +void UnboundVarChecker::visit(const Variable *op) { + if (!op->param.defined() && !op->image.defined() && !scope.contains(op->name)) { + offending_var = op->name; + } +} + +void UnboundVarChecker::visit(const Let *op) { + ScopedBinding<> bind(scope, op->name); + IRGraphVisitor::visit(op); +} + +void UnboundVarChecker::visit(const Call *op) { + IRGraphVisitor::visit(op); + if (check_func_calls && op->call_type == Call::Halide) { + offending_func = op->name; + } +} + +} // namespace Internal +} // namespace Halide diff --git a/src/FreeVariables.h b/src/FreeVariables.h new file mode 100644 index 000000000000..b1b69e063975 --- /dev/null +++ b/src/FreeVariables.h @@ -0,0 +1,88 @@ +#ifndef HALIDE_FREE_VARIABLES_H +#define HALIDE_FREE_VARIABLES_H + +/** \file + * Defines a visitor that collects the free variables of a piece of IR. */ + +#include +#include + +#include "Expr.h" +#include "IRVisitor.h" +#include "Scope.h" + +namespace Halide { +namespace Internal { + +/** Collects every Variable reference in an Expr/Stmt that isn't bound by an + * enclosing Let, LetStmt, or For within the visited subtree(s), mapped to + * its type. Variables carrying a Parameter or Buffer are included just + * like any other free reference. + * + * include() may be called more than once (on more than one Expr/Stmt) to + * accumulate the free variables of several pieces of IR into one result, + * e.g. an Expr and then the values of the Lets it may later be wrapped + * in. */ +class FreeVariables : public IRVisitor { +public: + std::map vars; + + void include(const Expr &e); + void include(const Stmt &s); + +protected: + using IRVisitor::visit; + void visit(const Variable *op) override; + void visit(const Let *op) override; + void visit(const LetStmt *op) override; + void visit(const For *op) override; + +private: + Scope<> scope; + + template + void visit_let(const LetOrLetStmt *op); +}; + +/** Convenience wrapper for the common case of finding the free variables of + * a single Expr or Stmt. */ +// @{ +std::map find_free_vars(const Expr &e); +std::map find_free_vars(const Stmt &s); +// @} + +/** Checks whether a subtree references any Var/RVar not bound by an + * enclosing Let (recorded in offending_var, if any), and, if + * check_func_calls is set, whether it calls any Halide Func (recorded in + * offending_func, if any). A Variable carrying a Parameter or Buffer is + * not considered offending. Correctly respects Let-shadowing via a Scope. + * + * This is graph-aware (IRGraphVisitor) rather than tree-based, because the + * Exprs it's used to validate are raw, user-authored conditions that may + * share subexpressions and haven't been through CSE. + * + * Used to validate conditions passed to Stage::specialize, + * Pipeline::add_requirement, and the region bounds of an RDom -- none of + * which may depend on a Var or RVar. */ +class UnboundVarChecker : public IRGraphVisitor { +public: + std::string offending_var; + std::string offending_func; + + explicit UnboundVarChecker(bool check_func_calls = false); + +protected: + using IRGraphVisitor::visit; + void visit(const Variable *op) override; + void visit(const Let *op) override; + void visit(const Call *op) override; + +private: + bool check_func_calls; + Scope<> scope; +}; + +} // namespace Internal +} // namespace Halide + +#endif diff --git a/src/Func.cpp b/src/Func.cpp index 468188530c67..f9ac1eb4bab5 100644 --- a/src/Func.cpp +++ b/src/Func.cpp @@ -22,6 +22,7 @@ #include "Debug.h" #include "ExprUsesVar.h" #include "FindCalls.h" +#include "FreeVariables.h" #include "Func.h" #include "Function.h" #include "IR.h" @@ -1418,28 +1419,13 @@ Stage &Stage::fuse(const VarOrRVar &inner, const VarOrRVar &outer, const VarOrRV return *this; } -namespace Internal { -class CheckForFreeVars : public IRGraphVisitor { -public: - string offending_var; - -protected: - using IRGraphVisitor::visit; - void visit(const Variable *var) override { - if (!var->param.defined() && !var->image.defined()) { - offending_var = var->name; - } - } -}; -} // namespace Internal - Stage Stage::specialize(const Expr &condition) { user_assert(condition.type().is_bool()) << "Argument passed to specialize must be of type bool\n"; definition.schedule().touched() = true; // The condition may not depend on Vars or RVars - Internal::CheckForFreeVars check; + Internal::UnboundVarChecker check; condition.accept(&check); if (!check.offending_var.empty()) { user_error << "Specialization condition " << condition << " for " << name() diff --git a/src/Pipeline.cpp b/src/Pipeline.cpp index e935fd4852df..7a6f31296d7d 100644 --- a/src/Pipeline.cpp +++ b/src/Pipeline.cpp @@ -6,6 +6,7 @@ #include "CodeGen_Internal.h" #include "Deserialization.h" #include "FindCalls.h" +#include "FreeVariables.h" #include "Func.h" #include "IRVisitor.h" #include "InferArguments.h" @@ -866,30 +867,12 @@ void Pipeline::add_requirement(const Expr &condition, const std::vector &e user_assert(defined()) << "Pipeline is undefined\n"; // It is an error for a requirement to reference a Func or a Var - class Checker : public IRGraphVisitor { - using IRGraphVisitor::visit; - - void visit(const Variable *op) override { - if (!op->param.defined()) { - user_error << "Requirement " << condition << " refers to Var or RVar " << op->name << "\n"; - } - } - - void visit(const Call *op) override { - if (op->call_type == Call::Halide) { - user_error << "Requirement " << condition << " calls Func " << op->name << "\n"; - } - IRGraphVisitor::visit(op); - } - - const Expr &condition; - - public: - Checker(const Expr &c) - : condition(c) { - c.accept(this); - } - } checker(condition); + Internal::UnboundVarChecker checker(/*check_func_calls=*/true); + condition.accept(&checker); + user_assert(checker.offending_func.empty()) + << "Requirement " << condition << " calls Func " << checker.offending_func << "\n"; + user_assert(checker.offending_var.empty()) + << "Requirement " << condition << " refers to Var or RVar " << checker.offending_var << "\n"; Expr error = Internal::requirement_failed_error(condition, error_args); contents->requirements.emplace_back(Internal::AssertStmt::make(condition, error)); diff --git a/src/RDom.cpp b/src/RDom.cpp index 713bae283866..a24d751b4759 100644 --- a/src/RDom.cpp +++ b/src/RDom.cpp @@ -1,6 +1,7 @@ #include #include +#include "FreeVariables.h" #include "IREquality.h" #include "IROperator.h" #include "IRPrinter.h" @@ -97,38 +98,6 @@ RDom::RDom(const ReductionDomain &d) } } -namespace { -class CheckRDomBounds : public IRGraphVisitor { - - using IRGraphVisitor::visit; - - void visit(const Call *op) override { - IRGraphVisitor::visit(op); - if (op->call_type == Call::Halide) { - offending_func = op->name; - } - } - - void visit(const Variable *op) override { - if (!op->param.defined() && - !op->image.defined() && - !internal_vars.contains(op->name)) { - offending_free_var = op->name; - } - } - - void visit(const Let *op) override { - ScopedBinding bind(internal_vars, op->name, 0); - IRGraphVisitor::visit(op); - } - Scope internal_vars; - -public: - string offending_func; - string offending_free_var; -}; -} // namespace - void RDom::validate_min_extent(const Expr &min, const Expr &extent) { user_assert(lossless_cast(Int(32), min).defined()) << "RDom min cannot be represented as an int32: " << min; @@ -143,7 +112,7 @@ void RDom::initialize_from_region(const Region ®ion, string name) { std::vector vars; for (size_t i = 0; i < region.size(); i++) { - CheckRDomBounds checker; + UnboundVarChecker checker(/*check_func_calls=*/true); user_assert(region[i].min.defined() && region[i].extent.defined()) << "The RDom " << name << " may not be constructed with undefined Exprs.\n"; region[i].min.accept(&checker); @@ -155,12 +124,12 @@ void RDom::initialize_from_region(const Region ®ion, string name) { << " " << region[i].min << " ... " << region[i].extent << "\n" << "These depend on a call to the Func " << checker.offending_func << ".\n" << "The bounds of an RDom may not depend on a call to a Func.\n"; - user_assert(checker.offending_free_var.empty()) + user_assert(checker.offending_var.empty()) << "The bounds of the RDom " << name << " in dimension " << i << " are:\n" << " " << region[i].min << " ... " << region[i].extent << "\n" - << "These depend on the variable " << checker.offending_free_var << ".\n" + << "These depend on the variable " << checker.offending_var << ".\n" << "The bounds of an RDom may not depend on a free variable.\n"; std::string rvar_uniquifier; diff --git a/src/Simplify.cpp b/src/Simplify.cpp index 97e1ca2b2629..5c89023939f7 100644 --- a/src/Simplify.cpp +++ b/src/Simplify.cpp @@ -2,6 +2,7 @@ #include "Simplify_Internal.h" #include "CSE.h" +#include "FreeVariables.h" #include "IRMutator.h" #include "Rename.h" #include "Substitute.h" @@ -455,34 +456,14 @@ bool can_prove(Expr e, const Scope &bounds) { }); // Collect the free variables of the renamed expression, for - // the random probing below. Needs its own scope-tracking, since - // renaming can (harmlessly) reuse a name across unrelated - // bindings that were never simultaneously live. - struct FindFreeVars : public IRVisitor { - using IRVisitor::visit; - - void visit(const Variable *op) override { - if (!lets.contains(op->name)) { - free_vars[op->name] = op->type; - } - } - - void visit(const Let *op) override { - op->value.accept(this); - ScopedBinding<> bind(lets, op->name); - op->body.accept(this); - } - - Scope<> lets; - map free_vars; - } finder; - renamed.accept(&finder); + // the random probing below. + map free_vars = find_free_vars(renamed); // Look for a concrete counter-example with random probing static std::mt19937 rng(0); for (int i = 0; i < 100; i++) { map s; - for (const auto &p : finder.free_vars) { + for (const auto &p : free_vars) { if (p.second.is_handle()) { // This aint gonna work return ""; diff --git a/src/SimplifyCorrelatedDifferences.cpp b/src/SimplifyCorrelatedDifferences.cpp index dc90aa6d7cef..94215875fd67 100644 --- a/src/SimplifyCorrelatedDifferences.cpp +++ b/src/SimplifyCorrelatedDifferences.cpp @@ -2,6 +2,7 @@ #include "CSE.h" #include "ExprUsesVar.h" +#include "FreeVariables.h" #include "IRMatch.h" #include "IRMutator.h" #include "IROperator.h" @@ -175,31 +176,6 @@ class SimplifyCorrelatedDifferences : public IRMutator { return s; } - // Add the names of any free variables in an expr to the provided set - void track_free_vars(const Expr &e, std::set *vars) { - class TrackFreeVars : public IRVisitor { - protected: - using IRVisitor::visit; - void visit(const Variable *op) override { - if (!scope.contains(op->name)) { - vars->insert(op->name); - } - } - void visit(const Let *op) override { - ScopedBinding<> bind(scope, op->name); - IRVisitor::visit(op); - } - - public: - std::set *vars; - Scope<> scope; - TrackFreeVars(std::set *vars) - : vars(vars) { - } - } tracker(vars); - tracker(e); - } - Expr cancel_correlated_subexpression(Expr e, const Expr &a, const Expr &b, bool correlated) { auto ma = is_monotonic(a, loop_var, monotonic); auto mb = is_monotonic(b, loop_var, monotonic); @@ -209,11 +185,11 @@ class SimplifyCorrelatedDifferences : public IRMutator { (ma == Monotonic::Increasing && mb == Monotonic::Decreasing && !correlated) || (ma == Monotonic::Decreasing && mb == Monotonic::Increasing && !correlated)) { - std::set vars; - track_free_vars(e, &vars); + FreeVariables free_vars; + free_vars.include(e); for (const auto &[var, value, may_substitute] : reverse_view(lets)) { - if (!may_substitute && vars.count(var)) { + if (!may_substitute && free_vars.vars.count(var)) { // We have to stop here. Can't continue // because there might be an outer let with // the same name that we *can* substitute in, @@ -221,7 +197,7 @@ class SimplifyCorrelatedDifferences : public IRMutator { // value. break; } - track_free_vars(value, &vars); + free_vars.include(value); e = Let::make(var, value, e); } e = common_subexpression_elimination(e); diff --git a/src/UniquifyVariableNames.cpp b/src/UniquifyVariableNames.cpp index c4485765862a..c7ff07c72f2b 100644 --- a/src/UniquifyVariableNames.cpp +++ b/src/UniquifyVariableNames.cpp @@ -1,7 +1,7 @@ #include "UniquifyVariableNames.h" +#include "FreeVariables.h" #include "IRMutator.h" #include "IROperator.h" -#include "IRVisitor.h" #include "Scope.h" namespace Halide { @@ -114,57 +114,15 @@ class UniquifyVariableNames : public IRMutator { } }; -class FindFreeVars : public IRVisitor { -protected: - using IRVisitor::visit; - - Scope<> scope; - - void visit(const Variable *op) override { - if (!scope.contains(op->name)) { - free_vars.push(op->name, op->name); - } - } - - template - void visit_let(const LetOrLetStmt *op) { - vector> frame; - decltype(op->body) body; - do { - op->value.accept(this); - frame.emplace_back(scope, op->name); - body = op->body; - op = body.template as(); - } while (op); - body.accept(this); - } - - void visit(const Let *op) override { - visit_let(op); - } - - void visit(const LetStmt *op) override { - visit_let(op); - } - - void visit(const For *op) override { - op->min.accept(this); - op->max.accept(this); - { - ScopedBinding<> bind(scope, op->name); - op->body.accept(this); - } - } - -public: - Scope free_vars; -}; } // namespace Stmt uniquify_variable_names(const Stmt &s) { - FindFreeVars finder; - s.accept(&finder); - return UniquifyVariableNames(&finder.free_vars)(s); + std::map free_vars = find_free_vars(s); + Scope free_var_names; + for (const auto &p : free_vars) { + free_var_names.push(p.first, p.first); + } + return UniquifyVariableNames(&free_var_names)(s); } } // namespace Internal diff --git a/test/correctness/CMakeLists.txt b/test/correctness/CMakeLists.txt index c82d5d5ca513..fbb318f48483 100644 --- a/test/correctness/CMakeLists.txt +++ b/test/correctness/CMakeLists.txt @@ -377,6 +377,7 @@ tests( tuple_update_ops.cpp two_vector_args.cpp typed_func.cpp + unbound_var_checker_let_shadowing.cpp undef.cpp uninitialized_read.cpp unique_func_image.cpp diff --git a/test/correctness/unbound_var_checker_let_shadowing.cpp b/test/correctness/unbound_var_checker_let_shadowing.cpp new file mode 100644 index 000000000000..5f80be902ac5 --- /dev/null +++ b/test/correctness/unbound_var_checker_let_shadowing.cpp @@ -0,0 +1,42 @@ +#include "Halide.h" +#include + +using namespace Halide; + +// Stage::specialize, Pipeline::add_requirement, and RDom's region bounds all +// reject conditions/bounds that depend on a free Var or RVar. The checks +// backing that must respect Let-shadowing: a Let that binds its own +// variable has no free variable at all, and must not be rejected. Prior to +// this test, the checks used by specialize() and add_requirement() did not +// track Let bindings, so they flagged the mere presence of any Variable +// node -- including one entirely bound by an enclosing Let -- as "free". +int main(int argc, char **argv) { + // A condition/bounds Expr with a Let that binds its own variable: it + // has no real free variable, since "t" never escapes this Let. + Expr self_contained_let = Internal::Let::make("t", 3, Internal::Variable::make(Int(32), "t") == 3); + + Var x; + Func f("f"); + f(x) = x; + + // Stage::specialize + f.specialize(self_contained_let); + + // Pipeline::add_requirement + Pipeline(f).add_requirement(self_contained_let, {}); + + // RDom region bounds + Expr let_min = Internal::Let::make("t", 0, Internal::Variable::make(Int(32), "t")); + RDom r(let_min, 10); + + Buffer result = f.realize({10}); + for (int i = 0; i < 10; i++) { + if (result(i) != i) { + printf("result(%d) = %d instead of %d\n", i, result(i), i); + return 1; + } + } + + printf("Success!\n"); + return 0; +}