Remove 'foo.buffer' from Closure entirely

This is a direct adaptation of what #6481 does for the old LLVM-based code, and allows elimination of one use of `get_pointer_or_null()`. PR is meant to be merged into factor_parallel_codegen, not master.

src/LowerParallelTasks.cpp

@@ -10,6 +10,7 @@
 #include "IROperator.h"
 #include "Module.h"
 #include "Param.h"
+#include "Simplify.h"
 
 namespace Halide {
 namespace Internal {
@@ -27,41 +28,49 @@ LoweredFunc generate_closure_ir(const std::string &name, const Closure &closure,
  Expr closure_arg = Variable::make(type_of<void *>(), closure_arg_name);
 
  Stmt wrapped_body = body;
- std::vector<Expr> type_args(closure.vars.size() + closure.buffers.size() * 2 + 1);
+ const size_t closure_entries_count = closure.vars.size() + closure.buffers.size();
+ std::vector<Expr> type_args(closure_entries_count + 1);
  std::string closure_type_name = unique_name("closure_struct_type");
  Expr struct_type = Variable::make(Handle(), closure_type_name);
  type_args[0] = StringImm::make(closure_type_name);
 
+ const auto make_load_typed_struct_member = [&closure_arg, &struct_type](const Type &t, int struct_index) -> Expr {
+ return Call::make(t, Call::load_typed_struct_member,
+ {closure_arg, struct_type, make_const(UInt(32), struct_index)},
+ Call::PureIntrinsic);
+ };
+
+ // Build the loads, then add them in reverse order, so that the final generated
+ // code is ascending by struct_index, which may provide marginally better code.
+ std::vector<std::pair<std::string, Expr>> load_lets;
+ load_lets.reserve(closure_entries_count);
+
  int struct_index = 0;
  for (const auto &v : closure.vars) {
- type_args[struct_index + 1] = make_zero(v.second);
- wrapped_body = LetStmt::make(v.first,
- Call::make(v.second, Call::load_typed_struct_member,
- {closure_arg, struct_type, make_const(UInt(32), struct_index)},
- Call::PureIntrinsic),
- wrapped_body);
+ const auto &name = v.first;
+ const auto &type = v.second;
+ type_args[struct_index + 1] = make_zero(type);
+ load_lets.emplace_back(name, make_load_typed_struct_member(type, struct_index));
  struct_index++;
  }
  for (const auto &b : closure.buffers) {
+ const auto &name = b.first;
  type_args[struct_index + 1] = make_zero(type_of<void *>());
- type_args[struct_index + 2] = make_zero(type_of<halide_buffer_t *>());
- wrapped_body = LetStmt::make(b.first,
- Call::make(type_of<void *>(), Call::load_typed_struct_member,
- {closure_arg, struct_type, make_const(UInt(32), struct_index)},
- Call::PureIntrinsic),
- wrapped_body);
- wrapped_body = LetStmt::make(b.first + ".buffer",
- Call::make(type_of<halide_buffer_t *>(), Call::load_typed_struct_member,
- {closure_arg, struct_type, make_const(UInt(32), struct_index + 1)},
- Call::PureIntrinsic),
- wrapped_body);
- struct_index += 2;
+ load_lets.emplace_back(name, make_load_typed_struct_member(type_of<void *>(), struct_index));
+ struct_index++;
+ }
+
+ for (auto it = load_lets.rbegin(); it != load_lets.rend(); ++it) {
+ wrapped_body = LetStmt::make(it->first, it->second, wrapped_body);
  }
 
  Expr struct_type_decl = Call::make(Handle(), Call::define_typed_struct, type_args, Call::PureIntrinsic);
  wrapped_body = Block::make(Evaluate::make(closure_arg), wrapped_body);
  wrapped_body = LetStmt::make(closure_type_name, struct_type_decl, wrapped_body);
 
+ // Call simplify() to remove any load_struct_member() calls we don't use.
+ wrapped_body = simplify(wrapped_body);
+
  // TODO(zvookin): Figure out how we want to handle name mangling of closures.
  // For now, the C++ backend makes them extern "C" so they have to be NameMangling::C.
  LoweredFunc result{name, args, wrapped_body, LinkageType::External, NameMangling::C};
@@ -72,15 +81,15 @@ LoweredFunc generate_closure_ir(const std::string &name, const Closure &closure,
 }
 
 Expr allocate_closure(const std::string &name, const Closure &closure) {
- std::vector<Expr> closure_elements;
- std::vector<Expr> closure_types;
+ const size_t closure_entries_count = closure.vars.size() + closure.buffers.size();
 
- std::string buffer_t_type_name = unique_name("buffer_t_type");
- Expr buffer_t_type = Variable::make(Handle(), buffer_t_type_name);
-
- // Place holder for closure struct type
- closure_elements.emplace_back(Expr());
+ std::vector<Expr> closure_elements;
+ closure_elements.reserve(closure_entries_count + 2);
+ closure_elements.emplace_back(Expr()); // Place holder for closure struct type
  closure_elements.emplace_back(1);
+
+ std::vector<Expr> closure_types;
+ closure_elements.reserve(closure_entries_count + 1);
  closure_types.emplace_back(unique_name(name + "_type"));
 
  for (const auto &v : closure.vars) {
@@ -91,10 +100,7 @@ Expr allocate_closure(const std::string &name, const Closure &closure) {
  for (const auto &b : closure.buffers) {
  Expr ptr_var = Variable::make(type_of<void *>(), b.first);
  closure_elements.emplace_back(ptr_var);
- closure_elements.emplace_back(Call::make(type_of<halide_buffer_t *>(), Call::get_pointer_symbol_or_null,
- {StringImm::make(b.first + ".buffer"), buffer_t_type}, Call::Intrinsic));
  closure_types.emplace_back(ptr_var);
- closure_types.emplace_back(buffer_t_type);
  }
  Expr closure_struct_type = Call::make(Handle(), Call::define_typed_struct, closure_types, Call::PureIntrinsic);
  std::string closure_type_name = unique_name("closure_struct_type");
@@ -104,10 +110,6 @@ Expr allocate_closure(const std::string &name, const Closure &closure) {
 
  Expr result = Let::make(closure_type_name, closure_struct_type,
  Call::make(type_of<void *>(), Call::make_typed_struct, closure_elements, Call::Intrinsic));
- if (!closure.buffers.empty()) {
- result = Let::make(buffer_t_type_name, Call::make(Handle(), Call::forward_declare_typed_struct, {StringImm::make("halide_buffer_t")}, Call::PureIntrinsic),
- result);
- }
  return result;
 }
 
@@ -509,4 +511,4 @@ Stmt lower_parallel_tasks(const Stmt &s, std::vector<LoweredFunc> &closure_imple
 }
 
 } // namespace Internal
-} // namespace Halide
+} // namespace Halide