fixup! [ModuloSchedule] Implement modulo variable expansion for pipelining

Author: ytmukai

llvm/include/llvm/CodeGen/TargetInstrInfo.h

@@ -765,19 +765,18 @@ class TargetInstrInfo : public MCInstrInfo {
     createTripCountGreaterCondition(int TC, MachineBasicBlock &MBB,
                                     SmallVectorImpl<MachineOperand> &Cond) = 0;
 
-    /// Create a condtion to determine if the remaining trip count represented
-    /// by the loop counter CounterReg is greater than TC. Some instructions
-    /// such as comparisons may be inserted at the bottom of MBB. CounterReg
-    /// must be accessible there.
+    /// Create a condition to determine if the remaining trip count for a phase
+    /// is greater than TC. Some instructions such as comparisons may be
+    /// inserted at the bottom of MBB. The all instructions expanded for the
+    /// phase must be inserted in MBB before calling this function. RegMap is
+    /// the map from the original registers to the expanded registers for the
+    /// phase.
     ///
-    /// The definition of the return value is the same as for the variant above.
-    virtual std::optional<bool>
-    createTripCountGreaterCondition(int TC, MachineBasicBlock &MBB,
-                                    SmallVectorImpl<MachineOperand> &Cond,
-                                    Register CounterReg) {
-      llvm_unreachable(
-          "Target didn't implement createTripCountGreaterCondition");
-    }
+    /// MBB can also be a predecessor of the prologue block. Then RegMap must be
+    /// empty and the compared value is the initial value of the trip count.
+    virtual void createRemainingIterationsGreaterCondition(
+        int TC, MachineBasicBlock &MBB, SmallVectorImpl<MachineOperand> &Cond,
+        DenseMap<unsigned, unsigned> RegMap) = 0;
 
     /// Modify the loop such that the trip count is
     /// OriginalTC + TripCountAdjust.
@@ -793,15 +792,9 @@ class TargetInstrInfo : public MCInstrInfo {
     /// valid; the loop has been removed.
     virtual void disposed() = 0;
 
-    /// Return the initial value of the loop counter.
-    virtual Register getCounterInitReg() {
-      llvm_unreachable("Target didn't implement getCounterInitReg");
-    }
-
-    /// Return the updated value of the loop counter in the original loop.
-    virtual Register getCounterUpdatedReg() {
-      llvm_unreachable("Target didn't implement getCounterUpdatedReg");
-    }
+    /// Return true if the target can expand pipelined schedule with modulo
+    /// variable expansion.
+    virtual bool isMVEExpanderSupported() = 0;
   };
 
   /// Analyze loop L, which must be a single-basic-block loop, and if the

llvm/lib/CodeGen/MachinePipeliner.cpp

@@ -681,9 +681,9 @@ void SwingSchedulerDAG::schedule() {
   if (ExperimentalCodeGen && NewInstrChanges.empty()) {
     PeelingModuloScheduleExpander MSE(MF, MS, &LIS);
     MSE.expand();
-  }
-  if (MVECodeGen && NewInstrChanges.empty() &&
-      ModuloScheduleExpanderMVE::canApply(Loop)) {
+  } else if (MVECodeGen && NewInstrChanges.empty() &&
+             LoopPipelinerInfo->isMVEExpanderSupported() &&
+             ModuloScheduleExpanderMVE::canApply(Loop)) {
     ModuloScheduleExpanderMVE MSE(MF, MS, LIS);
     MSE.expand();
   } else {

llvm/lib/CodeGen/ModuloSchedule.cpp

@@ -2262,9 +2262,8 @@ void ModuloScheduleExpanderMVE::generatePipelinedLoop() {
   Epilog->addSuccessor(NewExit);
 
   SmallVector<MachineOperand, 4> Cond;
-  LoopInfo->createTripCountGreaterCondition(
-      Schedule.getNumStages() + NumUnroll - 2, *Check, Cond,
-      LoopInfo->getCounterInitReg());
+  LoopInfo->createRemainingIterationsGreaterCondition(
+      Schedule.getNumStages() + NumUnroll - 2, *Check, Cond, ValueMapTy());
   TII->insertBranch(*Check, Prolog, NewPreheader, Cond, DebugLoc());
 
   // VRMaps map (prolog/kernel/epilog phase#, original register#) to new
@@ -2552,9 +2551,8 @@ void ModuloScheduleExpanderMVE::generateKernel(
 
   // If remaining trip count is greater than NumUnroll-1, loop continues
   SmallVector<MachineOperand, 4> Cond;
-  LoopInfo->createTripCountGreaterCondition(
-      NumUnroll - 1, *NewKernel, Cond,
-      KernelVRMap[NumUnroll - 1][LoopInfo->getCounterUpdatedReg()]);
+  LoopInfo->createRemainingIterationsGreaterCondition(
+      NumUnroll - 1, *NewKernel, Cond, KernelVRMap[NumUnroll - 1]);
   TII->insertBranch(*NewKernel, NewKernel, Epilog, Cond, DebugLoc());
 
   LLVM_DEBUG({
@@ -2591,11 +2589,13 @@ void ModuloScheduleExpanderMVE::generateEpilog(
     updateInstrUse(MI, StageNum, EpilogNum, EpilogVRMap, &KernelVRMap);
   }
 
-  // If there are remaining iterations, they are executed in the original loop
+  // If there are remaining iterations, they are executed in the original loop.
+  // Instructions related to loop control, such as loop counter comparison,
+  // are indicated by shouldIgnoreForPipelining() and are assumed to be placed
+  // in stage 0. Thus, the map is for the last one in the kernel.
   SmallVector<MachineOperand, 4> Cond;
-  LoopInfo->createTripCountGreaterCondition(
-      0, *Epilog, Cond,
-      KernelVRMap[NumUnroll - 1][LoopInfo->getCounterUpdatedReg()]);
+  LoopInfo->createRemainingIterationsGreaterCondition(
+      0, *Epilog, Cond, KernelVRMap[NumUnroll - 1]);
   TII->insertBranch(*Epilog, NewPreheader, NewExit, Cond, DebugLoc());
 
   LLVM_DEBUG({

llvm/lib/Target/ARM/ARMBaseInstrInfo.cpp

@@ -6823,11 +6823,20 @@ class ARMPipelinerLoopInfo : public TargetInstrInfo::PipelinerLoopInfo {
       llvm_unreachable("Unknown EndLoop");
   }
 
+  void createRemainingIterationsGreaterCondition(
+      int TC, MachineBasicBlock &MBB, SmallVectorImpl<MachineOperand> &Cond,
+      DenseMap<unsigned, unsigned> RegMap) override {
+    llvm_unreachable(
+        "Target didn't implement createRemainingIterationsGreaterCondition");
+  }
+
   void setPreheader(MachineBasicBlock *NewPreheader) override {}
 
   void adjustTripCount(int TripCountAdjust) override {}
 
   void disposed() override {}
+
+  bool isMVEExpanderSupported() override { return false; }
 };
 
 void ARMPipelinerLoopInfo::bumpCrossIterationPressure(RegPressureTracker &RPT,

llvm/lib/Target/Hexagon/HexagonInstrInfo.cpp

@@ -770,6 +770,13 @@ class HexagonPipelinerLoopInfo : public TargetInstrInfo::PipelinerLoopInfo {
     return TripCount > TC;
   }
 
+  void createRemainingIterationsGreaterCondition(
+      int TC, MachineBasicBlock &MBB, SmallVectorImpl<MachineOperand> &Cond,
+      DenseMap<unsigned, unsigned> RegMap) override {
+    llvm_unreachable(
+        "Target didn't implement createRemainingIterationsGreaterCondition");
+  }
+
   void setPreheader(MachineBasicBlock *NewPreheader) override {
     NewPreheader->splice(NewPreheader->getFirstTerminator(), Loop->getParent(),
                          Loop);
@@ -798,6 +805,8 @@ class HexagonPipelinerLoopInfo : public TargetInstrInfo::PipelinerLoopInfo {
   }
 
   void disposed() override { Loop->eraseFromParent(); }
+
+  bool isMVEExpanderSupported() override { return false; }
 };
 } // namespace
 

llvm/lib/Target/PowerPC/PPCInstrInfo.cpp

@@ -5439,6 +5439,13 @@ class PPCPipelinerLoopInfo : public TargetInstrInfo::PipelinerLoopInfo {
     return TripCount > TC;
   }
 
+  void createRemainingIterationsGreaterCondition(
+      int TC, MachineBasicBlock &MBB, SmallVectorImpl<MachineOperand> &Cond,
+      DenseMap<unsigned, unsigned> RegMap) override {
+    llvm_unreachable(
+        "Target didn't implement createRemainingIterationsGreaterCondition");
+  }
+
   void setPreheader(MachineBasicBlock *NewPreheader) override {
     // Do nothing. We want the LOOP setup instruction to stay in the *old*
     // preheader, so we can use BDZ in the prologs to adapt the loop trip count.
@@ -5463,6 +5470,8 @@ class PPCPipelinerLoopInfo : public TargetInstrInfo::PipelinerLoopInfo {
     // Ensure the loop setup instruction is deleted too.
     LoopCount->eraseFromParent();
   }
+
+  bool isMVEExpanderSupported() override { return false; }
 };
 } // namespace