[TableGen] Add support for emitting new function definition to return a range of results for Primary Key (#96174)

Garvit Gupta · web-flow · commit 6c903f05f301 · 2024-07-11T10:54:31.000-07:00
In the RISC-V architecture, multiple vendor-specific Control and Status
Registers (CSRs) share the same encoding. However, the existing lookup
function, which currently returns only a single result, falls short.
During disassembly, it consistently returns the first CSR encountered,
which may not be the correct CSR for the subtarget.

To address this issue, we modify the function definition to return a
range of results. These results can then be iterated upon to identify
the CSR that best fits the subtarget’s feature requirements. The
behavior of this new definition is controlled by a variable named
`ReturnRange`, which defaults to `false`.

Specifically, this patch introduces support for emitting a new lookup
function for the primary key. This function returns a pair of iterators
pointing to the first and last values, providing a comprehensive range
of values that satisfy the query
diff --git a/llvm/docs/TableGen/BackEnds.rst b/llvm/docs/TableGen/BackEnds.rst
@@ -717,6 +717,12 @@ This class provides six fields.
 
 * ``bit PrimaryKeyEarlyOut``. See the third example below.
 
+* ``bit PrimaryKeyReturnRange``. when set to 1, modifies the lookup function’s
+  definition to return a range of results rather than a single pointer to the
+  object. This feature proves useful when multiple objects meet the criteria
+  specified by the lookup function. Currently, it is supported only for primary
+  lookup functions. Refer to the second example below for further details.
+
 TableGen attempts to deduce the type of each of the table fields so that it
 can format the C++ initializers in the emitted table. It can deduce ``bit``,
 ``bits<n>``, ``string``, ``Intrinsic``, and ``Instruction``.  These can be
@@ -883,6 +889,84 @@ Here is the generated C++ code.
     return &*Idx;
   }
 
+In the above example, lets add one more record with encoding same as that of
+record ``CEntry<"Pear",  CBaz, 15>``.
+
+.. code-block:: text
+
+  def CFoobar : CEnum;
+  def : CEntry<"Banana", CFoobar, 15>;
+
+Below is the new generated ``CTable``
+
+.. code-block:: text
+
+  #ifdef GET_Table_IMPL
+  constexpr CEntry Table[] = {
+    { "Apple", CFoo, 0xA }, // 0
+    { "Apple", CBar, 0xD }, // 1
+    { "Banana", CFoobar, 0xF }, // 2
+    { "Pear", CBaz, 0xF }, // 3
+  };
+
+Since ``Banana`` lexicographically appears first, therefore in the ``CEntry``
+table, record with name ``Banana`` will come before the record with name
+``Pear``. Because of this, the ``lookupCEntryByEncoding`` function will always
+return a pointer to the record with name ``Banana`` even though in some cases
+the correct result can be the record with name ``Pear``. Such kind of scenario
+makes the exisitng lookup function insufficient because they always return a
+pointer to a single entry from the table, but instead it should return a range
+of results because multiple entries match the criteria sought by the lookup
+function. In this case, the definition of the lookup function needs to be
+modified to return a range of results which can be done by setting
+``PrimaryKeyReturnRange``.
+
+.. code-block:: text
+
+  def CTable : GenericTable {
+    let FilterClass = "CEntry";
+    let Fields = ["Name", "Kind", "Encoding"];
+    string TypeOf_Kind = "CEnum";
+    let PrimaryKey = ["Encoding"];
+    let PrimaryKeyName = "lookupCEntryByEncoding";
+    let PrimaryKeyReturnRange = true;
+  }
+
+Here is the modified lookup function.
+
+.. code-block:: text
+
+  llvm::iterator_range<const CEntry *> lookupCEntryByEncoding(uint16_t Encoding) {
+    struct KeyType {
+      uint16_t Encoding;
+    };
+    KeyType Key = {Encoding};
+    struct Comp {
+      bool operator()(const CEntry &LHS, const KeyType &RHS) const {
+        if (LHS.Encoding < RHS.Encoding)
+          return true;
+        if (LHS.Encoding > RHS.Encoding)
+          return false;
+        return false;
+      }
+      bool operator()(const KeyType &LHS, const CEntry &RHS) const {
+        if (LHS.Encoding < RHS.Encoding)
+          return true;
+        if (LHS.Encoding > RHS.Encoding)
+          return false;
+        return false;
+      }
+    };
+    auto Table = ArrayRef(Table);
+    auto It = std::equal_range(Table.begin(), Table.end(), Key, Comp());
+    return llvm::make_range(It.first, It.second);
+  }
+
+The new lookup function will return an iterator range with first pointer to the
+first result and the last pointer to the last matching result from the table.
+However, please note that the support for emitting modified definition exists
+for ``PrimaryKeyName`` only.
+
 The ``PrimaryKeyEarlyOut`` field, when set to 1, modifies the lookup
 function so that it tests the first field of the primary key to determine
 whether it is within the range of the collected records' primary keys. If
