[WIP] Print Loop* associated with isl_ast_for

[dpeixott]

Summary

First patch to keep Loop* associated with its corresponding
isl_ast_for_node.

[jdoerfert]

Minor comments and questions.

include/polly/CodeGen/IslAst.h

49

Any reason not to use a unordered small set? (Doesn't really matter, I'm just curious)

142

When do we need this specialization? Shouldn't we always handle the general (LoopSet) case?

lib/CodeGen/IslAst.cpp

274

drop the {}

[dpeixott]

include/polly/CodeGen/IslAst.h

49

I initially used an unordered set, but switched to the ordered set to keep the output deterministic. Otherwise, we could print out the loops in a different order depending on their pointer values when iterating through the set.

142

Good point. I will remove it. It was for lazy programmers that didn't want to deal with the reality of multiple loops.

lib/CodeGen/IslAst.cpp

274

Ah, is that what the linter was flagging? I got confused because it was pointing to the top of the file. Will fix it up.

[jdoerfert]

responses

include/polly/CodeGen/IslAst.h

49

I see the merrit of deterministic output here but it might not be needed here, however I do not want to argue against it.

lib/CodeGen/IslAst.cpp

274

It was not (some space in the setvector was), it is just our/llvm coding style for one statement conditionals/loops.

[jdoerfert]

Doesn't seem to be needed anymore.

include/polly/CodeGen/IslAst.h

 //
 //===----------------------------------------------------------------------===//
 
 #ifndef POLLY_ISL_AST_H
 #define POLLY_ISL_AST_H
 
 #include "polly/Config/config.h"
 #include "polly/ScopPass.h"
+#include "llvm/ADT/SetVector.h"
 
 #include "isl/ast.h"
 
 namespace llvm {
 class raw_ostream;
 }
 
 struct isl_ast_node;
 struct isl_ast_expr;
 struct isl_ast_build;
 struct isl_union_map;
 struct isl_pw_multi_aff;
 
 namespace polly {
 class Scop;
 class IslAst;
 class MemoryAccess;
 
 class IslAstInfo : public ScopPass {
 public:
   using MemoryAccessSet = SmallPtrSet<MemoryAccess *, 4>;
+  using LoopSet = SetVector<const Loop *, SmallVector<const Loop *, 2>,
+                            SmallPtrSet<const Loop *, 2> >;
 
   /// @brief Payload information used to annotate an AST node.
   struct IslAstUserPayload {
     /// @brief Construct and initialize the payload.
     IslAstUserPayload()
         : IsInnermost(false), IsInnermostParallel(false),
           IsOutermostParallel(false), IsReductionParallel(false),
           Build(nullptr) {}
struct IslAstUserPayload {
     /// @brief Flag to mark parallel loops which break reductions.
     bool IsReductionParallel;
 
     /// @brief The build environment at the time this node was constructed.
     isl_ast_build *Build;
 
     /// @brief Set of accesses which break reduction dependences.
     MemoryAccessSet BrokenReductions;
+
+    /// @brief Set of loops surrounding the node.
+    LoopSet Loops;
   };
 
 private:
   Scop *S;
   IslAst *Ast;
 
 public:
   static char ID;
public:
   static bool isReductionParallel(__isl_keep isl_ast_node *Node);
 
   /// @brief Get the nodes schedule or a nullptr if not available.
   static __isl_give isl_union_map *getSchedule(__isl_keep isl_ast_node *Node);
 
   /// @brief Get the nodes broken reductions or a nullptr if not available.
   static MemoryAccessSet *getBrokenReductions(__isl_keep isl_ast_node *Node);
 
+  /// @brief Get the original loops surrounding the new node or a nullptr
+  /// if not available.
+  static LoopSet *getLoops(__isl_keep isl_ast_node *Node);
+
+  /// @brief Get the original loop surrounding the new node or a nullptr
+  /// if not available.
+  static const Loop *getLoop(__isl_keep isl_ast_node *Node);
+
   /// @brief Get the nodes build context or a nullptr if not available.
   static __isl_give isl_ast_build *getBuild(__isl_keep isl_ast_node *Node);
 
   ///}
 
   virtual void getAnalysisUsage(AnalysisUsage &AU) const;
   virtual void releaseMemory();
 };
 }
 
 namespace llvm {
 class PassRegistry;
 void initializeIslAstInfoPass(llvm::PassRegistry &);
 }
 #endif /* POLLY_ISL_AST_H */

lib/CodeGen/IslAst.cpp

 //===----------------------------------------------------------------------===//
 
 #include "polly/CodeGen/CodeGeneration.h"
 #include "polly/CodeGen/IslAst.h"
 #include "polly/Dependences.h"
 #include "polly/LinkAllPasses.h"
 #include "polly/Options.h"
 #include "polly/ScopInfo.h"
+#include "llvm/ADT/SmallString.h"
 #include "llvm/Support/Debug.h"
+#include "llvm/Support/raw_ostream.h"
 
 #include "isl/union_map.h"
 #include "isl/list.h"
 #include "isl/ast_build.h"
 #include "isl/set.h"
 #include "isl/map.h"
 #include "isl/aff.h"
 
for (const auto &ReductionClause : Clauses) {
     str += MemoryAccess::getReductionOperatorStr(ReductionClause.first);
     // Remove the first two symbols (", ") to make the output look pretty.
     str += " : " + ReductionClause.second.substr(2) + ")";
   }
 
   return str;
 }
 
