Allow non-affine control flow -- Code Generation

[jdoerfert]

Summary

This is the code generation for region statements that are created
when non-affine control flow was present in the input. A new
generator, similar to the block or vector generator, for regions is
used to traverse and copy the region statement and to adjust the
control flow inside the new region in the end.

Depends on D7846

[jdoerfert]

The chain for non-affine subregions without loops passes all lnt tests and increses the number of SCoPs from ~1500 to ~1800

[sebpop]

LGTM.
Let's wait for Tobi's comments before committing the change.

lib/CodeGen/BlockGenerators.cpp

668

Update the comment to "only region statements can be copied by the region generator"

[sebpop]

Also, have you tried the non affine scop generator on llvm's test-suite?

[jdoerfert]

In D7847#128928, @jdoerfert wrote:

The chain for non-affine subregions without loops passes all lnt tests and increses the number of SCoPs from ~1500 to ~1800

When I talk about lnt I mean I run polly on the llvm-test-suite.

[grosser]

LGTM.

[jdoerfert]

commited in r230340.

include/polly/CodeGen/BlockGenerators.h

public:
   BlockGenerator(PollyIRBuilder &Builder, LoopInfo &LI, ScalarEvolution &SE,
                  DominatorTree &DT, IslExprBuilder *ExprBuilder = nullptr);
 
   /// @brief Copy the basic block.
   ///
   /// This copies the entire basic block and updates references to old values
   /// with references to new values, as defined by GlobalMap.
   ///
-  /// @param Stmt      The statement to code generate.
+  /// @param Stmt      The block statement to code generate.
   /// @param GlobalMap A mapping from old values to their new values
   ///                  (for values recalculated in the new ScoP, but not
   ///                  within this basic block).
   /// @param LTS       A map from old loops to new induction variables as SCEVs.
-  void copyBB(ScopStmt &Stmt, ValueMapT &GlobalMap, LoopToScevMapT &LTS);
+  void copyStmt(ScopStmt &Stmt, ValueMapT &GlobalMap, LoopToScevMapT &LTS);
 
 protected:
   PollyIRBuilder &Builder;
   LoopInfo &LI;
   ScalarEvolution &SE;
   IslExprBuilder *ExprBuilder;
 
   /// @brief The dominator tree of this function.
   DominatorTree &DT;
 
+  /// @brief Copy the given basic block.
+  ///
+  /// @param Stmt      The statement to code generate.
+  /// @param BB        The basic block to code generate.
+  /// @param BBMap     A mapping from old values to their new values in this
+  /// block.
+  /// @param GlobalMap A mapping from old values to their new values
+  ///                  (for values recalculated in the new ScoP, but not
+  ///                  within this basic block).
+  /// @param LTS       A map from old loops to new induction variables as SCEVs.
+  ///
+  /// @returns The copy of the basic block.
+  BasicBlock *copyBB(ScopStmt &Stmt, BasicBlock *BB, ValueMapT &BBMap,
+                     ValueMapT &GlobalMap, LoopToScevMapT &LTS);
+
   /// @brief Get the new version of a value.
   ///
   /// Given an old value, we first check if a new version of this value is
   /// available in the BBMap or GlobalMap. In case it is not and the value can
   /// be recomputed using SCEV, we do so. If we can not recompute a value
   /// using SCEV, but we understand that the value is constant within the scop,
   /// we return the old value.  If the value can still not be derived, this
   /// function will assert.
public:
   /// @param Schedule   A map from the statement to a schedule where the
   ///                   innermost dimension is the dimension of the innermost
   ///                   loop containing the statemenet.
   static void generate(BlockGenerator &BlockGen, ScopStmt &Stmt,
                        VectorValueMapT &GlobalMaps,
                        std::vector<LoopToScevMapT> &VLTS,
                        __isl_keep isl_map *Schedule) {
     VectorBlockGenerator Generator(BlockGen, GlobalMaps, VLTS, Schedule);
-    Generator.copyBB(Stmt);
+    Generator.copyStmt(Stmt);
   }
 
 private:
   // This is a vector of global value maps.  The first map is used for the first
   // vector lane, ...
   // Each map, contains information about Instructions in the old ScoP, which
   // are recalculated in the new SCoP. When copying the basic block, we replace
   // all referenes to the old instructions with their recalculated values.
private:
   bool extractScalarValues(const Instruction *Inst, ValueMapT &VectorMap,
                            VectorValueMapT &ScalarMaps);
 
   bool hasVectorOperands(const Instruction *Inst, ValueMapT &VectorMap);
 
   void copyInstruction(ScopStmt &Stmt, const Instruction *Inst,
                        ValueMapT &VectorMap, VectorValueMapT &ScalarMaps);
 
