[Polly] [ScopInfo] Do not use ScopStmt in Domain derivation of ScopInfo. NFC

[nandini12396]

Summary

ScopStmts were being used in the computation of the Domain of the SCoPs in ScopInfo. Once statements are split, there will not be a 1-to-1 correspondence between Stmts and Basic blocks. Thus this patch avoids the use of getStmtFor() by creating a map of BB to InvalidDomain and using it to compute the domain of the statements.

[nandini12396]

lib/Analysis/ScopInfo.cpp

2621–2622

@Meinersbur : I would like to clarify again what it is to be done here. You said that the easiest approach would be to iterate over all statements that represent this BB and set their Invalid Domain. But I got confused why that is inefficient? Could you please repeat what you meant by using the DomainMap which maps BasicBlocks to their domains?

[nandini12396]

lib/Analysis/ScopInfo.cpp

2621–2622

Ok, I think I understood what you meant. Create a map of BB to Invalid Domain. So there is no need of using getStmtFor(BB) by directly setting in the map and using it to get it.

[Meinersbur]

lib/Analysis/ScopInfo.cpp

2621–2622

Yes, I meant creating a map and in a first step, compute the domain only for that map.

InvalidDomainMap[EntryBB] = ...

Then, in a second step, apply the InvalidDomain to all statement for their BasicBlock.

See Scop::DomainMap for a similar approach.

[nandini12396]

@Meinersbur : Please have a look.

[Meinersbur]

Please don't remove an overload of getStmtFor if it is still called somewhere. It doesn't compile otherwise. It must also pass check-polly so we can see there are no unintended behavioral changes.

lib/Analysis/ScopInfo.cpp

2623–2624

Please only iterate and setInvalidDomain only once at the end of the function, when InvalidDomainMap contains all latest entries.

You cannot call getBasicBlock() on region statements.

2734

Note that there are also region statements, which do not have a 1-to-1 correlation between statements and a BasicBlock.

You might either normalize to a region's entry block, or use the RegionNodetype (the RN variable) instead.

2750–2755

Update InvalidDomainMap instead.

[nandini12396]

Hello Sir. I updated the diff to incorporate your comments. However, I don't know what I did wrong and there are 755 unexpected failures. Im still trying to understand that.

[nandini12396]

Sir, afaiu the problem seems to be that there maybe not be any ScopStmt, to which the first / last instruction of the BB belongs to; and hence the ScopStmt is null. So should I iterate over BB until I find the first/last instruction which corresponds to a ScopStmt ?

[Meinersbur]

include/polly/ScopInfo.h

2514–2517

Why did you move the location of this declaration?

lib/Analysis/ScopInfo.cpp

1327

As long as we did not remove getStmtFor(Instruction*), you can keep the assertion.

2654–2657

There is ScopStmt::getEntryBlock() which returns you the (Entry-)Block for the two cases.

2770

You must replace all occurrences of setInvalidDomain to update InvalidDomainMap instead.

2782

Shouldn't this changed to update InvalidDomainMap?

3021

The StmtMap does not contain all instructions. It never includes the terminator which PredBB->back() return. That is getStmtFor cannot any statement for its input.

