[eudsl-tblgen] add lots of stuff (support retrieving ArgumentInits from Record) (#318)

This is kind of a niche PR so I don't expect anyone to really closely
review but I'm describing it anyway just to capture some findings.

The basic feature added in this PR is the ability to recover from an
instantiated `def` which template args were used in the `class`
call/invocation/thing. E.g.,

```
class Test_Op<string mnemonic, list<Trait> traits = []>
    : Op<Test_Dialect, mnemonic, traits>;

def Test_AndOp : Test_Op<"and"> { ... }
```

With this PR you can recover that the `mnemonic` passed as the first
argument to `Test_Op` when instantiating `Test_AndOp` is `"and"`. This
sounds like something that should've been always possible but turns out
tablegen literally doesn't keep that information in any way/shape/form.
In order to make this work I had to hack the parser to keep it around
(note, this isn't completely straightforward because arguments can have
unevaluated [bang
operators](https://llvm.org/docs/TableGen/ProgRef.html#bang-operators)).

For context, I intended on using this functionality to translate LLVM
`Intrinsic`s into MLIR OpDefs but turns out it didn't really work out
(not because the feature doesn't work but because I have ~30k
`Intrinsic`s lol).

Author: makslevental

.github/workflows/build_llvm.yml

@@ -181,6 +181,9 @@ jobs:
           # Tar up MLIR/LLVM distro
           ###############################
           
+          # if timezone causes datetime to be ahead of local datetime then ninja will go into a loop... 
+          find $LLVM_INSTALL_DIR -exec touch -a -m -t 197001010000 {} \;
+          
           tar -czf "mlir_${{ matrix.name }}_$WHEEL_VERSION.tar.gz" -C "$LLVM_INSTALL_DIR/.." llvm-install
           rm -rf "$LLVM_BUILD_DIR" "$LLVM_SOURCE_DIR"
           

projects/common/eudsl/bind_vec_like.h

@@ -25,13 +25,13 @@ extern nanobind::class_<_SmallVector> smallVector;
 extern nanobind::class_<_ArrayRef> arrayRef;
 extern nanobind::class_<_MutableArrayRef> mutableArrayRef;
 
-template <typename Element, typename... Args>
-std::tuple<nanobind::class_<llvm::SmallVector<Element>>,
+template <typename Element, int Size = 4, typename... Args>
+std::tuple<nanobind::class_<llvm::SmallVector<Element, Size>>,
            nanobind::class_<llvm::ArrayRef<Element>>,
            nanobind::class_<llvm::MutableArrayRef<Element>>>
 bind_array_ref(nanobind::handle scope, Args &&...args) {
   using ArrayRef = llvm::ArrayRef<Element>;
-  using SmallVec = llvm::SmallVector<Element>;
+  using SmallVec = llvm::SmallVector<Element, Size>;
   using MutableArrayRef = llvm::MutableArrayRef<Element>;
   using ValueRef = Element &;
 

projects/eudsl-llvmpy/eudsl-llvmpy-generate.py

@@ -320,14 +320,15 @@ class LLVMMatchType(Generic[_T]):
         ),
         file=amdgcn_f,
     )
-    intrins = RecordKeeper().parse_td(
+    rk = RecordKeeper()
+    rk.parse_td(
         str(llvm_include_root / "llvm" / "IR" / "Intrinsics.td"),
         include_dirs=[str(llvm_include_root)],
     )
     int_regex = re.compile(r"_i(\d+)")
     fp_regex = re.compile(r"_f(\d+)")
 
-    defs = intrins.get_defs()
+    defs = rk.get_defs()
     for d in defs:
         intr = defs[d]
         if (

projects/eudsl-py/tests/test_bindings.py …eudsl-py/tests/test_eudsl_py_bindings.pyprojects/eudsl-py/tests/test_bindings.py renamed to projects/eudsl-py/tests/test_eudsl_py_bindings.py projects/eudsl-py/tests/test_bindings.py renamed to projects/eudsl-py/tests/test_eudsl_py_bindings.py

@@ -13,7 +13,7 @@ Homepage = "https://github.com/llvm/eudsl"
 
 [build-system]
 requires = [
-    "nanobind>=2.2.0",
+    "nanobind>=2.9.2",
     "scikit-build-core>=0.10.7",
     "typing_extensions>=4.12.2"
 ]
@@ -99,3 +99,6 @@ before-build = [
 before-build = [
     "ccache -z"
 ]
+
+[tool.pytest.ini_options]
+addopts = "-s"

projects/eudsl-tblgen/pyproject.toml

@@ -209,7 +209,6 @@ class TGLexer {
 public:
   typedef std::set<std::string> DependenciesSetTy;
 
-private:
   /// Dependencies - This is the list of all included files.
   DependenciesSetTy Dependencies;
 
@@ -240,7 +239,6 @@ class TGLexer {
   SMLoc getLoc() const;
   SMRange getLocRange() const;
 
-private:
   /// LexToken - Read the next token and return its code.
   tgtok::TokKind LexToken(bool FileOrLineStart = false);
 

projects/eudsl-tblgen/src/TGLexer.h

@@ -48,6 +48,7 @@ struct SubClassReference {
   SubClassReference() = default;
 
   bool isInvalid() const { return Rec == nullptr; }
+  void dump() const;
 };
 
 struct SubMultiClassReference {
@@ -72,6 +73,19 @@ LLVM_DUMP_METHOD void SubMultiClassReference::dump() const {
   for (const Init *TA : TemplateArgs)
     TA->dump();
 }
+
+LLVM_DUMP_METHOD void SubClassReference::dump() const {
+  errs() << "Subclass:\n";
+
+  Rec->dump();
+
+  errs() << "Template args:\n";
+  for (const Init *TA : TemplateArgs) {
+    TA->dump();
+    errs() << "\n";
+  }
+  errs() << "\n";
+}
 #endif
 
 static bool checkBitsConcrete(Record &R, const RecordVal &RV) {
@@ -307,6 +321,66 @@ bool TGParser::SetValue(Record *CurRec, SMLoc Loc, const Init *ValName,
   return false;
 }
 
+static std::string getCanonicalRecName(Record *CurRec) {
+  const Init *Name;
+  if (CurRec->isClass())
+    Name = VarInit::get(QualifiedNameOfImplicitName(*CurRec),
+                        StringRecTy::get(CurRec->getRecords()));
+  else
+    Name = CurRec->getNameInit();
+  return Name->getAsUnquotedString();
+}
+
+// https://nimrod.blog/posts/cpp-how-to-access-private-members-validly/
+template <typename Tag>
+struct Storage {
+  inline static typename Tag::type ptr;
+};
+
+template <typename Tag, typename Tag::type V>
+struct PtrTaker {
+  struct Transferer {
+    Transferer() { Storage<Tag>::ptr = V; }
+  };
+  inline static Transferer tr;
+};
+
+struct TemplateArgsTag {
+  using type = SmallVector<const Init *, 0> Record::*;
+};
+
+template struct PtrTaker<TemplateArgsTag, &Record::TemplateArgs>;
+
+static void
+addOrReplaceTemplateArgs(Record *rec,
+                         SmallVector<const ArgumentInit *, 4> newTemplateArgs) {
+  SmallVector<const Init *, 0> *mutableTemplateArgs =
+      &(*rec.*Storage<TemplateArgsTag>::ptr);
+  std::string recName = getCanonicalRecName(rec);
+  for (const ArgumentInit *newTemplateArg : newTemplateArgs) {
+    const Init **i = llvm::find_if(*mutableTemplateArgs, [&](const Init *tArg) {
+      const auto *oTArg = llvm::dyn_cast<ArgumentInit>(tArg);
+      if (oTArg->isNamed()) {
+        if (newTemplateArg->isNamed())
+          return oTArg->getName() == newTemplateArg->getName();
+        return false;
+      }
+      if (newTemplateArg->isPositional())
+        return oTArg->getIndex() == newTemplateArg->getIndex();
+      return false;
+    });
+    auto newVal = RecordVal(newTemplateArg, RecordRecTy::get(rec),
+                            RecordVal::FK_TemplateArg);
+    if (i == mutableTemplateArgs->end()) {
+      mutableTemplateArgs->push_back(newTemplateArg);
+    } else {
+      rec->removeValue(*i);
+      *i = newTemplateArg;
+    }
+    rec->addValue(newVal);
+  }
+}
+
 /// AddSubClass - Add SubClass as a subclass to CurRec, resolving its template
 /// args as SubClass's template arguments.
 bool TGParser::AddSubClass(Record *CurRec, SubClassReference &SubClass) {
@@ -336,9 +410,15 @@ bool TGParser::AddSubClass(Record *CurRec, SubClassReference &SubClass) {
                         StringRecTy::get(Records));
   else
     Name = CurRec->getNameInit();
+  auto canName = getCanonicalRecName(CurRec);
+  recordTemplateArgs[canName] = SubClass.TemplateArgs;
   R.set(QualifiedNameOfImplicitName(*SC), Name);
 
   CurRec->resolveReferences(R);
+  // we only want to keep args for defs not classes which pass
+  // args to other classes (i think so?)
+  if (!CurRec->isClass() && !SubClass.TemplateArgs.empty())
+    addOrReplaceTemplateArgs(CurRec, SubClass.TemplateArgs);
 
   // Since everything went well, we can now set the "superclass" list for the
   // current record.
@@ -531,7 +611,20 @@ bool TGParser::resolve(const std::vector<RecordsEntry> &Source,
       MapResolver R(Rec.get());
       for (const auto &S : Substs)
         R.set(S.first, S.second);
+      std::string eName = getCanonicalRecName(E.Rec.get());
+      assert(recordTemplateArgs.count(eName));
+      auto templateArgs = recordTemplateArgs[eName];
       Rec->resolveReferences(R);
+      SmallVector<const ArgumentInit *, 4> resolvedTemplateArgs;
+      for (const ArgumentInit *templateArg : templateArgs)
+        resolvedTemplateArgs.emplace_back(
+            llvm::cast<ArgumentInit>(templateArg->resolveReferences(R)));
+      std::string recName = getCanonicalRecName(Rec.get());
+      recordTemplateArgs[recName] = resolvedTemplateArgs;
+      // we only want to keep args for defs not classes which pass
+      // args to other classes (i think so?)
+      if (!Rec->isClass() && !resolvedTemplateArgs.empty())
+        addOrReplaceTemplateArgs(Rec.get(), resolvedTemplateArgs);
 
       if (Dest)
         Dest->push_back(std::move(Rec));
@@ -574,7 +667,7 @@ bool TGParser::addDefOne(std::unique_ptr<Record> Rec) {
   Rec->emitRecordDumps();
 
   // If ObjectBody has template arguments, it's an error.
-  assert(Rec->getTemplateArgs().empty() && "How'd this get template args?");
+  // assert(Rec->getTemplateArgs().empty() && "How'd this get template args?");
 
   for (DefsetRecord *Defset : Defsets) {
     DefInit *I = Rec->getDefInit();
@@ -2872,7 +2965,9 @@ const Init *TGParser::ParseSimpleValue(Record *CurRec, const RecTy *ItemType,
 
     if (TrackReferenceLocs)
       Class->appendReferenceLoc(NameLoc);
-    return VarDefInit::get(NameLoc.Start, Class, Args)->Fold();
+    const VarDefInit *v = VarDefInit::get(NameLoc.Start, Class, Args);
+    TheVarDefInitPool.insert(v);
+    return v->Fold();
   }
   case tgtok::l_brace: { // Value ::= '{' ValueList '}'
     SMLoc BraceLoc = Lex.getLoc();

projects/eudsl-tblgen/src/TGParser.cpp

@@ -90,7 +90,6 @@ class TGVarScope {
 public:
   enum ScopeKind { SK_Local, SK_Record, SK_ForeachLoop, SK_MultiClass };
 
-private:
   ScopeKind Kind;
   std::unique_ptr<TGVarScope> Parent;
   // A scope to hold variable definitions from defvar.
@@ -142,6 +141,9 @@ class TGParser {
   std::vector<SmallVector<LetRecord, 4>> LetStack;
   std::map<std::string, std::unique_ptr<MultiClass>> MultiClasses;
   std::map<std::string, const RecTy *> TypeAliases;
+  DenseSet<const VarDefInit *> TheVarDefInitPool;
+  std::unordered_map<std::string, SmallVector<const ArgumentInit *, 4>>
+      recordTemplateArgs;
 
   /// Loops - Keep track of any foreach loops we are within.
   ///
@@ -181,6 +183,12 @@ class TGParser {
         NoWarnOnUnusedTemplateArgs(NoWarnOnUnusedTemplateArgs),
         TrackReferenceLocs(TrackReferenceLocs) {}
 
+  DenseSet<const VarDefInit *> getVarDefInits() { return TheVarDefInitPool; }
+  std::unordered_map<std::string, SmallVector<const ArgumentInit *, 4>>
+  getRecordTemplateArgs() {
+    return recordTemplateArgs;
+  }
+
   /// ParseFile - Main entrypoint for parsing a tblgen file. These parser
   /// routines return true on error, or false on success.
   bool ParseFile();
@@ -219,7 +227,6 @@ class TGParser {
     CurScope = CurScope->extractParent();
   }
 
-private: // Semantic analysis methods.
   bool AddValue(Record *TheRec, SMLoc Loc, const RecordVal &RV);
   /// Set the value of a RecordVal within the given record. If `OverrideDefLoc`
   /// is set, the provided location overrides any existing location of the
@@ -253,7 +260,6 @@ class TGParser {
                                     ArrayRef<const ArgumentInit *> ArgValues,
                                     const Init *DefmName, SMLoc Loc);
 
-private: // Parser methods.
   bool consume(tgtok::TokKind K);
   bool ParseObjectList(MultiClass *MC = nullptr);
   bool ParseObject(MultiClass *MC);