+/// @brief Return a pragam string for annotating an isl ast.
+static std::string getLLVMLoopPragmaStr(const Loop *L) {
+  SmallString<128> Buf;
+  raw_svector_ostream OS(Buf);
+  OS << "#llvm loop(" << L->getHeader()->getName();
+  if (BasicBlock *Latch = L->getLoopLatch())
+    OS << ", " << Latch->getName();
+  OS << ")";
+  OS << " depth(" << L->getLoopDepth() << ")";
+  return OS.str();
+}
+
 /// @brief Callback executed for each for node in the ast in order to print it.
 static isl_printer *cbPrintFor(__isl_take isl_printer *Printer,
                                __isl_take isl_ast_print_options *Options,
                                __isl_keep isl_ast_node *Node, void *) {
 
   const std::string BrokenReductionsStr = getBrokenReductionsStr(Node);
   const std::string SimdPragmaStr = "#pragma simd";
   const std::string OmpPragmaStr = "#pragma omp parallel for";
 
+  if (IslAstInfo::LoopSet *Loops = IslAstInfo::getLoops(Node))
+    for (const Loop *L : *Loops)
+      Printer = printLine(Printer, getLLVMLoopPragmaStr(L));
+
   if (IslAstInfo::isInnermostParallel(Node))
     Printer = printLine(Printer, SimdPragmaStr + BrokenReductionsStr);
 
   if (IslAstInfo::isOutermostParallel(Node))
     Printer = printLine(Printer, OmpPragmaStr + BrokenReductionsStr);
 
   return isl_ast_node_for_print(Node, Printer, Options);
 }
static bool astScheduleDimIsParallel(__isl_keep isl_ast_build *Build,
   return true;
 }
 
 /// @brief Collect loop annotations from the orignal loops surrounding @p Body.
 ///
 /// We will look at all statements in @p Body and all loops formerly surrounding
 /// those statements and aggregate their loop annotations if they are invovled
 /// in the __new__ innermost dimension.
