Fix leaks in test/correctness/memoize.cpp

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;

test/correctness/memoize.cpp

@@ -59,15 +59,27 @@ extern "C" HALIDE_EXPORT_SYMBOL int computed_eviction_key(int a) {
 }
 HalideExtern_1(int, computed_eviction_key, int);
 
-void simple_free(JITUserContext *user_context, void *ptr) {
- free(ptr);
-}
-
-void *flakey_malloc(JITUserContext * /* user_context */, size_t x) {
+// A flaky allocator. Note that it has to be compatible with halide_free
+// though, because halide_free is going to be called by
+// memoization_cache_cleanup with a null user_context when we release the jit
+// shared runtimes at the endof this test. So it has to be aligned, and it has
+// to store the pointer to free just before the returned pointer.
+void *flaky_malloc(JITUserContext * /* user_context */, size_t x) {
  if ((rand() % 4) == 0) {
  return nullptr;
  } else {
- return malloc(x);
+ x = (x + 63) & (~31);
+ void *ptr = aligned_alloc(32, x);
+ void **ret = (void **)ptr;
+ ret += 4;
+ ret[-1] = ptr;
+ return ret;
+ }
+}
+
+void simple_free(JITUserContext *user_context, void *ptr) {
+ if (ptr != nullptr) {
+ free(((void **)ptr)[-1]);
  }
 }
 
@@ -600,7 +612,7 @@ int main(int argc, char **argv) {
 
  Pipeline pipe(g);
  pipe.jit_handlers().custom_error = record_error;
- pipe.jit_handlers().custom_malloc = flakey_malloc;
+ pipe.jit_handlers().custom_malloc = flaky_malloc;
  pipe.jit_handlers().custom_free = simple_free;
 
  int total_errors = 0;
@@ -644,7 +656,8 @@ int main(int argc, char **argv) {
  }
  }
 
- printf("In 100 attempts with flakey malloc, %d errors and %d full completions occured.\n", total_errors, completed);
+ printf("In 100 attempts with flaky malloc, %d errors and %d full completions occured.\n",
+ total_errors, completed);
  }
 
  {
@@ -733,6 +746,7 @@ int main(int argc, char **argv) {
 
  assert(call_count == 8);
  }
+ Internal::JITSharedRuntime::release_all();
 
  printf("Success!\n");
  return 0;