[llvm-exegesis] Add loop mode for repeating the snippet.

Summary:
Before this change the Executable function was made by duplicating the
snippet. This change adds a --repetion-mode={loop|duplicate} flag that
allows choosing between this behaviour and wrapping the snippet instructions
in a loop.

The new mode can help measurements when the snippet fits in the DSB by
short-cirtcuiting decoding. The loop adds a dec + jmp to the measurements, but
since these are not part of the critical path, they execute in parallel
with the measured code and do not impact measurements in practice.

Overview of the change:
 - New SnippetRepetitor abstraction that handles repeating the snippet.
   The assembler delegates repeating the instructions to this class.
 - ExegesisTarget learns how to decrement loop counter and jump.
 - Some refactoring of the assembler into FunctionFiller/BasicBlockFiller.

Reviewers: gchatelet

Subscribers: mgorny, tschuett, llvm-commits

Tags: #llvm

Differential Revision: https://reviews.llvm.org/D68125

llvm-svn: 373083

13 files changed, 376 insertions, 81 deletions

diff --git a/llvm/tools/llvm-exegesis/lib/Assembler.cpp b/llvm/tools/llvm-exegesis/lib/Assembler.cpp
index 9695c79b57a..4ee5f4bc2fd 100644
--- a/llvm/tools/llvm-exegesis/lib/Assembler.cpp
+++ b/llvm/tools/llvm-exegesis/lib/Assembler.cpp
@@ -8,6 +8,7 @@
 
 #include "Assembler.h"
 
+#include "SnippetRepetitor.h"
 #include "Target.h"
 #include "llvm/CodeGen/GlobalISel/CallLowering.h"
 #include "llvm/CodeGen/GlobalISel/MachineIRBuilder.h"
