Keep analyses live in unrolling (#1929)

Add code to keep the def-use manger and the inst-to-block mapping up-to-date. This means we do not have to rebuild them later.

To make this work, we will have to have to find places to update the
def-use manager. Updating the def-use manager is not straight forward
because we are unrolling loops, and we have circular references.

This forces one pass to register all of the definitions. A second one
to analyze the uses. Also because there will be references to the new
instructions in the old code, we want to register the definitions of the
new instructions early, so we can update the uses of the older code as
we go along.

The inst-to-block mapping is not too difficult. It can be done as instructions are created.

Fixes #1928.

source/opt/basic_block.cpp

@@ -35,9 +35,18 @@ const uint32_t kSelectionMergeMergeBlockIdInIdx = 0;
 BasicBlock* BasicBlock::Clone(IRContext* context) const {
   BasicBlock* clone = new BasicBlock(
       std::unique_ptr<Instruction>(GetLabelInst()->Clone(context)));
-  for (const auto& inst : insts_)
+  for (const auto& inst : insts_) {
     // Use the incoming context
     clone->AddInstruction(std::unique_ptr<Instruction>(inst.Clone(context)));
+  }
+
+  if (context->AreAnalysesValid(
+          IRContext::Analysis::kAnalysisInstrToBlockMapping)) {
+    for (auto& inst : *clone) {
+      context->set_instr_block(&inst, clone);
+    }
+  }
+
   return clone;
 }
 

source/opt/basic_block.h

@@ -54,6 +54,9 @@ class BasicBlock {
   //
   // The parent function will default to null and needs to be explicitly set by
   // the user.
+  //
+  // If the inst-to-block map in |context| is valid, then the new instructions
+  // will be inserted into the map.
   BasicBlock* Clone(IRContext*) const;
 
   // Sets the enclosing function for this basic block.

source/opt/loop_unroller.cpp

@@ -551,7 +551,10 @@ void LoopUnrollerUtilsImpl::RemoveDeadInstructions() {
 
 void LoopUnrollerUtilsImpl::ReplaceInductionUseWithFinalValue(Loop* loop) {
   context_->InvalidateAnalysesExceptFor(
-      IRContext::Analysis::kAnalysisLoopAnalysis);
+      IRContext::Analysis::kAnalysisLoopAnalysis |
+      IRContext::Analysis::kAnalysisDefUse |
+      IRContext::Analysis::kAnalysisInstrToBlockMapping);
+
   std::vector<Instruction*> inductions;
   loop->GetInductionVariables(inductions);
 
@@ -592,7 +595,8 @@ void LoopUnrollerUtilsImpl::FullyUnroll(Loop* loop) {
   RemoveDeadInstructions();
   // Invalidate all analyses.
   context_->InvalidateAnalysesExceptFor(
-      IRContext::Analysis::kAnalysisLoopAnalysis);
+      IRContext::Analysis::kAnalysisLoopAnalysis |
+      IRContext::Analysis::kAnalysisDefUse);
 }
 
 // Copy a given basic block, give it a new result_id, and store the new block
@@ -615,6 +619,7 @@ void LoopUnrollerUtilsImpl::CopyBasicBlock(Loop* loop, const BasicBlock* itr,
     if (!preserve_instructions) {
       Instruction* merge_inst = loop->GetHeaderBlock()->GetLoopMergeInst();
       merge_inst->SetInOperand(1, {basic_block->id()});
+      context_->UpdateDefUse(merge_inst);
     }
 
     state_.new_continue_block = basic_block;
@@ -655,15 +660,17 @@ void LoopUnrollerUtilsImpl::CopyBody(Loop* loop, bool eliminate_conditions) {
   }
 
   // Set the previous latch block to point to the new header.
-  Instruction& latch_branch = *state_.previous_latch_block_->tail();
-  latch_branch.SetInOperand(0, {state_.new_header_block->id()});
+  Instruction* latch_branch = state_.previous_latch_block_->terminator();
+  latch_branch->SetInOperand(0, {state_.new_header_block->id()});
+  context_->UpdateDefUse(latch_branch);
 
   // As the algorithm copies the original loop blocks exactly, the tail of the
   // latch block on iterations after the first one will be a branch to the new
   // header and not the actual loop header. The last continue block in the loop
   // should always be a backedge to the global header.
-  Instruction& new_latch_branch = *state_.new_latch_block->tail();
-  new_latch_branch.SetInOperand(0, {loop->GetHeaderBlock()->id()});
+  Instruction* new_latch_branch = state_.new_latch_block->terminator();
+  new_latch_branch->SetInOperand(0, {loop->GetHeaderBlock()->id()});
+  context_->AnalyzeUses(new_latch_branch);
 
   std::vector<Instruction*> inductions;
   loop->GetInductionVariables(inductions);
@@ -725,7 +732,10 @@ void LoopUnrollerUtilsImpl::FoldConditionBlock(BasicBlock* condition_block,
 
   context_->KillInst(&old_branch);
   // Add the new unconditional branch to the merge block.
-  InstructionBuilder builder{context_, condition_block};
+  InstructionBuilder builder(
+      context_, condition_block,
+      IRContext::Analysis::kAnalysisDefUse |
+          IRContext::Analysis::kAnalysisInstrToBlockMapping);
   builder.AddBranch(new_target);
 }
 
@@ -736,8 +746,9 @@ void LoopUnrollerUtilsImpl::CloseUnrolledLoop(Loop* loop) {
 
   // Remove the final backedge to the header and make it point instead to the
   // merge block.
-  state_.previous_latch_block_->tail()->SetInOperand(
-      0, {loop->GetMergeBlock()->id()});
+  Instruction* latch_instruction = state_.previous_latch_block_->terminator();
+  latch_instruction->SetInOperand(0, {loop->GetMergeBlock()->id()});
+  context_->UpdateDefUse(latch_instruction);
 
   // Remove all induction variables as the phis will now be invalid. Replace all
   // uses with the constant initializer value (all uses of phis will be in
@@ -821,6 +832,8 @@ void LoopUnrollerUtilsImpl::AddBlocksToFunction(
 // Assign all result_ids in |basic_block| instructions to new IDs and preserve
 // the mapping of new ids to old ones.
 void LoopUnrollerUtilsImpl::AssignNewResultIds(BasicBlock* basic_block) {
+  analysis::DefUseManager* def_use_mgr = context_->get_def_use_mgr();
+
   // Label instructions aren't covered by normal traversal of the
   // instructions.
   uint32_t new_label_id = context_->TakeNextId();
@@ -828,6 +841,7 @@ void LoopUnrollerUtilsImpl::AssignNewResultIds(BasicBlock* basic_block) {
   // Assign a new id to the label.
   state_.new_inst[basic_block->GetLabelInst()->result_id()] = new_label_id;
   basic_block->GetLabelInst()->SetResultId(new_label_id);
+  def_use_mgr->AnalyzeInstDef(basic_block->GetLabelInst());
 
   for (Instruction& inst : *basic_block) {
     uint32_t old_id = inst.result_id();
@@ -839,6 +853,7 @@ void LoopUnrollerUtilsImpl::AssignNewResultIds(BasicBlock* basic_block) {
 
     // Give the instruction a new id.
     inst.SetResultId(context_->TakeNextId());
+    def_use_mgr->AnalyzeInstDef(&inst);
 
     // Save the mapping of old_id -> new_id.
     state_.new_inst[old_id] = inst.result_id();
@@ -861,6 +876,7 @@ void LoopUnrollerUtilsImpl::RemapOperands(Instruction* inst) {
   };
 
   inst->ForEachInId(remap_operands_to_new_ids);
+  context_->AnalyzeUses(inst);
 }
 
 void LoopUnrollerUtilsImpl::RemapOperands(BasicBlock* basic_block) {

source/opt/loop_unroller.h

@@ -30,6 +30,13 @@ class LoopUnroller : public Pass {
 
   Status Process() override;
 
+  IRContext::Analysis GetPreservedAnalyses() override {
+    return IRContext::kAnalysisDefUse |
+           IRContext::kAnalysisInstrToBlockMapping |
+           IRContext::kAnalysisDecorations | IRContext::kAnalysisCombinators |
+           IRContext::kAnalysisNameMap;
+  }
+
  private:
   bool fully_unroll_;
   int unroll_factor_;