@@ -987,6 +1071,8 @@ function. This class provides three fields.
 
 * ``bit EarlyOut``. See the third example in `Generic Tables`_.
 
+* ``bit ReturnRange``. See the second example in `Generic Tables`_.
+
 Here is an example of a secondary key added to the ``CTable`` above. The
 generated function looks up entries based on the ``Name`` and ``Kind`` fields.
 
diff --git a/llvm/include/llvm/TableGen/SearchableTable.td b/llvm/include/llvm/TableGen/SearchableTable.td
@@ -114,6 +114,11 @@ class GenericTable {
 
   // See SearchIndex.EarlyOut
   bit PrimaryKeyEarlyOut = false;
+
+  // If true, will generate a different function signature which will return an
+  // iterator range of pointers giving the starting and end value of the range.
+  // e.g. lookupSysRegByEncoding returns multiple CSRs for same encoding.
+  bit PrimaryKeyReturnRange = false;
 }
 
 // Define a record derived from this class to generate an additional search
diff --git a/llvm/test/TableGen/generic-tables-return-range.td b/llvm/test/TableGen/generic-tables-return-range.td
@@ -0,0 +1,65 @@
+// RUN: llvm-tblgen -gen-searchable-tables -I %p/../../include %s | FileCheck %s
+
+include "llvm/TableGen/SearchableTable.td"
+
+class SysReg<string name, bits<12> op> {
+  string Name = name;
+  bits<12> Encoding = op;
+  code FeaturesRequired = [{ {} }];
+}
+
+def List1 : GenericTable {
+  let FilterClass = "SysReg";
+  let Fields = [
+     "Name", "Encoding", "FeaturesRequired",
+  ];
+
+  let PrimaryKey = [ "Encoding" ];
+  let PrimaryKeyName = "lookupSysRegByEncoding";
+  let PrimaryKeyReturnRange = true;
+}
+
+let FeaturesRequired = [{ {Feature1} }] in {
+def : SysReg<"csr1", 0x7C0>;
+}
+
+let FeaturesRequired = [{ {Feature2} }] in {
+def : SysReg<"csr2", 0x7C0>;
+}
+
+// CHECK: #ifdef GET_List1_DECL
+// CHECK-NEXT: llvm::iterator_range<const SysReg *> lookupSysRegByEncoding(uint16_t Encoding);
+// CHECK-NEXT: #endif
+
+// CHECK: #ifdef GET_List1_IMPL
+// CHECK-NEXT: constexpr SysReg List1[] = {
+// CHECK-NEXT:   { "csr1", 0x7C0,  {Feature1}  }, // 0
+// CHECK-NEXT:   { "csr2", 0x7C0,  {Feature2}  }, // 1
+// CHECK-NEXT:  };
+
+// CHECK: llvm::iterator_range<const SysReg *> lookupSysRegByEncoding(uint16_t Encoding) {
+// CHECK-NEXT: struct KeyType {
+// CHECK-NEXT:    uint16_t Encoding;
+// CHECK-NEXT:  };
+// CHECK-NEXT:  KeyType Key = {Encoding};
+// CHECK-NEXT:  struct Comp {
+// CHECK-NEXT:    bool operator()(const SysReg &LHS, const KeyType &RHS) const {
+// CHECK-NEXT:      if (LHS.Encoding < RHS.Encoding)
+// CHECK-NEXT:        return true;
+// CHECK-NEXT:      if (LHS.Encoding > RHS.Encoding)
+// CHECK-NEXT:        return false;
+// CHECK-NEXT:      return false;
+// CHECK-NEXT:    }
+// CHECK-NEXT:    bool operator()(const KeyType &LHS, const SysReg &RHS) const {
+// CHECK-NEXT:      if (LHS.Encoding < RHS.Encoding)
+// CHECK-NEXT:        return true;
+// CHECK-NEXT:      if (LHS.Encoding > RHS.Encoding)
+// CHECK-NEXT:        return false;
+// CHECK-NEXT:      return false;
+// CHECK-NEXT:    }
+// CHECK-NEXT:  };
+// CHECK-NEXT:  auto Table = ArrayRef(List1);
+// CHECK-NEXT:  auto It = std::equal_range(Table.begin(), Table.end(), Key, Comp());
+// CHECK-NEXT:  return llvm::make_range(It.first, It.second);
+// CHECK-NEXT: }
+// CHECK-NEXT: #endif
diff --git a/llvm/utils/TableGen/SearchableTableEmitter.cpp b/llvm/utils/TableGen/SearchableTableEmitter.cpp
@@ -68,6 +68,7 @@ struct SearchIndex {
   SMLoc Loc; // Source location of PrimaryKey or Key field definition.
   SmallVector<GenericField, 1> Fields;
   bool EarlyOut = false;
+  bool ReturnRange = false;
 };
 
 struct GenericTable {
@@ -198,7 +199,8 @@ class SearchableTableEmitter {
   bool parseFieldType(GenericField &Field, Init *II);
   std::unique_ptr<SearchIndex>
   parseSearchIndex(GenericTable &Table, const RecordVal *RecVal, StringRef Name,
-                   const std::vector<StringRef> &Key, bool EarlyOut);
+                   const std::vector<StringRef> &Key, bool EarlyOut,
+                   bool ReturnRange);
   void collectEnumEntries(GenericEnum &Enum, StringRef NameField,
                           StringRef ValueField,
                           const std::vector<Record *> &Items);
@@ -448,55 +450,81 @@ void SearchableTableEmitter::emitLookupFunction(const GenericTable &Table,
   }
   OS << "};\n";
 
-  OS << "  auto Table = ArrayRef(" << IndexName << ");\n";
-  OS << "  auto Idx = std::lower_bound(Table.begin(), Table.end(), Key,\n";
-  OS << "    [](const " << IndexTypeName << " &LHS, const KeyType &RHS) {\n";
+  OS << "  struct Comp {\n";
+  OS << "    bool operator()(const " << IndexTypeName
+     << " &LHS, const KeyType &RHS) const {\n";
 
-  for (const auto &Field : Index.Fields) {
-    if (isa<StringRecTy>(Field.RecType)) {
-      OS << "      int Cmp" << Field.Name << " = StringRef(LHS." << Field.Name
-         << ").compare(RHS." << Field.Name << ");\n";
-      OS << "      if (Cmp" << Field.Name << " < 0) return true;\n";
-      OS << "      if (Cmp" << Field.Name << " > 0) return false;\n";
-    } else if (Field.Enum) {
-      // Explicitly cast to unsigned, because the signedness of enums is
-      // compiler-dependent.
-      OS << "      if ((unsigned)LHS." << Field.Name << " < (unsigned)RHS."
-         << Field.Name << ")\n";
-      OS << "        return true;\n";
-      OS << "      if ((unsigned)LHS." << Field.Name << " > (unsigned)RHS."
-         << Field.Name << ")\n";
-      OS << "        return false;\n";
-    } else {
-      OS << "      if (LHS." << Field.Name << " < RHS." << Field.Name << ")\n";
-      OS << "        return true;\n";
-      OS << "      if (LHS." << Field.Name << " > RHS." << Field.Name << ")\n";
-      OS << "        return false;\n";
+  auto emitComparator = [&]() {
+    for (const auto &Field : Index.Fields) {
+      if (isa<StringRecTy>(Field.RecType)) {
+        OS << "      int Cmp" << Field.Name << " = StringRef(LHS." << Field.Name
+           << ").compare(RHS." << Field.Name << ");\n";
+        OS << "      if (Cmp" << Field.Name << " < 0) return true;\n";
+        OS << "      if (Cmp" << Field.Name << " > 0) return false;\n";
+      } else if (Field.Enum) {
+        // Explicitly cast to unsigned, because the signedness of enums is
+        // compiler-dependent.
+        OS << "      if ((unsigned)LHS." << Field.Name << " < (unsigned)RHS."
+           << Field.Name << ")\n";
+        OS << "        return true;\n";
+        OS << "      if ((unsigned)LHS." << Field.Name << " > (unsigned)RHS."
+           << Field.Name << ")\n";
+        OS << "        return false;\n";
+      } else {
+        OS << "      if (LHS." << Field.Name << " < RHS." << Field.Name
+           << ")\n";
+        OS << "        return true;\n";
+        OS << "      if (LHS." << Field.Name << " > RHS." << Field.Name
+           << ")\n";
+        OS << "        return false;\n";
+      }
     }
+    OS << "      return false;\n";
+    OS << "    }\n";
+  };
+  emitComparator();
+  bool ShouldReturnRange = Index.ReturnRange;
+  if (ShouldReturnRange) {
+    OS << "    bool operator()(const KeyType &LHS, const " << IndexTypeName
+       << " &RHS) const {\n";
+    emitComparator();
   }
 
-  OS << "      return false;\n";
-  OS << "    });\n\n";
-
-  OS << "  if (Idx == Table.end()";
+  OS << "  };\n";
+  OS << "  auto Table = ArrayRef(" << IndexName << ");\n";
+  if (ShouldReturnRange)
+    OS << "  auto It = std::equal_range(Table.begin(), Table.end(), Key, ";
+  else
+    OS << "  auto Idx = std::lower_bound(Table.begin(), Table.end(), Key, ";
+  OS << "Comp());\n";
 
-  for (const auto &Field : Index.Fields)
-    OS << " ||\n      Key." << Field.Name << " != Idx->" << Field.Name;
-  OS << ")\n    return nullptr;\n";
+  if (!ShouldReturnRange) {
+    OS << "  if (Idx == Table.end()";
+    for (const auto &Field : Index.Fields)
+      OS << " ||\n      Key." << Field.Name << " != Idx->" << Field.Name;
+  }
 
-  if (IsPrimary)
+  if (ShouldReturnRange)
+    OS << "  return llvm::make_range(It.first, It.second);\n";
+  else if (IsPrimary) {
+    OS << ")\n    return nullptr;\n\n";
     OS << "  return &*Idx;\n";
-  else
+  } else {
+    OS << ")\n    return nullptr;\n\n";
     OS << "  return &" << Table.Name << "[Idx->_index];\n";
+  }
 
   OS << "}\n";
 }
 
 void SearchableTableEmitter::emitLookupDeclaration(const GenericTable &Table,
                                                    const SearchIndex &Index,
                                                    raw_ostream &OS) {
-  OS << "const " << Table.CppTypeName << " *" << Index.Name << "(";
-
+  if (Index.ReturnRange)
+    OS << "llvm::iterator_range<const " << Table.CppTypeName << " *> ";
+  else
+    OS << "const " << Table.CppTypeName << " *";
+  OS << Index.Name << "(";
   ListSeparator LS;
   for (const auto &Field : Index.Fields)
     OS << LS << searchableFieldType(Table, Index, Field, TypeInArgument) << " "
@@ -541,9 +569,9 @@ void SearchableTableEmitter::emitGenericTable(const GenericTable &Table,
   // Indexes are sorted "{ Thing, PrimaryIdx }" arrays, so that a binary
   // search can be performed by "Thing".
   if (Table.PrimaryKey)
-    emitLookupFunction(Table, *Table.PrimaryKey, true, OS);
+    emitLookupFunction(Table, *Table.PrimaryKey, /*IsPrimary=*/true, OS);
   for (const auto &Index : Table.Indices)
-    emitLookupFunction(Table, *Index, false, OS);
+    emitLookupFunction(Table, *Index, /*IsPrimary=*/false, OS);
 
   OS << "#endif\n\n";
 }
@@ -569,11 +597,12 @@ bool SearchableTableEmitter::parseFieldType(GenericField &Field, Init *TypeOf) {
 
 std::unique_ptr<SearchIndex> SearchableTableEmitter::parseSearchIndex(
     GenericTable &Table, const RecordVal *KeyRecVal, StringRef Name,
-    const std::vector<StringRef> &Key, bool EarlyOut) {
+    const std::vector<StringRef> &Key, bool EarlyOut, bool ReturnRange) {
   auto Index = std::make_unique<SearchIndex>();
   Index->Name = std::string(Name);
   Index->Loc = KeyRecVal->getLoc();
   Index->EarlyOut = EarlyOut;
+  Index->ReturnRange = ReturnRange;
 
   for (const auto &FieldName : Key) {
     const GenericField *Field = Table.getFieldByName(FieldName);
@@ -769,7 +798,8 @@ void SearchableTableEmitter::run(raw_ostream &OS) {
           parseSearchIndex(*Table, TableRec->getValue("PrimaryKey"),
                            TableRec->getValueAsString("PrimaryKeyName"),
                            TableRec->getValueAsListOfStrings("PrimaryKey"),
-                           TableRec->getValueAsBit("PrimaryKeyEarlyOut"));
+                           TableRec->getValueAsBit("PrimaryKeyEarlyOut"),
+                           TableRec->getValueAsBit("PrimaryKeyReturnRange"));
 
       llvm::stable_sort(Table->Entries, [&](Record *LHS, Record *RHS) {
         return compareBy(LHS, RHS, *Table->PrimaryKey);
@@ -790,10 +820,10 @@ void SearchableTableEmitter::run(raw_ostream &OS) {
                           TableRec->getName());
 
     GenericTable &Table = *It->second;
-    Table.Indices.push_back(
-        parseSearchIndex(Table, IndexRec->getValue("Key"), IndexRec->getName(),
-                         IndexRec->getValueAsListOfStrings("Key"),
-                         IndexRec->getValueAsBit("EarlyOut")));
+    Table.Indices.push_back(parseSearchIndex(
+        Table, IndexRec->getValue("Key"), IndexRec->getName(),
+        IndexRec->getValueAsListOfStrings("Key"),
+        IndexRec->getValueAsBit("EarlyOut"), /*ReturnRange*/ false));
   }
 
   // Translate legacy tables.
@@ -847,8 +877,9 @@ void SearchableTableEmitter::run(raw_ostream &OS) {
          Class->getValueAsListOfStrings("SearchableFields")) {
       std::string Name =
           (Twine("lookup") + Table->CppTypeName + "By" + Field).str();
-      Table->Indices.push_back(parseSearchIndex(*Table, Class->getValue(Field),
-                                                Name, {Field}, false));
+      Table->Indices.push_back(
+          parseSearchIndex(*Table, Class->getValue(Field), Name, {Field},
+                           /*EarlyOut*/ false, /*ReturnRange*/ false));
     }
 
     Tables.emplace_back(std::move(Table));
