Move parallel/async lowering from LLVM codegen to a standard Halide IR lowering pass.

Makefile

@@ -491,6 +491,7 @@ SOURCE_FILES = \
  LLVM_Runtime_Linker.cpp \
  LoopCarry.cpp \
  Lower.cpp \
+ LowerParallelTasks.cpp \
  LowerWarpShuffles.cpp \
  MatlabWrapper.cpp \
  Memoization.cpp \
@@ -669,6 +670,7 @@ HEADER_FILES = \
  LLVM_Runtime_Linker.h \
  LoopCarry.h \
  Lower.h \
+ LowerParallelTasks.h \
  LowerWarpShuffles.h \
  MainPage.h \
  MatlabWrapper.h \

src/AsyncProducers.cpp

@@ -333,7 +333,7 @@ class ForkAsyncProducers : public IRMutator {
  vector<Expr> sema_vars;
  for (int i = 0; i < consumes.count; i++) {
  sema_names.push_back(op->name + ".semaphore_" + std::to_string(i));
- sema_vars.push_back(Variable::make(Handle(), sema_names.back()));
+ sema_vars.push_back(Variable::make(type_of<halide_semaphore_t *>(), sema_names.back()));
  }
 
  Stmt producer = GenerateProducerBody(op->name, sema_vars, cloned_acquires).mutate(body);

src/CMakeLists.txt

@@ -98,6 +98,7 @@ set(HEADER_FILES
  LLVM_Runtime_Linker.h
  LoopCarry.h
  Lower.h
+ LowerParallelTasks.h
  LowerWarpShuffles.h
  MainPage.h
  MatlabWrapper.h
@@ -258,6 +259,7 @@ set(SOURCE_FILES
  LLVM_Runtime_Linker.cpp
  LoopCarry.cpp
  Lower.cpp
+ LowerParallelTasks.cpp
  LowerWarpShuffles.cpp
  MatlabWrapper.cpp
  Memoization.cpp

src/Closure.cpp

@@ -1,11 +1,18 @@
 #include "Closure.h"
 #include "Debug.h"
+#include "ExprUsesVar.h"
+#include "IRMutator.h"
+#include "IROperator.h"
 
 namespace Halide {
 namespace Internal {
 
 using std::string;
 
+namespace {
+constexpr int DBG = 3;
+} // namespace
+
 void Closure::include(const Stmt &s, const string &loop_variable) {
  if (!loop_variable.empty()) {
  ignore.push(loop_variable);
@@ -38,7 +45,7 @@ void Closure::visit(const For *op) {
 void Closure::found_buffer_ref(const string &name, Type type,
  bool read, bool written, const Halide::Buffer<> &image) {
  if (!ignore.contains(name)) {
- debug(3) << "Adding buffer " << name << " to closure\n";
+ debug(DBG) << "Adding buffer " << name << " to closure:\n";
  Buffer &ref = buffers[name];
  ref.type = type.element_of(); // TODO: Validate type is the same as existing refs?
  ref.read = ref.read || read;
@@ -49,8 +56,15 @@ void Closure::found_buffer_ref(const string &name, Type type,
  ref.size = image.size_in_bytes();
  ref.dimensions = image.dimensions();
  }
+ debug(DBG) << " "
+ << " t=" << ref.type
+ << " d=" << (int)ref.dimensions
+ << " r=" << ref.read
+ << " w=" << ref.write
+ << " mt=" << (int)ref.memory_type
+ << " sz=" << ref.size << "\n";
  } else {
- debug(3) << "Not adding " << name << " to closure\n";
+ debug(DBG) << "Not adding buffer " << name << " to closure\n";
  }
 }
 
@@ -81,9 +95,9 @@ void Closure::visit(const Allocate *op) {
 
 void Closure::visit(const Variable *op) {
  if (ignore.contains(op->name)) {
- debug(3) << "Not adding " << op->name << " to closure\n";
+ debug(DBG) << "Not adding var " << op->name << " to closure\n";
  } else {
- debug(3) << "Adding " << op->name << " to closure\n";
+ debug(DBG) << "Adding var " << op->name << " to closure\n";
  vars[op->name] = op->type;
  }
 }
@@ -95,5 +109,71 @@ void Closure::visit(const Atomic *op) {
  op->body.accept(this);
 }
 
+Expr Closure::pack_into_struct() const {
+ std::vector<Expr> elements;
+
+ for (const auto &b : buffers) {
+ Expr ptr_var = Variable::make(type_of<void *>(), b.first);
+ elements.emplace_back(ptr_var);
+ }
+ for (const auto &v : vars) {
+ Expr var = Variable::make(v.second, v.first);
+ elements.emplace_back(var);
+ }
+
+ // Sort by decreasing size, to guarantee the struct is densely packed in
+ // memory. We don't actually rely on this, it's just nice to have.
+ std::stable_sort(elements.begin(), elements.end(),
+ [&](const Expr &a, const Expr &b) {
+ return a.type().bytes() > b.type().bytes();
+ });
+
+ Expr result = Call::make(Handle(),
+ Call::make_struct, elements, Call::Intrinsic);
+ return result;
+}
+
+Stmt Closure::unpack_from_struct(const Expr &e, const Stmt &s) const {
+ // Use the struct-packing code just to make sure the order of elements is
+ // the same.
+ Expr packed = pack_into_struct();
+
+ // Make a prototype of the packed struct
+ class ReplaceCallArgsWithZero : public IRMutator {
+ public:
+ using IRMutator::mutate;
+ Expr mutate(const Expr &e) override {
+ if (!e.as<Call>()) {
+ return make_zero(e.type());
+ } else {
+ return IRMutator::mutate(e);
+ }
+ }
+ } replacer;
+ string prototype_name = unique_name("closure_prototype");
+ Expr prototype = replacer.mutate(packed);
+ Expr prototype_var = Variable::make(Handle(), prototype_name);
+
+ const Call *c = packed.as<Call>();
+
+ Stmt result = s;
+ for (int idx = (int)c->args.size() - 1; idx >= 0; idx--) {
+ Expr arg = c->args[idx];
+ const Variable *var = arg.as<Variable>();
+ Expr val = Call::make(var->type,
+ Call::load_typed_struct_member,
+ {e, prototype_var, idx},
+ Call::Intrinsic);
+ if (stmt_uses_var(result, var->name)) {
+ // If a closure is generated for multiple consuming blocks of IR,
+ // then some of those blocks might only need some of the field.
+ result = LetStmt::make(var->name, val, result);
+ }
+ }
+ result = LetStmt::make(prototype_name, prototype, result);
+
+ return result;
+}
+
 } // namespace Internal
 } // namespace Halide

src/Closure.h

@@ -90,11 +90,19 @@ class Closure : public IRVisitor {
  **/
  void include(const Stmt &s, const std::string &loop_variable = "");
 
- /** External variables referenced. */
+ /** External variables referenced. There's code that assumes iterating over
+ * this repeatedly gives a consistent order, so don't swap out the data type
+ * for something non-deterministic. */
  std::map<std::string, Type> vars;
 
  /** External allocations referenced. */
  std::map<std::string, Buffer> buffers;
+
+ /** Pack a closure into a struct. */
+ Expr pack_into_struct() const;
+
+ /** Unpack a closure around a Stmt, putting all the names in scope. */
+ Stmt unpack_from_struct(const Expr &, const Stmt &) const;
 };
 
 } // namespace Internal

src/CodeGen_C.cpp

@@ -1820,7 +1820,7 @@ void CodeGen_C::compile(const LoweredFunc &f, const std::map<std::string, std::s
  set_name_mangling_mode(name_mangling);
 
  std::vector<std::string> namespaces;
- std::string simple_name = extract_namespaces(f.name, namespaces);
+ std::string simple_name = c_print_name(extract_namespaces(f.name, namespaces), false);
  if (!is_c_plus_plus_interface()) {
  user_assert(namespaces.empty()) << "Namespace qualifiers not allowed on function name if not compiling with Target::CPlusPlusNameMangling.\n";
  }
@@ -2024,7 +2024,13 @@ string CodeGen_C::print_assignment(Type t, const std::string &rhs) {
  if (cached == cache.end()) {
  id = unique_name('_');
  const char *const_flag = output_kind == CPlusPlusImplementation ? "const " : "";
- stream << get_indent() << print_type(t, AppendSpace) << const_flag << id << " = " << rhs << ";\n";
+ if (t.is_handle()) {
+ // Don't print void *, which might lose useful type information. just use auto.
+ stream << get_indent() << "auto *";
+ } else {
+ stream << get_indent() << print_type(t, AppendSpace);
+ }
+ stream << const_flag << id << " = " << rhs << ";\n";
  cache[rhs] = id;
  } else {
  id = cached->second;
@@ -2051,7 +2057,12 @@ void CodeGen_C::close_scope(const std::string &comment) {
 }
 
 void CodeGen_C::visit(const Variable *op) {
- id = print_name(op->name);
+ if (starts_with(op->name, "::")) {
+ // This is the name of a global, so we can't modify it.
+ id = op->name;
+ } else {
+ id = print_name(op->name);
+ }
 }
 
 void CodeGen_C::visit(const Cast *op) {
@@ -2369,12 +2380,19 @@ void CodeGen_C::visit(const Call *op) {
  } else if (op->is_intrinsic(Call::alloca)) {
  internal_assert(op->args.size() == 1);
  internal_assert(op->type.is_handle());
+ const int64_t *sz = as_const_int(op->args[0]);
  if (op->type == type_of<struct halide_buffer_t *>() &&
  Call::as_intrinsic(op->args[0], {Call::size_of_halide_buffer_t})) {
  stream << get_indent();
  string buf_name = unique_name('b');
  stream << "halide_buffer_t " << buf_name << ";\n";
  rhs << "&" << buf_name;
+ } else if (op->type == type_of<struct halide_semaphore_t *>() &&
+ sz && *sz == 16) {
+ stream << get_indent();
+ string semaphore_name = unique_name("sema");
+ stream << "halide_semaphore_t " << semaphore_name << ";\n";
+ rhs << "&" << semaphore_name;
  } else {
  // Make a stack of uint64_ts
  string size = print_expr(simplify((op->args[0] + 7) / 8));
@@ -2422,7 +2440,7 @@ void CodeGen_C::visit(const Call *op) {
  rhs << shape_name;
  } else {
  // Emit a declaration like:
- // struct {const int f_0, const char f_1, const int f_2} foo = {3, 'c', 4};
+ // struct {int f_0, int f_1, char f_2} foo = {3, 4, 'c'};
 
  // Get the args
  vector<string> values;
@@ -2433,7 +2451,7 @@ void CodeGen_C::visit(const Call *op) {
  // List the types.
  indent++;
  for (size_t i = 0; i < op->args.size(); i++) {
- stream << get_indent() << "const " << print_type(op->args[i].type()) << " f_" << i << ";\n";
+ stream << get_indent() << print_type(op->args[i].type()) << " f_" << i << ";\n";
  }
  indent--;
  string struct_name = unique_name('s');
@@ -2460,6 +2478,26 @@ void CodeGen_C::visit(const Call *op) {
  }
  rhs << "(&" << struct_name << ")";
  }
+ } else if (op->is_intrinsic(Call::load_typed_struct_member)) {
+ // Given a void * instance of a typed struct, an in-scope prototype
+ // struct of the same type, and the index of a slot, load the value of
+ // that slot.
+ //
+ // It is assumed that the slot index is valid for the given typed struct.
+ //
+ // TODO: this comment is replicated in CodeGen_LLVM and should be updated there too.
+ // TODO: https://github.com/halide/Halide/issues/6468
+
+ internal_assert(op->args.size() == 3);
+ std::string struct_instance = print_expr(op->args[0]);
+ std::string struct_prototype = print_expr(op->args[1]);
+ const int64_t *index = as_const_int(op->args[2]);
+ internal_assert(index != nullptr);
+ rhs << "((decltype(" << struct_prototype << "))"
+ << struct_instance << ")->f_" << *index;
+ } else if (op->is_intrinsic(Call::get_user_context)) {
+ internal_assert(op->args.empty());
+ rhs << "_ucon";
  } else if (op->is_intrinsic(Call::stringify)) {
  // Rewrite to an snprintf
  vector<string> printf_args;
@@ -2490,7 +2528,6 @@ void CodeGen_C::visit(const Call *op) {
  stream << get_indent() << "char " << buf_name << "[1024];\n";
  stream << get_indent() << "snprintf(" << buf_name << ", 1024, \"" << format_string << "\", " << with_commas(printf_args) << ");\n";
  rhs << buf_name;
-
  } else if (op->is_intrinsic(Call::register_destructor)) {
  internal_assert(op->args.size() == 2);
  const StringImm *fn = op->args[0].as<StringImm>();
@@ -2507,7 +2544,7 @@ void CodeGen_C::visit(const Call *op) {
  << "" << struct_name << "(void *ucon, void *a) : ucon(ucon), arg((void *)a) {} "
  << "~" << struct_name << "() { " << fn->value + "(ucon, arg); } "
  << "} " << instance_name << "(_ucon, " << arg << ");\n";
- rhs << print_expr(0);
+ rhs << "(void *)nullptr";
  } else if (op->is_intrinsic(Call::div_round_to_zero)) {
  rhs << print_expr(op->args[0]) << " / " << print_expr(op->args[1]);
  } else if (op->is_intrinsic(Call::mod_round_to_zero)) {
@@ -2710,9 +2747,10 @@ void CodeGen_C::visit(const Let *op) {
  if (op->value.type().is_handle()) {
  // The body might contain a Load that references this directly
  // by name, so we can't rewrite the name.
- stream << get_indent() << print_type(op->value.type())
- << " " << print_name(op->name)
- << " = " << id_value << ";\n";
+ std::string name = print_name(op->name);
+ stream << get_indent() << "auto "
+ << name << " = " << id_value << ";\n";
+ stream << get_indent() << "halide_unused(" << name << ");\n";
  } else {
  Expr new_var = Variable::make(op->value.type(), id_value);
  body = substitute(op->name, new_var, body);
@@ -2800,12 +2838,14 @@ void CodeGen_C::visit(const VectorReduce *op) {
 void CodeGen_C::visit(const LetStmt *op) {
  string id_value = print_expr(op->value);
  Stmt body = op->body;
+
  if (op->value.type().is_handle()) {
  // The body might contain a Load or Store that references this
  // directly by name, so we can't rewrite the name.
- stream << get_indent() << print_type(op->value.type())
- << " " << print_name(op->name)
- << " = " << id_value << ";\n";
+ std::string name = print_name(op->name);
+ stream << get_indent() << "auto "
+ << name << " = " << id_value << ";\n";
+ stream << get_indent() << "halide_unused(" << name << ");\n";
  } else {
  Expr new_var = Variable::make(op->value.type(), id_value);
  body = substitute(op->name, new_var, body);
@@ -3221,8 +3261,9 @@ extern "C" {
 HALIDE_FUNCTION_ATTRS
 int test1(struct halide_buffer_t *_buf_buffer, float _alpha, int32_t _beta, void const *__user_context) {
  void * const _ucon = const_cast<void *>(__user_context);
- void *_0 = _halide_buffer_get_host(_buf_buffer);
- void * _buf = _0;
+ auto *_0 = _halide_buffer_get_host(_buf_buffer);
+ auto _buf = _0;
+ halide_unused(_buf);
  {
  int64_t _1 = 43;
  int64_t _2 = _1 * _beta;
@@ -3248,7 +3289,7 @@ int test1(struct halide_buffer_t *_buf_buffer, float _alpha, int32_t _beta, void
  {
  char b0[1024];
  snprintf(b0, 1024, "%lld%s", (long long)(3), "\n");
- char const *_8 = b0;
+ auto *_8 = b0;
  halide_print(_ucon, _8);
  int32_t _9 = 0;
  int32_t _10 = return_second(_9, 3);

src/CodeGen_Hexagon.cpp

@@ -431,6 +431,30 @@ class InjectHVXLocks : public IRMutator {
  }
  Expr visit(const Call *op) override {
  uses_hvx = uses_hvx || op->type.is_vector();
+
+ if (op->name == "halide_do_par_for") {
+ // If we see a call to halide_do_par_for() at this point, it should mean that
+ // this statement was produced via HexagonOffload calling lower_parallel_tasks()
+ // explicitly; in this case, we won't see any parallel For statements, since they've
+ // all been transformed into closures already. To mirror the pattern above,
+ // we need to wrap the halide_do_par_for() call with an unlock/lock pair, but
+ // that's hard to do in Halide IR (we'd need to produce a Stmt to enforce the ordering,
+ // and the resulting Stmt can't easily be substituted for the Expr here). Rather than
+ // make fragile assumptions about the structure of the IR produced by lower_parallel_tasks(),
+ // we'll use a trick: we'll define a WEAK_INLINE function, _halide_hexagon_do_par_for,
+ // which simply encapsulates the unlock()/do_par_for()/lock() sequences, and swap out
+ // the call here. Since it is inlined, and since uses_hvx_var gets substituted at the end,
+ // we end up with LLVM IR that properly includes (or omits) the unlock/lock pair depending
+ // on the final value of uses_hvx_var in this scope.
+
+ internal_assert(op->call_type == Call::Extern);
+ internal_assert(op->args.size() == 4);
+
+ std::vector<Expr> args = op->args;
+ args.push_back(cast<int>(uses_hvx_var));
+
+ return Call::make(Int(32), "_halide_hexagon_do_par_for", args, Call::Extern);
+ }
  return op;
  }