Move utility for acting on each lane of ElementCount to common code [nfc]

This was first written for AddressSanitizer, but I'm about to reuse it for MemorySanitizer as well.

diff --git a/llvm/include/llvm/Transforms/Utils/BasicBlockUtils.h b/llvm/include/llvm/Transforms/Utils/BasicBlockUtils.h
index c97baaf..b657e2b 100644 (file)
--- a/llvm/include/llvm/Transforms/Utils/BasicBlockUtils.h
+++ b/llvm/include/llvm/Transforms/Utils/BasicBlockUtils.h
@@ -32,6 +32,7 @@ class BlockFrequencyInfo;
 class BranchProbabilityInfo;
 class DomTreeUpdater;
 class Function;
+class IRBuilderBase;
 class LoopInfo;
 class MDNode;
 class MemoryDependenceResults;
@@ -474,6 +475,18 @@ void SplitBlockAndInsertIfThenElse(Value *Cond, Instruction *SplitBefore,
                                    MDNode *BranchWeights = nullptr,
                                    DomTreeUpdater *DTU = nullptr);
 
+/// Utility function for performing a given action on each lane of a vector
+/// with \p EC elements.  To simplify porting legacy code, this defaults to
+/// unrolling the implied loop for non-scalable element counts, but this is
+/// not considered to be part of the contract of this routine, and is
+/// expected to change in the future. The callback takes as arguments an
+/// IRBuilder whose insert point is correctly set for instantiating the
+/// given index, and a value which is (at runtime) the index to access.
+/// This index *may* be a constant.
+void SplitBlockAndInsertForEachLane(ElementCount EC, Type *IndexTy,
+    Instruction *InsertBefore,
+    std::function<void(IRBuilderBase&, Value*)> Func);
+
 /// Check whether BB is the merge point of a if-region.
 /// If so, return the branch instruction that determines which entry into
 /// BB will be taken.  Also, return by references the block that will be

diff --git a/llvm/lib/Transforms/Instrumentation/AddressSanitizer.cpp b/llvm/lib/Transforms/Instrumentation/AddressSanitizer.cpp
index 3e1167d..65841f2 100644 (file)
--- a/llvm/lib/Transforms/Instrumentation/AddressSanitizer.cpp
+++ b/llvm/lib/Transforms/Instrumentation/AddressSanitizer.cpp
@@ -1439,69 +1439,6 @@ static void doInstrumentAddress(AddressSanitizer *Pass, Instruction *I,
                                          IsWrite, nullptr, UseCalls, Exp);
 }
 
-static void SplitBlockAndInsertSimpleForLoop(Value *End,
-                                             Instruction *SplitBefore,
-                                             Instruction *&BodyIP,
-                                             Value *&Index) {
-  BasicBlock *LoopPred = SplitBefore->getParent();
-  BasicBlock *LoopBody = SplitBlock(SplitBefore->getParent(), SplitBefore);
-  BasicBlock *LoopExit = SplitBlock(SplitBefore->getParent(), SplitBefore);
-
-  auto *Ty = End->getType();
-  auto &DL = SplitBefore->getModule()->getDataLayout();
-  const unsigned Bitwidth = DL.getTypeSizeInBits(Ty);
-
-  IRBuilder<> Builder(LoopBody->getTerminator());
-  auto *IV = Builder.CreatePHI(Ty, 2, "iv");
-  auto *IVNext =
-    Builder.CreateAdd(IV, ConstantInt::get(Ty, 1), IV->getName() + ".next",
-                      /*HasNUW=*/true, /*HasNSW=*/Bitwidth != 2);
-  auto *IVCheck = Builder.CreateICmpEQ(IVNext, End,
-                                       IV->getName() + ".check");
-  Builder.CreateCondBr(IVCheck, LoopExit, LoopBody);
-  LoopBody->getTerminator()->eraseFromParent();
-
-  // Populate the IV PHI.
-  IV->addIncoming(ConstantInt::get(Ty, 0), LoopPred);
-  IV->addIncoming(IVNext, LoopBody);
-
-  BodyIP = LoopBody->getFirstNonPHI();
-  Index = IV;
-}
-
-/// Utility function for performing a given action on each lane of a vector
-/// with \p EC elements.  To simplify porting legacy code, this defaults to
-/// unrolling the implied loop for non-scalable element counts, but this is
-/// not considered to be part of the contract of this routine, and is
-/// expected to change in the future.
-static void
-SplitBlockAndInsertForEachLane(ElementCount EC,
-                               Type *IndexTy,
-                               Instruction *InsertBefore,
-                               std::function<void(IRBuilderBase&, Value*)> Func) {
-
-  IRBuilder<> IRB(InsertBefore);
-
-  if (EC.isScalable()) {
-    Value *NumElements = IRB.CreateElementCount(IndexTy, EC);
-
-    Instruction *BodyIP;
-    Value *Index;
-    SplitBlockAndInsertSimpleForLoop(NumElements, InsertBefore,
-                                     BodyIP, Index);
-
-    IRB.SetInsertPoint(BodyIP);
-    Func(IRB, Index);
-    return;
-  }
-
-  unsigned Num = EC.getFixedValue();
-  for (unsigned Idx = 0; Idx < Num; ++Idx) {
-    IRB.SetInsertPoint(InsertBefore);
-    Func(IRB, ConstantInt::get(IndexTy, Idx));
-  }
-}
-
 static void instrumentMaskedLoadOrStore(AddressSanitizer *Pass,
                                         const DataLayout &DL, Type *IntptrTy,
                                         Value *Mask, Instruction *I,

diff --git a/llvm/lib/Transforms/Utils/BasicBlockUtils.cpp b/llvm/lib/Transforms/Utils/BasicBlockUtils.cpp
index 0dd3857..2d1be2c 100644 (file)
--- a/llvm/lib/Transforms/Utils/BasicBlockUtils.cpp
+++ b/llvm/lib/Transforms/Utils/BasicBlockUtils.cpp
@@ -32,6 +32,7 @@
 #include "llvm/IR/Instruction.h"
 #include "llvm/IR/Instructions.h"
 #include "llvm/IR/IntrinsicInst.h"
+#include "llvm/IR/IRBuilder.h"
 #include "llvm/IR/LLVMContext.h"
 #include "llvm/IR/Type.h"
 #include "llvm/IR/User.h"
@@ -1599,6 +1600,61 @@ void llvm::SplitBlockAndInsertIfThenElse(Value *Cond, Instruction *SplitBefore,
   }
 }
 
