[PDLL] Add support for tablegen includes and importing ODS information

This commit adds support for processing tablegen include files, and importing
various information from ODS. This includes operations, attribute+type constraints,
attribute/operation/type interfaces, etc. This will allow for much more robust tooling,
and also allows for referencing ODS constructs directly within PDLL (imported interfaces
can be used as constraints, operation result names can be used for member access, etc).

Differential Revision: https://reviews.llvm.org/D119900

llvm/include/llvm/Support/SourceMgr.h

@@ -100,6 +100,9 @@ class SourceMgr {
   SourceMgr &operator=(SourceMgr &&) = default;
   ~SourceMgr() = default;
 
+  /// Return the include directories of this source manager.
+  ArrayRef<std::string> getIncludeDirs() const { return IncludeDirectories; }
+
   void setIncludeDirs(const std::vector<std::string> &Dirs) {
     IncludeDirectories = Dirs;
   }
@@ -147,6 +150,14 @@ class SourceMgr {
     return Buffers.size();
   }
 
+  /// Takes the source buffers from the given source manager and append them to
+  /// the current manager.
+  void takeSourceBuffersFrom(SourceMgr &SrcMgr) {
+    std::move(SrcMgr.Buffers.begin(), SrcMgr.Buffers.end(),
+              std::back_inserter(Buffers));
+    SrcMgr.Buffers.clear();
+  }
+
   /// Search for a file with the specified name in the current directory or in
   /// one of the IncludeDirs.
   ///
@@ -156,6 +167,17 @@ class SourceMgr {
   unsigned AddIncludeFile(const std::string &Filename, SMLoc IncludeLoc,
                           std::string &IncludedFile);
 
+  /// Search for a file with the specified name in the current directory or in
+  /// one of the IncludeDirs, and try to open it **without** adding to the
+  /// SourceMgr. If the opened file is intended to be added to the source
+  /// manager, prefer `AddIncludeFile` instead.
+  ///
+  /// If no file is found, this returns an Error, otherwise it returns the
+  /// buffer of the stacked file. The full path to the included file can be
+  /// found in \p IncludedFile.
+  ErrorOr<std::unique_ptr<MemoryBuffer>>
+  OpenIncludeFile(const std::string &Filename, std::string &IncludedFile);
+
   /// Return the ID of the buffer containing the specified location.
   ///
   /// 0 is returned if the buffer is not found.

llvm/lib/Support/SourceMgr.cpp

@@ -40,6 +40,17 @@ static const size_t TabStop = 8;
 unsigned SourceMgr::AddIncludeFile(const std::string &Filename,
                                    SMLoc IncludeLoc,
                                    std::string &IncludedFile) {
+  ErrorOr<std::unique_ptr<MemoryBuffer>> NewBufOrErr =
+      OpenIncludeFile(Filename, IncludedFile);
+  if (!NewBufOrErr)
+    return 0;
+
+  return AddNewSourceBuffer(std::move(*NewBufOrErr), IncludeLoc);
+}
+
+ErrorOr<std::unique_ptr<MemoryBuffer>>
+SourceMgr::OpenIncludeFile(const std::string &Filename,
+                           std::string &IncludedFile) {
   IncludedFile = Filename;
   ErrorOr<std::unique_ptr<MemoryBuffer>> NewBufOrErr =
       MemoryBuffer::getFile(IncludedFile);
@@ -52,10 +63,7 @@ unsigned SourceMgr::AddIncludeFile(const std::string &Filename,
     NewBufOrErr = MemoryBuffer::getFile(IncludedFile);
   }
 
-  if (!NewBufOrErr)
-    return 0;
-
-  return AddNewSourceBuffer(std::move(*NewBufOrErr), IncludeLoc);
+  return NewBufOrErr;
 }
 
 unsigned SourceMgr::FindBufferContainingLoc(SMLoc Loc) const {

mlir/include/mlir/IR/OpBase.td

@@ -363,7 +363,8 @@ class DialectType<Dialect d, Pred condition, string descr = "",
 
 // A variadic type constraint. It expands to zero or more of the base type. This
 // class is used for supporting variadic operands/results.
-class Variadic<Type type> : TypeConstraint<type.predicate, type.summary> {
+class Variadic<Type type> : TypeConstraint<type.predicate, type.summary,
+                                           type.cppClassName> {
   Type baseType = type;
 }
 
@@ -379,7 +380,8 @@ class VariadicOfVariadic<Type type, string variadicSegmentAttrName>
 
 // An optional type constraint. It expands to either zero or one of the base
 // type. This class is used for supporting optional operands/results.
-class Optional<Type type> : TypeConstraint<type.predicate, type.summary> {
+class Optional<Type type> : TypeConstraint<type.predicate, type.summary,
+                                           type.cppClassName> {
   Type baseType = type;
 }
 

mlir/include/mlir/TableGen/Constraint.h

@@ -54,6 +54,11 @@ class Constraint {
   // description is not provided, returns the TableGen def name.
   StringRef getSummary() const;
 
+  /// Returns the name of the TablGen def of this constraint. In some cases
+  /// where the current def is anonymous, the name of the base def is used (e.g.
+  /// `Optional<>`/`Variadic<>` type constraints).
+  StringRef getDefName() const;
+
   Kind getKind() const { return kind; }
 
 protected:

mlir/include/mlir/Tools/PDLL/AST/Context.h

@@ -14,13 +14,17 @@
 
 namespace mlir {
 namespace pdll {
+namespace ods {
+class Context;
+} // namespace ods
+
 namespace ast {
 /// This class represents the main context of the PDLL AST. It handles
 /// allocating all of the AST constructs, and manages all state necessary for
 /// the AST.
 class Context {
 public:
-  Context();
+  explicit Context(ods::Context &odsContext);
   Context(const Context &) = delete;
   Context &operator=(const Context &) = delete;
 
@@ -30,13 +34,20 @@ class Context {
   /// Return the storage uniquer used for AST types.
   StorageUniquer &getTypeUniquer() { return typeUniquer; }
 
+  /// Return the ODS context used by the AST.
+  ods::Context &getODSContext() { return odsContext; }
+  const ods::Context &getODSContext() const { return odsContext; }
+
   /// Return the diagnostic engine of this context.
   DiagnosticEngine &getDiagEngine() { return diagEngine; }
 
 private:
   /// The diagnostic engine of this AST context.
   DiagnosticEngine diagEngine;
 
+  /// The ODS context used by the AST.
+  ods::Context &odsContext;
+
   /// The allocator used for AST nodes, and other entities allocated within the
   /// context.
   llvm::BumpPtrAllocator allocator;

mlir/include/mlir/Tools/PDLL/ODS/Constraint.h

@@ -0,0 +1,98 @@
+//===- Constraint.h - MLIR PDLL ODS Constraints -----------------*- C++ -*-===//
+//
+// 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
+//
+//===----------------------------------------------------------------------===//
+//
+// This file contains a PDLL description of ODS constraints. These are used to
+// support the import of constraints defined outside of PDLL.
+//
+//===----------------------------------------------------------------------===//
+
+#ifndef MLIR_TOOLS_PDLL_ODS_CONSTRAINT_H_
+#define MLIR_TOOLS_PDLL_ODS_CONSTRAINT_H_
+
+#include <string>
+
+#include "mlir/Support/LLVM.h"
+#include "llvm/ADT/SmallVector.h"
+#include "llvm/ADT/StringMap.h"
+
+namespace mlir {
+namespace pdll {
+namespace ods {
+
+//===----------------------------------------------------------------------===//
+// Constraint
+//===----------------------------------------------------------------------===//
+
+/// This class represents a generic ODS constraint.
+class Constraint {
+public:
+  /// Return the name of this constraint.
+  StringRef getName() const { return name; }
+
+  /// Return the summary of this constraint.
+  StringRef getSummary() const { return summary; }
+
+protected:
+  Constraint(StringRef name, StringRef summary)
+      : name(name.str()), summary(summary.str()) {}
+  Constraint(const Constraint &) = delete;
+
+private:
+  /// The name of the constraint.
+  std::string name;
+  /// A summary of the constraint.
+  std::string summary;
+};
+
+//===----------------------------------------------------------------------===//
+// AttributeConstraint
+//===----------------------------------------------------------------------===//
+
+/// This class represents a generic ODS Attribute constraint.
+class AttributeConstraint : public Constraint {
+public:
+  /// Return the name of the underlying c++ class of this constraint.
+  StringRef getCppClass() const { return cppClassName; }
+
+private:
+  AttributeConstraint(StringRef name, StringRef summary, StringRef cppClassName)
+      : Constraint(name, summary), cppClassName(cppClassName.str()) {}
+
+  /// The c++ class of the constraint.
+  std::string cppClassName;
+
+  /// Allow access to the constructor.
+  friend class Context;
+};
+
+//===----------------------------------------------------------------------===//
+// TypeConstraint
+//===----------------------------------------------------------------------===//
+
+/// This class represents a generic ODS Type constraint.
+class TypeConstraint : public Constraint {
+public:
+  /// Return the name of the underlying c++ class of this constraint.
+  StringRef getCppClass() const { return cppClassName; }
+
+private:
+  TypeConstraint(StringRef name, StringRef summary, StringRef cppClassName)
+      : Constraint(name, summary), cppClassName(cppClassName.str()) {}
+
+  /// The c++ class of the constraint.
+  std::string cppClassName;
+
+  /// Allow access to the constructor.
+  friend class Context;
+};
+
+} // namespace ods
+} // namespace pdll
+} // namespace mlir
+
+#endif // MLIR_TOOLS_PDLL_ODS_CONSTRAINT_H_

mlir/include/mlir/Tools/PDLL/ODS/Context.h

@@ -0,0 +1,78 @@
+//===- Context.h - MLIR PDLL ODS Context ------------------------*- C++ -*-===//
+//
+// 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
+//
+//===----------------------------------------------------------------------===//
+
+#ifndef MLIR_TOOLS_PDLL_ODS_CONTEXT_H_
+#define MLIR_TOOLS_PDLL_ODS_CONTEXT_H_
+
+#include <string>
+
+#include "mlir/Support/LLVM.h"
+#include "llvm/ADT/STLExtras.h"
+#include "llvm/ADT/SmallVector.h"
+#include "llvm/ADT/StringMap.h"
+
+namespace llvm {
+class SMLoc;
+} // namespace llvm
+
+namespace mlir {
+namespace pdll {
+namespace ods {
+class AttributeConstraint;
+class Dialect;
+class Operation;
+class TypeConstraint;
+
+/// This class contains all of the registered ODS operation classes.
+class Context {
+public:
+  Context();
+  ~Context();
+
+  /// Insert a new attribute constraint with the context. Returns the inserted
+  /// constraint, or a previously inserted constraint with the same name.
+  const AttributeConstraint &insertAttributeConstraint(StringRef name,
+                                                       StringRef summary,
+                                                       StringRef cppClass);
+
+  /// Insert a new type constraint with the context. Returns the inserted
+  /// constraint, or a previously inserted constraint with the same name.
+  const TypeConstraint &insertTypeConstraint(StringRef name, StringRef summary,
+                                             StringRef cppClass);
+
+  /// Insert a new dialect with the context. Returns the inserted dialect, or a
+  /// previously inserted dialect with the same name.
+  Dialect &insertDialect(StringRef name);
+
+  /// Lookup a dialect registered with the given name, or null if no dialect
+  /// with that name was inserted.
+  const Dialect *lookupDialect(StringRef name) const;
+
+  /// Insert a new operation with the context. Returns the inserted operation,
+  /// and a boolean indicating if the operation newly inserted (false if the
+  /// operation already existed).
+  std::pair<Operation *, bool>
+  insertOperation(StringRef name, StringRef summary, StringRef desc, SMLoc loc);
+
+  /// Lookup an operation registered with the given name, or null if no
+  /// operation with that name is registered.
+  const Operation *lookupOperation(StringRef name) const;
+
+  /// Print the contents of this context to the provided stream.
+  void print(raw_ostream &os) const;
+
+private:
+  llvm::StringMap<std::unique_ptr<AttributeConstraint>> attributeConstraints;
+  llvm::StringMap<std::unique_ptr<Dialect>> dialects;
+  llvm::StringMap<std::unique_ptr<TypeConstraint>> typeConstraints;
+};
+} // namespace ods
+} // namespace pdll
+} // namespace mlir
+
+#endif // MLIR_PDL_pdll_ODS_CONTEXT_H_

mlir/include/mlir/Tools/PDLL/ODS/Dialect.h

@@ -0,0 +1,64 @@
+//===- Dialect.h - PDLL ODS Dialect -----------------------------*- C++ -*-===//
+//
+// 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
+//
+//===----------------------------------------------------------------------===//
+
+#ifndef MLIR_TOOLS_PDLL_ODS_DIALECT_H_
+#define MLIR_TOOLS_PDLL_ODS_DIALECT_H_
+
+#include <string>
+
+#include "mlir/Support/LLVM.h"
+#include "llvm/ADT/STLExtras.h"
+#include "llvm/ADT/SmallVector.h"
+#include "llvm/ADT/StringMap.h"
+
+namespace mlir {
+namespace pdll {
+namespace ods {
+class Operation;
+
+/// This class represents an ODS dialect, and contains information on the
+/// constructs held within the dialect.
+class Dialect {
+public:
+  ~Dialect();
+
+  /// Return the name of this dialect.
+  StringRef getName() const { return name; }
+
+  /// Insert a new operation with the dialect. Returns the inserted operation,
+  /// and a boolean indicating if the operation newly inserted (false if the
+  /// operation already existed).
+  std::pair<Operation *, bool>
+  insertOperation(StringRef name, StringRef summary, StringRef desc, SMLoc loc);
+
+  /// Lookup an operation registered with the given name, or null if no
+  /// operation with that name is registered.
+  Operation *lookupOperation(StringRef name) const;
+
+  /// Return a map of all of the operations registered to this dialect.
+  const llvm::StringMap<std::unique_ptr<Operation>> &getOperations() const {
+    return operations;
+  }
+
+private:
+  explicit Dialect(StringRef name);
+
+  /// The name of the dialect.
+  std::string name;
+
+  /// The operations defined by the dialect.
+  llvm::StringMap<std::unique_ptr<Operation>> operations;
+
+  /// Allow access to the constructor.
+  friend class Context;
+};
+} // namespace ods
+} // namespace pdll
+} // namespace mlir
+
+#endif // MLIR_TOOLS_PDLL_ODS_DIALECT_H_