emitCudaKernel: avoid relying on isl set variable names in range of schedule

In particular, keep track of the AST build where each AST node was created
and use that to generate AST expressions.
Note that an AST build is a relatively large data structure, so
this may end up consuming some memory if many statements or AST nodes
are involved.
An alternative would be to generate the AST expressions while
the AST is being constructed, as is done in PPCG and in the LLVM
code generator, but this requires two passes over the Halide
data structures, once to generate the AST expressions and
once to print them and would be a more involved change overall.
A third alternative would be to generate AST expressions
for the individual statement indices during the AST generation
(using iteratorMap) and then to generate AST expressions in terms
of those statement indices during the printing of the AST,
plugging in the AST expressions for the statement indices.

Author: Sven Verdoolaege

include/tc/core/polyhedral/cuda/codegen.h

@@ -61,32 +61,52 @@ void emitMappedTensorAccess(
 
 } // namespace detail
 
-using IteratorMapsType =
-    std::unordered_map<isl::id, isl::pw_multi_aff, isl::IslIdIslHash>;
+/*
+ * Information attached to an AST node during printing of the AST.
+ * iteratorMap is the inverse schedule, mapping schedule dimensions
+ * to the indices of the statement corresponding to the AST node.
+ * build is the AST build at the point where the AST node is generated.
+ * It is used to generate AST expressions in that context.
+ */
+struct NodeInfo {
+  isl::pw_multi_aff iteratorMap;
+  isl::ast_build build;
+};
+/*
+ * Type used for mapping AST node identifier to the corresponding
+ * AST node information.
+ */
+using NodeInfoMapType =
+    std::unordered_map<isl::id, NodeInfo, isl::IslIdIslHash>;
 
 struct CodegenContext {
   CodegenContext(
       std::stringstream& ss_,
       const MappedScop& s,
-      const IteratorMapsType& i)
-      : ss(ss_), mappedScop(s), iteratorMaps(i) {}
+      const NodeInfoMapType& i)
+      : ss(ss_), mappedScop(s), nodeInfoMap(i) {}
   CodegenContext(const CodegenContext& c)
-      : ss(c.ss), mappedScop(c.mappedScop), iteratorMaps(c.iteratorMaps) {}
+      : ss(c.ss), mappedScop(c.mappedScop), nodeInfoMap(c.nodeInfoMap) {}
 
   const Scop& scop() const {
     return mappedScop.scop();
   }
 
   std::stringstream& ss;
   const MappedScop& mappedScop;
-  const IteratorMapsType& iteratorMaps;
+  const NodeInfoMapType& nodeInfoMap;
 };
 
 struct CodegenStatementContext : CodegenContext {
   CodegenStatementContext(const CodegenContext& c, isl::id astId)
       : CodegenContext(c), astNodeId(astId) {}
   isl::pw_multi_aff iteratorMap() const {
-    return this->iteratorMaps.at(astNodeId);
+    return this->nodeInfoMap.at(astNodeId).iteratorMap;
+  }
+  // Return the build where the AST node of this CodegenStatementContext
+  // was constructed.
+  isl::ast_build build() const {
+    return this->nodeInfoMap.at(astNodeId).build;
   }
   isl::id statementId() const {
     return this->iteratorMap().get_tuple_id(isl::dim_type::out);

src/core/polyhedral/cuda/codegen.cc

@@ -367,11 +367,7 @@ void emitReductionInit(
 namespace {
 template <typename AFF>
 void emitAccess(AFF access, const CodegenStatementContext& context) {
-  // Use a temporary isl::ast_build to print the expression.
-  // Ideally, this should use the build at the point
-  // where the user statement was created.
-  auto astBuild = isl::ast_build::from_context(access.domain());
-  context.ss << astBuild.access_from(access).to_C_str();
+  context.ss << context.build().access_from(access).to_C_str();
 }
 } // namespace
 
@@ -400,8 +396,8 @@ void AstPrinter::emitStmt(isl::ast_node_user node) {
   auto stmtId = usrExp.get_op_arg(0).get_id();
   auto nodeId = node.get_annotation();
   auto statementContext = CodegenStatementContext(context_, nodeId);
-  CHECK_EQ(context_.iteratorMaps.count(nodeId), 1)
-      << "no iterator remapping for op " << nodeId;
+  CHECK_EQ(context_.nodeInfoMap.count(nodeId), 1)
+      << "no info for node " << nodeId;
 
   WS ws;
   context_.ss << ws.tab();
@@ -514,11 +510,7 @@ void emitHalideExpr(
     void visit(const Halide::Internal::Variable* op) {
       auto pwAff = tc::polyhedral::detail::makeAffFromMappedExpr(
           Halide::Expr(op), context);
-      // Use a temporary isl::ast_build to print the expression.
-      // Ideally, this should use the build at the point
-      // where the user statement was created.
-      auto astBuild = isl::ast_build::from_context(pwAff.domain());
-      auto expr = astBuild.expr_from(pwAff);
+      auto expr = context.build().expr_from(pwAff);
       auto s = expr.to_C_str();
       if (!is_identifier_or_nonnegative_integer(expr)) {
         s = "(" + s + ")";
@@ -718,42 +710,32 @@ string emitCudaKernel(
   emitTensorViews(ss, scop.halide.inputs, paramValues);
   emitTmpDecl(ss, scop);
   emitPromotedArrayViewsHalide(ss, scop);
-  IteratorMapsType iteratorMaps;
-  auto collect = [&iteratorMaps](
+  NodeInfoMapType nodeInfoMap;
+  auto collect = [&nodeInfoMap](
                      isl::ast_node n, isl::ast_build b) -> isl::ast_node {
     auto collectIteratorMaps =
         [](isl::ast_node node,
            isl::ast_build build,
-           IteratorMapsType* iteratorMaps) -> isl::ast_node {
+           NodeInfoMapType* nodeInfoMap) -> isl::ast_node {
       auto user = node.as<isl::ast_node_user>();
       CHECK(user);
       auto expr = user.get_expr();
       auto stmtId = expr.get_op_arg(0).get_id();
-      // 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 nodeId = isl::id(
           node.get_ctx(),
           std::string(kAstNodeIdPrefix) + std::to_string(nAstNodes()++));
-      CHECK_EQ(0, iteratorMaps->count(nodeId)) << "entry exists: " << nodeId;
-      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));
-      }
+      CHECK_EQ(0, nodeInfoMap->count(nodeId)) << "entry exists: " << nodeId;
 
-      auto iteratorMap = isl::pw_multi_aff(scheduleMap.reverse());
-      iteratorMaps->emplace(nodeId, iteratorMap);
+      auto& nodeInfo = (*nodeInfoMap)[nodeId];
+      nodeInfo.iteratorMap = isl::pw_multi_aff(scheduleMap.reverse());
+      nodeInfo.build = build;
       return node.set_annotation(nodeId);
     };
 
-    return collectIteratorMaps(n, b, &iteratorMaps);
+    return collectIteratorMaps(n, b, &nodeInfoMap);
   };
 
   auto bands = detail::ScheduleTree::collect(
@@ -775,7 +757,7 @@ string emitCudaKernel(
   astBuild = astBuild.set_at_each_domain(collect);
   astBuild = astBuild.set_iterators(Codegen::makeLoopIterators(ctx, maxDepth));
   auto astNode = astBuild.node_from(schedule);
-  AstPrinter(CodegenContext(ss, mscop, iteratorMaps)).emit(astNode);
+  AstPrinter(CodegenContext(ss, mscop, nodeInfoMap)).emit(astNode);
   ss << "}" << endl;
 
   return ss.str();