Merge pull request #184 from facebookresearch/pr/codegen_llvm_names

emitLLVMKernel: avoid relying on isl set variable names

Author: skimo-openhub

include/tc/core/halide2isl.h

@@ -51,6 +51,8 @@ isl::aff makeIslAffFromInt(isl::space space, int64_t i);
 // does not correspond to a parameter or set dimension of the space.
 isl::aff makeIslAffFromExpr(isl::space space, const Halide::Expr& e);
 
+typedef std::unordered_map<isl::id, std::vector<std::string>, isl::IslIdIslHash>
+    IteratorMap;
 typedef std::unordered_map<isl::id, Halide::Internal::Stmt, isl::IslIdIslHash>
     StatementMap;
 typedef std::unordered_map<const Halide::Internal::IRNode*, isl::id> AccessMap;
@@ -73,6 +75,10 @@ struct ScheduleTreeAndAccesses {
   /// The correspondence between leaf Stmts and the statement ids
   /// refered to above.
   StatementMap statements;
+
+  /// The correspondence between statement ids and the outer loop iterators
+  /// of the corresponding leaf Stmt.
+  IteratorMap iterators;
 };
 
 /// Make a schedule tree from a Halide Stmt, along with auxiliary data

include/tc/core/polyhedral/scop.h

@@ -358,6 +358,17 @@ struct Scop {
   // Assumes such argument exists.
   const Halide::OutputImageParam& findArgument(isl::id id) const;
 
+  // Make an affine function from a Halide Expr that is defined
+  // over the instance set of the statement with identifier "stmtId" and
+  // with parameters specified by "paramSpace".  Return a
+  // null isl::aff if the expression is not affine.  Fail if any
+  // of the variables does not correspond to a parameter or
+  // an instance identifier of the statement.
+  isl::aff makeIslAffFromStmtExpr(
+      isl::id stmtId,
+      isl::space paramSpace,
+      const Halide::Expr& e) const;
+
   // Promote a tensor reference group to a storage of a given "kind",
   // inserting the copy
   // statements below the given node.  Inserts an Extension node below the give
@@ -419,9 +430,10 @@ struct Scop {
     std::unordered_map<isl::id, Halide::Internal::Stmt, isl::IslIdIslHash>
         statements;
     std::unordered_map<const Halide::Internal::IRNode*, isl::id> accesses;
+    halide2isl::IteratorMap iterators;
   } halide;
 
-  // Poyhedral IR
+  // Polyhedral IR
   //
   // The domain is collected from the root of the ScheduleTree; no redundant
   // state is kept.

src/core/halide2isl.cc

@@ -316,13 +316,33 @@ struct ScheduleTreeAndDomain {
   isl::union_set domain;
 };
 
+/*
+ * Helper function for extracting a schedule tree from a Halide Stmt,
+ * recursively descending over the Stmt.
+ * "s" is the current position in the recursive descent.
+ * "set" describes the bounds on the outer loop iterators.
+ * "outer" contains the names of the outer loop iterators
+ * from outermost to innermost.
+ * Return the schedule tree corresponding to the subtree at "s",
+ * along with a separated out domain.
+ *
+ * "reads" and "writes" collect the accesses found along the way.
+ * "accesses" collects the mapping from Call (for the reads) and Provide nodes
+ * (for the writes) to the corresponding tag in the access relations.
+ * "statements" collects the mapping from instance set tuple identifiers
+ * to the corresponding Provide node.
+ * "iterators" collects the mapping from instance set tuple identifiers
+ * to the corresponding outer loop iterator names, from outermost to innermost.
+ */
 ScheduleTreeAndDomain makeScheduleTreeHelper(
     const Stmt& s,
     isl::set set,
+    std::vector<std::string>& outer,
     isl::union_map* reads,
     isl::union_map* writes,
     AccessMap* accesses,
-    StatementMap* statements) {
+    StatementMap* statements,
+    IteratorMap* iterators) {
   ScheduleTreeAndDomain result;
   if (auto op = s.as<For>()) {
     // Add one additional dimension to our set of loop variables
@@ -358,8 +378,17 @@ ScheduleTreeAndDomain makeScheduleTreeHelper(
     }
 
     // Recursively descend.
+    auto outerNext = outer;
+    outerNext.push_back(op->name);
     auto body = makeScheduleTreeHelper(
-        op->body, set, reads, writes, accesses, statements);
+        op->body,
+        set,
+        outerNext,
+        reads,
+        writes,
+        accesses,
+        statements,
+        iterators);
 
     // Create an affine function that defines an ordering for all
     // the statements in the body of this loop over the values of
@@ -405,8 +434,8 @@ ScheduleTreeAndDomain makeScheduleTreeHelper(
     // children.
     std::vector<ScheduleTreeUPtr> trees;
     for (Stmt s : stmts) {
-      auto mem =
-          makeScheduleTreeHelper(s, set, reads, writes, accesses, statements);
+      auto mem = makeScheduleTreeHelper(
+          s, set, outer, reads, writes, accesses, statements, iterators);
       ScheduleTreeUPtr filter;
       if (mem.tree) {
         // No statement instances are shared between the blocks, so we
@@ -438,6 +467,7 @@ ScheduleTreeAndDomain makeScheduleTreeHelper(
     size_t stmtIndex = statements->size();
     isl::id id(set.get_ctx(), kStatementLabel + std::to_string(stmtIndex));
     statements->emplace(id, op);
+    iterators->emplace(id, outer);
     isl::set domain = set.set_tuple_id(id);
     result.domain = domain;
 
@@ -460,13 +490,16 @@ ScheduleTreeAndAccesses makeScheduleTree(isl::space paramSpace, const Stmt& s) {
   result.writes = result.reads = isl::union_map::empty(paramSpace);
 
   // Walk the IR building a schedule tree
+  std::vector<std::string> outer;
   auto treeAndDomain = makeScheduleTreeHelper(
       s,
       isl::set::universe(paramSpace),
+      outer,
       &result.reads,
       &result.writes,
       &result.accesses,
-      &result.statements);
+      &result.statements,
+      &result.iterators);
 
   // TODO: This fails if the stmt is just a Provide node, I'm not sure
   // what the schedule tree should look like in that case.

src/core/polyhedral/codegen_llvm.cc

@@ -72,8 +72,9 @@ isl::aff makeIslAffFromExpr(isl::space space, const Halide::Expr& e);
 
 namespace polyhedral {
 
+using IteratorMapType = std::unordered_map<std::string, isl::ast_expr>;
 using IteratorMapsType =
-    std::unordered_map<isl::id, isl::pw_multi_aff, isl::IslIdIslHash>;
+    std::unordered_map<isl::id, IteratorMapType, isl::IslIdIslHash>;
 
 using IteratorLLVMValueMapType =
     std::unordered_map<isl::id, llvm::Value*, isl::IslIdIslHash>;
@@ -96,14 +97,6 @@ llvm::Value* getLLVMConstantSignedInt64(int64_t v) {
   return llvm::ConstantInt::get(llvm::Type::getInt64Ty(llvmCtx), v, true);
 }
 
-isl::aff extractAff(isl::pw_multi_aff pma) {
-  isl::PMA pma_(pma);
-  CHECK_EQ(pma_.size(), 1);
-  isl::MA ma(pma_[0].second);
-  CHECK_EQ(ma.size(), 1);
-  return ma[0];
-}
-
 int64_t IslExprToSInt(isl::ast_expr e) {
   CHECK(isl_ast_expr_get_type(e.get()) == isl_ast_expr_type::isl_ast_expr_int);
   assert(sizeof(long) <= 8); // long is assumed to fit to 64bits
@@ -214,7 +207,7 @@ static constexpr int kOptLevel = 3;
 
 class CodeGen_TC : public Halide::Internal::CodeGen_X86 {
  public:
-  const isl::pw_multi_aff* iteratorMap_;
+  const IteratorMapType* iteratorMap_;
   CodeGen_TC(Target t) : CodeGen_X86(t) {}
 
   using CodeGen_X86::codegen;
@@ -249,6 +242,11 @@ class CodeGen_TC : public Halide::Internal::CodeGen_X86 {
     return std::move(module);
   }
 
+  // Convert an isl AST expression into an llvm::Value.
+  // Only expressions that consist of a pure identifier or
+  // a pure integer constant are currently supported.
+  llvm::Value* getValue(isl::ast_expr expr);
+
  protected:
   using CodeGen_X86::visit;
   void visit(const Halide::Internal::Call* call) override {
@@ -272,44 +270,7 @@ class CodeGen_TC : public Halide::Internal::CodeGen_X86 {
     }
   }
   void visit(const Halide::Internal::Variable* op) override {
-    auto aff = halide2isl::makeIslAffFromExpr(
-        iteratorMap_->get_space().range(), Halide::Expr(op));
-
-    auto subscriptPma = isl::pw_aff(aff).pullback(*iteratorMap_);
-    auto subscriptAff = extractAff(subscriptPma);
-
-    // sanity checks
-    CHECK_EQ(subscriptAff.dim(isl::dim_type::div), 0);
-    CHECK_EQ(subscriptAff.dim(isl::dim_type::out), 1);
-    for (int d = 0; d < subscriptAff.dim(isl::dim_type::param); ++d) {
-      auto v = subscriptAff.get_coefficient_val(isl::dim_type::param, d);
-      CHECK(v.is_zero());
-    }
-
-    llvm::Optional<int> posOne;
-    int sum = 0;
-    for (int d = 0; d < subscriptAff.dim(isl::dim_type::in); ++d) {
-      auto v = subscriptAff.get_coefficient_val(isl::dim_type::in, d);
-      CHECK(v.is_zero() or v.is_one());
-      if (v.is_zero()) {
-        continue;
-      }
-      ++sum;
-      posOne = d;
-    }
-    CHECK_LE(sum, 1);
-
-    if (sum == 0) {
-      value =
-          getLLVMConstantSignedInt64(toSInt(subscriptAff.get_constant_val()));
-      return;
-    }
-    CHECK(posOne);
-
-    std::string name(
-        isl_aff_get_dim_name(subscriptAff.get(), isl_dim_in, *posOne));
-
-    value = sym_get(name);
+    value = getValue(iteratorMap_->at(op->name));
   }
 
  public:
@@ -361,6 +322,21 @@ class CodeGen_TC : public Halide::Internal::CodeGen_X86 {
   }
 };
 
+llvm::Value* CodeGen_TC::getValue(isl::ast_expr expr) {
+  switch (isl_ast_expr_get_type(expr.get())) {
+    case isl_ast_expr_type::isl_ast_expr_id:
+      return sym_get(expr.get_id().get_name());
+    case isl_ast_expr_type::isl_ast_expr_int: {
+      auto val = isl::manage(isl_ast_expr_get_val(expr.get()));
+      CHECK(val.is_int());
+      return getLLVMConstantSignedInt64(val.get_num_si());
+    }
+    default:
+      LOG(FATAL) << "NYI";
+      return nullptr;
+  }
+}
+
 class LLVMCodegen {
   void collectTensor(const Halide::OutputImageParam& t) {
     auto sizes =
@@ -638,22 +614,7 @@ class LLVMCodegen {
     llvm::SmallVector<llvm::Value*, 5> subscriptValues;
 
     for (const auto& subscript : subscripts) {
-      switch (isl_ast_expr_get_type(subscript.get())) {
-        case isl_ast_expr_type::isl_ast_expr_id: {
-          subscriptValues.push_back(
-              halide_cg.sym_get(subscript.get_id().get_name()));
-          break;
-        }
-        case isl_ast_expr_type::isl_ast_expr_int: {
-          auto val = isl::manage(isl_ast_expr_get_val(subscript.get()));
-          CHECK_EQ(val.get_den_si(), 1);
-          subscriptValues.push_back(
-              getLLVMConstantSignedInt64(val.get_num_si()));
-          break;
-        }
-        default:
-          LOG(FATAL) << "NYI";
-      }
+      subscriptValues.push_back(halide_cg.getValue(subscript));
     }
 
     auto destAddr = halide_cg.get_builder().CreateInBoundsGEP(
@@ -703,34 +664,28 @@ IslCodegenRes codegenISL(const Scop& scop) {
            const Scop& scop,
            StmtSubscriptExprMapType& stmtSubscripts) -> isl::ast_node {
       auto expr = node.user_get_expr();
-      // We rename loop-related dimensions manually.
       auto schedule = build.get_schedule();
-      auto scheduleSpace = build.get_schedule_space();
       auto scheduleMap = isl::map::from_union_map(schedule);
 
       auto stmtId = expr.get_op_arg(0).get_id();
       // auto nodeId = isl::id(
       // node.get_ctx(),
       // std::string(kAstNodeIdPrefix) + std::to_string(nAstNodes()++));
       CHECK_EQ(0, iteratorMaps.count(stmtId)) << "entry exists: " << stmtId;
-      CHECK_EQ(
-          scheduleMap.dim(isl::dim_type::out),
-          scheduleSpace.dim(isl::dim_type::set));
-      for (int i = 0; i < scheduleSpace.dim(isl::dim_type::set); ++i) {
-        scheduleMap = scheduleMap.set_dim_id(
-            isl::dim_type::out,
-            i,
-            scheduleSpace.get_dim_id(isl::dim_type::set, i));
-      }
       auto iteratorMap = isl::pw_multi_aff(scheduleMap.reverse());
-      iteratorMaps.emplace(stmtId, iteratorMap);
+      auto iterators = scop.halide.iterators.at(stmtId);
+      auto& stmtIteratorMap = iteratorMaps[stmtId];
+      for (int i = 0; i < iterators.size(); ++i) {
+        auto expr = build.expr_from(iteratorMap.get_pw_aff(i));
+        stmtIteratorMap.emplace(iterators[i], expr);
+      }
       auto& subscripts = stmtSubscripts[stmtId];
       auto provide =
           scop.halide.statements.at(stmtId).as<Halide::Internal::Provide>();
       for (auto e : provide->args) {
         const auto& map = iteratorMap;
-        auto space = map.get_space().range();
-        auto aff = halide2isl::makeIslAffFromExpr(space, e);
+        auto space = map.get_space().params();
+        auto aff = scop.makeIslAffFromStmtExpr(stmtId, space, e);
         auto pulled = isl::pw_aff(aff).pullback(map);
         CHECK_EQ(pulled.n_piece(), 1);
         subscripts.push_back(build.expr_from(pulled));

src/core/polyhedral/scop.cc

@@ -70,6 +70,7 @@ ScopUPtr Scop::makeScop(
   scop->halide.statements = std::move(tree.statements);
   scop->halide.accesses = std::move(tree.accesses);
   scop->halide.reductions = halide2isl::findReductions(components.stmt);
+  scop->halide.iterators = std::move(tree.iterators);
 
   // Set partial schedule tuples for proper comparison with ISL
   // schedules (needs DFSPreorder numbering). Just for testing.
@@ -557,5 +558,23 @@ const Halide::OutputImageParam& Scop::findArgument(isl::id id) const {
   return *halide.inputs.begin();
 }
 
+isl::aff Scop::makeIslAffFromStmtExpr(
+    isl::id stmtId,
+    isl::space paramSpace,
+    const Halide::Expr& e) const {
+  auto ctx = stmtId.get_ctx();
+  auto iterators = halide.iterators.at(stmtId);
+  auto space = paramSpace.set_from_params();
+  space = space.add_dims(isl::dim_type::set, iterators.size());
+  // Set the names of the set dimensions of "space" for use
+  // by halide2isl::makeIslAffFromExpr.
+  for (int i = 0; i < iterators.size(); ++i) {
+    isl::id id(ctx, iterators[i]);
+    space = space.set_dim_id(isl::dim_type::set, i, id);
+  }
+  space = space.set_tuple_id(isl::dim_type::set, stmtId);
+  return halide2isl::makeIslAffFromExpr(space, e);
+}
+
 } // namespace polyhedral
 } // namespace tc