Fix leaks caused by self-referential parameter constraints

src/Dimension.cpp

@@ -50,42 +50,42 @@ Expr Dimension::stride() const {
     return Variable::make(Int(32), s.str(), param);
 }
 
-Dimension Dimension::set_extent(Expr extent) {
+Dimension Dimension::set_extent(const Expr &extent) {
     // Propagate constant bounds into estimates as well.
     if (is_const(extent)) {
         param.set_extent_constraint_estimate(d, extent);
     }
-    param.set_extent_constraint(d, std::move(extent));
+    param.set_extent_constraint(d, extent);
     return *this;
 }
 
-Dimension Dimension::set_min(Expr min) {
+Dimension Dimension::set_min(const Expr &min) {
     // Propagate constant bounds into estimates as well.
     if (is_const(min)) {
         param.set_min_constraint_estimate(d, min);
     }
-    param.set_min_constraint(d, std::move(min));
+    param.set_min_constraint(d, min);
     return *this;
 }
 
-Dimension Dimension::set_stride(Expr stride) {
-    param.set_stride_constraint(d, std::move(stride));
+Dimension Dimension::set_stride(const Expr &stride) {
+    param.set_stride_constraint(d, stride);
     return *this;
 }
 
-Dimension Dimension::set_bounds(Expr min, Expr extent) {
-    return set_min(std::move(min)).set_extent(std::move(extent));
+Dimension Dimension::set_bounds(const Expr &min, const Expr &extent) {
+    return set_min(min).set_extent(extent);
 }
 
-Dimension Dimension::set_estimate(Expr min, Expr extent) {
+Dimension Dimension::set_estimate(const Expr &min, const Expr &extent) {
     // Update the estimates on the linked Func as well.
     // (This matters mainly for OutputImageParams.)
     // Note that while it's possible/legal for a Dimension to have an undefined
     // Func, you shouldn't ever call set_estimate on such an instance.
     internal_assert(f.defined());
     f.set_estimate(f.args()[d], min, extent);
-    param.set_min_constraint_estimate(d, std::move(min));
-    param.set_extent_constraint_estimate(d, std::move(extent));
+    param.set_min_constraint_estimate(d, min);
+    param.set_extent_constraint_estimate(d, extent);
     return *this;
 }
 

src/Dimension.h

@@ -34,7 +34,7 @@ class Dimension {
     /** Set the min in a given dimension to equal the given
      * expression. Setting the mins to zero may simplify some
      * addressing math. */
-    Dimension set_min(Expr min);
+    Dimension set_min(const Expr &min);
 
     /** Set the extent in a given dimension to equal the given
      * expression. Images passed in that fail this check will generate
@@ -57,20 +57,20 @@ class Dimension {
      im.dim(0).set_extent((im.dim(0).extent()/32)*32);
      \endcode
      * tells the compiler that the extent is a multiple of 32. */
-    Dimension set_extent(Expr extent);
+    Dimension set_extent(const Expr &extent);
 
     /** Set the stride in a given dimension to equal the given
      * value. This is particularly helpful to set when
      * vectorizing. Known strides for the vectorized dimension
      * generate better code. */
-    Dimension set_stride(Expr stride);
+    Dimension set_stride(const Expr &stride);
 
     /** Set the min and extent in one call. */
-    Dimension set_bounds(Expr min, Expr extent);
+    Dimension set_bounds(const Expr &min, const Expr &extent);
 
     /** Set the min and extent estimates in one call. These values are only
      * used by the auto-scheduler and/or the RunGen tool/ */
-    Dimension set_estimate(Expr min, Expr extent);
+    Dimension set_estimate(const Expr &min, const Expr &extent);
 
     Expr min_estimate() const;
     Expr extent_estimate() const;

src/Parameter.cpp

@@ -3,6 +3,7 @@
 #include "Argument.h"
 #include "Float16.h"
 #include "IR.h"
+#include "IRMutator.h"
 #include "IROperator.h"
 
 namespace Halide {
@@ -194,34 +195,84 @@ bool Parameter::defined() const {
     return contents.defined();
 }
 
-void Parameter::set_min_constraint(int dim, Expr e) {
+// Helper function to remove any references in a parameter constraint to the
+// parameter itself, to avoid creating a reference count cycle and causing a
+// leak. Note that it's still possible to create a cycle by having two different
+// Parameters each have constraints that reference the other.
+Expr remove_self_references(const Parameter &p, const Expr &e) {
+    class RemoveSelfReferences : public IRMutator {
+        using IRMutator::visit;
+
+        Expr visit(const Variable *var) override {
+            if (var->param.same_as(p)) {
+                internal_assert(starts_with(var->name, p.name() + "."));
+                return Variable::make(var->type, var->name);
+            } else {
+                internal_assert(!starts_with(var->name, p.name() + "."));
+            }
+            return var;
+        }
+
+    public:
+        const Parameter &p;
+        RemoveSelfReferences(const Parameter &p)
+            : p(p) {
+        }
+    } mutator{p};
+    return mutator.mutate(e);
+}
+
+Expr restore_self_references(const Parameter &p, const Expr &e) {
+    class RestoreSelfReferences : public IRMutator {
+        using IRMutator::visit;
+
+        Expr visit(const Variable *var) override {
+            if (!var->image.defined() &&
+                !var->param.defined() &&
+                !var->reduction_domain.defined() &&
+                starts_with(var->name, p.name() + ".")) {
+                return Variable::make(var->type, var->name, p);
+            }
+            return var;
+        }
+
+    public:
+        const Parameter &p;
+        RestoreSelfReferences(const Parameter &p)
+            : p(p) {
+        }
+    } mutator{p};
+    return mutator.mutate(e);
+}
+
+void Parameter::set_min_constraint(int dim, const Expr &e) {
     check_is_buffer();
     check_dim_ok(dim);
-    contents->buffer_constraints[dim].min = std::move(e);
+    contents->buffer_constraints[dim].min = remove_self_references(*this, e);
 }
 
-void Parameter::set_extent_constraint(int dim, Expr e) {
+void Parameter::set_extent_constraint(int dim, const Expr &e) {
     check_is_buffer();
     check_dim_ok(dim);
-    contents->buffer_constraints[dim].extent = std::move(e);
+    contents->buffer_constraints[dim].extent = remove_self_references(*this, e);
 }
 
-void Parameter::set_stride_constraint(int dim, Expr e) {
+void Parameter::set_stride_constraint(int dim, const Expr &e) {
     check_is_buffer();
     check_dim_ok(dim);
-    contents->buffer_constraints[dim].stride = std::move(e);
+    contents->buffer_constraints[dim].stride = remove_self_references(*this, e);
 }
 
-void Parameter::set_min_constraint_estimate(int dim, Expr min) {
+void Parameter::set_min_constraint_estimate(int dim, const Expr &min) {
     check_is_buffer();
     check_dim_ok(dim);
-    contents->buffer_constraints[dim].min_estimate = std::move(min);
+    contents->buffer_constraints[dim].min_estimate = remove_self_references(*this, min);
 }
 
-void Parameter::set_extent_constraint_estimate(int dim, Expr extent) {
+void Parameter::set_extent_constraint_estimate(int dim, const Expr &extent) {
     check_is_buffer();
     check_dim_ok(dim);
-    contents->buffer_constraints[dim].extent_estimate = std::move(extent);
+    contents->buffer_constraints[dim].extent_estimate = remove_self_references(*this, extent);
 }
 
 void Parameter::set_host_alignment(int bytes) {
@@ -232,31 +283,31 @@ void Parameter::set_host_alignment(int bytes) {
 Expr Parameter::min_constraint(int dim) const {
     check_is_buffer();
     check_dim_ok(dim);
-    return contents->buffer_constraints[dim].min;
+    return restore_self_references(*this, contents->buffer_constraints[dim].min);
 }
 
 Expr Parameter::extent_constraint(int dim) const {
     check_is_buffer();
     check_dim_ok(dim);
-    return contents->buffer_constraints[dim].extent;
+    return restore_self_references(*this, contents->buffer_constraints[dim].extent);
 }
 
 Expr Parameter::stride_constraint(int dim) const {
     check_is_buffer();
     check_dim_ok(dim);
-    return contents->buffer_constraints[dim].stride;
+    return restore_self_references(*this, contents->buffer_constraints[dim].stride);
 }
 
 Expr Parameter::min_constraint_estimate(int dim) const {
     check_is_buffer();
     check_dim_ok(dim);
-    return contents->buffer_constraints[dim].min_estimate;
+    return restore_self_references(*this, contents->buffer_constraints[dim].min_estimate);
 }
 
 Expr Parameter::extent_constraint_estimate(int dim) const {
     check_is_buffer();
     check_dim_ok(dim);
-    return contents->buffer_constraints[dim].extent_estimate;
+    return restore_self_references(*this, contents->buffer_constraints[dim].extent_estimate);
 }
 
 int Parameter::host_alignment() const {

src/Parameter.h

@@ -124,11 +124,11 @@ class Parameter {
     /** Get and set constraints for the min, extent, stride, and estimates on
      * the min/extent. */
     //@{
-    void set_min_constraint(int dim, Expr e);
-    void set_extent_constraint(int dim, Expr e);
-    void set_stride_constraint(int dim, Expr e);
-    void set_min_constraint_estimate(int dim, Expr min);
-    void set_extent_constraint_estimate(int dim, Expr extent);
+    void set_min_constraint(int dim, const Expr &e);
+    void set_extent_constraint(int dim, const Expr &e);
+    void set_stride_constraint(int dim, const Expr &e);
+    void set_min_constraint_estimate(int dim, const Expr &min);
+    void set_extent_constraint_estimate(int dim, const Expr &extent);
     void set_host_alignment(int bytes);
     Expr min_constraint(int dim) const;
     Expr extent_constraint(int dim) const;