-static void collectLoopAnnotations(__isl_take isl_ast_node *Body) {
+static void collectLoopAnnotations(__isl_take isl_ast_node *Body,
+                                   IslAstUserPayload *Payload) {
   // Recurce for block and conditional statements but extract the annotations
   // once a user ast node was found.
   switch (isl_ast_node_get_type(Body)) {
   case isl_ast_node_block: {
     isl_ast_node_list *List = isl_ast_node_block_get_children(Body);
     for (int i = 0; i < isl_ast_node_list_n_ast_node(List); ++i)
-      collectLoopAnnotations(isl_ast_node_list_get_ast_node(List, i));
+      collectLoopAnnotations(isl_ast_node_list_get_ast_node(List, i), Payload);
     isl_ast_node_list_free(List);
     break;
   }
   case isl_ast_node_if: {
-    collectLoopAnnotations(isl_ast_node_if_get_then(Body));
+    collectLoopAnnotations(isl_ast_node_if_get_then(Body), Payload);
     if (isl_ast_node_if_has_else(Body))
-      collectLoopAnnotations(isl_ast_node_if_get_else(Body));
+      collectLoopAnnotations(isl_ast_node_if_get_else(Body), Payload);
     break;
   }
   case isl_ast_node_user: {
     isl_ast_expr *Expr, *UserExpr;
     isl_pw_multi_aff *ScatPMA;
     isl_pw_aff *ScatPA;
     isl_id *Id;
 
case isl_ast_node_user: {
     // Get rid of the constrains caused by the domain.
     ScatPA = isl_pw_aff_gist(ScatPA, Stmt->getDomain());
 
     // For each input dimension we check if it is actually used in the innermost
     // (now only) dimension. If so we can get the corresponding Loop and check
     // for annotations.
     for (unsigned u = 0, e = Stmt->getNumIterators(); u != e; u++)
       if (isl_pw_aff_involves_dims(ScatPA, isl_dim_in, u, 1))
-        if (const Loop *L = Stmt->getLoopForDimension(u))
-          /* TODO Actually check and extract annotations */ L->getLoopID();
+        if (const Loop *L = Stmt->getLoopForDimension(u)) {
+          Payload->Loops.insert(L);
+        }
 
     isl_pw_multi_aff_free(ScatPMA);
     isl_ast_expr_free(UserExpr);
     isl_ast_expr_free(Expr);
     isl_pw_aff_free(ScatPA);
     isl_id_free(Id);
     break;
   }
else
       Payload->IsInnermostParallel =
           astScheduleDimIsParallel(Build, BuildInfo->Deps, Payload);
   } else if (Payload->IsOutermostParallel)
     BuildInfo->InParallelFor = false;
 
   // For innermost loops collect all loop annotations from the orignal loop(s)
   // involved in this new innermost dimension.
   if (Payload->IsInnermost)
-    collectLoopAnnotations(isl_ast_node_for_get_body(Node));
+    collectLoopAnnotations(isl_ast_node_for_get_body(Node), Payload);
 
   isl_id_free(Id);
   return Node;
 }
 
 static __isl_give isl_ast_node *AtEachDomain(__isl_take isl_ast_node *Node,
                                              __isl_keep isl_ast_build *Build,
                                              void *User) {

 }
 
 IslAstInfo::MemoryAccessSet *
 IslAstInfo::getBrokenReductions(__isl_keep isl_ast_node *Node) {
   IslAstUserPayload *Payload = getNodePayload(Node);
   return Payload ? &Payload->BrokenReductions : nullptr;
 }
 
+IslAstInfo::LoopSet * IslAstInfo::getLoops(__isl_keep isl_ast_node *Node) {
+  IslAstUserPayload *Payload = getNodePayload(Node);
+  return Payload ? &Payload->Loops : nullptr;
+}
+
+const Loop * IslAstInfo::getLoop(__isl_keep isl_ast_node *Node) {
+  if (LoopSet *Loops = getLoops(Node))
+    if (Loops->size() == 1)
+      return *Loops->begin();
+  return nullptr;
+}
+
 isl_ast_build *IslAstInfo::getBuild(__isl_keep isl_ast_node *Node) {
   IslAstUserPayload *Payload = getNodePayload(Node);
   return Payload ? Payload->Build : nullptr;
 }
 
 void IslAstInfo::printScop(raw_ostream &OS) const {
   isl_ast_print_options *Options;
   isl_ast_node *RootNode = getAst();

lib/CodeGen/IslCodeGeneration.cpp

 // The Scop describes the high level memory behaviour of a control flow region.
 // Transformation passes can update the schedule (execution order) of statements
 // in the Scop. ISL is used to generate an abstract syntax tree that reflects
 // the updated execution order. This clast is used to create new LLVM-IR that is
 // computationally equivalent to the original control flow region, but executes
 // its code in the new execution order defined by the changed scattering.
 //
 //===----------------------------------------------------------------------===//
+#define DEBUG_TYPE "polly-codegen-isl"
 #include "polly/Config/config.h"
 #include "polly/CodeGen/IslExprBuilder.h"
 #include "polly/CodeGen/BlockGenerators.h"
 #include "polly/CodeGen/CodeGeneration.h"
 #include "polly/CodeGen/IslAst.h"
 #include "polly/CodeGen/IslExprBuilder.h"
 #include "polly/CodeGen/LoopGenerators.h"
 #include "polly/CodeGen/Utils.h"

 #include "isl/ast_build.h"
 #include "isl/set.h"
 #include "isl/map.h"
 #include "isl/aff.h"
 
 using namespace polly;
 using namespace llvm;
 
-#define DEBUG_TYPE "polly-codegen-isl"
 
 class IslNodeBuilder {
 public:
   IslNodeBuilder(PollyIRBuilder &Builder, LoopAnnotator &Annotator, Pass *P,
                  LoopInfo &LI, ScalarEvolution &SE, DominatorTree &DT)
       : Builder(Builder), Annotator(Annotator), ExprBuilder(Builder, IDToValue),
         P(P), LI(LI), SE(SE), DT(DT) {}
 
bool runOnScop(Scop &S) {
 
     assert(!S.getRegion().isTopLevelRegion() &&
            "Top level regions are not supported");
 
     simplifyRegion(&S, this);
 
     BasicBlock *StartBlock = executeScopConditionally(S, this);
     isl_ast_node *Ast = AstInfo.getAst();
+    DEBUG(dbgs() << "[IslCodeGen] "; AstInfo.printScop(dbgs()));
     LoopAnnotator Annotator;
     PollyIRBuilder Builder(StartBlock->getContext(), llvm::ConstantFolder(),
                            polly::IRInserter(Annotator));
     Builder.SetInsertPoint(StartBlock->begin());
 
     IslNodeBuilder NodeBuilder(Builder, Annotator, this, LI, SE, DT);
 
     Builder.SetInsertPoint(StartBlock->getSinglePredecessor()->begin());