[llvm] [RISCV][PoC] Schedule RVV instructions with same type first (PR #95924)

[Michael Maitland via llvm-commits]

================
@@ -0,0 +1,118 @@
+//===- RISCVMachineScheduler.cpp - MI Scheduler for RISC-V ----------------===//
+//
+// 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 "RISCVMachineScheduler.h"
+#include "MCTargetDesc/RISCVBaseInfo.h"
+#include "MCTargetDesc/RISCVMCTargetDesc.h"
+#include "RISCVInstrInfo.h"
+#include "RISCVSubtarget.h"
+#include "llvm/CodeGen/MachineOperand.h"
+#include "llvm/CodeGen/MachineScheduler.h"
+#include "llvm/CodeGen/ScheduleDAG.h"
+#include "llvm/MC/MCInstrDesc.h"
+#include "llvm/Support/Debug.h"
+#include "llvm/TargetParser/RISCVTargetParser.h"
+
+using namespace llvm;
+
+#define DEBUG_TYPE "riscv-prera-sched-strategy"
+
+static cl::opt<bool> EnableScheduleSameVType(
+    "riscv-enable-schedule-same-vtype", cl::init(false), cl::Hidden,
+    cl::desc("Enable scheduling RVV instructions with same vtype first"));
+
+SUnit *RISCVPreRAMachineSchedStrategy::pickNode(bool &IsTopNode) {
+  auto FindPotentialRVVInstructionInQueue =
+      [&](SchedBoundary &Boundary, ReadyQueue Q, bool ShouldBeTop) -> SUnit * {
+    for (SUnit *SU : Q) {
+      if (SU->isScheduled)
+        continue;
+
+      MachineInstr *MI = SU->getInstr();
+      const MCInstrDesc &Desc = MI->getDesc();
+      if (RISCVII::hasSEWOp(Desc.TSFlags)) {
+        unsigned CurVSEW = MI->getOperand(RISCVII::getSEWOpNum(Desc)).getImm();
+        RISCVII::VLMUL CurVLMUL = RISCVII::getLMul(Desc.TSFlags);
+        // FIXME: We should consider vl and policy here.
+        if (CurVSEW == PrevVSEW && CurVLMUL == PrevVLMUL) {
+          IsTopNode = ShouldBeTop;
+          // Boundary.removeReady(SU);
+          if (SU->isTopReady())
+            Top.removeReady(SU);
+          if (SU->isBottomReady())
+            Bot.removeReady(SU);
+          LLVM_DEBUG(dbgs() << "Scheduling SU(" << SU->NodeNum << ") "
+                            << *SU->getInstr());
+          return SU;
+        }
+      }
+    }
+    return nullptr;
+  };
+
+  auto FindPotentialRVVInstruction = [&](SchedBoundary &Boundary,
+                                         bool ShouldBeTop) -> SUnit * {
+    if (SUnit *Available = FindPotentialRVVInstructionInQueue(
+            Boundary, Boundary.Available, ShouldBeTop))
+      return Available;
+    if (SUnit *Pending = FindPotentialRVVInstructionInQueue(
+            Boundary, Boundary.Pending, ShouldBeTop))
+      return Pending;
+    return nullptr;
+  };
+
+  if (EnableScheduleSameVType) {
+    if (RegionPolicy.OnlyBottomUp) {
+      if (SUnit *SU = FindPotentialRVVInstruction(Bot, false))
+        return SU;
+    } else if (RegionPolicy.OnlyTopDown) {
+      if (SUnit *SU = FindPotentialRVVInstruction(Top, true))
+        return SU;
+    } else {
+      if (SUnit *SU =
----------------
michaelmaitland wrote:

I gave some more thought on this. I think we can reuse from `GenericScheduler::pickNode` and we should instead be overriding `pickNodeFromQueue`. `pickNodeFromQueue` is where the real picking of the node happens. `GenericScheduler::pickNode` just gets the correct queue based on direction and passes it to `pickNodeFromQueue`. Something like this:

```
SUnit *RISCVPreRAMachineSchedStrategy::pickNodeFromQueue(SchedBoundary &Zone,
                                                         const CandPolicy &ZonePolicy,
                                                         const RegPressureTracker &RPTracker,
                                                         SchedCandidate &Cand) {
  SchedCandidate RVVCand = FindRVVCandidate(Zone);
  GenericScheduler::pickNodeFromQueue(Zone, ZonePolicy, RPTracker, Cand);

  // Pass SchedBoundary only when comparing nodes from the same boundary.
  SchedBoundary *ZoneArg = Cand.AtTop == RVVCand.AtTop ? &Zone : nullptr;

  // TODO: we need to add our own heuristics here or inside an overriden
  // tryCandidate to make sure that we balance clustering RVV with same vtype
  // with the existing heuristics such as register pressure, latency, resource usage, etc.
  if (tryCandidate(RVVCand, Cand, ZoneArg)
    return RVVCand;
  return Cand;
}
```

https://github.com/llvm/llvm-project/pull/95924