[TableGen] Extend direct lookup to instruction values in generic tables. (#80486)

Currently, for some tables involving a single primary key field which is
integral and densely numbered, a direct lookup is generated rather than
a binary search. This patch extends the direct lookup function
generation to instructions, where the integral value corresponds to the
instruction's enum value.

While this isn't as common as for other tables, it does occur in at
least one downstream backend and one in-tree backend.

Added a unit test and minimally updated the documentation.

Author: nvjle

llvm/docs/TableGen/BackEnds.rst

@@ -817,7 +817,9 @@ The table entries in ``ATable`` are sorted in order by ``Val1``, and within
 each of those values, by ``Val2``. This allows a binary search of the table,
 which is performed in the lookup function by ``std::lower_bound``. The
 lookup function returns a reference to the found table entry, or the null
-pointer if no entry is found.
+pointer if no entry is found. If the table has a single primary key field
+which is integral and densely numbered, a direct lookup is generated rather
+than a binary search.
 
 This example includes a field whose type TableGen cannot deduce. The ``Kind``
 field uses the enumerated type ``CEnum`` defined above. To inform TableGen

llvm/test/TableGen/generic-tables-instruction.td

@@ -2,6 +2,10 @@
 // XFAIL: vg_leak
 
 include "llvm/TableGen/SearchableTable.td"
+include "llvm/Target/Target.td"
+
+def ArchInstrInfo : InstrInfo { }
+def Arch : Target { let InstructionSet = ArchInstrInfo; }
 
 // CHECK-LABEL: GET_InstrTable_IMPL
 // CHECK: constexpr MyInstr InstrTable[] = {
@@ -11,11 +15,20 @@ include "llvm/TableGen/SearchableTable.td"
 // CHECK:   { D, 0x8 },
 // CHECK: };
 
-class Instruction {
-  bit isPseudo = 0;
-}
+// A contiguous primary (Instruction) key should get a direct lookup instead of
+// binary search.
+// CHECK: const MyInstr *getCustomEncodingHelper(unsigned Opcode) {
+// CHECK:   if ((Opcode < B) ||
+// CHECK:       (Opcode > D))
+// CHECK:     return nullptr;
+// CHECK:   auto Table = ArrayRef(InstrTable);
+// CHECK:   size_t Idx = Opcode - B;
+// CHECK:   return &Table[Idx];
+
 
 class MyInstr<int op> : Instruction {
+  let OutOperandList = (outs);
+  let InOperandList = (ins);
   Instruction Opcode = !cast<Instruction>(NAME);
   bits<16> CustomEncoding = op;
 }
@@ -34,3 +47,45 @@ def InstrTable : GenericTable {
   let PrimaryKey = ["Opcode"];
   let PrimaryKeyName = "getCustomEncodingHelper";
 }
+
+
+// Non-contiguous instructions should get a binary search instead of direct
+// lookup.
+// CHECK: const MyInfoEntry *getTable2ByOpcode(unsigned Opcode) {
+// CHECK:   auto Idx = std::lower_bound(Table.begin(), Table.end(), Key,
+//
+// Verify contiguous check for SearchIndex.
+// const MyInfoEntry *getTable2ByValue(uint8_t Value) {
+// CHECK:   if ((Value < 0xB) ||
+// CHECK:      (Value > 0xD))
+// CHECK:    return nullptr;
+// CHECK:  auto Table = ArrayRef(Index);
+// CHECK:  size_t Idx = Value - 0xB;
+// CHECK:  return &InstrTable2[Table[Idx]._index];
+
+
+class MyInfoEntry<int V, string S> {
+  Instruction Opcode = !cast<Instruction>(NAME);
+  bits<4> Value = V;
+  string Name = S;
+}
+
+let OutOperandList = (outs), InOperandList = (ins) in {
+def W : Instruction, MyInfoEntry<12, "IW">;
+def X : Instruction;
+def Y : Instruction, MyInfoEntry<13, "IY">;
+def Z : Instruction, MyInfoEntry<11, "IZ">;
+}
+
+def InstrTable2 : GenericTable {
+  let FilterClass = "MyInfoEntry";
+  let Fields = ["Opcode", "Value", "Name"];
+
+  let PrimaryKey = ["Opcode"];
+  let PrimaryKeyName = "getTable2ByOpcode";
+}
+
+def getTable2ByValue : SearchIndex {
+  let Table = InstrTable2;
+  let Key = ["Value"];
+}

llvm/utils/TableGen/CodeGenTarget.cpp

@@ -538,6 +538,13 @@ void CodeGenTarget::ComputeInstrsByEnum() const {
         return std::make_tuple(!D1.getValueAsBit("isPseudo"), D1.getName()) <
                std::make_tuple(!D2.getValueAsBit("isPseudo"), D2.getName());
       });
+
+  // Build the instruction-to-int map using the same values emitted by
+  // InstrInfoEmitter::emitEnums.
+  assert(InstrToIntMap.empty());
+  unsigned Num = 0;
+  for (const CodeGenInstruction *Inst : InstrsByEnum)
+    InstrToIntMap[Inst->TheDef] = Num++;
 }
 
 

llvm/utils/TableGen/CodeGenTarget.h

