[NARY-REASSOCIATE] Simplify traversal logic by post deleting dead instructions

[ebrevnov]

Summary

Currently we delete optimized instructions as we go. That has several negative consequences. First it complicates traversal logic itself. Second if newly generated instruction has been deleted the traversal is repeated from scratch.

But real motivation for the change is upcoming change with support for min/max reassociation. Here we employ SCEV expander to generate code. As a result newly generated instructions may be inserted not right before original instruction (because SCEV may do hoisting) and there is no way to know 'next' instruction.

[ebrevnov]

ping

[mkazantsev]

llvm/lib/Transforms/Scalar/NaryReassociate.cpp

244

This is unused outside if, please move def closer to use.

276

nit: { } not needed.

llvm/test/Transforms/NaryReassociate/pr24301.ll

5

Not clear what has changed with your patch. Could you please auto-update the test, commit it as NFC and rebase on top of it to show the difference?

[ebrevnov]

llvm/lib/Transforms/Scalar/NaryReassociate.cpp

244

Actually, it is used in else part.

276

Ok

[ebrevnov]

Rebase

[ebrevnov]

Rebase

[mkazantsev]

LGTM

llvm/lib/Transforms/Scalar/NaryReassociate.cpp

 }
 
 bool NaryReassociatePass::doOneIteration(Function &F) {
   bool Changed = false;
   SeenExprs.clear();
   // Process the basic blocks in a depth first traversal of the dominator
   // tree. This order ensures that all bases of a candidate are in Candidates
   // when we process it.
+  SmallVector<WeakTrackingVH, 16> DeadInsts;
   for (const auto Node : depth_first(DT)) {
     BasicBlock *BB = Node->getBlock();
     for (auto I = BB->begin(); I != BB->end(); ++I) {
-      if (SE->isSCEVable(I->getType()) && isPotentiallyNaryReassociable(&*I)) {
-        const SCEV *OldSCEV = SE->getSCEV(&*I);
-        if (Instruction *NewI = tryReassociate(&*I)) {
-          Changed = true;
-          SE->forgetValue(&*I);
-          I->replaceAllUsesWith(NewI);
-          WeakVH NewIExist = NewI;
-          // If SeenExprs/NewIExist contains I's WeakTrackingVH/WeakVH, that
-          // entry will be replaced with nullptr if deleted.
-          RecursivelyDeleteTriviallyDeadInstructions(&*I, TLI);
-          if (!NewIExist) {
-            // Rare occation where the new instruction (NewI) have been removed,
-            // probably due to parts of the input code was dead from the
-            // beginning, reset the iterator and start over from the beginning
-            I = BB->begin();
-            continue;
-          }
-          I = NewI->getIterator();
-        }
-        // Add the rewritten instruction to SeenExprs; the original instruction
-        // is deleted.
-        const SCEV *NewSCEV = SE->getSCEV(&*I);
-        SeenExprs[NewSCEV].push_back(WeakTrackingVH(&*I));
+      Instruction *OrigI = &*I;
+
+      if (!SE->isSCEVable(OrigI->getType()) ||
+          !isPotentiallyNaryReassociable(OrigI))
+
+      const SCEV *OrigSCEV = SE->getSCEV(OrigI);
+      if (Instruction *NewI = tryReassociate(OrigI)) {
+        Changed = true;
+        OrigI->replaceAllUsesWith(NewI);
+
+        // Add 'OrigI' to the list of dead instructions.
+        DeadInsts.push_back(WeakTrackingVH(OrigI));
+        // Add the rewritten instruction to SeenExprs; the original
+        // instruction is deleted.
+        const SCEV *NewSCEV = SE->getSCEV(NewI);
+        SeenExprs[NewSCEV].push_back(WeakTrackingVH(NewI));
+
         // Ideally, NewSCEV should equal OldSCEV because tryReassociate(I)
         // is equivalent to I. However, ScalarEvolution::getSCEV may
-        // weaken nsw causing NewSCEV not to equal OldSCEV. For example, suppose
-        // we reassociate
+        // weaken nsw causing NewSCEV not to equal OldSCEV. For example,
+        // suppose we reassociate
         //   I = &a[sext(i +nsw j)] // assuming sizeof(a[0]) = 4
         // to
         //   NewI = &a[sext(i)] + sext(j).
         //
         // ScalarEvolution computes
         //   getSCEV(I)    = a + 4 * sext(i + j)
         //   getSCEV(newI) = a + 4 * sext(i) + 4 * sext(j)
         // which are different SCEVs.
         //
         // To alleviate this issue of ScalarEvolution not always capturing
         // equivalence, we add I to SeenExprs[OldSCEV] as well so that we can
         // map both SCEV before and after tryReassociate(I) to I.
         //
-        // This improvement is exercised in @reassociate_gep_nsw in nary-gep.ll.
-        if (NewSCEV != OldSCEV)
-          SeenExprs[OldSCEV].push_back(WeakTrackingVH(&*I));
-      }
+        // This improvement is exercised in @reassociate_gep_nsw in
+        // nary-gep.ll.
+        if (NewSCEV != OrigSCEV)
+          SeenExprs[OrigSCEV].push_back(WeakTrackingVH(NewI));
+      } else
+        SeenExprs[OrigSCEV].push_back(WeakTrackingVH(OrigI));
     }
   }
+  // Delete all dead instructions from 'DeadInsts'.
+  // Please note ScalarEvolution is updated along the way.
+  RecursivelyDeleteTriviallyDeadInstructionsPermissive(
+      DeadInsts, TLI, nullptr, [this](Value *V) { SE->forgetValue(V); });
+
   return Changed;
 }
 
 Instruction *NaryReassociatePass::tryReassociate(Instruction *I) {
   switch (I->getOpcode()) {
   case Instruction::Add:
   case Instruction::Mul:
     return tryReassociateBinaryOp(cast<BinaryOperator>(I));

llvm/test/Transforms/NaryReassociate/pr24301.ll

 ; NOTE: Assertions have been autogenerated by utils/update_test_checks.py
 ; RUN: opt < %s -nary-reassociate -S | FileCheck %s
 ; RUN: opt < %s -passes='nary-reassociate' -S | FileCheck %s
 
 define i32 @foo(i32 %t4) {
 ; CHECK-LABEL: @foo(
 ; CHECK-NEXT:  entry:
-; CHECK-NEXT:    [[T5:%.*]] = add i32 [[T4:%.*]], 8
-; CHECK-NEXT:    [[T14:%.*]] = add i32 [[T5]], -128
-; CHECK-NEXT:    [[T21:%.*]] = add i32 119, [[T4]]
-; CHECK-NEXT:    [[T23:%.*]] = add i32 [[T21]], -128
+; CHECK-NEXT:    [[T13:%.*]] = add i32 [[T4:%.*]], -128
+; CHECK-NEXT:    [[T14:%.*]] = add i32 [[T13]], 8
+; CHECK-NEXT:    [[T23:%.*]] = add i32 [[T13]], 119
 ; CHECK-NEXT:    ret i32 [[T23]]
 ;
 entry:
   %t5 = add i32 %t4, 8
   %t13 = add i32 %t4, -128  ; deleted
   %t14 = add i32 %t13, 8    ; => %t5 + -128
   %t21 = add i32 119, %t4
   ; do not rewrite %t23 against %t13 because %t13 is already deleted
   %t23 = add i32 %t21, -128
   ret i32 %t23
 }
 
 ; This is essentially the same test as the previous one but intermidiate result (t14) has a use.
 define i32 @foo2(i32 %t4) {
 ; CHECK-LABEL: @foo2(
 ; CHECK-NEXT:  entry:
-; CHECK-NEXT:    [[T5:%.*]] = add i32 [[T4:%.*]], 8
-; CHECK-NEXT:    [[T14:%.*]] = add i32 [[T5]], -128
-; CHECK-NEXT:    [[T21:%.*]] = add i32 119, [[T4]]
-; CHECK-NEXT:    [[T23:%.*]] = add i32 [[T21]], -128
+; CHECK-NEXT:    [[T13:%.*]] = add i32 [[T4:%.*]], -128
+; CHECK-NEXT:    [[T14:%.*]] = add i32 [[T13]], 8
+; CHECK-NEXT:    [[T23:%.*]] = add i32 [[T13]], 119
 ; CHECK-NEXT:    [[RES:%.*]] = add i32 [[T14]], [[T23]]
 ; CHECK-NEXT:    ret i32 [[RES]]
 ;
 entry:
-  %t5 = add i32 %t4, 8
-  %t13 = add i32 %t4, -128  ; deleted
-  %t14 = add i32 %t13, 8    ; => %t5 + -128
-  %t21 = add i32 119, %t4
-  %t23 = add i32 %t21, -128
+  %t5 = add i32 %t4, 8     ; initially dead
+  %t13 = add i32 %t4, -128
+  %t14 = add i32 %t13, 8
+  %t21 = add i32 119, %t4  ; => dead after reassociation
+  %t23 = add i32 %t21, -128; => reassociated to t13 + 119
   %res = add i32 %t14, %t23
   ret i32 %res
 }