+/// Insert a for (int i = 0; i < End; i++) loop structure (with the exception
+/// that \p End is assumed > 0, and thus not checked on entry) at \p
+/// SplitBefore.  Returns the first insert point in the loop body, and the
+/// PHINode for the induction variable (i.e. "i" above).
+static std::pair<Instruction*, Value*>
+SplitBlockAndInsertSimpleForLoop(Value *End, Instruction *SplitBefore) {
+  BasicBlock *LoopPred = SplitBefore->getParent();
+  BasicBlock *LoopBody = SplitBlock(SplitBefore->getParent(), SplitBefore);
+  BasicBlock *LoopExit = SplitBlock(SplitBefore->getParent(), SplitBefore);
+
+  auto *Ty = End->getType();
+  auto &DL = SplitBefore->getModule()->getDataLayout();
+  const unsigned Bitwidth = DL.getTypeSizeInBits(Ty);
+
+  IRBuilder<> Builder(LoopBody->getTerminator());
+  auto *IV = Builder.CreatePHI(Ty, 2, "iv");
+  auto *IVNext =
+    Builder.CreateAdd(IV, ConstantInt::get(Ty, 1), IV->getName() + ".next",
+                      /*HasNUW=*/true, /*HasNSW=*/Bitwidth != 2);
+  auto *IVCheck = Builder.CreateICmpEQ(IVNext, End,
+                                       IV->getName() + ".check");
+  Builder.CreateCondBr(IVCheck, LoopExit, LoopBody);
+  LoopBody->getTerminator()->eraseFromParent();
+
+  // Populate the IV PHI.
+  IV->addIncoming(ConstantInt::get(Ty, 0), LoopPred);
+  IV->addIncoming(IVNext, LoopBody);
+
+  return std::make_pair(LoopBody->getFirstNonPHI(), IV);
+}
+
+void llvm::SplitBlockAndInsertForEachLane(ElementCount EC,
+     Type *IndexTy, Instruction *InsertBefore,
+     std::function<void(IRBuilderBase&, Value*)> Func) {
+
+  IRBuilder<> IRB(InsertBefore);
+
+  if (EC.isScalable()) {
+    Value *NumElements = IRB.CreateElementCount(IndexTy, EC);
+
+    auto [BodyIP, Index] =
+      SplitBlockAndInsertSimpleForLoop(NumElements, InsertBefore);
+
+    IRB.SetInsertPoint(BodyIP);
+    Func(IRB, Index);
+    return;
+  }
+
+  unsigned Num = EC.getFixedValue();
+  for (unsigned Idx = 0; Idx < Num; ++Idx) {
+    IRB.SetInsertPoint(InsertBefore);
+    Func(IRB, ConstantInt::get(IndexTy, Idx));
+  }
+}
+
 BranchInst *llvm::GetIfCondition(BasicBlock *BB, BasicBlock *&IfTrue,
                                  BasicBlock *&IfFalse) {
   PHINode *SomePHI = dyn_cast<PHINode>(BB->begin());