@@ -74,6 +74,7 @@ class CodeGenTarget {
 
   mutable StringRef InstNamespace;
   mutable std::vector<const CodeGenInstruction*> InstrsByEnum;
+  mutable DenseMap<const Record *, unsigned> InstrToIntMap;
   mutable unsigned NumPseudoInstructions = 0;
 public:
   CodeGenTarget(RecordKeeper &Records);
@@ -192,6 +193,15 @@ class CodeGenTarget {
     return InstrsByEnum;
   }
 
+  /// Return the integer enum value corresponding to this instruction record.
+  unsigned getInstrIntValue(const Record *R) const {
+    if (InstrsByEnum.empty())
+      ComputeInstrsByEnum();
+    auto V = InstrToIntMap.find(R);
+    assert(V != InstrToIntMap.end() && "instr not in InstrToIntMap");
+    return V->second;
+  }
+
   typedef ArrayRef<const CodeGenInstruction *>::const_iterator inst_iterator;
   inst_iterator inst_begin() const{return getInstructionsByEnumValue().begin();}
   inst_iterator inst_end() const { return getInstructionsByEnumValue().end(); }

llvm/utils/TableGen/SearchableTableEmitter.cpp

@@ -7,12 +7,14 @@
 //===----------------------------------------------------------------------===//
 //
 // This tablegen backend emits a generic array initialized by specified fields,
-// together with companion index tables and lookup functions (binary search,
-// currently).
+// together with companion index tables and lookup functions. The lookup
+// function generated is either a direct lookup (when a single primary key field
+// is integral and densely numbered) or a binary search otherwise.
 //
 //===----------------------------------------------------------------------===//
 
 #include "CodeGenIntrinsics.h"
+#include "CodeGenTarget.h"
 #include "llvm/ADT/ArrayRef.h"
 #include "llvm/ADT/DenseMap.h"
 #include "llvm/ADT/STLExtras.h"
@@ -90,6 +92,7 @@ struct GenericTable {
 
 class SearchableTableEmitter {
   RecordKeeper &Records;
+  std::unique_ptr<CodeGenTarget> Target;
   DenseMap<Init *, std::unique_ptr<CodeGenIntrinsic>> Intrinsics;
   std::vector<std::unique_ptr<GenericEnum>> Enums;
   DenseMap<Record *, GenericEnum *> EnumMap;
@@ -201,18 +204,23 @@ class SearchableTableEmitter {
                           const std::vector<Record *> &Items);
   void collectTableEntries(GenericTable &Table,
                            const std::vector<Record *> &Items);
+  int64_t getNumericKey(const SearchIndex &Index, Record *Rec);
 };
 
 } // End anonymous namespace.
 
 // For search indices that consists of a single field whose numeric value is
 // known, return that numeric value.
-static int64_t getNumericKey(const SearchIndex &Index, Record *Rec) {
+int64_t SearchableTableEmitter::getNumericKey(const SearchIndex &Index,
+                                              Record *Rec) {
   assert(Index.Fields.size() == 1);
 
   if (Index.Fields[0].Enum) {
     Record *EnumEntry = Rec->getValueAsDef(Index.Fields[0].Name);
     return Index.Fields[0].Enum->EntryMap[EnumEntry]->second;
+  } else if (Index.Fields[0].IsInstruction) {
+    Record *TheDef = Rec->getValueAsDef(Index.Fields[0].Name);
+    return Target->getInstrIntValue(TheDef);
   }
 
   return getInt(Rec, Index.Fields[0].Name);
@@ -370,20 +378,31 @@ void SearchableTableEmitter::emitLookupFunction(const GenericTable &Table,
   bool IsContiguous = false;
 
   if (Index.Fields.size() == 1 &&
-      (Index.Fields[0].Enum || isa<BitsRecTy>(Index.Fields[0].RecType))) {
+      (Index.Fields[0].Enum || isa<BitsRecTy>(Index.Fields[0].RecType) ||
+       Index.Fields[0].IsInstruction)) {
+    int64_t FirstKeyVal = getNumericKey(Index, IndexRows[0]);
     IsContiguous = true;
     for (unsigned i = 0; i < IndexRows.size(); ++i) {
-      if (getNumericKey(Index, IndexRows[i]) != i) {
+      if (getNumericKey(Index, IndexRows[i]) != (FirstKeyVal + i)) {
         IsContiguous = false;
         break;
       }
     }
   }
 
   if (IsContiguous) {
+    const GenericField &Field = Index.Fields[0];
+    std::string FirstRepr = primaryRepresentation(
+        Index.Loc, Field, IndexRows[0]->getValueInit(Field.Name));
+    std::string LastRepr = primaryRepresentation(
+        Index.Loc, Field, IndexRows.back()->getValueInit(Field.Name));
+    OS << "  if ((" << Field.Name << " < " << FirstRepr << ") ||\n";
+    OS << "      (" << Field.Name << " > " << LastRepr << "))\n";
+    OS << "    return nullptr;\n";
     OS << "  auto Table = ArrayRef(" << IndexName << ");\n";
-    OS << "  size_t Idx = " << Index.Fields[0].Name << ";\n";
-    OS << "  return Idx >= Table.size() ? nullptr : ";
+    OS << "  size_t Idx = " << Index.Fields[0].Name << " - " << FirstRepr
+       << ";\n";
+    OS << "  return ";
     if (IsPrimary)
       OS << "&Table[Idx]";
     else
@@ -663,6 +682,8 @@ void SearchableTableEmitter::run(raw_ostream &OS) {
   // Emit tables in a deterministic order to avoid needless rebuilds.
   SmallVector<std::unique_ptr<GenericTable>, 4> Tables;
   DenseMap<Record *, GenericTable *> TableMap;
+  if (!Records.getAllDerivedDefinitionsIfDefined("Instruction").empty())
+    Target = std::make_unique<CodeGenTarget>(Records);
 
   // Collect all definitions first.
   for (auto *EnumRec : Records.getAllDerivedDefinitions("GenericEnum")) {