This applies similarly to getStmtFor(&BB->front() as well which can be a synthesizable value.

[nandini12396]

The most appearing cause of failure of the test cases is:

opt: /home/nandini/llvm/tools/polly/lib/Analysis/ScopInfo.cpp:3058: bool polly::Scop::propagateDomainConstraints(llvm::Region*, llvm::DominatorTree&, llvm::LoopInfo&): Assertion `Domain' failed.

But I do not seem to understand the reason for this assertion fail. Can you please give me some pointers?

[nandini12396]

lib/Analysis/ScopInfo.cpp

1473

@Meinersbur : I think I am doing something wrong here. Can you please help.

[Meinersbur]

lib/Analysis/ScopBuilder.cpp

956 ↗

(On Diff #103408)

The ScopStmt's invalid domains have already set in buildDomains. That is, the elements of InvalidDomainMap have been consumed already and setting them has no effect.

lib/Analysis/ScopInfo.cpp

2130

getBasicBlock() returns nullptr if Stmt is a region statement.

2653–2654

When we have multiple statements per BasicBlock, this would consume the same InvalidDomainMap[Stmt.getEntryBlock()] twice, giving us an use-after-free sometime later.

How are you sure that all entries of InvalidDomainMap are an entry block of some statement? There could be an BB from the middle of a region statement. Such a block's invalid domain would not be consumed, therefore leaking.

2838

You are leaking isl objects here: InvalidDomainMap[ExitBB] may already contain a value.

This is probably not the only leak. You can find some leak by using valgrind --leak-check=full on a failing test case with -polly-on-isl-error-abort=0 added could be useful to find those.

[nandini12396]

I think it will be better to change the isl functions to their corresponding c++ version in another patch, since the change is large.

[Meinersbur]

Could you make use of getFirstNonBoxedLoopFor? Then it LGTM.

lib/Analysis/ScopInfo.cpp

2784

Please try to avoid unrelated and unnecessary changes.

2824–2825

Could you extract this into its own function?

We once had the following code in ScopHelper.cpp, you could revive them:

llvm::Loop *polly::getFirstNonBoxedLoopFor(llvm::Loop *L, llvm::LoopInfo &LI,
                                           const BoxedLoopsSetTy &BoxedLoops) {
  while (BoxedLoops.count(L))
    L = L->getParentLoop();
  return L;
}

llvm::Loop *polly::getFirstNonBoxedLoopFor(llvm::BasicBlock *BB,
                                           llvm::LoopInfo &LI,
                                           const BoxedLoopsSetTy &BoxedLoops) {
  Loop *L = LI.getLoopFor(BB);
  return getFirstNonBoxedLoopFor(L, LI, BoxedLoops);
}

Please don't forget to also the change the other place where getFirstNonBoxedLoopFor could be used.

[nandini12396]

Addressed comments.

[Meinersbur]

LGTM, going to commit...

(When submitting a diff, please ensure that formatting is correct by running "make check-polly")

include/polly/ScopInfo.h

private:
   ///
   /// We need a shared_ptr with reference counter to delete the context when all
   /// isl objects are deleted. We will distribute the shared_ptr to all objects
   /// that use the context to create isl objects, and increase the reference
   /// counter. By doing this, we guarantee that the context is deleted when we
   /// delete the last object that creates isl objects with the context.
   std::shared_ptr<isl_ctx> IslCtx;
 
-  /// A map from basic blocks to SCoP statements.
-  DenseMap<BasicBlock *, ScopStmt *> StmtMap;
+  /// A map from instructions to SCoP statements.
+  DenseMap<Instruction *, ScopStmt *> StmtMap;
 
   /// A map from basic blocks to their domains.
   DenseMap<BasicBlock *, isl_set *> DomainMap;
 
   /// Constraints on parameters.
   isl_set *Context;
 
   /// The affinator used to translate SCEVs to isl expressions.
private:
 
   /// Propagate invalid domains of statements through @p R.
   ///
   /// This method will propagate invalid statement domains through @p R and at
   /// the same time add error block domains to them. Additionally, the domains
   /// of error statements and those only reachable via error statements will be
   /// replaced by an empty set. Later those will be removed completely.
   ///
-  /// @param R  The currently traversed region.
-  /// @param DT The DominatorTree for the current function.
-  /// @param LI The LoopInfo for the current function.
-  ///
+  /// @param R                The currently traversed region.
+  /// @param DT               The DominatorTree for the current function.
+  /// @param LI               The LoopInfo for the current function.
+  /// @param InvalidDomainMap BB to InvalidDomain map for the BB of current
+  ///                         region.
+  //
   /// @returns True if there was no problem and false otherwise.
-  bool propagateInvalidStmtDomains(Region *R, DominatorTree &DT, LoopInfo &LI);
+  bool propagateInvalidStmtDomains(
+      Region *R, DominatorTree &DT, LoopInfo &LI,
+      DenseMap<BasicBlock *, isl_set *> &InvalidDomainMap);
 
   /// Compute the domain for each basic block in @p R.
   ///
   /// @param R  The region we currently traverse.
   /// @param DT The DominatorTree for the current function.
   /// @param LI The LoopInfo for the current function.
   ///
   /// @returns True if there was no problem and false otherwise.
public:
   std::string getContextStr() const;
 
   /// Get an isl string representing the assumed context.
   std::string getAssumedContextStr() const;
 
   /// Get an isl string representing the invalid context.
   std::string getInvalidContextStr() const;
 
+  /// Return the ScopStmt an instruction belongs to, or nullptr if it
+  /// does not belong to any statement in this Scop.
+  ScopStmt *getStmtFor(Instruction *Inst) const;
+
   /// Return the ScopStmt for the given @p BB or nullptr if there is
   ///        none.
   ScopStmt *getStmtFor(BasicBlock *BB) const;
 
   /// Return the ScopStmt that represents the Region @p R, or nullptr if
   ///        it is not represented by any statement in this Scop.
   ScopStmt *getStmtFor(Region *R) const;
 
   /// Return the ScopStmt that represents @p RN; can return nullptr if
   ///        the RegionNode is not within the SCoP or has been removed due to
   ///        simplifications.
   ScopStmt *getStmtFor(RegionNode *RN) const;
 
-  /// Return the ScopStmt an instruction belongs to, or nullptr if it
-  ///        does not belong to any statement in this Scop.
-  ScopStmt *getStmtFor(Instruction *Inst) const {
-    return getStmtFor(Inst->getParent());
-  }
-
-  /// Return the number of statements in the SCoP.
   size_t getSize() const { return Stmts.size(); }
 
   /// @name Statements Iterators
   ///
   /// These iterators iterate over all statements of this Scop.
   //@{
   typedef StmtSet::iterator iterator;
   typedef StmtSet::const_iterator const_iterator;

lib/Analysis/ScopInfo.cpp

 void ScopStmt::addAccess(MemoryAccess *Access) {
   Instruction *AccessInst = Access->getAccessInstruction();
 
   if (Access->isArrayKind()) {
     MemoryAccessList &MAL = InstructionToAccess[AccessInst];
     MAL.emplace_front(Access);
   } else if (Access->isValueKind() && Access->isWrite()) {
     Instruction *AccessVal = cast<Instruction>(Access->getAccessValue());
-    assert(Parent.getStmtFor(AccessVal) == this);
     assert(!ValueWrites.lookup(AccessVal));
 
     ValueWrites[AccessVal] = Access;
   } else if (Access->isValueKind() && Access->isRead()) {
     Value *AccessVal = Access->getAccessValue();
     assert(!ValueReads.lookup(AccessVal));
 
     ValueReads[AccessVal] = Access;
static __isl_give isl_set *buildConditionSet(ICmpInst::Predicate Pred,
   return setDimensionIds(Domain, ConsequenceCondSet);
 }
 
 /// Build the conditions sets for the switch @p SI in the @p Domain.
 ///
 /// This will fill @p ConditionSets with the conditions under which control
 /// will be moved from @p SI to its successors. Hence, @p ConditionSets will
 /// have as many elements as @p SI has successors.
 static bool
 buildConditionSets(ScopStmt &Stmt, SwitchInst *SI, Loop *L,
                    __isl_keep isl_set *Domain,
                    SmallVectorImpl<__isl_give isl_set *> &ConditionSets) {
 
   Value *Condition = getConditionFromTerminator(SI);
   assert(Condition && "No condition for switch");
 
   Scop &S = *Stmt.getParent();
for (auto *Param : DetectedParams) {
         continue;
       NewParams.insert(Param);
     }
 
     SmallVector<isl_set *, 2> ConditionSets;
     auto *TI = InScop ? CI->getParent()->getTerminator() : nullptr;
     auto &Stmt = InScop ? *getStmtFor(CI->getParent()) : *Stmts.begin();
     auto *Dom = InScop ? getDomainConditions(&Stmt) : isl_set_copy(Context);
     bool Valid = buildConditionSets(Stmt, Val, TI, L, Dom, ConditionSets);
     isl_set_free(Dom);
 
     if (!Valid)
       continue;
 
     isl_set *AssumptionCtx = nullptr;
     if (InScop) {
       AssumptionCtx = isl_set_complement(isl_set_params(ConditionSets[1]));
bool Scop::buildDomains(Region *R, DominatorTree &DT, LoopInfo &LI) {
   auto *L = IsOnlyNonAffineRegion ? nullptr : LI.getLoopFor(EntryBB);
   int LD = getRelativeLoopDepth(L);
   auto *S = isl_set_universe(isl_space_set_alloc(getIslCtx(), 0, LD + 1));
 
   while (LD-- >= 0) {
     S = addDomainDimId(S, LD + 1, L);
     L = L->getParentLoop();
   }
 
-  // Initialize the invalid domain.
-  auto *EntryStmt = getStmtFor(EntryBB);
-  EntryStmt->setInvalidDomain(isl_set_empty(isl_set_get_space(S)));
+  /// A map from basic blocks to their invalid domains.
+  DenseMap<BasicBlock *, isl_set *> InvalidDomainMap;
 
+  InvalidDomainMap[EntryBB] = isl_set_empty(isl_set_get_space(S));
   DomainMap[EntryBB] = S;
 
   if (IsOnlyNonAffineRegion)
     return !containsErrorBlock(R->getNode(), *R, LI, DT);
 
   if (!buildDomainsWithBranchConstraints(R, DT, LI))
     return false;
 
   if (!propagateDomainConstraints(R, DT, LI))
     return false;
 
   // Error blocks and blocks dominated by them have been assumed to never be
   // executed. Representing them in the Scop does not add any value. In fact,
   // it is likely to cause issues during construction of the ScopStmts. The
   // contents of error blocks have not been verified to be expressible and
   // will cause problems when building up a ScopStmt for them.
   // Furthermore, basic blocks dominated by error blocks may reference
   // instructions in the error block which, if the error block is not modeled,
   // can themselves not be constructed properly. To this end we will replace
   // the domains of error blocks and those only reachable via error blocks
   // with an empty set. Additionally, we will record for each block under which
   // parameter combination it would be reached via an error block in its
   // InvalidDomain. This information is needed during load hoisting.
-  if (!propagateInvalidStmtDomains(R, DT, LI))
+  if (!propagateInvalidStmtDomains(R, DT, LI, InvalidDomainMap))
     return false;
 
+  // Initialize the invalid domain.
+  for (ScopStmt &Stmt : Stmts)
+    if (Stmt.isRegionStmt())
+      Stmt.setInvalidDomain(InvalidDomainMap[Stmt.getRegion()->getEntry()]);
+    else
+      Stmt.setInvalidDomain(InvalidDomainMap[Stmt.getBasicBlock()]);
+
   return true;
 }
 
 /// Adjust the dimensions of @p Dom that was constructed for @p OldL
 ///        to be compatible to domains constructed for loop @p NewL.
 ///
 /// This function assumes @p NewL and @p OldL are equal or there is a CFG
 /// edge from @p OldL to @p NewL.
if (OldDepth == NewDepth) {
     int NumDim = isl_set_n_dim(Dom);
     assert(NumDim >= Diff);
     Dom = isl_set_project_out(Dom, isl_dim_set, NumDim - Diff, Diff);
   }
 
   return Dom;
 }
 
-bool Scop::propagateInvalidStmtDomains(Region *R, DominatorTree &DT,
-                                       LoopInfo &LI) {
+bool Scop::propagateInvalidStmtDomains(
+    Region *R, DominatorTree &DT, LoopInfo &LI,
+    DenseMap<BasicBlock *, isl_set *> &InvalidDomainMap) {
+
   ReversePostOrderTraversal<Region *> RTraversal(R);
   for (auto *RN : RTraversal) {
 
     // Recurse for affine subregions but go on for basic blocks and non-affine
     // subregions.
     if (RN->isSubRegion()) {
       Region *SubRegion = RN->getNodeAs<Region>();
       if (!isNonAffineSubRegion(SubRegion)) {
-        propagateInvalidStmtDomains(SubRegion, DT, LI);
+        propagateInvalidStmtDomains(SubRegion, DT, LI, InvalidDomainMap);
         continue;
       }
     }
 
     bool ContainsErrorBlock = containsErrorBlock(RN, getRegion(), LI, DT);
     BasicBlock *BB = getRegionNodeBasicBlock(RN);
-    ScopStmt *Stmt = getStmtFor(BB);
     isl_set *&Domain = DomainMap[BB];
     assert(Domain && "Cannot propagate a nullptr");
 
-    auto *InvalidDomain = Stmt->getInvalidDomain();
+    auto *InvalidDomain = InvalidDomainMap[BB];
     bool IsInvalidBlock =
         ContainsErrorBlock || isl_set_is_subset(Domain, InvalidDomain);
 
     if (!IsInvalidBlock) {
       InvalidDomain = isl_set_intersect(InvalidDomain, isl_set_copy(Domain));
     } else {
       isl_set_free(InvalidDomain);
       InvalidDomain = Domain;
       isl_set *DomPar = isl_set_params(isl_set_copy(Domain));
       recordAssumption(ERRORBLOCK, DomPar, BB->getTerminator()->getDebugLoc(),
                        AS_RESTRICTION);
       Domain = nullptr;
     }
 
     if (isl_set_is_empty(InvalidDomain)) {
-      Stmt->setInvalidDomain(InvalidDomain);
+      for (ScopStmt &Stmt : Stmts)
+        if (Stmt.isRegionStmt())
+          InvalidDomainMap[Stmt.getRegion()->getEntry()] = InvalidDomain;
+        else
+          InvalidDomainMap[Stmt.getBasicBlock()] = InvalidDomain;
       continue;
     }
 
     auto *BBLoop = getRegionNodeLoop(RN, LI);
     auto *TI = BB->getTerminator();
     unsigned NumSuccs = RN->isSubRegion() ? 1 : TI->getNumSuccessors();
     for (unsigned u = 0; u < NumSuccs; u++) {
       auto *SuccBB = getRegionNodeSuccessor(RN, TI, u);
-      auto *SuccStmt = getStmtFor(SuccBB);
+      auto *SuccStmt = getStmtFor(&SuccBB->front());
 
       // Skip successors outside the SCoP.
       if (!SuccStmt)
         continue;
 
       // Skip backedges.
       if (DT.dominates(SuccBB, BB))
         continue;
 
       auto *SuccBBLoop = SuccStmt->getSurroundingLoop();
       auto *AdjustedInvalidDomain = adjustDomainDimensions(
           *this, isl_set_copy(InvalidDomain), BBLoop, SuccBBLoop);
       auto *SuccInvalidDomain = SuccStmt->getInvalidDomain();
       SuccInvalidDomain =
           isl_set_union(SuccInvalidDomain, AdjustedInvalidDomain);
       SuccInvalidDomain = isl_set_coalesce(SuccInvalidDomain);
       unsigned NumConjucts = isl_set_n_basic_set(SuccInvalidDomain);
       SuccStmt->setInvalidDomain(SuccInvalidDomain);
 
       // Check if the maximal number of domain disjunctions was reached.
       // In case this happens we will bail.
       if (NumConjucts < MaxDisjunctsInDomain)
         continue;
 
       isl_set_free(InvalidDomain);
       invalidate(COMPLEXITY, TI->getDebugLoc());
       return false;
     }
-
-    Stmt->setInvalidDomain(InvalidDomain);
   }
-
   return true;
 }
 
 void Scop::propagateDomainConstraintsToRegionExit(
     BasicBlock *BB, Loop *BBLoop,
     SmallPtrSetImpl<BasicBlock *> &FinishedExitBlocks, LoopInfo &LI) {
 
   // Check if the block @p BB is the entry of a region. If so we propagate it's
for (auto *LatchBB : LatchBBs)
       if (BB != LatchBB && BBReg->contains(LatchBB))
         return;
     L = L->getParentLoop();
   }
 
   auto *Domain = DomainMap[BB];
   assert(Domain && "Cannot propagate a nullptr");
 
-  auto *ExitStmt = getStmtFor(ExitBB);
+  auto *ExitStmt = getStmtFor(&ExitBB->back());
   auto *ExitBBLoop = ExitStmt->getSurroundingLoop();
 
   // Since the dimensions of @p BB and @p ExitBB might be different we have to
   // adjust the domain before we can propagate it.
   auto *AdjustedDomain =
       adjustDomainDimensions(*this, isl_set_copy(Domain), BBLoop, ExitBBLoop);
   auto *&ExitDomain = DomainMap[ExitBB];
 
   // If the exit domain is not yet created we set it otherwise we "add" the
   // current domain.
   ExitDomain =
       ExitDomain ? isl_set_union(AdjustedDomain, ExitDomain) : AdjustedDomain;
 
   // Initialize the invalid domain.
   ExitStmt->setInvalidDomain(isl_set_empty(isl_set_get_space(ExitDomain)));
 
   FinishedExitBlocks.insert(ExitBB);
 }
 
 bool Scop::buildDomainsWithBranchConstraints(Region *R, DominatorTree &DT,
                                              LoopInfo &LI) {
   // To create the domain for each block in R we iterate over all blocks and
   // subregions in R and propagate the conditions under which the current region
for (auto *RN : RTraversal) {
     // Now iterate over the successors and set their initial domain based on
     // their condition set. We skip back edges here and have to be careful when
     // we leave a loop not to keep constraints over a dimension that doesn't
     // exist anymore.
     assert(RN->isSubRegion() || TI->getNumSuccessors() == ConditionSets.size());
     for (unsigned u = 0, e = ConditionSets.size(); u < e; u++) {
       isl_set *CondSet = ConditionSets[u];
       BasicBlock *SuccBB = getRegionNodeSuccessor(RN, TI, u);
+      auto *SuccStmt = getStmtFor(&SuccBB->front());
 
-      auto *SuccStmt = getStmtFor(SuccBB);
       // Skip blocks outside the region.
       if (!SuccStmt) {
         isl_set_free(CondSet);
         continue;
       }
 
       // If we propagate the domain of some block to "SuccBB" we do not have to
       // adjust the domain.
Scop::getPredecessorDomainConstraints(BasicBlock *BB,
                                       DominatorTree &DT, LoopInfo &LI) {
   // If @p BB is the ScopEntry we are done
   if (R.getEntry() == BB)
     return isl_set_universe(isl_set_get_space(Domain));
 
   // The region info of this function.
   auto &RI = *R.getRegionInfo();
 
-  auto *BBLoop = getStmtFor(BB)->getSurroundingLoop();
+  auto *BBLoop = getStmtFor(&BB->front())->getSurroundingLoop();
 
   // A domain to collect all predecessor domains, thus all conditions under
   // which the block is executed. To this end we start with the empty domain.
   isl_set *PredDom = isl_set_empty(isl_set_get_space(Domain));
 
   // Set of regions of which the entry block domain has been propagated to BB.
   // all predecessors inside any of the regions can be skipped.
   SmallSet<Region *, 8> PropagatedRegions;
for (auto *PredBB : predecessors(BB)) {
     // If a valid region for propagation was found use the entry of that region
     // for propagation, otherwise the PredBB directly.
     if (PredR->getExit() == BB) {
       PredBB = PredR->getEntry();
       PropagatedRegions.insert(PredR);
     }
 
     auto *PredBBDom = getDomainConditions(PredBB);
-    auto *PredBBLoop = getStmtFor(PredBB)->getSurroundingLoop();
+    auto *PredBBLoop = getStmtFor(&PredBB->back())->getSurroundingLoop();
+
     PredBBDom = adjustDomainDimensions(*this, PredBBDom, PredBBLoop, BBLoop);
 
     PredDom = isl_set_union(PredDom, PredBBDom);
   }
 
   return PredDom;
 }
 
for (auto StmtIt = Stmts.begin(), StmtEnd = Stmts.end(); StmtIt != StmtEnd;) {
     if (!RemoveStmt) {
       StmtIt++;
       continue;
     }
 
     // Remove the statement because it is unnecessary.
     if (Stmt.isRegionStmt())
       for (BasicBlock *BB : Stmt.getRegion()->blocks())
-        StmtMap.erase(BB);
+        for (Instruction &Inst : *BB)
+          StmtMap.erase(&Inst);
     else
-      StmtMap.erase(Stmt.getBasicBlock());
+      for (Instruction &Inst : *Stmt.getBasicBlock())
+        StmtMap.erase(&Inst);
 
     StmtIt = Stmts.erase(StmtIt);
   }
 }
 
 InvariantEquivClassTy *Scop::lookupInvariantEquivClass(Value *Val) {
   LoadInst *LInst = dyn_cast<LoadInst>(Val);
   if (!LInst)
static isl::multi_union_pw_aff mapToDimension(isl::union_set USet, int N) {
   return isl::multi_union_pw_aff(isl::union_pw_multi_aff(Result));
 }
 
 void Scop::addScopStmt(BasicBlock *BB, Loop *SurroundingLoop,
                        std::vector<Instruction *> Instructions) {
   assert(BB && "Unexpected nullptr!");
   Stmts.emplace_back(*this, *BB, SurroundingLoop, Instructions);
   auto *Stmt = &Stmts.back();
-  StmtMap[BB] = Stmt;
+  for (Instruction *Inst : Instructions)
+    StmtMap[Inst] = Stmt;
 }
 
 void Scop::addScopStmt(Region *R, Loop *SurroundingLoop) {
   assert(R && "Unexpected nullptr!");
   Stmts.emplace_back(*this, *R, SurroundingLoop);
   auto *Stmt = &Stmts.back();
   for (BasicBlock *BB : R->blocks())
-    StmtMap[BB] = Stmt;
+    for (Instruction &Inst : *BB)
+      StmtMap[&Inst] = Stmt;
 }
 
 ScopStmt *Scop::addScopStmt(__isl_take isl_map *SourceRel,
                             __isl_take isl_map *TargetRel,
                             __isl_take isl_set *Domain) {
 #ifndef NDEBUG
   isl_set *SourceDomain = isl_map_domain(isl_map_copy(SourceRel));
   isl_set *TargetDomain = isl_map_domain(isl_map_copy(TargetRel));
if (Schedule) {
           combineInSequence(NextLoopData.Schedule, Schedule);
     }
 
     NextLoopData.NumBlocksProcessed += NumBlocksProcessed;
     LoopData = NextLoopData;
   }
 }
 
+ScopStmt *Scop::getStmtFor(Instruction *Inst) const {
+  auto StmtMapIt = StmtMap.find(Inst);
+  if (StmtMapIt == StmtMap.end())
+    return nullptr;
+  return StmtMapIt->second;
+}
+
 ScopStmt *Scop::getStmtFor(BasicBlock *BB) const {
-  auto StmtMapIt = StmtMap.find(BB);
+  auto StmtMapIt = StmtMap.find(&BB->front());
   if (StmtMapIt == StmtMap.end())
     return nullptr;
   return StmtMapIt->second;
 }
 
 ScopStmt *Scop::getStmtFor(RegionNode *RN) const {
   if (RN->isSubRegion())
     return getStmtFor(RN->getNodeAs<Region>());