emitLLVMKernel: avoid relying on isl set variable names

codegenISL relied on set variable names in two ways.
First, it assumed that the variable names of the domain
of the schedule were preserved through the entire
AST generation process.
Explicitly set the names of the space sent to
halide2isl::makeIslAffFromExpr instead of assuming
they are still available in the schedule.
halide2isl::makeIslAffFromExpr itself still relies on
the names in the space itself being preserved, but
the space does not get modified in between,
so the risk of the names disappearing is reduced.

Second, codegenISL would manually set the names of the range
of the schedule, expecting them to be preserved across subsequent operations,
but this is again not guaranteed.
Convert the corresponding expressions to isl::ast_expr objects instead.
As a nice bonus, the code generation for subscripts and variables
now share the same mechanism, while before one used isl::ast_expr
objects, while the other used an isl::pw_multi_aff.
As a result, a lot of code can be removed.

Author: Sven Verdoolaege

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

src/core/polyhedral/codegen_llvm.cc

@@ -72,7 +72,7 @@ isl::aff makeIslAffFromExpr(isl::space space, const Halide::Expr& e);
 
 namespace polyhedral {
 
-using IteratorMapType = isl::pw_multi_aff;
+using IteratorMapType = std::unordered_map<std::string, isl::ast_expr>;
 using IteratorMapsType =
     std::unordered_map<isl::id, IteratorMapType, isl::IslIdIslHash>;
 
@@ -97,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
@@ -278,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:
@@ -709,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

@@ -558,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