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

@@ -6,6 +6,10 @@ 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 +42,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 +53,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 +92,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;
  }
 }

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";
  }
@@ -2369,12 +2369,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));
@@ -2460,6 +2467,110 @@ void CodeGen_C::visit(const Call *op) {
  }
  rhs << "(&" << struct_name << ")";
  }
+ } else if (op->is_intrinsic(Call::declare_struct_type)) {
+ // Emit a declaration like:
+ // struct {const int f_0, const char f_1, const int f_2} foo;
+
+ // Declares a struct type. Returns a null pointer of the new type.
+ internal_assert(op->args.size() >= 2);
+
+ const StringImm *str_imm = op->args[0].as<StringImm>();
+ internal_assert(str_imm != nullptr);
+ std::string name = c_print_name(str_imm->value, false);
+
+ const int64_t *mode = as_const_int(op->args[1]);
+ internal_assert(mode != nullptr);
+
+ if (*mode == 1 || !starts_with(name, "halide_")) {
+ stream << get_indent() << "struct " << name;
+ if (*mode == 1) {
+ stream << " {\n";
+ // List the types.
+ indent++;
+ for (size_t i = 2; i < op->args.size(); i++) {
+ stream << get_indent() << print_type(op->args[i].type(), AppendSpace) << "f_" << i - 2 << ";\n";
+ }
+ indent--;
+ stream << get_indent() << "};\n";
+ } else {
+ stream << ";\n";
+ }
+ }
+ rhs << "((" << name << " *)nullptr)";
+ } else if (op->is_intrinsic(Call::make_typed_struct)) {
+ internal_assert(op->args.size() >= 2);
+
+ string type_carrier = print_expr(op->args[0]);
+ const int64_t *count_ptr = as_const_int(op->args[1]);
+ internal_assert(count_ptr != nullptr);
+ int64_t count = *count_ptr;
+
+ // Get the args
+ vector<string> values;
+ for (size_t i = 2; i < op->args.size(); i++) {
+ values.push_back(print_expr(op->args[i]));
+ }
+ string struct_name = unique_name('s');
+ string array_specifier;
+ if (count > 0) {
+ array_specifier = "[]";
+ }
+ stream << get_indent() << "std::remove_pointer<decltype(" << type_carrier << ")>::type " << struct_name << array_specifier << " = {\n";
+ // List the values.
+ indent++;
+ int item = 0;
+ size_t initializer_count = values.size() / count;
+ internal_assert((values.size() % count) == 0);
+ do {
+ if (count > 0) {
+ stream << get_indent() << "{\n";
+ indent++;
+ }
+ for (size_t i = 0; i < initializer_count; i++) {
+ stream << get_indent() << values[i + item * initializer_count];
+ if (i < op->args.size() - 1) {
+ stream << ",";
+ }
+ stream << "\n";
+ }
+ if (count > 0) {
+ indent--;
+ stream << get_indent() << "},\n";
+ }
+ item += 1;
+ } while (item < count);
+ indent--;
+ stream << get_indent() << "};\n";
+ // Return a pointer to it of the appropriate type
+
+ // TODO: This is dubious type-punning. We really need to
+ // find a better way to do this. We dodge the problem for
+ // the specific case of buffer shapes in the case above.
+ if (op->type.handle_type) {
+ rhs << "(" << print_type(op->type) << ")";
+ }
+ rhs << "(&" << struct_name << ((count > 0) ? "[0])" : ")");
+ } else if (op->is_intrinsic(Call::load_struct_member)) {
+ internal_assert(op->args.size() == 3);
+ const uint64_t *index = as_const_uint(op->args[2]);
+ internal_assert(index != nullptr);
+ std::string struct_base = print_expr(op->args[0]);
+ std::string type_carrier = print_expr(op->args[1]);
+ rhs << "((decltype(" << type_carrier << "))" << struct_base << ")->"
+ << "f_" << *index;
+ } else if (op->is_intrinsic(Call::resolve_function_name)) {
+ internal_assert(op->args.size() == 1);
+ const Call *decl_call = op->args[0].as<Call>();
+ internal_assert(decl_call != nullptr);
+ std::string c_name = c_print_name(decl_call->name, false);
+ rhs << "(&" << c_name << ")";
+ } else if (op->is_intrinsic(Call::get_user_context)) {
+ internal_assert(op->args.empty());
+ rhs << "_ucon";
+ } else if (op->is_intrinsic(Call::get_pointer_symbol_or_null)) {
+ internal_assert(op->args.size() == 2);
+ // TODO(zalman|abadams): Figure out how to get rid of the maybe foo.buffer exists, maybe it doesn't thing in closures.
+ rhs << "(nullptr)";
  } else if (op->is_intrinsic(Call::stringify)) {
  // Rewrite to an snprintf
  vector<string> printf_args;
@@ -2490,7 +2601,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>();
@@ -2560,6 +2670,21 @@ void CodeGen_C::visit(const Call *op) {
  stream << get_indent() << rhs.str() << ";\n";
  // Make an innocuous assignment value for our caller (probably an Evaluate node) to ignore.
  print_assignment(op->type, "0");
+ } else if (op->is_intrinsic(Call::declare_struct_type) ||
+ op->is_intrinsic(Call::make_typed_struct) ||
+ op->is_intrinsic(Call::load_struct_member) ||
+ op->is_intrinsic(Call::resolve_function_name)) {
+ // print_assigment will get the type info wrong for this case.
+ std::string rhs_str = rhs.str();
+ auto cached = cache.find(rhs_str);
+ if (cached == cache.end()) {
+ id = unique_name('_');
+ stream << get_indent() << "auto " << id << " = " << rhs_str << ";\n";
+ } else {
+ id = cached->second;
+ }
+ // Avoid unused variable warnings.
+ stream << get_indent() << "halide_unused(" << id << ");\n";
  } else {
  print_assignment(op->type, rhs.str());
  }
@@ -2710,9 +2835,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 +2926,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);
@@ -3222,7 +3350,8 @@ 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 _buf = _0;
+ halide_unused(_buf);
  {
  int64_t _1 = 43;
  int64_t _2 = _1 * _beta;

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;
  }