@@ -28,21 +29,22 @@ namespace exegesis {
 static constexpr const char ModuleID[] = "ExegesisInfoTest";
 static constexpr const char FunctionID[] = "foo";
 
-static std::vector<llvm::MCInst>
+// Fills the given basic block with register setup code, and returns true if
+// all registers could be setup correctly.
+static bool
 generateSnippetSetupCode(const ExegesisTarget &ET,
                          const llvm::MCSubtargetInfo *const MSI,
                          llvm::ArrayRef<RegisterValue> RegisterInitialValues,
-                         bool &IsSnippetSetupComplete) {
-  IsSnippetSetupComplete = true;
-  std::vector<llvm::MCInst> Result;
+                         BasicBlockFiller &BBF) {
+  bool IsSnippetSetupComplete = true;
   for (const RegisterValue &RV : RegisterInitialValues) {
     // Load a constant in the register.
     const auto SetRegisterCode = ET.setRegTo(*MSI, RV.Register, RV.Value);
     if (SetRegisterCode.empty())
       IsSnippetSetupComplete = false;
-    Result.insert(Result.end(), SetRegisterCode.begin(), SetRegisterCode.end());
+    BBF.addInstructions(SetRegisterCode);
   }
-  return Result;
+  return IsSnippetSetupComplete;
 }
 
 // Small utility function to add named passes.
@@ -67,8 +69,7 @@ static bool addPass(llvm::PassManagerBase &PM, llvm::StringRef PassName,
   return false;
 }
 
-// Creates a void(int8*) MachineFunction.
-static llvm::MachineFunction &
+llvm::MachineFunction &
 createVoidVoidPtrMachineFunction(llvm::StringRef FunctionID,
                                  llvm::Module *Module,
                                  llvm::MachineModuleInfo *MMI) {
@@ -85,38 +86,43 @@ createVoidVoidPtrMachineFunction(llvm::StringRef FunctionID,
   return MMI->getOrCreateMachineFunction(*F);
 }
 
-static void fillMachineFunction(llvm::MachineFunction &MF,
-                                llvm::ArrayRef<unsigned> LiveIns,
-                                llvm::ArrayRef<llvm::MCInst> Instructions) {
-  llvm::MachineBasicBlock *MBB = MF.CreateMachineBasicBlock();
-  MF.push_back(MBB);
-  for (const unsigned Reg : LiveIns)
-    MBB->addLiveIn(Reg);
-  const llvm::MCInstrInfo *MCII = MF.getTarget().getMCInstrInfo();
-  llvm::DebugLoc DL;
-  for (const llvm::MCInst &Inst : Instructions) {
-    const unsigned Opcode = Inst.getOpcode();
-    const llvm::MCInstrDesc &MCID = MCII->get(Opcode);
-    llvm::MachineInstrBuilder Builder = llvm::BuildMI(MBB, DL, MCID);
-    for (unsigned OpIndex = 0, E = Inst.getNumOperands(); OpIndex < E;
-         ++OpIndex) {
-      const llvm::MCOperand &Op = Inst.getOperand(OpIndex);
-      if (Op.isReg()) {
-        const bool IsDef = OpIndex < MCID.getNumDefs();
-        unsigned Flags = 0;
-        const llvm::MCOperandInfo &OpInfo = MCID.operands().begin()[OpIndex];
-        if (IsDef && !OpInfo.isOptionalDef())
-          Flags |= llvm::RegState::Define;
-        Builder.addReg(Op.getReg(), Flags);
-      } else if (Op.isImm()) {
-        Builder.addImm(Op.getImm());
-      } else if (!Op.isValid()) {
-        llvm_unreachable("Operand is not set");
-      } else {
-        llvm_unreachable("Not yet implemented");
-      }
+BasicBlockFiller::BasicBlockFiller(llvm::MachineFunction &MF,
+                                   llvm::MachineBasicBlock *MBB,
+                                   const llvm::MCInstrInfo *MCII)
+    : MF(MF), MBB(MBB), MCII(MCII) {}
+
+void BasicBlockFiller::addInstruction(const llvm::MCInst &Inst,
+                                      const llvm::DebugLoc &DL) {
+  const unsigned Opcode = Inst.getOpcode();
+  const llvm::MCInstrDesc &MCID = MCII->get(Opcode);
+  llvm::MachineInstrBuilder Builder = llvm::BuildMI(MBB, DL, MCID);
+  for (unsigned OpIndex = 0, E = Inst.getNumOperands(); OpIndex < E;
+       ++OpIndex) {
+    const llvm::MCOperand &Op = Inst.getOperand(OpIndex);
+    if (Op.isReg()) {
+      const bool IsDef = OpIndex < MCID.getNumDefs();
+      unsigned Flags = 0;
+      const llvm::MCOperandInfo &OpInfo = MCID.operands().begin()[OpIndex];
+      if (IsDef && !OpInfo.isOptionalDef())
+        Flags |= llvm::RegState::Define;
+      Builder.addReg(Op.getReg(), Flags);
+    } else if (Op.isImm()) {
+      Builder.addImm(Op.getImm());
+    } else if (!Op.isValid()) {
+      llvm_unreachable("Operand is not set");
+    } else {
+      llvm_unreachable("Not yet implemented");
     }
   }
+}
+
+void BasicBlockFiller::addInstructions(ArrayRef<llvm::MCInst> Insts,
+                                       const llvm::DebugLoc &DL) {
+  for (const MCInst &Inst : Insts)
+    addInstruction(Inst, DL);
+}
+
+void BasicBlockFiller::addReturn(const llvm::DebugLoc &DL) {
   // Insert the return code.
   const llvm::TargetInstrInfo *TII = MF.getSubtarget().getInstrInfo();
   if (TII->getReturnOpcode() < TII->getNumOpcodes()) {
@@ -128,6 +134,21 @@ static void fillMachineFunction(llvm::MachineFunction &MF,
   }
 }
 
+FunctionFiller::FunctionFiller(llvm::MachineFunction &MF,
+                               std::vector<unsigned> RegistersSetUp)
+    : MF(MF), MCII(MF.getTarget().getMCInstrInfo()), Entry(addBasicBlock()),
+      RegistersSetUp(std::move(RegistersSetUp)) {}
+
+BasicBlockFiller FunctionFiller::addBasicBlock() {
+  llvm::MachineBasicBlock *MBB = MF.CreateMachineBasicBlock();
+  MF.push_back(MBB);
+  return BasicBlockFiller(MF, MBB, MCII);
+}
+
+ArrayRef<unsigned> FunctionFiller::getRegistersSetUp() const {
+  return RegistersSetUp;
+}
+
 static std::unique_ptr<llvm::Module>
 createModule(const std::unique_ptr<llvm::LLVMContext> &Context,
              const llvm::DataLayout DL) {
@@ -155,7 +176,7 @@ void assembleToStream(const ExegesisTarget &ET,
                       std::unique_ptr<llvm::LLVMTargetMachine> TM,
                       llvm::ArrayRef<unsigned> LiveIns,
                       llvm::ArrayRef<RegisterValue> RegisterInitialValues,
-                      llvm::ArrayRef<llvm::MCInst> Instructions,
+                      const FillFunction &Fill,
                       llvm::raw_pwrite_stream &AsmStream) {
   std::unique_ptr<llvm::LLVMContext> Context =
       std::make_unique<llvm::LLVMContext>();
@@ -171,29 +192,34 @@ void assembleToStream(const ExegesisTarget &ET,
   auto &Properties = MF.getProperties();
   Properties.set(llvm::MachineFunctionProperties::Property::NoVRegs);
   Properties.reset(llvm::MachineFunctionProperties::Property::IsSSA);
+  Properties.set(llvm::MachineFunctionProperties::Property::NoPHIs);
 
   for (const unsigned Reg : LiveIns)
     MF.getRegInfo().addLiveIn(Reg);
 
-  bool IsSnippetSetupComplete;
-  std::vector<llvm::MCInst> Code =
-      generateSnippetSetupCode(ET, TM->getMCSubtargetInfo(),
-                               RegisterInitialValues, IsSnippetSetupComplete);
+  std::vector<unsigned> RegistersSetUp;
+  for (const auto &InitValue : RegisterInitialValues) {
+    RegistersSetUp.push_back(InitValue.Register);
+  }
+  FunctionFiller Sink(MF, std::move(RegistersSetUp));
+  auto Entry = Sink.getEntry();
+  for (const unsigned Reg : LiveIns)
+    Entry.MBB->addLiveIn(Reg);
 
-  Code.insert(Code.end(), Instructions.begin(), Instructions.end());
+  const bool IsSnippetSetupComplete = generateSnippetSetupCode(
+      ET, TM->getMCSubtargetInfo(), RegisterInitialValues, Entry);
 
   // If the snippet setup is not complete, we disable liveliness tracking. This
   // means that we won't know what values are in the registers.
   if (!IsSnippetSetupComplete)
     Properties.reset(llvm::MachineFunctionProperties::Property::TracksLiveness);
 
+  Fill(Sink);
+
   // prologue/epilogue pass needs the reserved registers to be frozen, this
   // is usually done by the SelectionDAGISel pass.
   MF.getRegInfo().freezeReservedRegs(MF);
 
-  // Fill the MachineFunction from the instructions.
-  fillMachineFunction(MF, LiveIns, Code);
-
   // We create the pass manager, run the passes to populate AsmBuffer.
   llvm::MCContext &MCContext = MMI->getContext();
   llvm::legacy::PassManager PM;
diff --git a/llvm/tools/llvm-exegesis/lib/Assembler.h b/llvm/tools/llvm-exegesis/lib/Assembler.h
index ff1e8db1a0c..5937de9f053 100644
--- a/llvm/tools/llvm-exegesis/lib/Assembler.h
+++ b/llvm/tools/llvm-exegesis/lib/Assembler.h
@@ -40,6 +40,49 @@ class ExegesisTarget;
 // convention and target machine).
 llvm::BitVector getFunctionReservedRegs(const llvm::TargetMachine &TM);
 
+// Helper to fill in a basic block.
+class BasicBlockFiller {
+public:
+  BasicBlockFiller(MachineFunction &MF, MachineBasicBlock *MBB,
+                   const MCInstrInfo *MCII);
+
+  void addInstruction(const MCInst &Inst, const DebugLoc &DL = DebugLoc());
+  void addInstructions(ArrayRef<llvm::MCInst> Insts,
+                       const DebugLoc &DL = DebugLoc());
+
+  void addReturn(const DebugLoc &DL = DebugLoc());
+
+  MachineFunction &MF;
+  MachineBasicBlock *const MBB;
+  const MCInstrInfo *const MCII;
+};
+
+// Helper to fill in a function.
+class FunctionFiller {
+public:
+  FunctionFiller(MachineFunction &MF, std::vector<unsigned> RegistersSetUp);
+
+  // Adds a basic block to the function.
+  BasicBlockFiller addBasicBlock();
+
+  // Returns the function entry point.
+  BasicBlockFiller getEntry() { return Entry; }
+
+  MachineFunction &MF;
+  const MCInstrInfo *const MCII;
+
+  // Returns the set of registers in the snippet setup code.
+  ArrayRef<unsigned> getRegistersSetUp() const;
+
+private:
+  BasicBlockFiller Entry;
+  // The set of registers that are set up in the basic block.
+  std::vector<unsigned> RegistersSetUp;
+};
+
+// A callback that fills a function.
+using FillFunction = std::function<void(FunctionFiller &)>;
+
 // Creates a temporary `void foo(char*)` function containing the provided
 // Instructions. Runs a set of llvm Passes to provide correct prologue and
 // epilogue. Once the MachineFunction is ready, it is assembled for TM to
@@ -48,7 +91,7 @@ void assembleToStream(const ExegesisTarget &ET,
                       std::unique_ptr<llvm::LLVMTargetMachine> TM,
                       llvm::ArrayRef<unsigned> LiveIns,
                       llvm::ArrayRef<RegisterValue> RegisterInitialValues,
-                      llvm::ArrayRef<llvm::MCInst> Instructions,
+                      const FillFunction &Fill,
                       llvm::raw_pwrite_stream &AsmStream);
 
 // Creates an ObjectFile in the format understood by the host.
@@ -81,6 +124,12 @@ struct ExecutableFunction {
   llvm::StringRef FunctionBytes;
 };
 
+// Creates a void(int8*) MachineFunction.
+llvm::MachineFunction &
+createVoidVoidPtrMachineFunction(llvm::StringRef FunctionID,
+                                 llvm::Module *Module,
+                                 llvm::MachineModuleInfo *MMI);
+
 } // namespace exegesis
 } // namespace llvm
 
diff --git a/llvm/tools/llvm-exegesis/lib/BenchmarkResult.h b/llvm/tools/llvm-exegesis/lib/BenchmarkResult.h
index e54e8c2a8da..17ffd0a8c87 100644
--- a/llvm/tools/llvm-exegesis/lib/BenchmarkResult.h
+++ b/llvm/tools/llvm-exegesis/lib/BenchmarkResult.h
@@ -66,6 +66,8 @@ struct InstructionBenchmark {
   // The number of instructions inside the repeated snippet. For example, if a
   // snippet of 3 instructions is repeated 4 times, this is 12.
   int NumRepetitions = 0;
+  enum RepetitionModeE { Duplicate, Loop };
+  RepetitionModeE RepetitionMode;
   // Note that measurements are per instruction.
   std::vector<BenchmarkMeasure> Measurements;
   std::string Error;
diff --git a/llvm/tools/llvm-exegesis/lib/BenchmarkRunner.cpp b/llvm/tools/llvm-exegesis/lib/BenchmarkRunner.cpp
index 38abf3eefa3..a08983817ad 100644
--- a/llvm/tools/llvm-exegesis/lib/BenchmarkRunner.cpp
+++ b/llvm/tools/llvm-exegesis/lib/BenchmarkRunner.cpp
@@ -33,17 +33,6 @@ BenchmarkRunner::BenchmarkRunner(const LLVMState &State,
 
 BenchmarkRunner::~BenchmarkRunner() = default;
 
-// Repeat the snippet until there are at least MinInstructions in the resulting
-// code.
-static std::vector<llvm::MCInst>
-GenerateInstructions(const BenchmarkCode &BC, const size_t MinInstructions) {
-  if (BC.Instructions.empty())
-    return {};
-  std::vector<llvm::MCInst> Code = BC.Instructions;
-  for (int I = 0; Code.size() < MinInstructions; ++I)
-    Code.push_back(BC.Instructions[I % BC.Instructions.size()]);
-  return Code;
-}
 
 namespace {
 class FunctionExecutorImpl : public BenchmarkRunner::FunctionExecutor {
@@ -95,10 +84,9 @@ private:
 };
 } // namespace
 
-InstructionBenchmark
-BenchmarkRunner::runConfiguration(const BenchmarkCode &BC,
-                                  unsigned NumRepetitions,
-                                  bool DumpObjectToDisk) const {
+InstructionBenchmark BenchmarkRunner::runConfiguration(
+    const BenchmarkCode &BC, unsigned NumRepetitions,
+    const SnippetRepetitor &Repetitor, bool DumpObjectToDisk) const {
   InstructionBenchmark InstrBenchmark;
   InstrBenchmark.Mode = Mode;
   InstrBenchmark.CpuName = State.getTargetMachine().getTargetCPU();
@@ -119,9 +107,10 @@ BenchmarkRunner::runConfiguration(const BenchmarkCode &BC,
   {
     llvm::SmallString<0> Buffer;
     llvm::raw_svector_ostream OS(Buffer);
-    assembleToStream(State.getExegesisTarget(), State.createTargetMachine(),
-                     BC.LiveIns, BC.RegisterInitialValues,
-                     GenerateInstructions(BC, kMinInstructionsForSnippet), OS);
+    assembleToStream(
+        State.getExegesisTarget(), State.createTargetMachine(), BC.LiveIns,
+        BC.RegisterInitialValues,
+        Repetitor.Repeat(BC.Instructions, kMinInstructionsForSnippet), OS);
     const ExecutableFunction EF(State.createTargetMachine(),
                                 getObjectFromBuffer(OS.str()));
     const auto FnBytes = EF.getFunctionBytes();
@@ -130,11 +119,12 @@ BenchmarkRunner::runConfiguration(const BenchmarkCode &BC,
 
   // Assemble NumRepetitions instructions repetitions of the snippet for
   // measurements.
-  const auto Code = GenerateInstructions(BC, InstrBenchmark.NumRepetitions);
+  const auto Filler =
+      Repetitor.Repeat(BC.Instructions, InstrBenchmark.NumRepetitions);
 
   llvm::object::OwningBinary<llvm::object::ObjectFile> ObjectFile;
   if (DumpObjectToDisk) {
-    auto ObjectFilePath = writeObjectFile(BC, Code);
+    auto ObjectFilePath = writeObjectFile(BC, Filler);
     if (llvm::Error E = ObjectFilePath.takeError()) {
       InstrBenchmark.Error = llvm::toString(std::move(E));
       return InstrBenchmark;
@@ -146,7 +136,7 @@ BenchmarkRunner::runConfiguration(const BenchmarkCode &BC,
     llvm::SmallString<0> Buffer;
     llvm::raw_svector_ostream OS(Buffer);
     assembleToStream(State.getExegesisTarget(), State.createTargetMachine(),
-                     BC.LiveIns, BC.RegisterInitialValues, Code, OS);
+                     BC.LiveIns, BC.RegisterInitialValues, Filler, OS);
     ObjectFile = getObjectFromBuffer(OS.str());
   }
 
@@ -172,7 +162,7 @@ BenchmarkRunner::runConfiguration(const BenchmarkCode &BC,
 
 llvm::Expected<std::string>
 BenchmarkRunner::writeObjectFile(const BenchmarkCode &BC,
-                                 llvm::ArrayRef<llvm::MCInst> Code) const {
+                                 const FillFunction &FillFunction) const {
   int ResultFD = 0;
   llvm::SmallString<256> ResultPath;
   if (llvm::Error E = llvm::errorCodeToError(llvm::sys::fs::createTemporaryFile(
@@ -180,7 +170,7 @@ BenchmarkRunner::writeObjectFile(const BenchmarkCode &BC,
     return std::move(E);
   llvm::raw_fd_ostream OFS(ResultFD, true /*ShouldClose*/);
   assembleToStream(State.getExegesisTarget(), State.createTargetMachine(),
-                   BC.LiveIns, BC.RegisterInitialValues, Code, OFS);
+                   BC.LiveIns, BC.RegisterInitialValues, FillFunction, OFS);
   return ResultPath.str();
 }
 
diff --git a/llvm/tools/llvm-exegesis/lib/BenchmarkRunner.h b/llvm/tools/llvm-exegesis/lib/BenchmarkRunner.h
index 29e98619301..13227a7bfb3 100644
--- a/llvm/tools/llvm-exegesis/lib/BenchmarkRunner.h
+++ b/llvm/tools/llvm-exegesis/lib/BenchmarkRunner.h
@@ -20,6 +20,7 @@
 #include "BenchmarkResult.h"
 #include "LlvmState.h"
 #include "MCInstrDescView.h"
+#include "SnippetRepetitor.h"
 #include "llvm/MC/MCInst.h"
 #include "llvm/Support/Error.h"
 #include <cstdlib>
@@ -46,6 +47,7 @@ public:
 
   InstructionBenchmark runConfiguration(const BenchmarkCode &Configuration,
                                         unsigned NumRepetitions,
+                                        const SnippetRepetitor &Repetitor,
                                         bool DumpObjectToDisk) const;
 
   // Scratch space to run instructions that touch memory.
@@ -84,7 +86,7 @@ private:
 
   llvm::Expected<std::string>
   writeObjectFile(const BenchmarkCode &Configuration,
-                  llvm::ArrayRef<llvm::MCInst> Code) const;
+                  const FillFunction &Fill) const;
 
   const std::unique_ptr<ScratchSpace> Scratch;
 };
diff --git a/llvm/tools/llvm-exegesis/lib/CMakeLists.txt b/llvm/tools/llvm-exegesis/lib/CMakeLists.txt
index ef1a411d9ef..1cd6c8f8418 100644
--- a/llvm/tools/llvm-exegesis/lib/CMakeLists.txt
+++ b/llvm/tools/llvm-exegesis/lib/CMakeLists.txt
@@ -31,6 +31,7 @@ add_library(LLVMExegesis
   RegisterValue.cpp
   SchedClassResolution.cpp
   SnippetGenerator.cpp
+  SnippetRepetitor.cpp
   Target.cpp
   Uops.cpp
   )
diff --git a/llvm/tools/llvm-exegesis/lib/SnippetGenerator.cpp b/llvm/tools/llvm-exegesis/lib/SnippetGenerator.cpp
index cd93589d34d..21fbe583636 100644
--- a/llvm/tools/llvm-exegesis/lib/SnippetGenerator.cpp
+++ b/llvm/tools/llvm-exegesis/lib/SnippetGenerator.cpp
@@ -37,9 +37,10 @@ SnippetGenerator::SnippetGenerator(const LLVMState &State) : State(State) {}
 SnippetGenerator::~SnippetGenerator() = default;
 
 llvm::Expected<std::vector<BenchmarkCode>>
-SnippetGenerator::generateConfigurations(const Instruction &Instr) const {
+SnippetGenerator::generateConfigurations(
+    const Instruction &Instr, const llvm::BitVector &ExtraForbiddenRegs) const {
   llvm::BitVector ForbiddenRegs = State.getRATC().reservedRegisters();
-
+  ForbiddenRegs |= ExtraForbiddenRegs;
   // If the instruction has memory registers, prevent the generator from
   // using the scratch register and its aliasing registers.
   if (Instr.hasMemoryOperands()) {
diff --git a/llvm/tools/llvm-exegesis/lib/SnippetGenerator.h b/llvm/tools/llvm-exegesis/lib/SnippetGenerator.h
index 966da2a8c41..c2ea1c12473 100644
--- a/llvm/tools/llvm-exegesis/lib/SnippetGenerator.h
+++ b/llvm/tools/llvm-exegesis/lib/SnippetGenerator.h
@@ -57,7 +57,8 @@ public:
 
   // Calls generateCodeTemplate and expands it into one or more BenchmarkCode.
   llvm::Expected<std::vector<BenchmarkCode>>
-  generateConfigurations(const Instruction &Instr) const;
+  generateConfigurations(const Instruction &Instr,
+                         const llvm::BitVector &ExtraForbiddenRegs) const;
 
   // Given a snippet, computes which registers the setup code needs to define.
   std::vector<RegisterValue> computeRegisterInitialValues(
diff --git a/llvm/tools/llvm-exegesis/lib/SnippetRepetitor.cpp b/llvm/tools/llvm-exegesis/lib/SnippetRepetitor.cpp
new file mode 100644
index 00000000000..551ee5341a0
--- /dev/null
+++ b/llvm/tools/llvm-exegesis/lib/SnippetRepetitor.cpp
@@ -0,0 +1,116 @@
+//===-- SnippetRepetitor.cpp ------------------------------------*- C++ -*-===//
+//
+// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
+// See https://llvm.org/LICENSE.txt for license information.
+// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+//
+//===----------------------------------------------------------------------===//
+
+#include <array>
+#include <string>
+
+#include "SnippetRepetitor.h"
+#include "Target.h"
+#include "llvm/CodeGen/TargetInstrInfo.h"
+#include "llvm/CodeGen/TargetSubtargetInfo.h"
+
+namespace llvm {
+namespace exegesis {
+namespace {
+
+class DuplicateSnippetRepetitor : public SnippetRepetitor {
+public:
+  using SnippetRepetitor::SnippetRepetitor;
+
+  // Repeats the snippet until there are at least MinInstructions in the
+  // resulting code.
+  FillFunction Repeat(ArrayRef<MCInst> Instructions,
+                      unsigned MinInstructions) const override {
+    return [Instructions, MinInstructions](FunctionFiller &Filler) {
+      auto Entry = Filler.getEntry();
+      if (!Instructions.empty()) {
+        // Add the whole snippet at least once.
+        Entry.addInstructions(Instructions);
+        for (unsigned I = Instructions.size(); I < MinInstructions; ++I) {
+          Entry.addInstruction(Instructions[I % Instructions.size()]);
+        }
+      }
+      Entry.addReturn();
+    };
+  }
+
+  BitVector getReservedRegs() const override {
+    // We're using no additional registers.
+    return State.getRATC().emptyRegisters();
+  }
+};
+
+class LoopSnippetRepetitor : public SnippetRepetitor {
+public:
+  explicit LoopSnippetRepetitor(const LLVMState &State)
+      : SnippetRepetitor(State),
+        LoopCounter(State.getExegesisTarget().getLoopCounterRegister(
+            State.getTargetMachine().getTargetTriple())) {}
+
+  // Loop over the snippet ceil(MinInstructions / Instructions.Size()) times.
+  FillFunction Repeat(ArrayRef<MCInst> Instructions,
+                      unsigned MinInstructions) const override {
+    return [this, Instructions, MinInstructions](FunctionFiller &Filler) {
+      const auto &ET = State.getExegesisTarget();
+      auto Entry = Filler.getEntry();
+      auto Loop = Filler.addBasicBlock();
+      auto Exit = Filler.addBasicBlock();
+
+      // Set loop counter to the right value:
+      const APInt LoopCount(32, (MinInstructions + Instructions.size() - 1) /
+                                    Instructions.size());
+      for (const MCInst &Inst :
+           ET.setRegTo(State.getSubtargetInfo(), LoopCounter, LoopCount))
+        Entry.addInstruction(Inst);
+
+      // Set up the loop basic block.
+      Entry.MBB->addSuccessor(Loop.MBB, llvm::BranchProbability::getOne());
+      Loop.MBB->addSuccessor(Loop.MBB, llvm::BranchProbability::getOne());
+      // The live ins are: the loop counter, the registers that were setup by
+      // the entry block, and entry block live ins.
+      Loop.MBB->addLiveIn(LoopCounter);
+      for (unsigned Reg : Filler.getRegistersSetUp())
+        Loop.MBB->addLiveIn(Reg);
+      for (const auto &LiveIn : Entry.MBB->liveins())
+        Loop.MBB->addLiveIn(LiveIn);
+      Loop.addInstructions(Instructions);
+      ET.decrementLoopCounterAndLoop(*Loop.MBB, State.getInstrInfo());
+
+      // Set up the exit basic block.
+      Loop.MBB->addSuccessor(Exit.MBB, llvm::BranchProbability::getZero());
+      Exit.addReturn();
+    };
+  }
+
+  BitVector getReservedRegs() const override {
+    // We're using a single loop counter, but we have to reserve all aliasing
+    // registers.
+    return State.getRATC().getRegister(LoopCounter).aliasedBits();
+  }
+
+private:
+  const unsigned LoopCounter;
+};
+
+} // namespace
+
+SnippetRepetitor::~SnippetRepetitor() {}
+
+std::unique_ptr<const SnippetRepetitor>
+SnippetRepetitor::Create(InstructionBenchmark::RepetitionModeE Mode,
+                         const LLVMState &State) {
+  switch (Mode) {
+  case InstructionBenchmark::Duplicate:
+    return std::make_unique<DuplicateSnippetRepetitor>(State);
+  case InstructionBenchmark::Loop:
+    return std::make_unique<LoopSnippetRepetitor>(State);
+  }
+}
+
+} // namespace exegesis
+} // namespace llvm
diff --git a/llvm/tools/llvm-exegesis/lib/SnippetRepetitor.h b/llvm/tools/llvm-exegesis/lib/SnippetRepetitor.h
new file mode 100644
index 00000000000..038dcac1e16
--- /dev/null
+++ b/llvm/tools/llvm-exegesis/lib/SnippetRepetitor.h
@@ -0,0 +1,53 @@
+//===-- SnippetRepetitor.h --------------------------------------*- C++ -*-===//
+//
+// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
+// See https://llvm.org/LICENSE.txt for license information.
+// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+//
+//===----------------------------------------------------------------------===//
+///
+/// \file
+/// Defines helpers to fill functions with repetitions of a snippet.
+///
+//===----------------------------------------------------------------------===//
+
+#ifndef LLVM_TOOLS_LLVM_EXEGESIS_FUNCTIONFILLER_H
+#define LLVM_TOOLS_LLVM_EXEGESIS_FUNCTIONFILLER_H
+
+#include "Assembler.h"
+#include "BenchmarkResult.h"
+#include "LlvmState.h"
+#include "llvm/ADT/BitVector.h"
+#include "llvm/CodeGen/MachineFunction.h"
+#include "llvm/MC/MCInst.h"
+#include "llvm/MC/MCInstrInfo.h"
+#include "llvm/Object/Binary.h"
+
+namespace llvm {
+namespace exegesis {
+
+class SnippetRepetitor {
+public:
+  static std::unique_ptr<const SnippetRepetitor>
+  Create(InstructionBenchmark::RepetitionModeE Mode, const LLVMState &State);
+
+  virtual ~SnippetRepetitor();
+
+  // Returns the set of registers that are reserved by the repetitor.
+  virtual BitVector getReservedRegs() const = 0;
+
+  // Returns a functor that repeats `Instructions` so that the function executes
+  // at least `MinInstructions` instructions.
+  virtual FillFunction Repeat(ArrayRef<MCInst> Instructions,
+                              unsigned MinInstructions) const = 0;
+
+  explicit SnippetRepetitor(const LLVMState &State) : State(State) {}
+
+protected:
+  const LLVMState &State;
+};
+
+} // namespace exegesis
+} // namespace llvm
+
+#endif
diff --git a/llvm/tools/llvm-exegesis/lib/Target.h b/llvm/tools/llvm-exegesis/lib/Target.h
index 0c574780865..4b0c9d17dd7 100644
--- a/llvm/tools/llvm-exegesis/lib/Target.h
+++ b/llvm/tools/llvm-exegesis/lib/Target.h
@@ -85,11 +85,22 @@ public:
   // Fills memory operands with references to the address at [Reg] + Offset.
   virtual void fillMemoryOperands(InstructionTemplate &IT, unsigned Reg,
                                   unsigned Offset) const {
-
     llvm_unreachable(
         "fillMemoryOperands() requires getScratchMemoryRegister() > 0");
   }
 
+  // Returns a counter usable as a loop counter.
+  virtual unsigned getLoopCounterRegister(const llvm::Triple &) const {
+    return 0;
+  }
+
+  // Adds the code to decrement the loop counter and
+  virtual void decrementLoopCounterAndLoop(MachineBasicBlock &MBB,
+                                           const llvm::MCInstrInfo &MII) const {
+    llvm_unreachable("decrementLoopCounterAndBranch() requires "
+                     "getLoopCounterRegister() > 0");
+  }
+
   // Returns a list of unavailable registers.
   // Targets can use this to prevent some registers to be automatically selected
   // for use in snippets.
diff --git a/llvm/tools/llvm-exegesis/lib/X86/Target.cpp b/llvm/tools/llvm-exegesis/lib/X86/Target.cpp
index 45e36117693..ca86ebcf2a2 100644
--- a/llvm/tools/llvm-exegesis/lib/X86/Target.cpp
+++ b/llvm/tools/llvm-exegesis/lib/X86/Target.cpp
@@ -436,6 +436,8 @@ private:
 
   unsigned getScratchMemoryRegister(const llvm::Triple &TT) const override;
 
+  unsigned getLoopCounterRegister(const llvm::Triple &) const override;
+
   unsigned getMaxMemoryAccessSize() const override { return 64; }
 
   void randomizeMCOperand(const Instruction &Instr, const Variable &Var,
@@ -445,6 +447,9 @@ private:
   void fillMemoryOperands(InstructionTemplate &IT, unsigned Reg,
                           unsigned Offset) const override;
 
+  void decrementLoopCounterAndLoop(MachineBasicBlock &MBB,
+                                   const llvm::MCInstrInfo &MII) const override;
+
   std::vector<llvm::MCInst> setRegTo(const llvm::MCSubtargetInfo &STI,
                                      unsigned Reg,
                                      const llvm::APInt &Value) const override;
@@ -476,6 +481,12 @@ private:
 // prefix.
 const unsigned ExegesisX86Target::kUnavailableRegisters[4] = {X86::AH, X86::BH,
                                                               X86::CH, X86::DH};
+
+// We're using one of R8-R15 because these registers are never hardcoded in
+// instructions (e.g. MOVS writes to EDI, ESI, EDX), so they have less
+// conflicts.
+constexpr const unsigned kLoopCounterReg = X86::R8;
+
 } // namespace
 
 void ExegesisX86Target::addTargetSpecificPasses(
@@ -494,6 +505,14 @@ ExegesisX86Target::getScratchMemoryRegister(const llvm::Triple &TT) const {
   return TT.isOSWindows() ? llvm::X86::RCX : llvm::X86::RDI;
 }
 
+unsigned
+ExegesisX86Target::getLoopCounterRegister(const llvm::Triple &TT) const {
+  if (!TT.isArch64Bit()) {
+    return 0;
+  }
+  return kLoopCounterReg;
+}
+
 void ExegesisX86Target::randomizeMCOperand(
     const Instruction &Instr, const Variable &Var,
     llvm::MCOperand &AssignedValue,
@@ -538,6 +557,17 @@ void ExegesisX86Target::fillMemoryOperands(InstructionTemplate &IT,
   SetOp(MemOpIdx + 4, MCOperand::createReg(0));      // Segment
 }
 
+void ExegesisX86Target::decrementLoopCounterAndLoop(
+    MachineBasicBlock &MBB, const llvm::MCInstrInfo &MII) const {
+  BuildMI(&MBB, DebugLoc(), MII.get(X86::ADD64ri8))
+      .addDef(kLoopCounterReg)
+      .addUse(kLoopCounterReg)
+      .addImm(-1);
+  BuildMI(&MBB, DebugLoc(), MII.get(X86::JCC_1))
+      .addMBB(&MBB)
+      .addImm(X86::COND_NE);
+}
+
 std::vector<llvm::MCInst>
 ExegesisX86Target::setRegTo(const llvm::MCSubtargetInfo &STI, unsigned Reg,
                             const llvm::APInt &Value) const {
diff --git a/llvm/tools/llvm-exegesis/llvm-exegesis.cpp b/llvm/tools/llvm-exegesis/llvm-exegesis.cpp
index 91bb30703a9..6c4567870d2 100644
--- a/llvm/tools/llvm-exegesis/llvm-exegesis.cpp
+++ b/llvm/tools/llvm-exegesis/llvm-exegesis.cpp
@@ -17,6 +17,7 @@
 #include "lib/Clustering.h"
 #include "lib/LlvmState.h"
 #include "lib/PerfHelper.h"
+#include "lib/SnippetRepetitor.h"
 #include "lib/Target.h"
 #include "lib/TargetSelect.h"
 #include "llvm/ADT/StringExtras.h"
@@ -80,6 +81,14 @@ static cl::opt<exegesis::InstructionBenchmark::ModeE> BenchmarkMode(
                clEnumValN(exegesis::InstructionBenchmark::Unknown, "analysis",
                           "Analysis")));
 
+static cl::opt<exegesis::InstructionBenchmark::RepetitionModeE> RepetitionMode(
+    "repetition-mode", cl::desc("how to repeat the instruction snippet"),
+    cl::cat(BenchmarkOptions),
+    cl::values(clEnumValN(exegesis::InstructionBenchmark::Duplicate,
+                          "duplicate", "Duplicate the snippet"),
+               clEnumValN(exegesis::InstructionBenchmark::Loop, "loop",
+                          "Loop over the snippet")));
+
 static cl::opt<unsigned>
     NumRepetitions("num-repetitions",
                    cl::desc("number of time to repeat the asm snippet"),
@@ -192,7 +201,8 @@ getOpcodesOrDie(const llvm::MCInstrInfo &MCInstrInfo) {
 
 // Generates code snippets for opcode `Opcode`.
 static llvm::Expected<std::vector<BenchmarkCode>>
-generateSnippets(const LLVMState &State, unsigned Opcode) {
+generateSnippets(const LLVMState &State, unsigned Opcode,
+                 const llvm::BitVector &ForbiddenRegs) {
   const Instruction &Instr = State.getIC().getInstr(Opcode);
   const llvm::MCInstrDesc &InstrDesc = *Instr.Description;
   // Ignore instructions that we cannot run.
@@ -209,7 +219,7 @@ generateSnippets(const LLVMState &State, unsigned Opcode) {
       State.getExegesisTarget().createSnippetGenerator(BenchmarkMode, State);
   if (!Generator)
     llvm::report_fatal_error("cannot create snippet generator");
-  return Generator->generateConfigurations(Instr);
+  return Generator->generateConfigurations(Instr, ForbiddenRegs);
 }
 
 namespace {
@@ -372,6 +382,8 @@ void benchmarkMain() {
   const LLVMState State(CpuName);
   const auto Opcodes = getOpcodesOrDie(State.getInstrInfo());
 
+  const auto Repetitor = SnippetRepetitor::Create(RepetitionMode, State);
+
   std::vector<BenchmarkCode> Configurations;
   if (!Opcodes.empty()) {
     for (const unsigned Opcode : Opcodes) {
@@ -383,7 +395,8 @@ void benchmarkMain() {
                      << ": ignoring instruction without sched class\n";
         continue;
       }
-      auto ConfigsForInstr = generateSnippets(State, Opcode);
+      auto ConfigsForInstr =
+          generateSnippets(State, Opcode, Repetitor->getReservedRegs());
       if (!ConfigsForInstr) {
         llvm::logAllUnhandledErrors(
             ConfigsForInstr.takeError(), llvm::errs(),
@@ -411,8 +424,8 @@ void benchmarkMain() {
     BenchmarkFile = "-";
 
   for (const BenchmarkCode &Conf : Configurations) {
-    InstructionBenchmark Result =
-        Runner->runConfiguration(Conf, NumRepetitions, DumpObjectToDisk);
+    InstructionBenchmark Result = Runner->runConfiguration(
+        Conf, NumRepetitions, *Repetitor, DumpObjectToDisk);
     ExitOnErr(Result.writeYaml(State, BenchmarkFile));
   }
   exegesis::pfm::pfmTerminate();