[ExpandReductions] Don't push all intrinsics to the worklist. Just push reductions.

[craig.topper]

Summary

We were previously pushing all intrinsics used in a function to the
worklist. This is wasteful for memory in a function with a lot of
intrinsics.

We also ask TTI if we should expand every intrinsic, but we only
have expansion support for the reduction intrinsics. This just
wastes time for the non-reduction intrinsics.

This patch only pushes reduction intrinsics into the worklist and
skips other intrinsics.

[RKSimon]

llvm/lib/CodeGen/ExpandReductions.cpp

99

Do we gain anything by doing the if TTI->shouldExpandReduction(II) here to decide whether to add to WorkList?

[spatel]

llvm/lib/CodeGen/ExpandReductions.cpp

82–105

Could reduce with something like:

for (auto &I : instructions(F)) {

Note - either way, this pass may be in danger of dying if used on unreachable blocks that contain weird IR.
See for example: D67766

[craig.topper]

llvm/lib/CodeGen/ExpandReductions.cpp

82–105

Does the UnreachableBlockEliminationPass that we run earlier in the codegen pipeline help prevent that?

[craig.topper]

Call shouldExpandReduction before putting in the worklist

[aemerson]

LGTM.

[spatel]

llvm/lib/CodeGen/ExpandReductions.cpp

82–105

Nice - I didn't know that existed. It's just a wrapper around a util function: llvm::EliminateUnreachableBlocks().

So yes, that should make it pretty safe. Still a chance that fuzzers will target this pass in a different pipeline or that something between that running and this running would create a dead block with bogus code, but the odds are low.

llvm/lib/CodeGen/ExpandReductions.cpp

RecurrenceDescriptor::MinMaxRecurrenceKind getMRK(Intrinsic::ID ID) {
   default:
     return RecurrenceDescriptor::MRK_Invalid;
   }
 }
 
 bool expandReductions(Function &F, const TargetTransformInfo *TTI) {
   bool Changed = false;
   SmallVector<IntrinsicInst *, 4> Worklist;
-  for (inst_iterator I = inst_begin(F), E = inst_end(F); I != E; ++I)
-    if (auto II = dyn_cast<IntrinsicInst>(&*I))
+  for (auto &I : instructions(F)) {
+    if (auto *II = dyn_cast<IntrinsicInst>(&I)) {
+      switch (II->getIntrinsicID()) {
+      default: break;
+      case Intrinsic::experimental_vector_reduce_v2_fadd:
+      case Intrinsic::experimental_vector_reduce_v2_fmul:
+      case Intrinsic::experimental_vector_reduce_add:
+      case Intrinsic::experimental_vector_reduce_mul:
+      case Intrinsic::experimental_vector_reduce_and:
+      case Intrinsic::experimental_vector_reduce_or:
+      case Intrinsic::experimental_vector_reduce_xor:
+      case Intrinsic::experimental_vector_reduce_smax:
+      case Intrinsic::experimental_vector_reduce_smin:
+      case Intrinsic::experimental_vector_reduce_umax:
+      case Intrinsic::experimental_vector_reduce_umin:
+      case Intrinsic::experimental_vector_reduce_fmax:
+      case Intrinsic::experimental_vector_reduce_fmin:
+        if (TTI->shouldExpandReduction(II))
           Worklist.push_back(II);
 
-  for (auto *II : Worklist) {
-    if (!TTI->shouldExpandReduction(II))
-      continue;
+        break;
+      }
+    }
+  }
 
+  for (auto *II : Worklist) {
     FastMathFlags FMF =
         isa<FPMathOperator>(II) ? II->getFastMathFlags() : FastMathFlags{};
     Intrinsic::ID ID = II->getIntrinsicID();
     RecurrenceDescriptor::MinMaxRecurrenceKind MRK = getMRK(ID);
 
     Value *Rdx = nullptr;
     IRBuilder<> Builder(II);
     IRBuilder<>::FastMathFlagGuard FMFGuard(Builder);
     Builder.setFastMathFlags(FMF);
     switch (ID) {
+    default: llvm_unreachable("Unexpected intrinsic!");
     case Intrinsic::experimental_vector_reduce_v2_fadd:
     case Intrinsic::experimental_vector_reduce_v2_fmul: {
       // FMFs must be attached to the call, otherwise it's an ordered reduction
       // and it can't be handled by generating a shuffle sequence.
       Value *Acc = II->getArgOperand(0);
       Value *Vec = II->getArgOperand(1);
       if (!FMF.allowReassoc())
         Rdx = getOrderedReduction(Builder, Acc, Vec, getOpcode(ID), MRK);
       else {
         if (!isPowerOf2_32(Vec->getType()->getVectorNumElements()))
           continue;
 
         Rdx = getShuffleReduction(Builder, Vec, getOpcode(ID), MRK);
         Rdx = Builder.CreateBinOp((Instruction::BinaryOps)getOpcode(ID),
                                   Acc, Rdx, "bin.rdx");
       }
-    } break;
+      break;
+    }
     case Intrinsic::experimental_vector_reduce_add:
     case Intrinsic::experimental_vector_reduce_mul:
     case Intrinsic::experimental_vector_reduce_and:
     case Intrinsic::experimental_vector_reduce_or:
     case Intrinsic::experimental_vector_reduce_xor:
     case Intrinsic::experimental_vector_reduce_smax:
     case Intrinsic::experimental_vector_reduce_smin:
     case Intrinsic::experimental_vector_reduce_umax:
     case Intrinsic::experimental_vector_reduce_umin:
     case Intrinsic::experimental_vector_reduce_fmax:
     case Intrinsic::experimental_vector_reduce_fmin: {
       Value *Vec = II->getArgOperand(0);
       if (!isPowerOf2_32(Vec->getType()->getVectorNumElements()))
         continue;
 
       Rdx = getShuffleReduction(Builder, Vec, getOpcode(ID), MRK);
-    } break;
-    default:
-      continue;
+      break;
+    }
     }
     II->replaceAllUsesWith(Rdx);
     II->eraseFromParent();
     Changed = true;
   }
   return Changed;
 }