[Polly] [BlockGen] Support partial writes in regions

[grosser]

Summary

The RegionGenerator traditionally kept a BlockMap that mapped from original
basic blocks to newly generated basic blocks. With the introduction of partial
writes such a 1:1 mapping is not possible any more, as a single basic block
can be code generated into multiple basic blocks. Hence, depending on the use
case we need to either use the first basic block or the last basic block.

This is intended to address the last four cases of incorrect code generation
in our AOSP buildbot and hopefully should turn it green.

[Meinersbur]

We should think a bit more about when to use StartBlock and when to use EndBlock. Test cases to find the right choices could be helpful.

lib/CodeGen/BlockGenerators.cpp

1272

Should this be EndBlockMap? Because it is the end of the previous BB that dominates BBCopy.

I changed it in the code and nothing changes. Is there a test case such that we could test this?

1354

I'd expect only the end block ever by used (to repair dominance or PHI incoming blocks), but I also think it is safer to just apply it.

1439

Call it RegionMapEnd?

Do we even need it? I removed it and tests still pass. Having only ValueMap per BB is easier to maintain (otherwise why shouldn't the *.partial blocks have one as well.). The call to copyInstScalar below will only update one of them: RegionMap, i.e. the EndBlock's value map will necessarily out of date. We could also keep the EndBlock's value map, because only in that one all values would be available.

1467–1469

Should these be EndBlockMap? Since we look for the predecessor of an original block, which is the end of the copy. Test case?

test/Isl/CodeGen/partial_write_in_region.ll

13

Can you remove the references do metadata (!3 etc) that we are not checking for? I observed metadata numbering sometimes being unstable.

[grosser]

Address Michael's comments.

[grosser]

Hi Michael,

thank you for the comments. I believe this is now (close to) commit ready, but I would appreciate a final OK. Hope with this change, we can turn the buildbots green!

lib/CodeGen/BlockGenerators.cpp

1272

I expanded the test case to cover this case. In fact, we had to lookup the dominator in EndBlockMap and had to return the lookup from StartBlockMap to make sure the correct definitions are provided.

1354

OK, I leave it then.

1439

No, we do not need it. I dropped it. The only place where we would have needed is was in add addOperandToPHI, which I now changed to lookup the necessary values in RegionMaps[BBCopyStart].

1467–1469

It should. Due to https://bugs.llvm.org//show_bug.cgi?id=33265, I had some issues to create a working test case, but now added one with
test/Isl/CodeGen/partial_write_in_region_with_loop.ll.

[Meinersbur]

lib/CodeGen/BlockGenerators.cpp

1518

If we try consistly have the BBMaps of the StartBlock store the correct value maps: This looks up the EndBlock, which will probably be always empty.

[Meinersbur]

Otherwise, LGTM.

include/polly/CodeGen/BlockGenerators.h

public:
   /// GlobalMap.
   ///
   /// @param Stmt      The statement to code generate.
   /// @param LTS       A map from old loops to new induction variables as SCEVs.
   void copyStmt(ScopStmt &Stmt, LoopToScevMapT &LTS,
                 __isl_keep isl_id_to_ast_expr *IdToAstExp);
 
 private:
-  /// A map from old to new blocks in the region.
-  DenseMap<BasicBlock *, BasicBlock *> BlockMap;
+  /// A map from old to the first new block in the region, that was created to
+  /// model the old basic block.
+  DenseMap<BasicBlock *, BasicBlock *> StartBlockMap;
+
+  /// A map from old to the last new block in the region, that was created to
+  /// model the old basic block.
+  DenseMap<BasicBlock *, BasicBlock *> EndBlockMap;
 
   /// The "BBMaps" for the whole region (one for each block).
   DenseMap<BasicBlock *, ValueMapT> RegionMaps;
 
   /// Mapping to remember PHI nodes that still need incoming values.
   using PHINodePairTy = std::pair<PHINode *, PHINode *>;
   DenseMap<BasicBlock *, SmallVector<PHINodePairTy, 4>> IncompletePHINodeMap;
 

lib/CodeGen/BlockGenerators.cpp

void VectorBlockGenerator::copyStmt(
 
   verifyNoScalarStores(Stmt);
 }
 
 BasicBlock *RegionGenerator::repairDominance(BasicBlock *BB,
                                              BasicBlock *BBCopy) {
 
   BasicBlock *BBIDom = DT.getNode(BB)->getIDom()->getBlock();
-  BasicBlock *BBCopyIDom = BlockMap.lookup(BBIDom);
+  BasicBlock *BBCopyIDom = StartBlockMap.lookup(BBIDom);
 
   if (BBCopyIDom)
     DT.changeImmediateDominator(BBCopy, BBCopyIDom);
 
   return BBCopyIDom;
 }
 
 // This is to determine whether an llvm::Value (defined in @p BB) is usable when

 }
 
 void RegionGenerator::copyStmt(ScopStmt &Stmt, LoopToScevMapT &LTS,
                                isl_id_to_ast_expr *IdToAstExp) {
   assert(Stmt.isRegionStmt() &&
          "Only region statements can be copied by the region generator");
 
   // Forget all old mappings.
-  BlockMap.clear();
+  StartBlockMap.clear();
+  EndBlockMap.clear();
   RegionMaps.clear();
   IncompletePHINodeMap.clear();
 
   // Collection of all values related to this subregion.
   ValueMapT ValueMap;
 
   // The region represented by the statement.
   Region *R = Stmt.getRegion();
 
   // Create a dedicated entry for the region where we can reload all demoted
   // inputs.
   BasicBlock *EntryBB = R->getEntry();
   BasicBlock *EntryBBCopy = SplitBlock(Builder.GetInsertBlock(),
                                        &*Builder.GetInsertPoint(), &DT, &LI);
   EntryBBCopy->setName("polly.stmt." + EntryBB->getName() + ".entry");
   Builder.SetInsertPoint(&EntryBBCopy->front());
 
   ValueMapT &EntryBBMap = RegionMaps[EntryBBCopy];
   generateScalarLoads(Stmt, LTS, EntryBBMap, IdToAstExp);
 
   for (auto PI = pred_begin(EntryBB), PE = pred_end(EntryBB); PI != PE; ++PI)
-    if (!R->contains(*PI))
-      BlockMap[*PI] = EntryBBCopy;
+    if (!R->contains(*PI)) {
+      StartBlockMap[*PI] = EntryBBCopy;
+      EndBlockMap[*PI] = EntryBBCopy;
+    }
 
   // Iterate over all blocks in the region in a breadth-first search.
   std::deque<BasicBlock *> Blocks;
   SmallSetVector<BasicBlock *, 8> SeenBlocks;
   Blocks.push_back(EntryBB);
   SeenBlocks.insert(EntryBB);
 
   while (!Blocks.empty()) {
while (!Blocks.empty()) {
     auto Inserted = RegionMaps.insert(std::make_pair(BBCopy, *InitBBMap));
     ValueMapT &RegionMap = Inserted.first->second;
 
     // Copy the block with the BlockGenerator.
     Builder.SetInsertPoint(&BBCopy->front());
     copyBB(Stmt, BB, BBCopy, RegionMap, LTS, IdToAstExp);
 
     // In order to remap PHI nodes we store also basic block mappings.
-    BlockMap[BB] = BBCopy;
+    StartBlockMap[BB] = BBCopy;
+    EndBlockMap[BB] = Builder.GetInsertBlock();
 
     // Add values to incomplete PHI nodes waiting for this block to be copied.
     for (const PHINodePairTy &PHINodePair : IncompletePHINodeMap[BB])
       addOperandToPHI(Stmt, PHINodePair.first, PHINodePair.second, BB, LTS);
     IncompletePHINodeMap[BB].clear();
 
     // 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))
         Blocks.push_back(*SI);
 
     // Remember value in case it is visible after this subregion.
     if (isDominatingSubregionExit(DT, R, BB))
       ValueMap.insert(RegionMap.begin(), RegionMap.end());
   }
 
   // Now create a new dedicated region exit block and add it to the region map.
   BasicBlock *ExitBBCopy = SplitBlock(Builder.GetInsertBlock(),
                                       &*Builder.GetInsertPoint(), &DT, &LI);
   ExitBBCopy->setName("polly.stmt." + R->getExit()->getName() + ".exit");
-  BlockMap[R->getExit()] = ExitBBCopy;
+  StartBlockMap[R->getExit()] = ExitBBCopy;
+  EndBlockMap[R->getExit()] = ExitBBCopy;
 
-  BasicBlock *ExitDomBBCopy = BlockMap.lookup(findExitDominator(DT, R));
+  BasicBlock *ExitDomBBCopy = EndBlockMap.lookup(findExitDominator(DT, R));
   assert(ExitDomBBCopy &&
          "Common exit dominator must be within region; at least the entry node "
          "must match");
   DT.changeImmediateDominator(ExitBBCopy, ExitDomBBCopy);
 
   // 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) {
 
-    BasicBlock *BBCopy = BlockMap[BB];
+    BasicBlock *BBCopyStart = StartBlockMap[BB];
+    BasicBlock *BBCopyEnd = EndBlockMap[BB];
     TerminatorInst *TI = BB->getTerminator();
     if (isa<UnreachableInst>(TI)) {
-      while (!BBCopy->empty())
-        BBCopy->begin()->eraseFromParent();
-      new UnreachableInst(BBCopy->getContext(), BBCopy);
+      while (!BBCopyEnd->empty())
+        BBCopyEnd->begin()->eraseFromParent();
+      new UnreachableInst(BBCopyEnd->getContext(), BBCopyEnd);
       continue;
     }
 
-    Instruction *BICopy = BBCopy->getTerminator();
+    Instruction *BICopy = BBCopyEnd->getTerminator();
 
-    ValueMapT &RegionMap = RegionMaps[BBCopy];
-    RegionMap.insert(BlockMap.begin(), BlockMap.end());
+    ValueMapT &RegionMap = RegionMaps[BBCopyStart];
+    RegionMap.insert(StartBlockMap.begin(), StartBlockMap.end());
+    ValueMapT &RegionMap2 = RegionMaps[BBCopyEnd];
+    RegionMap2.insert(StartBlockMap.begin(), StartBlockMap.end());
 
     Builder.SetInsertPoint(BICopy);
     copyInstScalar(Stmt, TI, RegionMap, LTS);
     BICopy->eraseFromParent();
   }
 
   // Add counting PHI nodes to all loops in the region that can be used as
   // replacement for SCEVs refering to the old loop.
   for (BasicBlock *BB : SeenBlocks) {
     Loop *L = LI.getLoopFor(BB);
     if (L == nullptr || L->getHeader() != BB || !R->contains(L))
       continue;
 
-    BasicBlock *BBCopy = BlockMap[BB];
+    BasicBlock *BBCopy = StartBlockMap[BB];
     Value *NullVal = Builder.getInt32(0);
     PHINode *LoopPHI =
         PHINode::Create(Builder.getInt32Ty(), 2, "polly.subregion.iv");
     Instruction *LoopPHIInc = BinaryOperator::CreateAdd(
         LoopPHI, Builder.getInt32(1), "polly.subregion.iv.inc");
     LoopPHI->insertBefore(&BBCopy->front());
     LoopPHIInc->insertBefore(BBCopy->getTerminator());
 
     for (auto *PredBB : make_range(pred_begin(BB), pred_end(BB))) {
       if (!R->contains(PredBB))
         continue;
       if (L->contains(PredBB))
-        LoopPHI->addIncoming(LoopPHIInc, BlockMap[PredBB]);
+        LoopPHI->addIncoming(LoopPHIInc, StartBlockMap[PredBB]);
       else
-        LoopPHI->addIncoming(NullVal, BlockMap[PredBB]);
+        LoopPHI->addIncoming(NullVal, StartBlockMap[PredBB]);
     }
 
     for (auto *PredBBCopy : make_range(pred_begin(BBCopy), pred_end(BBCopy)))
       if (LoopPHI->getBasicBlockIndex(PredBBCopy) < 0)
         LoopPHI->addIncoming(NullVal, PredBBCopy);
 
     LTS[L] = SE.getUnknown(LoopPHI);
   }
 
   // Continue generating code in the exit block.
   Builder.SetInsertPoint(&*ExitBBCopy->getFirstInsertionPt());
 
   // Write values visible to other statements.
   generateScalarStores(Stmt, LTS, ValueMap, IdToAstExp);
-  BlockMap.clear();
+  StartBlockMap.clear();
+  EndBlockMap.clear();
   RegionMaps.clear();
   IncompletePHINodeMap.clear();
 }
 
 PHINode *RegionGenerator::buildExitPHI(MemoryAccess *MA, LoopToScevMapT &LTS,
                                        ValueMapT &BBMap, Loop *L) {
   ScopStmt *Stmt = MA->getStatement();
   Region *SubR = Stmt->getRegion();
   auto Incoming = MA->getIncoming();
 
   PollyIRBuilder::InsertPointGuard IPGuard(Builder);
   PHINode *OrigPHI = cast<PHINode>(MA->getAccessInstruction());
   BasicBlock *NewSubregionExit = Builder.GetInsertBlock();
 
   // This can happen if the subregion is simplified after the ScopStmts
   // have been created; simplification happens as part of CodeGeneration.
   if (OrigPHI->getParent() != SubR->getExit()) {
     BasicBlock *FormerExit = SubR->getExitingBlock();
     if (FormerExit)
-      NewSubregionExit = BlockMap.lookup(FormerExit);
+      NewSubregionExit = StartBlockMap.lookup(FormerExit);
   }
 
   PHINode *NewPHI = PHINode::Create(OrigPHI->getType(), Incoming.size(),
                                     "polly." + OrigPHI->getName(),
                                     NewSubregionExit->getFirstNonPHI());
 
   // Add the incoming values to the PHI.
   for (auto &Pair : Incoming) {
     BasicBlock *OrigIncomingBlock = Pair.first;
-    BasicBlock *NewIncomingBlock = BlockMap.lookup(OrigIncomingBlock);
+    BasicBlock *NewIncomingBlock = EndBlockMap.lookup(OrigIncomingBlock);
     Builder.SetInsertPoint(NewIncomingBlock->getTerminator());
     assert(RegionMaps.count(NewIncomingBlock));
     ValueMapT *LocalBBMap = &RegionMaps[NewIncomingBlock];
 
     Value *OrigIncomingValue = Pair.second;
     Value *NewIncomingValue =
         getNewValue(*Stmt, OrigIncomingValue, *LocalBBMap, LTS, L);
     NewPHI->addIncoming(NewIncomingValue, NewIncomingBlock);
   }
 
   return NewPHI;
void RegionGenerator::generateScalarStores(
   }
 }
 
 void RegionGenerator::addOperandToPHI(ScopStmt &Stmt, PHINode *PHI,
                                       PHINode *PHICopy, BasicBlock *IncomingBB,
                                       LoopToScevMapT &LTS) {
   // If the incoming block was not yet copied mark this PHI as incomplete.
   // Once the block will be copied the incoming value will be added.
-  BasicBlock *BBCopy = BlockMap[IncomingBB];
+  BasicBlock *BBCopy = EndBlockMap[IncomingBB];
   if (!BBCopy) {
     assert(Stmt.contains(IncomingBB) &&
            "Bad incoming block for PHI in non-affine region");
     IncompletePHINodeMap[IncomingBB].push_back(std::make_pair(PHI, PHICopy));
     return;
   }
 
   assert(RegionMaps.count(BBCopy) && "Incoming PHI block did not have a BBMap");

test/Isl/CodeGen/partial_write_in_region.ll

+; RUN: opt %loadPolly -polly-import-jscop -polly-import-jscop-dir=%S \
+; RUN: -polly-import-jscop-postfix=transformed -polly-codegen \
+; RUN: -S < %s | FileCheck %s
+;
+;    void foo(long A[], float B[]) {
+;      for (long i = 0; i < 1024; i++)
+;        if (A[i])
+;          B[i]++;
+;    }
+
+; CHECK: polly.stmt.bb5:                                   ; preds = %polly.stmt.bb2
+; CHECK-NEXT:   %scevgep3 = getelementptr float, float* %B, i64 %polly.indvar
+; CHECK-NEXT:   %tmp7_p_scalar_ = load float, float* %scevgep3, align 4, !alias.scope !3, !noalias !4
+; CHECK-NEXT:   %p_tmp8 = fadd float %tmp7_p_scalar_, 1.000000e+00
+; CHECK-NEXT:   %8 = icmp sle i64 %polly.indvar, 9
+; CHECK-NEXT:   %polly.Stmt_bb2__TO__bb9_MayWrite2.cond = icmp ne i1 %8, false
+; CHECK-NEXT:   br i1 %polly.Stmt_bb2__TO__bb9_MayWrite2.cond, label %polly.stmt.bb5.Stmt_bb2__TO__bb9_MayWrite2.partial, label %polly.stmt.bb5.cont
+
+; CHECK: polly.stmt.bb5.Stmt_bb2__TO__bb9_MayWrite2.partial: ; preds = %polly.stmt.bb5
+; CHECK-NEXT:   %polly.access.B4 = getelementptr float, float* %B, i64 %polly.indvar
+; CHECK-NEXT:   store float %p_tmp8, float* %polly.access.B4, align 4, !alias.scope !3, !noalias !4
+; CHECK-NEXT:   br label %polly.stmt.bb5.cont
+
+; CHECK: polly.stmt.bb5.cont:                              ; preds = %polly.stmt.bb5, %polly.stmt.bb5.Stmt_bb2__TO__bb9_MayWrite2.partial
+; CHECK-NEXT:   br label %polly.stmt.bb9.exit
+
+target datalayout = "e-m:e-i64:64-f80:128-n8:16:32:64-S128"
+
+define void @partial_write_in_region(i64* %A, float* %B) {
+bb:
+  br label %bb1
+
+bb1:                                              ; preds = %bb10, %bb
+  %i.0 = phi i64 [ 0, %bb ], [ %tmp11, %bb10 ]
+  %exitcond = icmp ne i64 %i.0, 1024
+  br i1 %exitcond, label %bb2, label %bb12
+
+bb2:                                              ; preds = %bb1
+  %tmp = getelementptr inbounds i64, i64* %A, i64 %i.0
+  %tmp3 = load i64, i64* %tmp, align 8
+  %tmp4 = icmp eq i64 %tmp3, 0
+  br i1 %tmp4, label %bb9, label %bb5
+
+bb5:                                              ; preds = %bb2
+  %tmp6 = getelementptr inbounds float, float* %B, i64 %i.0
+  %tmp7 = load float, float* %tmp6, align 4
+  %tmp8 = fadd float %tmp7, 1.000000e+00
+  store float %tmp8, float* %tmp6, align 4
+  br label %bb9
+
+bb9:                                              ; preds = %bb2, %bb5
+  br label %bb10
+
+bb10:                                             ; preds = %bb9
+  %tmp11 = add nuw nsw i64 %i.0, 1
+  br label %bb1
+
+bb12:                                             ; preds = %bb1
+  ret void
+}

test/Isl/CodeGen/partial_write_in_region___%bb1---%bb12.jscop

+{
+   "arrays" : [
+      {
+         "name" : "MemRef_A",
+         "sizes" : [ "*" ],
+         "type" : "i64"
+      },
+      {
+         "name" : "MemRef_B",
+         "sizes" : [ "*" ],
+         "type" : "float"
+      }
+   ],
+   "context" : "{  :  }",
+   "name" : "%bb1---%bb12",
+   "statements" : [
+      {
+         "accesses" : [
+            {
+               "kind" : "read",
+               "relation" : "{ Stmt_bb2__TO__bb9[i0] -> MemRef_A[i0] }"
+            },
+            {
+               "kind" : "read",
+               "relation" : "{ Stmt_bb2__TO__bb9[i0] -> MemRef_B[i0] }"
+            },
+            {
+               "kind" : "write",
+               "relation" : "{ Stmt_bb2__TO__bb9[i0] -> MemRef_B[i0] }"
+            }
+         ],
+         "domain" : "{ Stmt_bb2__TO__bb9[i0] : 0 <= i0 <= 1023 }",
+         "name" : "Stmt_bb2__TO__bb9",
+         "schedule" : "{ Stmt_bb2__TO__bb9[i0] -> [i0] }"
+      }
+   ]
+}

test/Isl/CodeGen/partial_write_in_region___%bb1---%bb12.jscop.transformed

+{
+   "arrays" : [
+      {
+         "name" : "MemRef_A",
+         "sizes" : [ "*" ],
+         "type" : "i64"
+      },
+      {
+         "name" : "MemRef_B",
+         "sizes" : [ "*" ],
+         "type" : "float"
+      }
+   ],
+   "context" : "{  :  }",
+   "name" : "%bb1---%bb12",
+   "statements" : [
+      {
+         "accesses" : [
+            {
+               "kind" : "read",
+               "relation" : "{ Stmt_bb2__TO__bb9[i0] -> MemRef_A[i0] }"
+            },
+            {
+               "kind" : "read",
+               "relation" : "{ Stmt_bb2__TO__bb9[i0] -> MemRef_B[i0] }"
+            },
+            {
+               "kind" : "write",
+               "relation" : "{ Stmt_bb2__TO__bb9[i0] -> MemRef_B[i0] : i0 < 10}"
+            }
+         ],
+         "domain" : "{ Stmt_bb2__TO__bb9[i0] : 0 <= i0 <= 1023 }",
+         "name" : "Stmt_bb2__TO__bb9",
+         "schedule" : "{ Stmt_bb2__TO__bb9[i0] -> [i0] }"
+      }
+   ]
+}