-  void copyBB(ScopStmt &Stmt);
+  void copyStmt(ScopStmt &Stmt);
+};
+
+/// @brief Generator for new versions of polyhedral region statements.
+class RegionGenerator : public BlockGenerator {
+public:
+  /// @brief Create a generator for regions.
+  ///
+  /// @param BlockGen A generator for basic blocks.
+  RegionGenerator(BlockGenerator &BlockGen) : BlockGenerator(BlockGen) {}
+
+  /// @brief Copy the region statement @p Stmt.
+  ///
+  /// This copies the entire region represented by @p Stmt and updates
+  /// references to old values with references to new values, as defined by
+  /// GlobalMap.
+  ///
+  /// @param Stmt      The statement to code generate.
+  /// @param GlobalMap A mapping from old values to their new values
+  ///                  (for values recalculated in the new ScoP, but not
+  ///                  within this basic block).
+  /// @param LTS       A map from old loops to new induction variables as SCEVs.
+  void copyStmt(ScopStmt &Stmt, ValueMapT &GlobalMap, LoopToScevMapT &LTS);
 };
 }
 #endif

lib/CodeGen/BlockGenerators.cpp

default:
       // Other intrinsics are copied.
       break;
     }
   }
 
   copyInstScalar(Stmt, Inst, BBMap, GlobalMap, LTS);
 }
 
