text: #include "StorageFolding.h"
  
  #include "Bounds.h"
  #include "CSE.h"
  #include "Debug.h"
  #include "ExprUsesVar.h"
  #include "IRMutator.h"
  #include "IROperator.h"
  #include "IRPrinter.h"
  #include "Monotonic.h"
  #include "Simplify.h"
  #include "Substitute.h"
  #include <utility>
  
  namespace Halide {
  namespace Internal {
  
  namespace {
  
  int64_t next_power_of_two(int64_t x) {
      return static_cast<int64_t>(1) << static_cast<int64_t>(std::ceil(std::log2(x)));
  }
  
  using std::map;
  using std::string;
  using std::vector;
  
  // Count the number of producers of a particular func.
  class CountProducers : public IRVisitor {
      const std::string &name;
  
      void visit(const ProducerConsumer *op) override {
          if (op->is_producer && (op->name == name)) {
              count++;
          } else {
              IRVisitor::visit(op);
          }
      }
  
      using IRVisitor::visit;
  
  public:
      int count = 0;
  
      CountProducers(const std::string &name)
          : name(name) {
      }
  };
  
  int count_producers(const Stmt &in, const std::string &name) {
      CountProducers counter(name);
      in.accept(&counter);
      return counter.count;
  }
  
  // Fold the storage of a function in a particular dimension by a particular factor
  class FoldStorageOfFunction : public IRMutator {
      string func;
      int dim;
      Expr factor;
      string dynamic_footprint;
  
      using IRMutator::visit;
  
      Expr visit(const Call *op) override {
          Expr expr = IRMutator::visit(op);
          op = expr.as<Call>();
          internal_assert(op);
          if (op->name == func && op->call_type == Call::Halide) {
              vector<Expr> args = op->args;
              internal_assert(dim < (int)args.size());
              args[dim] = is_const_one(factor) ? 0 : (args[dim] % factor);
              expr = Call::make(op->type, op->name, args, op->call_type,
                                op->func, op->value_index, op->image, op->param);
          } else if (op->name == Call::buffer_crop) {
              Expr source = op->args[2];
              const Variable *buf_var = source.as<Variable>();
              if (buf_var &&
                  starts_with(buf_var->name, func + ".") &&
                  ends_with(buf_var->name, ".buffer")) {
                  // We are taking a crop of a folded buffer. For now
                  // we'll just assert that the crop doesn't wrap
                  // around, so that the crop doesn't need to be treated
                  // as a folded buffer too. But to take the crop, we
                  // need to use folded coordinates, and then restore
                  // the non-folded min after the crop operation.
  
                  // Pull out the expressions we need
                  internal_assert(op->args.size() >= 5);
                  Expr mins_arg = op->args[3];
                  Expr extents_arg = op->args[4];
                  const Call *mins_call = mins_arg.as<Call>();
                  const Call *extents_call = extents_arg.as<Call>();
                  internal_assert(mins_call && extents_call);
                  vector<Expr> mins = mins_call->args;
                  const vector<Expr> &extents = extents_call->args;
                  internal_assert(dim < (int)mins.size() && dim < (int)extents.size());
                  Expr old_min = mins[dim];
                  Expr old_extent = extents[dim];
  
                  // Rewrite the crop args
                  mins[dim] = old_min % factor;
                  Expr new_mins = Call::make(type_of<int *>(), Call::make_struct, mins, Call::Intrinsic);
                  vector<Expr> new_args = op->args;
                  new_args[3] = new_mins;
                  expr = Call::make(op->type, op->name, new_args, op->call_type);
  
                  // Inject the assertion
                  Expr no_wraparound = mins[dim] + extents[dim] <= factor;
  
                  Expr valid_min = old_min;
                  if (!dynamic_footprint.empty()) {
                      // If the footprint is being tracked dynamically, it's
                      // not enough to just check we don't overlap a
                      // fold. We also need to check the min against the
                      // valid min.
  
                      // TODO: dynamic footprint is no longer the min, and may be tracked separately on producer and consumer sides (head vs tail)
                      valid_min =
                          Load::make(Int(32), dynamic_footprint, 0, Buffer<>(), Parameter(), const_true(), ModulusRemainder());
                      Expr check = (old_min >= valid_min &&
                                    (old_min + old_extent - 1) < valid_min + factor);
                      no_wraparound = no_wraparound && check;
                  }
  
                  Expr error = Call::make(Int(32), "halide_error_bad_extern_fold",
                                          {Expr(func), Expr(dim), old_min, old_extent, valid_min, factor},
                                          Call::Extern);
                  expr = Call::make(op->type, Call::require,
                                    {no_wraparound, expr, error}, Call::Intrinsic);
  
                  // Restore the correct min coordinate
                  expr = Call::make(op->type, Call::buffer_set_bounds,
                                    {expr, dim, old_min, old_extent}, Call::Extern);
              }
          }
          return expr;
      }
  
      Stmt visit(const Provide *op) override {
          Stmt stmt = IRMutator::visit(op);
          op = stmt.as<Provide>();
          internal_assert(op);
          if (op->name == func) {
              vector<Expr> args = op->args;
              args[dim] = is_const_one(factor) ? 0 : (args[dim] % factor);
              stmt = Provide::make(op->name, op->values, args, op->predicate);
          }
          return stmt;
      }
  
  public:
      FoldStorageOfFunction(string f, int d, Expr e, string p)
          : func(std::move(f)), dim(d), factor(std::move(e)), dynamic_footprint(std::move(p)) {
      }
  };
  
  // Inject dynamic folding checks against a tracked live range.
  class InjectFoldingCheck : public IRMutator {
      Function func;
      string head, tail, loop_var;
      Expr sema_var;
      int dim;
      bool in_produce;
      const StorageDim &storage_dim;
      using IRMutator::visit;
  
      Stmt visit(const ProducerConsumer *op) override {
          if (op->name == func.name()) {
              Stmt body = op->body;
              if (op->is_producer) {
                  if (func.has_extern_definition()) {
                      // We'll update the valid min at the buffer_crop call.
                      in_produce = true;
                      body = mutate(op->body);
                  } else {
                      // Update valid range based on bounds written to.
                      Box b = box_provided(body, func.name());
                      Expr old_leading_edge =
                          Load::make(Int(32), head + "_next", 0, Buffer<>(), Parameter(), const_true(), ModulusRemainder());
  
                      internal_assert(!b.empty());
  
                      // Track the logical address range the memory
                      // currently represents.
                      Expr new_leading_edge;
                      if (storage_dim.fold_forward) {
                          new_leading_edge = max(b[dim].max, old_leading_edge);
                      } else {
                          new_leading_edge = min(b[dim].min, old_leading_edge);
                      }
  
                      string new_leading_edge_var_name = unique_name('t');
                      Expr new_leading_edge_var = Variable::make(Int(32), new_leading_edge_var_name);
  
                      Stmt update_leading_edge =
                          Store::make(head, new_leading_edge_var, 0, Parameter(), const_true(), ModulusRemainder());
                      Stmt update_next_leading_edge =
                          Store::make(head + "_next", new_leading_edge_var, 0, Parameter(), const_true(), ModulusRemainder());
  
                      // Check the region being written to in this
                      // iteration lies within the range of coordinates
                      // currently represented.
                      Expr fold_non_monotonic_error =
                          Call::make(Int(32), "halide_error_bad_fold",
                                     {func.name(), storage_dim.var, loop_var},
                                     Call::Extern);
  
                      Expr in_valid_range;
                      if (storage_dim.fold_forward) {
                          in_valid_range = b[dim].min > new_leading_edge - storage_dim.fold_factor;
                      } else {
                          in_valid_range = b[dim].max < new_leading_edge + storage_dim.fold_factor;
                      }
                      Stmt check_in_valid_range =
                          AssertStmt::make(in_valid_range, fold_non_monotonic_error);
  
                      Expr extent = b[dim].max - b[dim].min + 1;
  
                      // Separately check the extent for *this* loop iteration fits.
                      Expr fold_too_small_error =
                          Call::make(Int(32), "halide_error_fold_factor_too_small",
                                     {func.name(), storage_dim.var, storage_dim.fold_factor, loop_var, extent},
                                     Call::Extern);
  
                      Stmt check_extent =
                          AssertStmt::make(extent <= storage_dim.fold_factor, fold_too_small_error);
  
                      Stmt checks = Block::make({check_extent, check_in_valid_range,
                                                 update_leading_edge, update_next_leading_edge});
                      if (func.schedule().async()) {
                          Expr to_acquire;
                          if (storage_dim.fold_forward) {
                              to_acquire = new_leading_edge_var - old_leading_edge;
                          } else {
                              to_acquire = old_leading_edge - new_leading_edge_var;
                          }
                          body = Block::make(checks, body);
                          body = Acquire::make(sema_var, to_acquire, body);
                          body = LetStmt::make(new_leading_edge_var_name, new_leading_edge, body);
                      } else {
                          checks = LetStmt::make(new_leading_edge_var_name, new_leading_edge, checks);
                          body = Block::make(checks, body);
                      }
                  }
  
              } else {
                  // Check the accessed range against the valid range.
                  Box b = box_required(body, func.name());
                  if (b.empty()) {
                      // Must be used in an extern call (TODO:
                      // assert this, TODO: What if it's used in an
                      // extern call and native Halide). We'll
                      // update the valid min at the buffer_crop
                      // call.
                      in_produce = false;
                      body = mutate(op->body);
                  } else {
                      Expr leading_edge =
                          Load::make(Int(32), tail + "_next", 0, Buffer<>(), Parameter(), const_true(), ModulusRemainder());
  
                      if (func.schedule().async()) {
                          Expr new_leading_edge;
                          if (storage_dim.fold_forward) {
                              new_leading_edge = b[dim].min - 1 + storage_dim.fold_factor;
                          } else {
                              new_leading_edge = b[dim].max + 1 - storage_dim.fold_factor;
                          }
                          string new_leading_edge_name = unique_name('t');
                          Expr new_leading_edge_var = Variable::make(Int(32), new_leading_edge_name);
                          Expr to_release;
                          if (storage_dim.fold_forward) {
                              to_release = new_leading_edge_var - leading_edge;
                          } else {
                              to_release = leading_edge - new_leading_edge_var;
                          }
                          Expr release_producer =
                              Call::make(Int(32), "halide_semaphore_release", {sema_var, to_release}, Call::Extern);
                          // The consumer is going to get its own forked copy of the footprint, so it needs to update it too.
                          Stmt update_leading_edge = Store::make(tail, new_leading_edge_var, 0, Parameter(), const_true(), ModulusRemainder());
                          update_leading_edge = Block::make(Store::make(tail + "_next", new_leading_edge_var, 0, Parameter(), const_true(), ModulusRemainder()),
                                                            update_leading_edge);
                          update_leading_edge = Block::make(Evaluate::make(release_producer), update_leading_edge);
                          update_leading_edge = LetStmt::make(new_leading_edge_name, new_leading_edge, update_leading_edge);
                          body = Block::make(update_leading_edge, body);
                      } else {
                          Expr check;
                          if (storage_dim.fold_forward) {
                              check = (b[dim].min > leading_edge - storage_dim.fold_factor && b[dim].max <= leading_edge);
                          } else {
                              check = (b[dim].max < leading_edge + storage_dim.fold_factor && b[dim].min >= leading_edge);
                          }
                          Expr bad_fold_error = Call::make(Int(32), "halide_error_bad_fold",
                                                           {func.name(), storage_dim.var, loop_var},
                                                           Call::Extern);
                          body = Block::make(AssertStmt::make(check, bad_fold_error), body);
                      }
                  }
              }
  
              return ProducerConsumer::make(op->name, op->is_producer, body);
          } else {
              return IRMutator::visit(op);
          }
      }
  
      Stmt visit(const LetStmt *op) override {
          if (starts_with(op->name, func.name() + ".") &&
              ends_with(op->name, ".tmp_buffer")) {
  
              Stmt body = op->body;
              Expr buf = Variable::make(type_of<halide_buffer_t *>(), op->name);
  
              if (in_produce) {
                  // We're taking a crop of the buffer to act as an output
                  // to an extern stage. Update the valid min or max
                  // coordinate accordingly.
  
                  Expr leading_edge;
                  if (storage_dim.fold_forward) {
                      leading_edge =
                          Call::make(Int(32), Call::buffer_get_max, {buf, dim}, Call::Extern);
                  } else {
                      leading_edge =
                          Call::make(Int(32), Call::buffer_get_min, {buf, dim}, Call::Extern);
                  }
  
                  Stmt update_leading_edge =
                      Store::make(head, leading_edge, 0, Parameter(), const_true(), ModulusRemainder());
                  body = Block::make(update_leading_edge, body);
  
                  // We don't need to make sure the min is moving
                  // monotonically, because we can't do sliding window on
                  // extern stages, so we don't have to worry about whether
                  // we're preserving valid values from previous loop
                  // iterations.
  
                  if (func.schedule().async()) {
                      Expr old_leading_edge =
                          Load::make(Int(32), head, 0, Buffer<>(), Parameter(), const_true(), ModulusRemainder());
                      Expr to_acquire;
                      if (storage_dim.fold_forward) {
                          to_acquire = leading_edge - old_leading_edge;
                      } else {
                          to_acquire = old_leading_edge - leading_edge;
                      }
                      body = Acquire::make(sema_var, to_acquire, body);
                  }
              } else {
                  // We're taking a crop of the buffer to act as an input
                  // to an extern stage. Update the valid min or max
                  // coordinate accordingly.
  
                  Expr leading_edge;
                  if (storage_dim.fold_forward) {
                      leading_edge =
                          Call::make(Int(32), Call::buffer_get_min, {buf, dim}, Call::Extern) - 1 + storage_dim.fold_factor;
                  } else {
                      leading_edge =
                          Call::make(Int(32), Call::buffer_get_max, {buf, dim}, Call::Extern) + 1 - storage_dim.fold_factor;
                  }
  
                  Stmt update_leading_edge =
                      Store::make(tail, leading_edge, 0, Parameter(), const_true(), ModulusRemainder());
                  body = Block::make(update_leading_edge, body);
  
                  if (func.schedule().async()) {
                      Expr old_leading_edge =
                          Load::make(Int(32), tail, 0, Buffer<>(), Parameter(), const_true(), ModulusRemainder());
                      Expr to_release;
                      if (storage_dim.fold_forward) {
                          to_release = leading_edge - old_leading_edge;
                      } else {
                          to_release = old_leading_edge - leading_edge;
                      }
                      Expr release_producer =
                          Call::make(Int(32), "halide_semaphore_release", {sema_var, to_release}, Call::Extern);
                      body = Block::make(Evaluate::make(release_producer), body);
                  }
              }
  
              return LetStmt::make(op->name, op->value, body);
          } else {
              return LetStmt::make(op->name, op->value, mutate(op->body));
          }
      }
  
  public:
      InjectFoldingCheck(Function func,
                         string head, string tail,
                         string loop_var, Expr sema_var,
                         int dim, const StorageDim &storage_dim)
          : func(std::move(func)),
            head(std::move(head)), tail(std::move(tail)), loop_var(std::move(loop_var)), sema_var(std::move(sema_var)),
            dim(dim), storage_dim(storage_dim) {
      }
  };
  
  struct Semaphore {
      string name;
      Expr var;
      Expr init;
  };
  
  class HasExternConsumer : public IRVisitor {
  
      using IRVisitor::visit;
  
      void visit(const Variable *op) override {
          if (op->name == func + ".buffer") {
              result = true;
          }
      }
  
      const std::string &func;
  
  public:
      HasExternConsumer(const std::string &func)
          : func(func) {
      }
      bool result = false;
  };
  
  class VectorAccessOfFoldedDim : public IRVisitor {
      using IRVisitor::visit;
  
      void visit(const Provide *op) override {
          if (op->name == func) {
              internal_assert(dim < (int)op->args.size());
              if (expr_uses_vars(op->args[dim], vector_vars)) {
                  result = true;
              }
          } else {
              IRVisitor::visit(op);
          }
      }
  
      void visit(const Call *op) override {
          if (op->name == func &&
              op->call_type == Call::Halide) {
              internal_assert(dim < (int)op->args.size());
              if (expr_uses_vars(op->args[dim], vector_vars)) {
                  result = true;
              }
          } else {
              IRVisitor::visit(op);
          }
      }
  
      template<typename LetOrLetStmt>
      void visit_let(const LetOrLetStmt *op) {
          op->value.accept(this);
          bool is_vec = expr_uses_vars(op->value, vector_vars);
          ScopedBinding<> bind(is_vec, vector_vars, op->name);
          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 {
          ScopedBinding<> bind(op->for_type == ForType::Vectorized,
                               vector_vars, op->name);
          IRVisitor::visit(op);
      }
  
      Scope<> vector_vars;
      const string &func;
      int dim;
  
  public:
      bool result = false;
      VectorAccessOfFoldedDim(const string &func, int dim)
          : func(func), dim(dim) {
      }
  };
  
  // Attempt to fold the storage of a particular function in a statement
  class AttemptStorageFoldingOfFunction : public IRMutator {
      Function func;
      bool explicit_only;
  
      using IRMutator::visit;
  
      Stmt visit(const ProducerConsumer *op) override {
          if (op->name == func.name()) {
              // Can't proceed into the pipeline for this func
              return op;
          } else {
              return IRMutator::visit(op);
          }
      }
  
      Stmt visit(const For *op) override {
          if (op->for_type != ForType::Serial && op->for_type != ForType::Unrolled) {
              // We can't proceed into a parallel for loop.
  
              // TODO: If there's no overlap between the region touched
              // by the threads as this loop counter varies
              // (i.e. there's no cross-talk between threads), then it's
              // safe to proceed.
              return op;
          }
  
          Stmt stmt;
          Stmt body = op->body;
  
          Box provided = box_provided(body, func.name());
          Box required = box_required(body, func.name());
          // For storage folding, we don't care about conditional reads.
          required.used = Expr();
          Box box = box_union(provided, required);
  
          Expr loop_var = Variable::make(Int(32), op->name);
          Expr loop_min = Variable::make(Int(32), op->name + ".loop_min");
          Expr loop_max = Variable::make(Int(32), op->name + ".loop_max");
  
          string dynamic_footprint;
  
          Scope<Interval> bounds;
          bounds.push(op->name, Interval(op->min, simplify(op->min + op->extent - 1)));
  
          Scope<Interval> steady_bounds;
          steady_bounds.push(op->name, Interval(simplify(op->min + 1), simplify(op->min + op->extent - 1)));
  
          HasExternConsumer has_extern_consumer(func.name());
          body.accept(&has_extern_consumer);
  
          // Try each dimension in turn from outermost in
          for (size_t i = box.size(); i > 0; i--) {
              int dim = (int)(i - 1);
  
              if (!box[dim].is_bounded()) {
                  continue;
              }
  
              Expr min = simplify(common_subexpression_elimination(box[dim].min));
              Expr max = simplify(common_subexpression_elimination(box[dim].max));
  
              if (is_const(min) || is_const(max)) {
                  debug(3) << "\nNot considering folding " << func.name()
                           << " over for loop over " << op->name
                           << " dimension " << i - 1 << "\n"
                           << " because the min or max are constants."
                           << "Min: " << min << "\n"
                           << "Max: " << max << "\n";
                  continue;
              }
  
              Expr min_provided, max_provided, min_required, max_required;
              if (func.schedule().async() && !explicit_only) {
                  if (!provided.empty()) {
                      min_provided = simplify(provided[dim].min);
                      max_provided = simplify(provided[dim].max);
                  }
                  if (!required.empty()) {
                      min_required = simplify(required[dim].min);
                      max_required = simplify(required[dim].max);
                  }
              }
              string sema_name = func.name() + ".folding_semaphore." + unique_name('_');
              Expr sema_var = Variable::make(type_of<halide_semaphore_t *>(), sema_name);
  
              // Consider the initial iteration and steady state
              // separately for all these proofs.
              Expr loop_var = Variable::make(Int(32), op->name);
              Expr steady_state = (op->min < loop_var);
  
              Expr min_steady = simplify(substitute(steady_state, const_true(), min), true, steady_bounds);
              Expr max_steady = simplify(substitute(steady_state, const_true(), max), true, steady_bounds);
              Expr min_initial = simplify(substitute(steady_state, const_false(), min), true, bounds);
              Expr max_initial = simplify(substitute(steady_state, const_false(), max), true, bounds);
              Expr extent_initial = simplify(substitute(loop_var, op->min, max_initial - min_initial + 1), true, bounds);
              Expr extent_steady = simplify(max_steady - min_steady + 1, true, steady_bounds);
              Expr extent = Max::make(extent_initial, extent_steady);
              extent = simplify(common_subexpression_elimination(extent), true, bounds);
  
              // Find the StorageDim corresponding to dim.
              const std::vector<StorageDim> &storage_dims = func.schedule().storage_dims();
              auto storage_dim_i = std::find_if(storage_dims.begin(), storage_dims.end(),
                                                [&](const StorageDim &i) { return i.var == func.args()[dim]; });
              internal_assert(storage_dim_i != storage_dims.end());
              const StorageDim &storage_dim = *storage_dim_i;
  
              Expr explicit_factor;
              if (!is_pure(min) ||
                  !is_pure(max) ||
                  has_extern_consumer.result ||
                  expr_uses_var(min, op->name) ||
                  expr_uses_var(max, op->name)) {
                  // We only use the explicit fold factor if the fold is
                  // relevant for this loop. If the fold isn't relevant
                  // for this loop, the added asserts will be too
                  // conservative.
                  explicit_factor = storage_dim.fold_factor;
              }
  
              debug(3) << "\nConsidering folding " << func.name()
                       << " over for loop over " << op->name
                       << " dimension " << i - 1 << "\n"
                       << "Min: " << min << "\n"
                       << "Max: " << max << "\n"
                       << "Extent: " << extent << "\n"
                       << "explicit_factor: " << explicit_factor << "\n";
  
              // First, attempt to detect if the loop is monotonically
              // increasing or decreasing (if we allow automatic folding).
              bool can_fold_forwards = false, can_fold_backwards = false;
  
              if (!explicit_only) {
                  // We can't clobber data that will be read later. If
                  // async, the producer can't un-release slots in the
                  // circular buffer.
                  can_fold_forwards = (is_monotonic(min, op->name) == Monotonic::Increasing);
                  can_fold_backwards = (is_monotonic(max, op->name) == Monotonic::Decreasing);
                  if (func.schedule().async()) {
                      // Our semaphore acquire primitive can't take
                      // negative values, so we can't un-acquire slots
                      // in the circular buffer.
                      can_fold_forwards &= (is_monotonic(max_provided, op->name) == Monotonic::Increasing);
                      can_fold_backwards &= (is_monotonic(min_provided, op->name) == Monotonic::Decreasing);
                      // We need to be able to analyze the required footprint to know how much to release
                      can_fold_forwards &= min_required.defined();
                      can_fold_backwards &= max_required.defined();
                  }
              }
  
              // Uncomment to pretend that static analysis always fails (for testing)
              // can_fold_forwards = can_fold_backwards = false;
  
              if (!can_fold_forwards && !can_fold_backwards) {
                  if (explicit_factor.defined()) {
                      // If we didn't find a monotonic dimension, and we
                      // have an explicit fold factor, we need to
                      // dynamically check that the min/max do in fact
                      // monotonically increase/decrease. We'll allocate
                      // some stack space to store the valid footprint,
                      // update it outside produce nodes, and check it
                      // outside consume nodes.
  
                      string head, tail;
                      if (func.schedule().async()) {
                          // If we're async, we need to keep a separate
                          // counter for the producer and consumer. They
                          // are coupled by a semaphore. The counter
                          // represents the max index the producer may
                          // write to. The invariant is that the
                          // semaphore count is the difference between
                          // the counters. So...
                          //
                          // when folding forwards,  semaphore == head - tail
                          // when folding backwards, semaphore == tail - head
                          //
                          // We'll initialize to head = tail, and
                          // semaphore = 0. Every time the producer or
                          // consumer wants to move the counter, it must
                          // also acquire or release the semaphore to
                          // prevent them from diverging too far.
                          dynamic_footprint = func.name() + ".folding_semaphore." + op->name + unique_name('_');
                          head = dynamic_footprint + ".head";
                          tail = dynamic_footprint + ".tail";
                      } else {
                          dynamic_footprint = func.name() + "." + op->name + unique_name('_') + ".head";
                          head = tail = dynamic_footprint;
                      }
  
                      body = InjectFoldingCheck(func,
                                                head, tail,
                                                op->name,
                                                sema_var,
                                                dim,
                                                storage_dim)
                                 .mutate(body);
  
                      if (storage_dim.fold_forward) {
                          can_fold_forwards = true;
                      } else {
                          can_fold_backwards = true;
                      }
                  } else {
                      // Can't do much with this dimension
                      if (!explicit_only) {
                          debug(3) << "Not folding because loop min or max not monotonic in the loop variable\n"
                                   << "min_initial = " << min_initial << "\n"
                                   << "min_steady = " << min_steady << "\n"
                                   << "max_initial = " << max_initial << "\n"
                                   << "max_steady = " << max_steady << "\n";
                      } else {
                          debug(3) << "Not folding because there is no explicit storage folding factor\n";
                      }
                      continue;
                  }
              }
  
              internal_assert(can_fold_forwards || can_fold_backwards);
  
              Expr factor;
              if (explicit_factor.defined()) {
                  if (dynamic_footprint.empty() && !func.schedule().async()) {
                      // We were able to prove monotonicity
                      // statically, but we may need a runtime
                      // assertion for maximum extent. In many cases
                      // it will simplify away. For async schedules
                      // it gets dynamically tracked anyway.
                      Expr error = Call::make(Int(32), "halide_error_fold_factor_too_small",
                                              {func.name(), storage_dim.var, explicit_factor, op->name, extent},
                                              Call::Extern);
                      body = Block::make(AssertStmt::make(extent <= explicit_factor, error), body);
                  }
                  factor = explicit_factor;
              } else {
                  // The max of the extent over all values of the loop variable must be a constant
                  Scope<Interval> scope;
                  scope.push(op->name, Interval(loop_min, loop_max));
                  Expr max_extent = find_constant_bound(extent, Direction::Upper, scope);
                  scope.pop(op->name);
  
                  const int max_fold = 1024;
                  const int64_t *const_max_extent = as_const_int(max_extent);
                  if (const_max_extent && *const_max_extent <= max_fold) {
                      factor = static_cast<int>(next_power_of_two(*const_max_extent));
                  } else {
                      // Try a little harder to find a bounding power of two
                      int e = max_fold * 2;
                      bool success = false;
                      while (e > 0 && can_prove(extent <= e / 2)) {
                          success = true;
                          e /= 2;
                      }
                      if (success) {
                          factor = e;
                      } else {
                          debug(3) << "Not folding because extent not bounded by a constant not greater than " << max_fold << "\n"
                                   << "extent = " << extent << "\n"
                                   << "max extent = " << max_extent << "\n";
                          // Try the next dimension
                          continue;
                      }
                  }
              }
  
              internal_assert(factor.defined());
  
              if (!explicit_factor.defined()) {
                  VectorAccessOfFoldedDim vector_access_of_folded_dim{func.name(), dim};
                  body.accept(&vector_access_of_folded_dim);
                  if (vector_access_of_folded_dim.result) {
                      user_warning
                          << "Not folding Func " << func.name() << " along dimension " << func.args()[dim]
                          << " because there is vectorized access to that Func in that dimension and "
                          << "storage folding was not explicitly requested in the schedule. In previous "
                          << "versions of Halide this would have folded with factor " << factor
                          << ". To restore the old behavior add " << func.name()
                          << ".fold_storage(" << func.args()[dim] << ", " << factor
                          << ") to your schedule.\n";
                      // Try the next dimension
                      continue;
                  }
              }
  
              debug(3) << "Proceeding with factor " << factor << "\n";
  
              Fold fold = {(int)i - 1, factor};
              dims_folded.push_back(fold);
              {
                  string head;
                  if (!dynamic_footprint.empty() && func.schedule().async()) {
                      head = dynamic_footprint + ".head";
                  } else {
                      head = dynamic_footprint;
                  }
                  body = FoldStorageOfFunction(func.name(), (int)i - 1, factor, head).mutate(body);
              }
  
              // If the producer is async, it can run ahead by
              // some amount controlled by a semaphore.
              if (func.schedule().async()) {
                  Semaphore sema;
                  sema.name = sema_name;
                  sema.var = sema_var;
                  sema.init = 0;
  
                  if (dynamic_footprint.empty()) {
                      // We are going to do the sem acquires and releases using static analysis of the boxes accessed.
                      sema.init = factor;
                      // Do the analysis of how much to acquire and release statically
                      Expr to_acquire, to_release;
                      if (can_fold_forwards) {
                          Expr max_provided_prev = substitute(op->name, loop_var - 1, max_provided);
                          Expr min_required_next = substitute(op->name, loop_var + 1, min_required);
                          to_acquire = max_provided - max_provided_prev;  // This is the first time we use these entries
                          to_release = min_required_next - min_required;  // This is the last time we use these entries
                      } else {
                          internal_assert(can_fold_backwards);
                          Expr min_provided_prev = substitute(op->name, loop_var - 1, min_provided);
                          Expr max_required_next = substitute(op->name, loop_var + 1, max_required);
                          to_acquire = min_provided_prev - min_provided;  // This is the first time we use these entries
                          to_release = max_required - max_required_next;  // This is the last time we use these entries
                      }
  
                      if (provided.used.defined()) {
                          to_acquire = select(provided.used, to_acquire, 0);
                      }
                      // We should always release the required region, even if we don't use it.
  
                      // On the first iteration, we need to acquire the extent of the region shared
                      // between the producer and consumer, and we need to release it on the last
                      // iteration.
                      to_acquire = select(loop_var > loop_min, to_acquire, extent);
                      to_release = select(loop_var < loop_max, to_release, extent);
  
                      // We may need dynamic assertions that a positive
                      // amount of the semaphore is acquired/released,
                      // and that the semaphore is initialized to a
                      // positive value. If we are able to prove it,
                      // these checks will simplify away.
                      string to_release_name = unique_name('t');
                      Expr to_release_var = Variable::make(Int(32), to_release_name);
                      string to_acquire_name = unique_name('t');
                      Expr to_acquire_var = Variable::make(Int(32), to_acquire_name);
  
                      Expr bad_fold_error =
                          Call::make(Int(32), "halide_error_bad_fold",
                                     {func.name(), storage_dim.var, op->name},
                                     Call::Extern);
  
                      Expr release_producer =
                          Call::make(Int(32), "halide_semaphore_release", {sema.var, to_release_var}, Call::Extern);
                      Stmt release = Evaluate::make(release_producer);
                      Stmt check_release = AssertStmt::make(to_release_var >= 0 && to_release <= factor, bad_fold_error);
                      release = Block::make(check_release, release);
                      release = LetStmt::make(to_release_name, to_release, release);
  
                      Stmt check_acquire = AssertStmt::make(to_acquire_var >= 0 && to_acquire_var <= factor, bad_fold_error);
  
                      body = Block::make(body, release);
                      body = Acquire::make(sema.var, to_acquire_var, body);
                      body = Block::make(check_acquire, body);
                      body = LetStmt::make(to_acquire_name, to_acquire, body);
                  } else {
                      // We injected runtime tracking and semaphore logic already
                  }
                  dims_folded.back().semaphore = sema;
              }
  
              if (!dynamic_footprint.empty()) {
                  if (func.schedule().async()) {
                      dims_folded.back().head = dynamic_footprint + ".head";
                      dims_folded.back().tail = dynamic_footprint + ".tail";
                  } else {
                      dims_folded.back().head = dynamic_footprint;
                      dims_folded.back().tail.clear();
                  }
                  dims_folded.back().fold_forward = storage_dim.fold_forward;
              }
  
              Expr min_next = substitute(op->name, loop_var + 1, min);
  
              if (can_prove(max < min_next)) {
                  // There's no overlapping usage between loop
                  // iterations, so we can continue to search
                  // for further folding opportunities
                  // recursively.
              } else if (!body.same_as(op->body)) {
                  stmt = For::make(op->name, op->min, op->extent, op->for_type, op->device_api, body);
                  break;
              } else {
                  stmt = op;
                  debug(3) << "Not folding because loop min or max not monotonic in the loop variable\n"
                           << "min = " << min << "\n"
                           << "max = " << max << "\n";
                  break;
              }
          }
  
          // If there's no communication of values from one loop
          // iteration to the next (which may happen due to sliding),
          // then we're safe to fold an inner loop.
          if (box_contains(provided, required)) {
              body = mutate(body);
          }
  
          if (body.same_as(op->body)) {
              stmt = op;
          } else {
              stmt = For::make(op->name, op->min, op->extent, op->for_type, op->device_api, body);
          }
  
          if (func.schedule().async() && !dynamic_footprint.empty()) {
              // Step the counters backwards over the entire extent of
              // the realization, in case we're in an inner loop and are
              // going to run this loop again with the same
              // semaphore. Our invariant is that the difference between
              // the two counters is the semaphore.
              //
              // Doing this instead of synchronizing and resetting the
              // counters and semaphores lets producers advance to the
              // next scanline while a consumer is still on the last few
              // pixels of the previous scanline.
  
              Expr head = Load::make(Int(32), dynamic_footprint + ".head", 0, Buffer<>(), Parameter(), const_true(), ModulusRemainder());
              Expr tail = Load::make(Int(32), dynamic_footprint + ".tail", 0, Buffer<>(), Parameter(), const_true(), ModulusRemainder());
              Expr step = Variable::make(Int(32), func.name() + ".extent." + std::to_string(dims_folded.back().dim)) + dims_folded.back().factor;
              Stmt reset_head = Store::make(dynamic_footprint + ".head_next", head - step, 0, Parameter(), const_true(), ModulusRemainder());
              Stmt reset_tail = Store::make(dynamic_footprint + ".tail_next", tail - step, 0, Parameter(), const_true(), ModulusRemainder());
              stmt = Block::make({stmt, reset_head, reset_tail});
          }
          return stmt;
      }
  
  public:
      struct Fold {
          int dim;
          Expr factor;
          Semaphore semaphore;
          string head, tail;
          bool fold_forward;
      };
      vector<Fold> dims_folded;
  
      AttemptStorageFoldingOfFunction(Function f, bool explicit_only)
          : func(std::move(f)), explicit_only(explicit_only) {
      }
  };
  
  // Look for opportunities for storage folding in a statement
  class StorageFolding : public IRMutator {
      const map<string, Function> &env;
  
      using IRMutator::visit;
  
      Stmt visit(const Realize *op) override {
          Stmt body = mutate(op->body);
  
          // Get the function associated with this realization, which
          // contains the explicit fold directives from the schedule.
          auto func_it = env.find(op->name);
          Function func = func_it != env.end() ? func_it->second : Function();
  
          // Don't attempt automatic storage folding if there is
          // more than one produce node for this func.
          bool explicit_only = count_producers(body, op->name) != 1;
          AttemptStorageFoldingOfFunction folder(func, explicit_only);
          if (explicit_only) {
              debug(3) << "Attempting to fold " << op->name << " explicitly\n";
          } else {
              debug(3) << "Attempting to fold " << op->name << " automatically or explicitly\n";
          }
          body = folder.mutate(body);
  
          if (body.same_as(op->body)) {
              return op;
          } else if (folder.dims_folded.empty()) {
              return Realize::make(op->name, op->types, op->memory_type, op->bounds, op->condition, body);
          } else {
              Region bounds = op->bounds;
  
              // Collapse down the extent in the folded dimension
              for (const auto &dim : folder.dims_folded) {
                  int d = dim.dim;
                  Expr f = dim.factor;
                  internal_assert(d >= 0 &&
                                  d < (int)bounds.size());
                  bounds[d] = Range(0, f);
              }
  
              Stmt stmt = Realize::make(op->name, op->types, op->memory_type, bounds, op->condition, body);
  
              // Each fold may have an associated semaphore that needs initialization, along with some counters
              for (const auto &fold : folder.dims_folded) {
                  auto sema = fold.semaphore;
                  if (sema.var.defined()) {
                      Expr sema_space = Call::make(type_of<halide_semaphore_t *>(), "halide_make_semaphore",
                                                   {sema.init}, Call::Extern);
                      stmt = LetStmt::make(sema.name, sema_space, stmt);
                  }
                  Expr init;
                  if (fold.fold_forward) {
                      init = op->bounds[fold.dim].min;
                  } else {
                      init = op->bounds[fold.dim].min + op->bounds[fold.dim].extent - 1;
                  }
                  if (!fold.head.empty()) {
                      stmt = Block::make(Store::make(fold.head + "_next", init, 0, Parameter(), const_true(), ModulusRemainder()), stmt);
                      stmt = Allocate::make(fold.head + "_next", Int(32), MemoryType::Stack, {}, const_true(), stmt);
                      stmt = Block::make(Store::make(fold.head, init, 0, Parameter(), const_true(), ModulusRemainder()), stmt);
                      stmt = Allocate::make(fold.head, Int(32), MemoryType::Stack, {}, const_true(), stmt);
                  }
                  if (!fold.tail.empty()) {
                      internal_assert(func.schedule().async()) << "Expected a single counter for synchronous folding";
                      stmt = Block::make(Store::make(fold.tail + "_next", init, 0, Parameter(), const_true(), ModulusRemainder()), stmt);
                      stmt = Allocate::make(fold.tail + "_next", Int(32), MemoryType::Stack, {}, const_true(), stmt);
                      stmt = Block::make(Store::make(fold.tail, init, 0, Parameter(), const_true(), ModulusRemainder()), stmt);
                      stmt = Allocate::make(fold.tail, Int(32), MemoryType::Stack, {}, const_true(), stmt);
                  }
              }
  
              return stmt;
          }
      }
  
  public:
      StorageFolding(const map<string, Function> &env)
          : env(env) {
      }
  };
  
  }  // namespace
  
  Stmt storage_folding(const Stmt &s, const std::map<std::string, Function> &env) {
      return StorageFolding(env).mutate(s);
  }
  
  }  // namespace Internal
  }  // namespace Halide
  