[RISCV][MCA] Pick the correct VPseudo sched class for indexed memory operation (#128978)

It seems like we had been picking the wrong VPseudo scheduling class for
indexed memory operations in RISCVMCACustomBehavior: the VPseudo opcode
of indexed memory ops encode two EMULs, one for index and the other for
data. However, in RISCVInversePseudoTable, we're only able to look up
against one of them, yielding an incorrect VPseudo opcode with the wrong
data EEW (index EEW is encoded in the opcode). Since scheduling classes
for indexed memory ops uses data EMUL / EEW in their scheduling class,
we would eventually fetch the wrong scheduling classes with faulty data
EEW.

This patch fixes this issue by deducting the correct index EMUL with
LMUL (data EMUL), SEW (data EEW), and index EEW. With these parameters
we can thus fetch the correct VPseudo opcode with `getVLXPseudo` /
`getVLXSEGPseudo` and friends.

The new search table, RISCVBaseVXMemOpTable, is created to extract the
NF and index EEW info from MC opcode. Otherwise we need to write a
gigantic switch statement to decode this info.

Author: mshockwave

llvm/lib/Target/RISCV/MCA/RISCVCustomBehaviour.cpp

@@ -14,12 +14,26 @@
 #include "RISCVCustomBehaviour.h"
 #include "MCTargetDesc/RISCVMCTargetDesc.h"
 #include "RISCV.h"
+#include "RISCVISelDAGToDAG.h"
 #include "TargetInfo/RISCVTargetInfo.h"
 #include "llvm/MC/TargetRegistry.h"
 #include "llvm/Support/Debug.h"
 
 #define DEBUG_TYPE "llvm-mca-riscv-custombehaviour"
 
+namespace llvm::RISCV {
+struct VXMemOpInfo {
+  unsigned Log2IdxEEW : 3;
+  unsigned IsOrdered : 1;
+  unsigned IsStore : 1;
+  unsigned NF : 4;
+  unsigned BaseInstr;
+};
+
+#define GET_RISCVBaseVXMemOpTable_IMPL
+#include "RISCVGenSearchableTables.inc"
+} // namespace llvm::RISCV
+
 namespace llvm {
 namespace mca {
 
@@ -247,21 +261,58 @@ unsigned RISCVInstrumentManager::getSchedClassID(
   // and SEW, or (Opcode, LMUL, 0) if does not depend on SEW.
   uint8_t SEW = SI ? SI->getSEW() : 0;
 
-  const RISCVVInversePseudosTable::PseudoInfo *RVV = nullptr;
-  if (opcodeHasEEWAndEMULInfo(Opcode)) {
+  std::optional<unsigned> VPOpcode;
+  if (const auto *VXMO = RISCV::getVXMemOpInfo(Opcode)) {
+    // Calculate the expected index EMUL. For indexed operations,
+    // the DataEEW and DataEMUL are equal to SEW and LMUL, respectively.
+    unsigned IndexEMUL = ((1 << VXMO->Log2IdxEEW) * LMUL) / SEW;
+
+    if (!VXMO->NF) {
+      // Indexed Load / Store.
+      if (VXMO->IsStore) {
+        if (const auto *VXP = RISCV::getVSXPseudo(
+                /*Masked=*/0, VXMO->IsOrdered, VXMO->Log2IdxEEW, LMUL,
+                IndexEMUL))
+          VPOpcode = VXP->Pseudo;
+      } else {
+        if (const auto *VXP = RISCV::getVLXPseudo(
+                /*Masked=*/0, VXMO->IsOrdered, VXMO->Log2IdxEEW, LMUL,
+                IndexEMUL))
+          VPOpcode = VXP->Pseudo;
+      }
+    } else {
+      // Segmented Indexed Load / Store.
+      if (VXMO->IsStore) {
+        if (const auto *VXP =
+                RISCV::getVSXSEGPseudo(VXMO->NF, /*Masked=*/0, VXMO->IsOrdered,
+                                       VXMO->Log2IdxEEW, LMUL, IndexEMUL))
+          VPOpcode = VXP->Pseudo;
+      } else {
+        if (const auto *VXP =
+                RISCV::getVLXSEGPseudo(VXMO->NF, /*Masked=*/0, VXMO->IsOrdered,
+                                       VXMO->Log2IdxEEW, LMUL, IndexEMUL))
+          VPOpcode = VXP->Pseudo;
+      }
+    }
+  } else if (opcodeHasEEWAndEMULInfo(Opcode)) {
     RISCVVType::VLMUL VLMUL = static_cast<RISCVVType::VLMUL>(LMUL);
     auto [EEW, EMUL] = getEEWAndEMUL(Opcode, VLMUL, SEW);
-    RVV = RISCVVInversePseudosTable::getBaseInfo(Opcode, EMUL, EEW);
+    if (const auto *RVV =
+            RISCVVInversePseudosTable::getBaseInfo(Opcode, EMUL, EEW))
+      VPOpcode = RVV->Pseudo;
   } else {
     // Check if it depends on LMUL and SEW
-    RVV = RISCVVInversePseudosTable::getBaseInfo(Opcode, LMUL, SEW);
+    const auto *RVV = RISCVVInversePseudosTable::getBaseInfo(Opcode, LMUL, SEW);
     // Check if it depends only on LMUL
     if (!RVV)
       RVV = RISCVVInversePseudosTable::getBaseInfo(Opcode, LMUL, 0);
+
+    if (RVV)
+      VPOpcode = RVV->Pseudo;
   }
 
   // Not a RVV instr
-  if (!RVV) {
+  if (!VPOpcode) {
     LLVM_DEBUG(
         dbgs() << "RVCB: Could not find PseudoInstruction for Opcode "
                << MCII.getName(Opcode)
@@ -277,8 +328,8 @@ unsigned RISCVInstrumentManager::getSchedClassID(
                     << MCII.getName(Opcode) << ", LMUL=" << LI->getData()
                     << ", SEW=" << (SI ? SI->getData() : "Unspecified")
                     << ". Overriding original SchedClassID=" << SchedClassID
-                    << " with " << MCII.getName(RVV->Pseudo) << '\n');
-  return MCII.get(RVV->Pseudo).getSchedClass();
+                    << " with " << MCII.getName(*VPOpcode) << '\n');
+  return MCII.get(*VPOpcode).getSchedClass();
 }
 
 } // namespace mca

llvm/lib/Target/RISCV/RISCVInstrInfoV.td

@@ -316,6 +316,22 @@ class VSXSEGSched<int nf, int eew, bit isOrdered, string emul,
 class VSXSEGSchedMC<int nf, int eew, bit isOrdered>:
   VSXSEGSched<nf, eew, isOrdered, "WorstCase", forceMasked=1>;
 
+class RISCVVXMemOpMC<bits<3> E, bit Ordered, bit Store, bits<4> N = 0> {
+  bits<3> Log2EEW = E;
+  bits<1> IsOrdered = Ordered;
+  bits<1> IsStore = Store;
+  bits<4> NF = N;
+  Instruction BaseInstr = !cast<Instruction>(NAME);
+}
+
+def RISCVBaseVXMemOpTable : GenericTable {
+  let FilterClass = "RISCVVXMemOpMC";
+  let CppTypeName = "VXMemOpInfo";
+  let Fields = ["Log2EEW", "IsOrdered", "IsStore", "NF", "BaseInstr"];
+  let PrimaryKey = ["BaseInstr"];
+  let PrimaryKeyName = "getVXMemOpInfo";
+}
+
 //===----------------------------------------------------------------------===//
 // Instruction class templates
 //===----------------------------------------------------------------------===//
@@ -560,16 +576,20 @@ multiclass VIndexLoadStore<int eew> {
 
   def VLUXEI # eew # _V :
     VIndexedLoad<MOPLDIndexedUnord, w, "vluxei" # eew # ".v">,
+    RISCVVXMemOpMC<!logtwo(eew), Ordered=false, Store=false>,
     VLXSchedMC<eew, isOrdered=0>;
   def VLOXEI # eew # _V :
     VIndexedLoad<MOPLDIndexedOrder, w, "vloxei" # eew # ".v">,
+    RISCVVXMemOpMC<!logtwo(eew), Ordered=true, Store=false>,
     VLXSchedMC<eew, isOrdered=1>;
 
   def VSUXEI # eew # _V :
     VIndexedStore<MOPSTIndexedUnord, w, "vsuxei" # eew # ".v">,
+    RISCVVXMemOpMC<!logtwo(eew), Ordered=false, Store=true>,
     VSXSchedMC<eew, isOrdered=0>;
   def VSOXEI # eew # _V :
     VIndexedStore<MOPSTIndexedOrder, w, "vsoxei" # eew # ".v">,
+    RISCVVXMemOpMC<!logtwo(eew), Ordered=true, Store=true>,
     VSXSchedMC<eew, isOrdered=1>;
 }
 
@@ -1760,18 +1780,22 @@ let Predicates = [HasVInstructions] in {
       def VLUXSEG#nf#EI#eew#_V :
         VIndexedSegmentLoad<!add(nf, -1), MOPLDIndexedUnord, w,
                             "vluxseg"#nf#"ei"#eew#".v">,
+        RISCVVXMemOpMC<!logtwo(eew), Ordered=false, Store=false, N=nf>,
         VLXSEGSchedMC<nf, eew, isOrdered=0>;
       def VLOXSEG#nf#EI#eew#_V :
         VIndexedSegmentLoad<!add(nf, -1), MOPLDIndexedOrder, w,
                             "vloxseg"#nf#"ei"#eew#".v">,
+        RISCVVXMemOpMC<!logtwo(eew), Ordered=true, Store=false, N=nf>,
         VLXSEGSchedMC<nf, eew, isOrdered=1>;
       def VSUXSEG#nf#EI#eew#_V :
         VIndexedSegmentStore<!add(nf, -1), MOPSTIndexedUnord, w,
                              "vsuxseg"#nf#"ei"#eew#".v">,
+        RISCVVXMemOpMC<!logtwo(eew), Ordered=false, Store=true, N=nf>,
         VSXSEGSchedMC<nf, eew, isOrdered=0>;
       def VSOXSEG#nf#EI#eew#_V :
         VIndexedSegmentStore<!add(nf, -1), MOPSTIndexedOrder, w,
                              "vsoxseg"#nf#"ei"#eew#".v">,
+        RISCVVXMemOpMC<!logtwo(eew), Ordered=true, Store=true, N=nf>,
         VSXSEGSchedMC<nf, eew, isOrdered=1>;
     }
   }