Fix Closure API

src/Closure.cpp

@@ -6,11 +6,14 @@ namespace Internal {
 
 using std::string;
 
-Closure::Closure(const Stmt &s, const string &loop_variable) {
+void Closure::include(const Stmt &s, const string &loop_variable) {
     if (!loop_variable.empty()) {
         ignore.push(loop_variable);
     }
     s.accept(this);
+    if (!loop_variable.empty()) {
+        ignore.pop(loop_variable);
+    }
 }
 
 void Closure::visit(const Let *op) {

src/Closure.h

@@ -71,14 +71,24 @@ class Closure : public IRVisitor {
 public:
     Closure() = default;
 
+    // Movable but not copyable.
+    Closure(const Closure &) = delete;
+    Closure &operator=(const Closure &) = delete;
+    Closure(Closure &&) = default;
+    Closure &operator=(Closure &&) = default;
+
     /** Traverse a statement and find all references to external
      * symbols.
      *
      * When the closure encounters a read or write to 'foo', it
      * assumes that the host pointer is found in the symbol table as
      * 'foo.host', and any halide_buffer_t pointer is found under
-     * 'foo.buffer'. */
-    Closure(const Stmt &s, const std::string &loop_variable = "");
+     * 'foo.buffer'.
+     *
+     * Calling this multiple times (on multiple statements) is legal
+     * (and will produce a unified closure).
+     **/
+    void include(const Stmt &s, const std::string &loop_variable = "");
 
     /** External variables referenced. */
     std::map<std::string, Type> vars;

src/DeviceArgument.cpp

@@ -5,18 +5,6 @@
 namespace Halide {
 namespace Internal {
 
-HostClosure::HostClosure(const Stmt &s, const std::string &loop_variable) {
-    // It looks tempting to just collapse this into a call to the Closure ctor,
-    // since the bodies are the same, but there is a subtle trap: if we do that,
-    // the vtable for 'this' will only be that of a Closure, not a HostClosure,
-    // and so our overrides for visit() won't be triggered, and we'll get
-    // incorrect results.
-    if (!loop_variable.empty()) {
-        ignore.push(loop_variable);
-    }
-    s.accept(this);
-}
-
 std::vector<DeviceArgument> HostClosure::arguments() {
     if (debug::debug_level() >= 2) {
         debug(2) << *this;

src/DeviceArgument.h

@@ -94,13 +94,14 @@ struct DeviceArgument {
  * produce a vector of DeviceArgument objects. */
 class HostClosure : public Closure {
 public:
-    HostClosure(const Stmt &s, const std::string &loop_variable = "");
+    HostClosure() = default;
 
     /** Get a description of the captured arguments. */
     std::vector<DeviceArgument> arguments();
 
 protected:
     using Internal::Closure::visit;
+
     void visit(const For *loop) override;
     void visit(const Call *op) override;
 };

src/HexagonOffload.cpp

@@ -755,7 +755,8 @@ class InjectHexagonRpc : public IRMutator {
         // Build a closure for the device code.
         // TODO: Should this move the body of the loop to Hexagon,
         // or the loop itself? Currently, this moves the loop itself.
-        Closure c(body);
+        Closure c;
+        c.include(body);
 
         // A buffer parameter potentially generates 3 scalar parameters (min,
         // extent, stride) per dimension. Pipelines with many buffers may

src/OffloadGPULoops.cpp

@@ -149,7 +149,8 @@ class InjectGpuOffload : public IRMutator {
                  << bounds.num_blocks[3] << ") blocks\n";
 
         // compute a closure over the state passed into the kernel
-        HostClosure c(loop->body, loop->name);
+        HostClosure c;
+        c.include(loop->body, loop->name);
 
         // Determine the arguments that must be passed into the halide function
         vector<DeviceArgument> closure_args = c.arguments();