-void BlockGenerator::copyBB(ScopStmt &Stmt, ValueMapT &GlobalMap,
+void BlockGenerator::copyStmt(ScopStmt &Stmt, ValueMapT &GlobalMap,
                               LoopToScevMapT &LTS) {
+  assert(Stmt.isBlockStmt() &&
+         "Only block statements can be copied by the block generator");
+
+  ValueMapT BBMap;
+
   BasicBlock *BB = Stmt.getBasicBlock();
+  copyBB(Stmt, BB, BBMap, GlobalMap, LTS);
+}
+
+BasicBlock *BlockGenerator::copyBB(ScopStmt &Stmt, BasicBlock *BB,
+                                   ValueMapT &BBMap, ValueMapT &GlobalMap,
+                                   LoopToScevMapT &LTS) {
   BasicBlock *CopyBB =
       SplitBlock(Builder.GetInsertBlock(), Builder.GetInsertPoint(), &DT, &LI);
   CopyBB->setName("polly.stmt." + BB->getName());
   Builder.SetInsertPoint(CopyBB->begin());
 
-  ValueMapT BBMap;
-
   for (Instruction &Inst : *BB)
     copyInstruction(Stmt, &Inst, BBMap, GlobalMap, LTS);
+
+  return CopyBB;
 }
 
 VectorBlockGenerator::VectorBlockGenerator(BlockGenerator &BlockGen,
                                            VectorValueMapT &GlobalMaps,
                                            std::vector<LoopToScevMapT> &VLTS,
                                            isl_map *Schedule)
     : BlockGenerator(BlockGen), GlobalMaps(GlobalMaps), VLTS(VLTS),
       Schedule(Schedule) {
if (hasVectorOperands(Inst, VectorMap)) {
 
     // Falltrough: We generate scalar instructions, if we don't know how to
     // generate vector code.
   }
 
   copyInstScalarized(Stmt, Inst, VectorMap, ScalarMaps);
 }
 
-void VectorBlockGenerator::copyBB(ScopStmt &Stmt) {
+void VectorBlockGenerator::copyStmt(ScopStmt &Stmt) {
+  assert(Stmt.isBlockStmt() && "TODO: Only block statements can be copied by "
+                               "the vector block generator");
+
   BasicBlock *BB = Stmt.getBasicBlock();
   BasicBlock *CopyBB =
       SplitBlock(Builder.GetInsertBlock(), Builder.GetInsertPoint(), &DT, &LI);
   CopyBB->setName("polly.stmt." + BB->getName());
   Builder.SetInsertPoint(CopyBB->begin());
 
   // Create two maps that store the mapping from the original instructions of
   // the old basic block to their copies in the new basic block. Those maps
void VectorBlockGenerator::copyStmt(ScopStmt &Stmt) {
   // time. Or (if the values are calculated using scalar operations), it
   // appears once in every dimension of the scalarMap.
   VectorValueMapT ScalarBlockMap(getVectorWidth());
   ValueMapT VectorBlockMap;
 
   for (Instruction &Inst : *BB)
     copyInstruction(Stmt, &Inst, VectorBlockMap, ScalarBlockMap);
 }
+
+void RegionGenerator::copyStmt(ScopStmt &Stmt, ValueMapT &GlobalMap,
+                               LoopToScevMapT &LTS) {
+  assert(Stmt.isRegionStmt() &&
+         "Only block statements can be copied by the block generator");
+
+  // The region represented by the statement.
+  Region *R = Stmt.getRegion();
+
+  // The "BBMap" for the whole region.
+  ValueMapT RegionMap;
+
+  // Iterate over all blocks in the region in a breadth-first search.
+  std::deque<BasicBlock *> Blocks;
+  SmallPtrSet<BasicBlock *, 8> SeenBlocks;
+  Blocks.push_back(R->getEntry());
+  SeenBlocks.insert(R->getEntry());
+
+  while (!Blocks.empty()) {
+    BasicBlock *BB = Blocks.front();
+    Blocks.pop_front();
+
+    // Copy the block with the BlockGenerator.
+    BasicBlock *BBCopy = copyBB(Stmt, BB, RegionMap, GlobalMap, LTS);
+
+    // And continue with new successors inside the region.
+    for (auto SI = succ_begin(BB), SE = succ_end(BB); SI != SE; SI++)
+      if (R->contains(*SI) && SeenBlocks.insert(*SI).second)
+        Blocks.push_back(*SI);
+
+    // In order to remap PHI nodes we store also basic block mappings.
+    RegionMap[BB] = BBCopy;
+  }
+
+  // Now create a new dedicated region exit block and add it to the region map.
+  BasicBlock *RegionExit =
+      SplitBlock(Builder.GetInsertBlock(), Builder.GetInsertPoint(), &DT, &LI);
+  RegionExit->setName("polly.stmt." + R->getExit()->getName() + ".pre");
+  RegionMap[R->getExit()] = RegionExit;
+
+  // As the block generator doesn't handle control flow we need to add the
+  // region control flow by hand after all blocks have been copied.
+  for (BasicBlock *BB : SeenBlocks) {
+
+    BranchInst *BI = cast<BranchInst>(BB->getTerminator());
+
+    BasicBlock *BBCopy = cast<BasicBlock>(RegionMap[BB]);
+    Instruction *BICopy = BBCopy->getTerminator();
+
+    Builder.SetInsertPoint(BBCopy);
+    copyInstScalar(Stmt, BI, RegionMap, GlobalMap, LTS);
+    BICopy->eraseFromParent();
+  }
+
+  // Reset the old insert point for the build.
+  Builder.SetInsertPoint(RegionExit->begin());
+}

lib/CodeGen/IslCodeGeneration.cpp

 class IslNodeBuilder {
 public:
   IslNodeBuilder(PollyIRBuilder &Builder, ScopAnnotator &Annotator, Pass *P,
                  const DataLayout &DL, LoopInfo &LI, ScalarEvolution &SE,
                  DominatorTree &DT, Scop &S)
       : S(S), Builder(Builder), Annotator(Annotator),
         Rewriter(new SCEVExpander(SE, "polly")),
         ExprBuilder(Builder, IDToValue, *Rewriter),
-        BlockGen(Builder, LI, SE, DT, &ExprBuilder), P(P), DL(DL), LI(LI),
-        SE(SE), DT(DT) {}
+        BlockGen(Builder, LI, SE, DT, &ExprBuilder), RegionGen(BlockGen), P(P),
+        DL(DL), LI(LI), SE(SE), DT(DT) {}
 
   ~IslNodeBuilder() { delete Rewriter; }
 
   void addParameters(__isl_take isl_set *Context);
   void create(__isl_take isl_ast_node *Node);
   IslExprBuilder &getExprBuilder() { return ExprBuilder; }
 
 private:
   Scop &S;
   PollyIRBuilder &Builder;
   ScopAnnotator &Annotator;
 
   /// @brief A SCEVExpander to create llvm values from SCEVs.
   SCEVExpander *Rewriter;
 
   IslExprBuilder ExprBuilder;
   BlockGenerator BlockGen;
+
+  /// @brief Generator for region statements.
+  RegionGenerator RegionGen;
+
   Pass *P;
   const DataLayout &DL;
   LoopInfo &LI;
   ScalarEvolution &SE;
   DominatorTree &DT;
 
   /// @brief The current iteration of out-of-scop loops
   ///
void IslNodeBuilder::createUser(__isl_take isl_ast_node *User) {
   isl_ast_expr_free(StmtExpr);
 
   LTS.insert(OutsideLoopIterations.begin(), OutsideLoopIterations.end());
 
   Stmt = (ScopStmt *)isl_id_get_user(Id);
   Stmt->setAstBuild(IslAstInfo::getBuild(User));
 
   createSubstitutions(Expr, Stmt, VMap, LTS);
-  BlockGen.copyBB(*Stmt, VMap, LTS);
+  if (Stmt->isBlockStmt())
+    BlockGen.copyStmt(*Stmt, VMap, LTS);
+  else
+    RegionGen.copyStmt(*Stmt, VMap, LTS);
 
   isl_ast_node_free(User);
   isl_id_free(Id);
 }
 
 void IslNodeBuilder::createBlock(__isl_take isl_ast_node *Block) {
   isl_ast_node_list *List = isl_ast_node_block_get_children(Block);
 

test/Isl/CodeGen/non_affine_float_compare.ll

+; RUN: opt %loadPolly -polly-codegen-isl -polly-no-early-exit -polly-allow-nonaffine-branches -S < %s | FileCheck %s
+;
+;    void f(float *A) {
+;      for (int i = 0; i < 1024; i++)
+;        if (A[i] == A[i - 1])
+;          A[i]++;
+;       A[i]++;
+;    }
+;
+;
+; CHECK: polly.stmt.bb2:
+; CHECK:   %scevgep[[R0:[0-9]*]] = getelementptr float* %A, i64 %polly.indvar
+; CHECK:   %tmp3_p_scalar_ = load float* %scevgep[[R0]], align 4, !alias.scope !0, !noalias !2
+; CHECK:   %scevgep[[R2:[0-9]*]] = getelementptr float* %scevgep{{[0-9]*}}, i64 %polly.indvar
+; CHECK:   %tmp6_p_scalar_ = load float* %scevgep[[R2]], align 4, !alias.scope !0, !noalias !2
+; CHECK:   %p_tmp7 = fcmp oeq float %tmp3_p_scalar_, %tmp6_p_scalar_
+; CHECK:   br i1 %p_tmp7, label %polly.stmt.bb8, label %polly.stmt.bb12.pre
+
+; CHECK: polly.stmt.bb8:
+; CHECK:   %scevgep[[R3:[0-9]*]] = getelementptr float* %A, i64 %polly.indvar
+; CHECK:   %tmp10_p_scalar_ = load float* %scevgep[[R3]], align 4, !alias.scope !0, !noalias !2
+; CHECK:   %p_tmp11 = fadd float %tmp10_p_scalar_, 1.000000e+00
+; CHECK:   store float %p_tmp11, float* %scevgep[[R3]], align 4, !alias.scope !0, !noalias !2
+; CHECK:   br label %polly.stmt.bb12.pre
+
+; CHECK: polly.stmt.bb12.pre:
+; CHECK:   br label %polly.stmt.bb12
+
+; CHECK: polly.stmt.bb12:
+; CHECK:   %scevgep[[R4:[0-9]*]] = getelementptr float* %A, i64 %polly.indvar
+; CHECK:   %tmp10b_p_scalar_ = load float* %scevgep[[R4]], align 4, !alias.scope !0, !noalias !2
+; CHECK:   %p_tmp11b = fadd float %tmp10b_p_scalar_, 1.000000e+00
+; CHECK:   store float %p_tmp11b, float* %scevgep[[R4]], align 4, !alias.scope !0, !noalias !2
+; CHECK:   %polly.indvar_next = add nsw i64 %polly.indvar, 1
+; CHECK:   %polly.loop_cond = icmp sle i64 %polly.indvar, 1022
+; CHECK:   br i1 %polly.loop_cond, label %polly.loop_header, label %polly.loop_exit
+
+target datalayout = "e-m:e-i64:64-f80:128-n8:16:32:64-S128"
+
+define void @f(float* %A) {
+bb:
+  br label %bb1
+
+bb1:                                              ; preds = %bb13, %bb
+  %indvars.iv = phi i64 [ %indvars.iv.next, %bb13 ], [ 0, %bb ]
+  %exitcond = icmp ne i64 %indvars.iv, 1024
+  br i1 %exitcond, label %bb2, label %bb14
+
+bb2:                                              ; preds = %bb1
+  %tmp = getelementptr inbounds float* %A, i64 %indvars.iv
+  %tmp3 = load float* %tmp, align 4
+  %tmp4 = add nsw i64 %indvars.iv, -1
+  %tmp5 = getelementptr inbounds float* %A, i64 %tmp4
+  %tmp6 = load float* %tmp5, align 4
+  %tmp7 = fcmp oeq float %tmp3, %tmp6
+  br i1 %tmp7, label %bb8, label %bb12
+
+bb8:                                              ; preds = %bb2
+  %tmp9 = getelementptr inbounds float* %A, i64 %indvars.iv
+  %tmp10 = load float* %tmp9, align 4
+  %tmp11 = fadd float %tmp10, 1.000000e+00
+  store float %tmp11, float* %tmp9, align 4
+  br label %bb12
+
+bb12:                                             ; preds = %bb8, %bb2
+  %tmp9b = getelementptr inbounds float* %A, i64 %indvars.iv
+  %tmp10b = load float* %tmp9b, align 4
+  %tmp11b = fadd float %tmp10b, 1.000000e+00
+  store float %tmp11b, float* %tmp9b, align 4
+  br label %bb13
+
+bb13:                                             ; preds = %bb12
+  %indvars.iv.next = add nuw nsw i64 %indvars.iv, 1
+  br label %bb1
+
+bb14:                                             ; preds = %bb1
+  ret void
+}