[mlir:PDLL] Rework the C++ generation of native Constraint/Rewrite arguments and results

[rriddle]

Summary

The current translation uses the old "ugly"/"raw" form which used PDLValue for the arguments
and results. This commit updates the C++ generation to use the recently added sugar that
allows for directly using the desired types for the arguments and result of PDL functions.
In addition, this commit also properly imports the C++ class for ODS operations, constraints,
and interfaces. This allows for a much more convienent C++ API than previously granted
with the raw/low-level types.

[jpienaar]

Nice clean-up

mlir/docs/PDLL.md

1157

able allowed? (Guessing s/able//)

mlir/lib/Tools/PDLL/CodeGen/CPPGen.cpp

176

Why 0?

[rriddle]

mlir/lib/Tools/PDLL/CodeGen/CPPGen.cpp

176

This is called for constraints on VariableDecls, which have only one input. I merged this function into the VariableDecl overload so that the single input expectation is more obvious.

mlir/docs/PDLL.md

Pattern {
   let value: Value = ...;
   let value2: Value = ...;
   HasOneUse(value);
   HasSameElementType(value, value2);
 }
 ```
 
 The arguments of the constraint are accessible within the code block via the
-same name. The type of these native variables are mapped directly to the
-corresponding MLIR type of the [core constraint](#core-constraints) used. For
-example, an `Op` corresponds to a variable of type `Operation *`.
-
-The results of the constraint can be populated using the provided `results`
-variable. This variable is a `PDLResultList`, and expects results to be
-populated in the order that they are defined within the result list of the
-constraint declaration.
+same name. See the ["type translation"](#native-constraint-type-translations) below for
+detailed information on how PDLL types are converted to native types. In addition to the
+PDLL arguments, the code block may also access the current `PatternRewriter` using
+`rewriter`. The result type of the native constraint function is implicitly defined
+as a `::mlir::LogicalResult`.
 
-In addition to the above, the code block may also access the current
-`PatternRewriter` using `rewriter`.
+Taking the constraints defined above as an example, these function would roughly be
+translated into:
+
+```c++
+LogicalResult HasOneUse(PatternRewriter &rewriter, Value value) {
+  return success(value.hasOneUse());
+}
+LogicalResult HasSameElementType(Value value1, Value value2) {
+  return success(value1.getType().cast<ShapedType>().getElementType() ==
+                 value2.getType().cast<ShapedType>().getElementType());
+}
+```
+
+TODO: Native constraints should also be able allowed to return values in certain cases.
+
+###### Native Constraint Type Translations
+
+The types of argument and result variables are generally mapped to the corresponding
+MLIR type of the [constraint](#constraints) used. Below is a detailed description
+of how the mapped type of a variable is determined for the various different types of
+constraints.
+
+* Attr, Op, Type, TypeRange, Value, ValueRange:
+
+These are all core constraints, and are mapped directly to the MLIR equivalent
+(that their names suggest), namely:
+
+  * `Attr`       -> "::mlir::Attribute"
+  * `Op`         -> "::mlir::Operation *"
+  * `Type`       -> "::mlir::Type"
+  * `TypeRange`  -> "::mlir::TypeRange"
+  * `Value`      -> "::mlir::Value"
+  * `ValueRange` -> "::mlir::ValueRange"
+
+* Op<dialect.name>
+
+A named operation constraint has a unique translation. If the ODS registration of the
+referenced operation has been included, the qualified C++ is used. If the ODS information
+is not available, this constraint maps to "::mlir::Operation *", similarly to the unnamed
+variant. For example, given the following:
+
+```pdll
+// `my_ops.td` provides the ODS definition of the `my_dialect` operations, such as
+// `my_dialect.bar` used below.
+#include "my_ops.td"
+
+Constraint Cst(op: Op<my_dialect.bar>) [{
+  return success(op ... );
+}];
+```
+
+The native type used for `op` may be of the form `my_dialect::BarOp`, as opposed to the
+default `::mlir::Operation *`. Below is a sample translation of the above constraint:
+
+```c++
+LogicalResult Cst(my_dialect::BarOp op) {
+  return success(op ... );
+}
+```
+
+* Imported ODS Constraints
+
+Aside from the core constraints, certain constraints imported from ODS may use a unique
+native type. How to enable this unique type depends on the ODS constraint construct that
+was imported:
+
+  * `Attr` constraints
+    - Imported `Attr` constraints utilize the `storageType` field for native type translation.
+  
+  * `Type` constraints
+    - Imported `Type` constraints utilize the `cppClassName` field for native type translation.
+
+  * `AttrInterface`/`OpInterface`/`TypeInterface` constraints
+    - Imported interfaces utilize the `cppClassName` field for native type translation.
 
 #### Defining Constraints Inline
 
 In addition to global scope, PDLL Constraints and Native Constraints defined in
 PDLL may be specified *inline* at any level of nesting. This means that they may
 be defined in Patterns, other Constraints, Rewrites, etc:
 
 ```pdll

 ##### Defining Native Rewrites in PDLL
 
 In addition to importing native rewrites, PDLL also supports defining native
 rewrites directly when compiling ahead-of-time (AOT) for C++. These rewrites can
 be defined by specifying a string code block after the rewrite declaration:
 
 ```pdll
 Rewrite BuildOp(value: Value) -> (foo: Op<my_dialect.foo>, bar: Op<my_dialect.bar>) [{
-  // We push back the results into the `results` variable in the order defined
-  // by the result list of the rewrite declaration.
-  results.push_back(rewriter.create<my_dialect::FooOp>(value));
-  results.push_back(rewriter.create<my_dialect::BarOp>());
+  return {rewriter.create<my_dialect::FooOp>(value), rewriter.create<my_dialect::BarOp>()};
 }];
 
 Pattern {
   let root = op<my_dialect.foo>(input: Value);
   rewrite root with {
     // Invoke the native rewrite and use the results when replacing the root.
     let results = BuildOp(input);
     replace root with (results.foo, results.bar);
   }
 }
 ```
 
 The arguments of the rewrite are accessible within the code block via the
-same name. The type of these native variables are mapped directly to the
-corresponding MLIR type of the [core constraint](#core-constraints) used. For
-example, an `Op` corresponds to a variable of type `Operation *`.
-
-The results of the rewrite can be populated using the provided `results`
-variable. This variable is a `PDLResultList`, and expects results to be
-populated in the order that they are defined within the result list of the
-rewrite declaration.
+same name. See the ["type translation"](#native-rewrite-type-translations) below for
+detailed information on how PDLL types are converted to native types. In addition to the
+PDLL arguments, the code block may also access the current `PatternRewriter` using
+`rewriter`. See the ["result translation"](#native-rewrite-result-translation) section
+for detailed information on how the result type of the native function is determined.
+
+Taking the rewrite defined above as an example, this function would roughly be
+translated into:
+
+```c++
+std::tuple<my_dialect::FooOp, my_dialect::BarOp> BuildOp(Value value) {
+  return {rewriter.create<my_dialect::FooOp>(value), rewriter.create<my_dialect::BarOp>()};
+}
+```
 
-In addition to the above, the code block may also access the current
-`PatternRewriter` using `rewriter`.
+###### Native Rewrite Type Translations
+
+The types of argument and result variables are generally mapped to the corresponding
+MLIR type of the [constraint](#constraints) used. The rules of native `Rewrite` type translation
+are identical to those of native `Constraint`s, please view the corresponding
+[native `Constraint` type translation](#native-constraint-type-translations) section for a
+detailed description of how the mapped type of a variable is determined.
+
+###### Native Rewrite Result Translation
+
+The results of a native rewrite are directly translated to the results of the native function,
+using the type translation rules [described above](#native-rewrite-type-translations). The section
+below describes the various result translation scenarios:
+
+* Zero Result
+
+```pdll
+Rewrite createOp() [{
+  rewriter.create<my_dialect::FooOp>();
+}];
+```
+
+In the case where a native `Rewrite` has no results, the native function returns `void`:
+
+```c++
+void createOp(PatternRewriter &rewriter) {
+  rewriter.create<my_dialect::FooOp>();
+}
+```
+
+* Single Result
+
+```pdll
+Rewrite createOp() -> Op<my_dialect.foo> [{
+  return rewriter.create<my_dialect::FooOp>();
+}];
+```
+
+In the case where a native `Rewrite` has a single result, the native function returns the corresponding
+native type for that single result:
+
+```c++
+my_dialect::FooOp createOp(PatternRewriter &rewriter) {
+  return rewriter.create<my_dialect::FooOp>();
+}
+```
+
+* Multi Result
+
+```pdll
+Rewrite complexRewrite(value: Value) -> (Op<my_dialect.foo>, FunctionOpInterface) [{
+  ...
+}];
+```
+
+In the case where a native `Rewrite` has multiple results, the native function returns a `std::tuple<...>`
+containing the corresponding native types for each of the results:
+
+```c++
+std::tuple<my_dialect::FooOp, FunctionOpInterface>
+complexRewrite(PatternRewriter &rewriter, Value value) {
+  ...
+}
+```
 
 #### Defining Rewrites Inline
 
 In addition to global scope, PDLL Rewrites and Native Rewrites defined in PDLL
 may be specified *inline* at any level of nesting. This means that they may be
 defined in Patterns, other Rewrites, etc:
 
 ```pdll

mlir/include/mlir/IR/PatternMatch.h

verifyAsArg(function_ref<LogicalResult(const Twine &)> errorFn,
               BaseT baseValue, size_t argIdx) {
     return TypeSwitch<BaseT, LogicalResult>(baseValue)
         .Case([&](T) { return success(); })
         .Default([&](BaseT) {
           return errorFn("expected argument " + Twine(argIdx) +
                          " to be of type: " + llvm::getTypeName<T>());
         });
   }
+  using ProcessPDLValueBasedOn<T, BaseT>::verifyAsArg;
+
   static T processAsArg(BaseT baseValue) {
     return baseValue.template cast<T>();
   }
+  using ProcessPDLValueBasedOn<T, BaseT>::processAsArg;
+
   static void processAsResult(PatternRewriter &, PDLResultList &results,
                               T value) {
     results.push_back(value);
   }
 };
 
 //===----------------------------------------------------------------------===//
 // Attribute
 
 template <>
 struct ProcessPDLValue<Attribute> : public ProcessBuiltinPDLValue<Attribute> {};
 template <typename T>
 struct ProcessPDLValue<T,
                        std::enable_if_t<std::is_base_of<Attribute, T>::value>>
     : public ProcessDerivedPDLValue<T, Attribute> {};
 
 /// Handling for various Attribute value types.
 template <>
 struct ProcessPDLValue<StringRef>
     : public ProcessPDLValueBasedOn<StringRef, StringAttr> {
   static StringRef processAsArg(StringAttr value) { return value.getValue(); }
+  using ProcessPDLValueBasedOn<StringRef, StringAttr>::processAsArg;
+
   static void processAsResult(PatternRewriter &rewriter, PDLResultList &results,
                               StringRef value) {
     results.push_back(rewriter.getStringAttr(value));
   }
 };
 template <>
 struct ProcessPDLValue<std::string>
     : public ProcessPDLValueBasedOn<std::string, StringAttr> {

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

 /// represents either an input operation to match, or an operation to create
 /// within a rewrite.
 class OperationExpr final
     : public Node::NodeBase<OperationExpr, Expr>,
       private llvm::TrailingObjects<OperationExpr, Expr *,
                                     NamedAttributeDecl *> {
 public:
   static OperationExpr *create(Context &ctx, SMRange loc,
+                               const ods::Operation *odsOp,
                                const OpNameDecl *nameDecl,
                                ArrayRef<Expr *> operands,
                                ArrayRef<Expr *> resultTypes,
                                ArrayRef<NamedAttributeDecl *> attributes);
 
   /// Return the name of the operation, or None if there isn't one.
   Optional<StringRef> getName() const;
 

 ///     - This is a native constraint, i.e. a constraint whose implementation is
 ///       defined in C++(or potentially some other non-PDLL language). The
 ///       implementation of this constraint is specified as a string code block
 ///       in PDLL.
 ///   * a PDLL constraint
 ///     - This is a constraint which is defined using only PDLL constructs.
 class UserConstraintDecl final
     : public Node::NodeBase<UserConstraintDecl, ConstraintDecl>,
-      llvm::TrailingObjects<UserConstraintDecl, VariableDecl *> {
+      llvm::TrailingObjects<UserConstraintDecl, VariableDecl *, StringRef> {
 public:
   /// Create a native constraint with the given optional code block.
-  static UserConstraintDecl *createNative(Context &ctx, const Name &name,
-                                          ArrayRef<VariableDecl *> inputs,
-                                          ArrayRef<VariableDecl *> results,
-                                          Optional<StringRef> codeBlock,
-                                          Type resultType) {
-    return createImpl(ctx, name, inputs, results, codeBlock, /*body=*/nullptr,
-                      resultType);
+  static UserConstraintDecl *
+  createNative(Context &ctx, const Name &name, ArrayRef<VariableDecl *> inputs,
+               ArrayRef<VariableDecl *> results, Optional<StringRef> codeBlock,
+               Type resultType, ArrayRef<StringRef> nativeInputTypes = {}) {
+    return createImpl(ctx, name, inputs, nativeInputTypes, results, codeBlock,
+                      /*body=*/nullptr, resultType);
   }
 
   /// Create a PDLL constraint with the given body.
   static UserConstraintDecl *createPDLL(Context &ctx, const Name &name,
                                         ArrayRef<VariableDecl *> inputs,
                                         ArrayRef<VariableDecl *> results,
                                         const CompoundStmt *body,
                                         Type resultType) {
-    return createImpl(ctx, name, inputs, results, /*codeBlock=*/llvm::None,
-                      body, resultType);
+    return createImpl(ctx, name, inputs, /*nativeInputTypes=*/llvm::None,
+                      results, /*codeBlock=*/llvm::None, body, resultType);
   }
 
   /// Return the name of the constraint.
   const Name &getName() const { return *Decl::getName(); }
 
   /// Return the input arguments of this constraint.
   MutableArrayRef<VariableDecl *> getInputs() {
     return {getTrailingObjects<VariableDecl *>(), numInputs};
   }
   ArrayRef<VariableDecl *> getInputs() const {
     return const_cast<UserConstraintDecl *>(this)->getInputs();
   }
 
+  /// Return the explicit native type to use for the given input. Returns None
+  /// if no explicit type was set.
+  Optional<StringRef> getNativeInputType(unsigned index) const;
+
   /// Return the explicit results of the constraint declaration. May be empty,
   /// even if the constraint has results (e.g. in the case of inferred results).
   MutableArrayRef<VariableDecl *> getResults() {
     return {getTrailingObjects<VariableDecl *>() + numInputs, numResults};
   }
   ArrayRef<VariableDecl *> getResults() const {
     return const_cast<UserConstraintDecl *>(this)->getResults();
   }
public:
   /// Returns true if this constraint is external.
   bool isExternal() const { return !constraintBody && !codeBlock; }
 
 private:
   /// Create either a PDLL constraint or a native constraint with the given
   /// components.
   static UserConstraintDecl *
   createImpl(Context &ctx, const Name &name, ArrayRef<VariableDecl *> inputs,
+             ArrayRef<StringRef> nativeInputTypes,
              ArrayRef<VariableDecl *> results, Optional<StringRef> codeBlock,
              const CompoundStmt *body, Type resultType);
 
-  UserConstraintDecl(const Name &name, unsigned numInputs, unsigned numResults,
+  UserConstraintDecl(const Name &name, unsigned numInputs,
+                     bool hasNativeInputTypes, unsigned numResults,
                      Optional<StringRef> codeBlock, const CompoundStmt *body,
                      Type resultType)
       : Base(name.getLoc(), &name), numInputs(numInputs),
         numResults(numResults), codeBlock(codeBlock), constraintBody(body),
         resultType(resultType) {}
 
   /// The number of inputs to this constraint.
   unsigned numInputs;
 
   /// The number of explicit results to this constraint.
   unsigned numResults;
 
   /// The optional code block of this constraint.
   Optional<StringRef> codeBlock;
 
   /// The optional body of this constraint.
   const CompoundStmt *constraintBody;
 
   /// The result type of the constraint.
   Type resultType;
 
+  /// Flag indicating if this constraint has explicit native input types.
+  bool hasNativeInputTypes;
+
   /// Allow access to various internals.
-  friend llvm::TrailingObjects<UserConstraintDecl, VariableDecl *>;
+  friend llvm::TrailingObjects<UserConstraintDecl, VariableDecl *, StringRef>;
+  size_t numTrailingObjects(OverloadToken<VariableDecl *>) const {
+    return numInputs + numResults;
+  }
 };
 
 //===----------------------------------------------------------------------===//
 // NamedAttributeDecl
 //===----------------------------------------------------------------------===//
 
 /// This Decl represents a NamedAttribute, and contains a string name and
 /// attribute value.
public:
 
   /// Return the result type of this decl.
   Type getResultType() const {
     if (const auto *cst = dyn_cast<UserConstraintDecl>(this))
       return cst->getResultType();
     return cast<UserRewriteDecl>(this)->getResultType();
   }
 
+  /// Return the explicit results of the declaration. Note that these may be
+  /// empty, even if the callable has results (e.g. in the case of inferred
+  /// results).
+  ArrayRef<VariableDecl *> getResults() const {
+    if (const auto *cst = dyn_cast<UserConstraintDecl>(this))
+      return cst->getResults();
+    return cast<UserRewriteDecl>(this)->getResults();
+  }
+
+  /// Return the optional code block of this callable, if this is a native
+  /// callable with a provided implementation.
+  Optional<StringRef> getCodeBlock() const {
+    if (const auto *cst = dyn_cast<UserConstraintDecl>(this))
+      return cst->getCodeBlock();
+    return cast<UserRewriteDecl>(this)->getCodeBlock();
+  }
+
   /// Support LLVM type casting facilities.
   static bool classof(const Node *decl) {
     return isa<UserConstraintDecl, UserRewriteDecl>(decl);
   }
 };
 
 //===----------------------------------------------------------------------===//
 // VariableDecl

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

 //===- Types.h --------------------------------------------------*- 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_AST_TYPES_H_
 #define MLIR_TOOLS_PDLL_AST_TYPES_H_
 
 #include "mlir/Support/LLVM.h"
 #include "mlir/Support/StorageUniquer.h"
 
 namespace mlir {
 namespace pdll {
+namespace ods {
+class Operation;
+} // namespace ods
+
 namespace ast {
 class Context;
 
 namespace detail {
 struct AttributeTypeStorage;
 struct ConstraintTypeStorage;
 struct OperationTypeStorage;
 struct RangeTypeStorage;

 /// This class represents a PDLL type that corresponds to an mlir::Operation.
 class OperationType : public Type::TypeBase<detail::OperationTypeStorage> {
 public:
   using Base::Base;
 
   /// Return an instance of the Operation type with an optional operation name.
   /// If no name is provided, this type may refer to any operation.
   static OperationType get(Context &context,
-                           Optional<StringRef> name = llvm::None);
+                           Optional<StringRef> name = llvm::None,
+                           const ods::Operation *odsOp = nullptr);
 
   /// Return the name of this operation type, or None if it doesn't have on.
   Optional<StringRef> getName() const;
+
+  /// Return the ODS operation that this type refers to, or nullptr if the ODS
+  /// operation is unknown.
+  const ods::Operation *getODSOperation() const;
 };
 
 //===----------------------------------------------------------------------===//
 // RangeType
 //===----------------------------------------------------------------------===//
 
 /// This class represents a PDLL type that corresponds to a range of elements
 /// with a given element type.

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

auto getDialects() const {
     return llvm::make_pointee_range(llvm::make_second_range(dialects));
   }
 
   /// 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,
-                  bool supportsResultTypeInferrence, SMLoc loc);
+                  StringRef nativeClassName, bool supportsResultTypeInferrence,
+                  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;
 

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

public:
   /// 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,
-                  bool supportsResultTypeInferrence, SMLoc loc);
+                  StringRef nativeClassName, bool supportsResultTypeInferrence,
+                  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;

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

public:
   StringRef getName() const { return name; }
 
   /// Returns the summary of the operation.
   StringRef getSummary() const { return summary; }
 
   /// Returns the description of the operation.
   StringRef getDescription() const { return description; }
 
+  /// Returns the native class name of the operation.
+  StringRef getNativeClassName() const { return nativeClassName; }
+
   /// Returns the attributes of this operation.
   ArrayRef<Attribute> getAttributes() const { return attributes; }
 
   /// Returns the operands of this operation.
   ArrayRef<OperandOrResult> getOperands() const { return operands; }
 
   /// Returns the results of this operation.
   ArrayRef<OperandOrResult> getResults() const { return results; }
 
   /// Return if the operation is known to support result type inferrence.
   bool hasResultTypeInferrence() const { return supportsTypeInferrence; }
 
 private:
   Operation(StringRef name, StringRef summary, StringRef desc,
-            bool supportsTypeInferrence, SMLoc loc);
+            StringRef nativeClassName, bool supportsTypeInferrence, SMLoc loc);
 
   /// The name of the operation.
   std::string name;
 
   /// The documentation of the operation.
   std::string summary;
   std::string description;
 
+  /// The native class name of the operation, used when generating native code.
+  std::string nativeClassName;
+
   /// Flag indicating if the operation is known to support type inferrence.
   bool supportsTypeInferrence;
 
   /// The source location of this operation.
   SMRange location;
 
   /// The operands of the operation.
   SmallVector<OperandOrResult> operands;

mlir/lib/Tools/PDLL/AST/Nodes.cpp

MemberAccessExpr *MemberAccessExpr::create(Context &ctx, SMRange loc,
   return new (ctx.getAllocator().Allocate<MemberAccessExpr>()) MemberAccessExpr(
       loc, parentExpr, memberName.copy(ctx.getAllocator()), type);
 }
 
 //===----------------------------------------------------------------------===//
 // OperationExpr
 //===----------------------------------------------------------------------===//
 
-OperationExpr *OperationExpr::create(
-    Context &ctx, SMRange loc, const OpNameDecl *name,
-    ArrayRef<Expr *> operands, ArrayRef<Expr *> resultTypes,
+OperationExpr *
+OperationExpr::create(Context &ctx, SMRange loc, const ods::Operation *odsOp,
+                      const OpNameDecl *name, ArrayRef<Expr *> operands,
+                      ArrayRef<Expr *> resultTypes,
                       ArrayRef<NamedAttributeDecl *> attributes) {
   unsigned allocSize =
       OperationExpr::totalSizeToAlloc<Expr *, NamedAttributeDecl *>(
           operands.size() + resultTypes.size(), attributes.size());
   void *rawData =
       ctx.getAllocator().Allocate(allocSize, alignof(OperationExpr));
 
-  Type resultType = OperationType::get(ctx, name->getName());
+  Type resultType = OperationType::get(ctx, name->getName(), odsOp);
   OperationExpr *opExpr = new (rawData)
       OperationExpr(loc, resultType, name, operands.size(), resultTypes.size(),
                     attributes.size(), name->getLoc());
   std::uninitialized_copy(operands.begin(), operands.end(),
                           opExpr->getOperands().begin());
   std::uninitialized_copy(resultTypes.begin(), resultTypes.end(),
                           opExpr->getResultTypes().begin());
   std::uninitialized_copy(attributes.begin(), attributes.end(),
ValueRangeConstraintDecl *ValueRangeConstraintDecl::create(Context &ctx,
   return new (ctx.getAllocator().Allocate<ValueRangeConstraintDecl>())
       ValueRangeConstraintDecl(loc, typeExpr);
 }
 
 //===----------------------------------------------------------------------===//
 // UserConstraintDecl
 //===----------------------------------------------------------------------===//
 
+Optional<StringRef>
+UserConstraintDecl::getNativeInputType(unsigned index) const {
+  return hasNativeInputTypes ? getTrailingObjects<StringRef>()[index]
+                             : Optional<StringRef>();
+}
+
 UserConstraintDecl *UserConstraintDecl::createImpl(
     Context &ctx, const Name &name, ArrayRef<VariableDecl *> inputs,
-    ArrayRef<VariableDecl *> results, Optional<StringRef> codeBlock,
-    const CompoundStmt *body, Type resultType) {
-  unsigned allocSize = UserConstraintDecl::totalSizeToAlloc<VariableDecl *>(
-      inputs.size() + results.size());
+    ArrayRef<StringRef> nativeInputTypes, ArrayRef<VariableDecl *> results,
+    Optional<StringRef> codeBlock, const CompoundStmt *body, Type resultType) {
+  bool hasNativeInputTypes = !nativeInputTypes.empty();
+  assert(!hasNativeInputTypes || nativeInputTypes.size() == inputs.size());
+
+  unsigned allocSize =
+      UserConstraintDecl::totalSizeToAlloc<VariableDecl *, StringRef>(
+          inputs.size() + results.size(),
+          hasNativeInputTypes ? inputs.size() : 0);
   void *rawData =
       ctx.getAllocator().Allocate(allocSize, alignof(UserConstraintDecl));
   if (codeBlock)
     codeBlock = codeBlock->copy(ctx.getAllocator());
 
-  UserConstraintDecl *decl = new (rawData) UserConstraintDecl(
-      name, inputs.size(), results.size(), codeBlock, body, resultType);
+  UserConstraintDecl *decl = new (rawData)
+      UserConstraintDecl(name, inputs.size(), hasNativeInputTypes,
+                         results.size(), codeBlock, body, resultType);
   std::uninitialized_copy(inputs.begin(), inputs.end(),
                           decl->getInputs().begin());
   std::uninitialized_copy(results.begin(), results.end(),
                           decl->getResults().begin());
+  if (hasNativeInputTypes) {
+    StringRef *nativeInputTypesPtr = decl->getTrailingObjects<StringRef>();
+    for (unsigned i = 0, e = inputs.size(); i < e; ++i)
+      nativeInputTypesPtr[i] = nativeInputTypes[i].copy(ctx.getAllocator());
+  }
+
   return decl;
 }
 
 //===----------------------------------------------------------------------===//
 // NamedAttributeDecl
 //===----------------------------------------------------------------------===//
 
 NamedAttributeDecl *NamedAttributeDecl::create(Context &ctx, const Name &name,

mlir/lib/Tools/PDLL/AST/TypeDetail.h

 
 struct ConstraintTypeStorage : public TypeStorageBase<ConstraintTypeStorage> {};
 
 //===----------------------------------------------------------------------===//
 // OperationType
 //===----------------------------------------------------------------------===//
 
 struct OperationTypeStorage
-    : public TypeStorageBase<OperationTypeStorage, StringRef> {
+    : public TypeStorageBase<OperationTypeStorage,
+                             std::pair<StringRef, const ods::Operation *>> {
   using Base::Base;
 
   static OperationTypeStorage *
-  construct(StorageUniquer::StorageAllocator &alloc, StringRef key) {
-    return new (alloc.allocate<OperationTypeStorage>())
-        OperationTypeStorage(alloc.copyInto(key));
+  construct(StorageUniquer::StorageAllocator &alloc,
+            const std::pair<StringRef, const ods::Operation *> &key) {
+    return new (alloc.allocate<OperationTypeStorage>()) OperationTypeStorage(
+        std::make_pair(alloc.copyInto(key.first), key.second));
   }
 };
 
 //===----------------------------------------------------------------------===//
 // RangeType
 //===----------------------------------------------------------------------===//
 
 struct RangeTypeStorage : public TypeStorageBase<RangeTypeStorage, Type> {

mlir/lib/Tools/PDLL/AST/Types.cpp

 //===- Types.cpp ----------------------------------------------------------===//
 //
 // 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
 //
 //===----------------------------------------------------------------------===//
 
 #include "mlir/Tools/PDLL/AST/Types.h"
 #include "TypeDetail.h"
 #include "mlir/Tools/PDLL/AST/Context.h"
 
 using namespace mlir;
+using namespace mlir::pdll;
 using namespace mlir::pdll::ast;
 
 MLIR_DEFINE_EXPLICIT_TYPE_ID(mlir::pdll::ast::detail::AttributeTypeStorage)
 MLIR_DEFINE_EXPLICIT_TYPE_ID(mlir::pdll::ast::detail::ConstraintTypeStorage)
 MLIR_DEFINE_EXPLICIT_TYPE_ID(mlir::pdll::ast::detail::OperationTypeStorage)
 MLIR_DEFINE_EXPLICIT_TYPE_ID(mlir::pdll::ast::detail::RangeTypeStorage)
 MLIR_DEFINE_EXPLICIT_TYPE_ID(mlir::pdll::ast::detail::RewriteTypeStorage)
 MLIR_DEFINE_EXPLICIT_TYPE_ID(mlir::pdll::ast::detail::TupleTypeStorage)

 ConstraintType ConstraintType::get(Context &context) {
   return context.getTypeUniquer().get<ImplTy>();
 }
 
 //===----------------------------------------------------------------------===//
 // OperationType
 //===----------------------------------------------------------------------===//
 
-OperationType OperationType::get(Context &context, Optional<StringRef> name) {
+OperationType OperationType::get(Context &context, Optional<StringRef> name,
+                                 const ods::Operation *odsOp) {
   return context.getTypeUniquer().get<ImplTy>(
-      /*initFn=*/function_ref<void(ImplTy *)>(), name.getValueOr(""));
+      /*initFn=*/function_ref<void(ImplTy *)>(),
+      std::make_pair(name.getValueOr(""), odsOp));
 }
 
 Optional<StringRef> OperationType::getName() const {
-  StringRef name = getImplAs<ImplTy>()->getValue();
+  StringRef name = getImplAs<ImplTy>()->getValue().first;
   return name.empty() ? Optional<StringRef>() : Optional<StringRef>(name);
 }
 
+const ods::Operation *OperationType::getODSOperation() const {
+  return getImplAs<ImplTy>()->getValue().second;
+}
+
 //===----------------------------------------------------------------------===//
 // RangeType
 //===----------------------------------------------------------------------===//
 
 RangeType RangeType::get(Context &context, Type elementType) {
   return context.getTypeUniquer().get<ImplTy>(
       /*initFn=*/function_ref<void(ImplTy *)>(), elementType);
 }

mlir/lib/Tools/PDLL/CodeGen/CPPGen.cpp

 //
 //===----------------------------------------------------------------------===//
 
 #include "mlir/Tools/PDLL/CodeGen/CPPGen.h"
 #include "mlir/Dialect/PDL/IR/PDL.h"
 #include "mlir/Dialect/PDL/IR/PDLOps.h"
 #include "mlir/IR/BuiltinOps.h"
 #include "mlir/Tools/PDLL/AST/Nodes.h"
+#include "mlir/Tools/PDLL/ODS/Operation.h"
 #include "llvm/ADT/SmallString.h"
 #include "llvm/ADT/StringExtras.h"
 #include "llvm/ADT/StringSet.h"
 #include "llvm/ADT/TypeSwitch.h"
 #include "llvm/Support/ErrorHandling.h"
 #include "llvm/Support/FormatVariadic.h"
 
 using namespace mlir;
private:
   /// Generate C++ code for all user defined constraints and rewrites with
   /// native code.
   void generateConstraintAndRewrites(const ast::Module &astModule,
                                      ModuleOp module,
                                      StringSet<> &nativeFunctions);
   void generate(const ast::UserConstraintDecl *decl,
                 StringSet<> &nativeFunctions);
   void generate(const ast::UserRewriteDecl *decl, StringSet<> &nativeFunctions);
-  void generateConstraintOrRewrite(StringRef name, bool isConstraint,
-                                   ArrayRef<ast::VariableDecl *> inputs,
-                                   StringRef codeBlock,
+  void generateConstraintOrRewrite(const ast::CallableDecl *decl,
+                                   bool isConstraint,
                                    StringSet<> &nativeFunctions);
 
+  /// Return the native name of the type to use for the single input argument of
+  /// the given constraint.
+  StringRef getNativeTypeName(const ast::UserConstraintDecl *cst);
+
+  /// Return the native name for the type of the given type.
+  StringRef getNativeTypeName(ast::Type type);
+
+  /// Return the native name for the type of the given variable decl.
+  StringRef getNativeTypeName(ast::VariableDecl *decl);
+
   /// The stream to output to.
   raw_ostream &os;
 };
 } // namespace
 
 void CodeGen::generate(const ast::Module &astModule, ModuleOp module) {
   SetVector<std::string, SmallVector<std::string>, StringSet<>> patternNames;
   StringSet<> nativeFunctions;
TypeSwitch<const ast::Decl *>(decl)
                   usedFns.contains(decl->getName().getName()))
                 this->generate(decl, nativeFunctions);
             });
   }
 }
 
 void CodeGen::generate(const ast::UserConstraintDecl *decl,
                        StringSet<> &nativeFunctions) {
-  return generateConstraintOrRewrite(decl->getName().getName(),
-                                     /*isConstraint=*/true, decl->getInputs(),
-                                     *decl->getCodeBlock(), nativeFunctions);
+  return generateConstraintOrRewrite(cast<ast::CallableDecl>(decl),
+                                     /*isConstraint=*/true, nativeFunctions);
 }
 
 void CodeGen::generate(const ast::UserRewriteDecl *decl,
                        StringSet<> &nativeFunctions) {
-  return generateConstraintOrRewrite(decl->getName().getName(),
-                                     /*isConstraint=*/false, decl->getInputs(),
-                                     *decl->getCodeBlock(), nativeFunctions);
+  return generateConstraintOrRewrite(cast<ast::CallableDecl>(decl),
+                                     /*isConstraint=*/false, nativeFunctions);
 }
 
-void CodeGen::generateConstraintOrRewrite(StringRef name, bool isConstraint,
-                                          ArrayRef<ast::VariableDecl *> inputs,
-                                          StringRef codeBlock,
-                                          StringSet<> &nativeFunctions) {
-  nativeFunctions.insert(name);
+StringRef CodeGen::getNativeTypeName(const ast::UserConstraintDecl *cst) {
+  if (Optional<StringRef> name = cst->getNativeInputType(0))
+    return *name;
+  return getNativeTypeName(cst->getInputs()[0]);
+}
 
-  // TODO: Should there be something explicit for handling optionality?
-  auto getCppType = [&](ast::Type type) -> StringRef {
-    return llvm::TypeSwitch<ast::Type, StringRef>(type)
-        .Case([&](ast::AttributeType) { return "::mlir::Attribute"; })
-        .Case([&](ast::OperationType) {
-          // TODO: Allow using the derived Op class when possible.
-          return "::mlir::Operation *";
-        })
-        .Case([&](ast::TypeType) { return "::mlir::Type"; })
-        .Case([&](ast::ValueType) { return "::mlir::Value"; })
-        .Case([&](ast::TypeRangeType) { return "::mlir::TypeRange"; })
-        .Case([&](ast::ValueRangeType) { return "::mlir::ValueRange"; });
-  };
-  os << "static " << (isConstraint ? "::mlir::LogicalResult " : "void ") << name
-     << "PDLFn(::mlir::PatternRewriter &rewriter, "
-     << (isConstraint ? "" : "::mlir::PDLResultList &results, ")
-     << "::llvm::ArrayRef<::mlir::PDLValue> values) {\n";
-
-  const char *argumentInitStr = R"(
-  {0} {1} = {{};
-  if (values[{2}])
-    {1} = values[{2}].cast<{0}>();
-  (void){1};
-)";
-  for (const auto &it : llvm::enumerate(inputs)) {
-    const ast::VariableDecl *input = it.value();
-    os << llvm::formatv(argumentInitStr, getCppType(input->getType()),
-                        input->getName().getName(), it.index());
-  }
+StringRef CodeGen::getNativeTypeName(ast::Type type) {
+      .Case([&](ast::OperationType opType) -> StringRef {
+        // Use the derived Op class when available.
+        if (const auto *odsOp = opType.getODSOperation())
+          return odsOp->getNativeClassName();
 
-  os << "  " << codeBlock.trim() << "\n}\n";
+StringRef CodeGen::getNativeTypeName(ast::VariableDecl *decl) {
+  // Try to extract a type name from the variable's constraints.
+  for (ast::ConstraintRef &cst : decl->getConstraints()) {
+    if (auto *userCst = dyn_cast<ast::UserConstraintDecl>(cst.constraint))
+      return getNativeTypeName(userCst);
+  }
+
+  // Otherwise, use the type of the variable.
+  return getNativeTypeName(decl->getType());
+}
+
+void CodeGen::generateConstraintOrRewrite(const ast::CallableDecl *decl,
+                                          bool isConstraint,
+                                          StringSet<> &nativeFunctions) {
+  StringRef name = decl->getName()->getName();
+  nativeFunctions.insert(name);
+
+  os << "static ";
+
+  // TODO: Work out a proper modeling for "optionality".
+
+  // Emit the result type.
+  // If this is a constraint, we always return a LogicalResult.
+  // TODO: This will need to change if we allow Constraints to return values as
+  // well.
+  if (isConstraint) {
+    os << "::mlir::LogicalResult";
+  } else {
+    // Otherwise, generate a type based on the results of the callable.
+    // If the callable has explicit results, use those to build the result.
+    // Otherwise, use the type of the callable.
+    ArrayRef<ast::VariableDecl *> results = decl->getResults();
+    if (results.empty()) {
+      os << "void";
+    } else if (results.size() == 1) {
+      os << getNativeTypeName(results[0]);
+    } else {
+      os << "std::tuple<";
+      llvm::interleaveComma(results, os, [&](ast::VariableDecl *result) {
+        os << getNativeTypeName(result);
+      });
+      os << ">";
+    }
+  }
+
+  os << " " << name << "PDLFn(::mlir::PatternRewriter &rewriter";
+  if (!decl->getInputs().empty()) {
+    os << ", ";
+    llvm::interleaveComma(decl->getInputs(), os, [&](ast::VariableDecl *input) {
+      os << getNativeTypeName(input) << " " << input->getName().getName();
+    });
+  }
+  os << ") {\n";
+  os << "  " << decl->getCodeBlock()->trim() << "\n}\n\n";
 }
 
 //===----------------------------------------------------------------------===//
 // CPPGen
 //===----------------------------------------------------------------------===//
 
 void mlir::pdll::codegenPDLLToCPP(const ast::Module &astModule, ModuleOp module,
                                   raw_ostream &os) {
   CodeGen codegen(os);
   codegen.generate(astModule, module);
 }

mlir/lib/Tools/PDLL/CodeGen/MLIRGen.cpp

if (ast::OperationType opType = parentType.dyn_cast<ast::OperationType>()) {
     if (isa<ast::AllResultsMemberAccessExpr>(expr)) {
       Type mlirType = genType(expr->getType());
       if (mlirType.isa<pdl::ValueType>())
         return builder.create<pdl::ResultOp>(loc, mlirType, parentExprs[0],
                                              builder.getI32IntegerAttr(0));
       return builder.create<pdl::ResultsOp>(loc, mlirType, parentExprs[0]);
     }
 
-    assert(opType.getName() && "expected valid operation name");
-    const ods::Operation *odsOp = odsContext.lookupOperation(*opType.getName());
+    const ods::Operation *odsOp = opType.getODSOperation();
     assert(odsOp && "expected valid ODS operation information");
 
     // Find the result with the member name or by index.
     ArrayRef<ods::OperandOrResult> results = odsOp->getResults();
     unsigned resultIndex = results.size();
     if (llvm::isDigit(name[0])) {
       name.getAsInteger(/*Radix=*/10, resultIndex);
     } else {

mlir/lib/Tools/PDLL/ODS/Context.cpp

 
 const Dialect *Context::lookupDialect(StringRef name) const {
   auto it = dialects.find(name);
   return it == dialects.end() ? nullptr : &*it->second;
 }
 
 std::pair<Operation *, bool>
 Context::insertOperation(StringRef name, StringRef summary, StringRef desc,
+                         StringRef nativeClassName,
                          bool supportsResultTypeInferrence, SMLoc loc) {
   std::pair<StringRef, StringRef> dialectAndName = name.split('.');
   return insertDialect(dialectAndName.first)
-      .insertOperation(name, summary, desc, supportsResultTypeInferrence, loc);
+      .insertOperation(name, summary, desc, nativeClassName,
+                       supportsResultTypeInferrence, loc);
 }
 
 const Operation *Context::lookupOperation(StringRef name) const {
   std::pair<StringRef, StringRef> dialectAndName = name.split('.');
   if (const Dialect *dialect = lookupDialect(dialectAndName.first))
     return dialect->lookupOperation(name);
   return nullptr;
 }

mlir/lib/Tools/PDLL/ODS/Dialect.cpp

 // Dialect
 //===----------------------------------------------------------------------===//
 
 Dialect::Dialect(StringRef name) : name(name.str()) {}
 Dialect::~Dialect() = default;
 
 std::pair<Operation *, bool>
 Dialect::insertOperation(StringRef name, StringRef summary, StringRef desc,
+                         StringRef nativeClassName,
                          bool supportsResultTypeInferrence, llvm::SMLoc loc) {
   std::unique_ptr<Operation> &operation = operations[name];
   if (operation)
     return std::make_pair(&*operation, /*wasInserted*/ false);
 
-  operation.reset(
-      new Operation(name, summary, desc, supportsResultTypeInferrence, loc));
+  operation.reset(new Operation(name, summary, desc, nativeClassName,
+                                supportsResultTypeInferrence, loc));
   return std::make_pair(&*operation, /*wasInserted*/ true);
 }
 
 Operation *Dialect::lookupOperation(StringRef name) const {
   auto it = operations.find(name);
   return it != operations.end() ? it->second.get() : nullptr;
 }

mlir/lib/Tools/PDLL/ODS/Operation.cpp

 using namespace mlir;
 using namespace mlir::pdll::ods;
 
 //===----------------------------------------------------------------------===//
 // Operation
 //===----------------------------------------------------------------------===//
 
 Operation::Operation(StringRef name, StringRef summary, StringRef desc,
-                     bool supportsTypeInferrence, llvm::SMLoc loc)
+                     StringRef nativeClassName, bool supportsTypeInferrence,
+                     llvm::SMLoc loc)
     : name(name.str()), summary(summary.str()),
+      nativeClassName(nativeClassName.str()),
       supportsTypeInferrence(supportsTypeInferrence),
       location(loc, llvm::SMLoc::getFromPointer(loc.getPointer() + 1)) {
   llvm::raw_string_ostream descOS(description);
   raw_indented_ostream(descOS).printReindented(desc.rtrim(" \t"));
 }

mlir/lib/Tools/PDLL/Parser/Parser.cpp

private:
   void processTdIncludeRecords(llvm::RecordKeeper &tdRecords,
                                SmallVectorImpl<ast::Decl *> &decls);
 
   /// Create a user defined native constraint for a constraint imported from
   /// ODS.
   template <typename ConstraintT>
   ast::Decl *createODSNativePDLLConstraintDecl(StringRef name,
                                                StringRef codeBlock, SMRange loc,
-                                               ast::Type type);
+                                               ast::Type type,
+                                               StringRef nativeType);
   template <typename ConstraintT>
   ast::Decl *
   createODSNativePDLLConstraintDecl(const tblgen::Constraint &constraint,
-                                    SMRange loc, ast::Type type);
+                                    SMRange loc, ast::Type type,
+                                    StringRef nativeType);
 
   //===--------------------------------------------------------------------===//
   // Decls
 
   /// This structure contains the set of pattern metadata that may be parsed.
   struct ParsedPatternMetadata {
     Optional<uint16_t> benefit;
     bool hasBoundedRecursion = false;
if (type == valueRangeTy) {
                                                      valueRangeTy);
       return success();
     }
 
     // Allow conversion to a single value by constraining the result range.
     if (type == valueTy) {
       // If the operation is registered, we can verify if it can ever have a
       // single result.
-      Optional<StringRef> opName = exprOpType.getName();
-      if (const ods::Operation *odsOp = lookupODSOperation(opName)) {
+      if (const ods::Operation *odsOp = exprOpType.getODSOperation()) {
         if (odsOp->getResults().empty()) {
           return emitConvertError()->attachNote(
               llvm::formatv("see the definition of `{0}`, which was defined "
                             "with zero results",
                             odsOp->getName()),
               odsOp->getLoc());
         }
 
for (llvm::Record *def : tdRecords.getAllDerivedDefinitions("Op")) {
     // Check to see if this operation is known to support type inferrence.
     bool supportsResultTypeInferrence =
         op.getTrait("::mlir::InferTypeOpInterface::Trait");
 
     bool inserted = false;
     ods::Operation *odsOp = nullptr;
     std::tie(odsOp, inserted) = odsContext.insertOperation(
         op.getOperationName(), op.getSummary(), op.getDescription(),
-        supportsResultTypeInferrence, op.getLoc().front());
+        op.getQualCppClassName(), supportsResultTypeInferrence,
+        op.getLoc().front());
 
     // Ignore operations that have already been added.
     if (!inserted)
       continue;
 
     for (const tblgen::NamedAttribute &attr : op.getAttributes()) {
       odsOp->appendAttribute(
           attr.name, attr.attr.isOptional(),
           odsContext.insertAttributeConstraint(attr.attr.getUniqueDefName(),
                                                attr.attr.getSummary(),
                                                attr.attr.getStorageType()));
     }
     for (const tblgen::NamedTypeConstraint &operand : op.getOperands()) {
       odsOp->appendOperand(operand.name, getLengthKind(operand),
                            addTypeConstraint(operand));
     }
     for (const tblgen::NamedTypeConstraint &result : op.getResults()) {
       odsOp->appendResult(result.name, getLengthKind(result),
                           addTypeConstraint(result));
     }
   }
   /// Attr constraints.
   for (llvm::Record *def : tdRecords.getAllDerivedDefinitions("Attr")) {
     if (!def->isAnonymous() && !curDeclScope->lookup(def->getName())) {
+      tblgen::Attribute constraint(def);
       decls.push_back(
           createODSNativePDLLConstraintDecl<ast::AttrConstraintDecl>(
-              tblgen::AttrConstraint(def),
-              convertLocToRange(def->getLoc().front()), attrTy));
+              constraint, convertLocToRange(def->getLoc().front()), attrTy,
+              constraint.getStorageType()));
     }
   }
   /// Type constraints.
   for (llvm::Record *def : tdRecords.getAllDerivedDefinitions("Type")) {
     if (!def->isAnonymous() && !curDeclScope->lookup(def->getName())) {
+      tblgen::TypeConstraint constraint(def);
       decls.push_back(
           createODSNativePDLLConstraintDecl<ast::TypeConstraintDecl>(
-              tblgen::TypeConstraint(def),
-              convertLocToRange(def->getLoc().front()), typeTy));
+              constraint, convertLocToRange(def->getLoc().front()), typeTy,
+              constraint.getCPPClassName()));
     }
   }
   /// Interfaces.
   ast::Type opTy = ast::OperationType::get(ctx);
   for (llvm::Record *def : tdRecords.getAllDerivedDefinitions("Interface")) {
     StringRef name = def->getName();
     if (def->isAnonymous() || curDeclScope->lookup(name) ||
         def->isSubClassOf("DeclareInterfaceMethods"))
       continue;
     SMRange loc = convertLocToRange(def->getLoc().front());
 
-    StringRef className = def->getValueAsString("cppClassName");
-    StringRef cppNamespace = def->getValueAsString("cppNamespace");
+    std::string cppClassName =
+        llvm::formatv("{0}::{1}", def->getValueAsString("cppNamespace"),
+                      def->getValueAsString("cppClassName"))
+            .str();
     std::string codeBlock =
-        llvm::formatv("return ::mlir::success(llvm::isa<{0}::{1}>(self));",
-                      cppNamespace, className)
+        llvm::formatv("return ::mlir::success(llvm::isa<{0}>(self));",
+                      cppClassName)
             .str();
 
     if (def->isSubClassOf("OpInterface")) {
       decls.push_back(createODSNativePDLLConstraintDecl<ast::OpConstraintDecl>(
-          name, codeBlock, loc, opTy));
+          name, codeBlock, loc, opTy, cppClassName));
     } else if (def->isSubClassOf("AttrInterface")) {
       decls.push_back(
           createODSNativePDLLConstraintDecl<ast::AttrConstraintDecl>(
-              name, codeBlock, loc, attrTy));
+              name, codeBlock, loc, attrTy, cppClassName));
     } else if (def->isSubClassOf("TypeInterface")) {
       decls.push_back(
           createODSNativePDLLConstraintDecl<ast::TypeConstraintDecl>(
-              name, codeBlock, loc, typeTy));
+              name, codeBlock, loc, typeTy, cppClassName));
     }
   }
 }
 
 template <typename ConstraintT>
 ast::Decl *
 Parser::createODSNativePDLLConstraintDecl(StringRef name, StringRef codeBlock,
-                                          SMRange loc, ast::Type type) {
+                                          SMRange loc, ast::Type type,
+                                          StringRef nativeType) {
   // Build the single input parameter.
   ast::DeclScope *argScope = pushDeclScope();
   auto *paramVar = ast::VariableDecl::create(
       ctx, ast::Name::create(ctx, "self", loc), type,
       /*initExpr=*/nullptr, ast::ConstraintRef(ConstraintT::create(ctx, loc)));
   argScope->add(paramVar);
   popDeclScope();
 
   // Build the native constraint.
   auto *constraintDecl = ast::UserConstraintDecl::createNative(
       ctx, ast::Name::create(ctx, name, loc), paramVar,
-      /*results=*/llvm::None, codeBlock, ast::TupleType::get(ctx));
+      /*results=*/llvm::None, codeBlock, ast::TupleType::get(ctx), nativeType);
   curDeclScope->add(constraintDecl);
   return constraintDecl;
 }
 
 template <typename ConstraintT>
 ast::Decl *
 Parser::createODSNativePDLLConstraintDecl(const tblgen::Constraint &constraint,
-                                          SMRange loc, ast::Type type) {
+                                          SMRange loc, ast::Type type,
+                                          StringRef nativeType) {
   // Format the condition template.
   tblgen::FmtContext fmtContext;
   fmtContext.withSelf("self");
   std::string codeBlock = tblgen::tgfmt(
       "return ::mlir::success(" + constraint.getConditionTemplate() + ");",
       &fmtContext);
 
   return createODSNativePDLLConstraintDecl<ConstraintT>(
-      constraint.getUniqueDefName(), codeBlock, loc, type);
+      constraint.getUniqueDefName(), codeBlock, loc, type, nativeType);
 }
 
 //===----------------------------------------------------------------------===//
 // Decls
 
 FailureOr<ast::Decl *> Parser::parseTopLevelDecl() {
   FailureOr<ast::Decl *> decl;
   switch (curToken.getKind()) {
LogicalResult Parser::validateVariableConstraint(const ast::ConstraintRef &ref,
   if (const auto *cst = dyn_cast<ast::AttrConstraintDecl>(ref.constraint)) {
     if (const ast::Expr *typeExpr = cst->getTypeExpr()) {
       if (failed(validateTypeConstraintExpr(typeExpr)))
         return failure();
     }
     constraintType = ast::AttributeType::get(ctx);
   } else if (const auto *cst =
                  dyn_cast<ast::OpConstraintDecl>(ref.constraint)) {
-    constraintType = ast::OperationType::get(ctx, cst->getName());
+    constraintType = ast::OperationType::get(
+        ctx, cst->getName(), lookupODSOperation(cst->getName()));
   } else if (isa<ast::TypeConstraintDecl>(ref.constraint)) {
     constraintType = typeTy;
   } else if (isa<ast::TypeRangeConstraintDecl>(ref.constraint)) {
     constraintType = typeRangeTy;
   } else if (const auto *cst =
                  dyn_cast<ast::ValueConstraintDecl>(ref.constraint)) {
     if (const ast::Expr *typeExpr = cst->getTypeExpr()) {
       if (failed(validateTypeConstraintExpr(typeExpr)))

 FailureOr<ast::Type> Parser::validateMemberAccess(ast::Expr *parentExpr,
                                                   StringRef name, SMRange loc) {
   ast::Type parentType = parentExpr->getType();
   if (ast::OperationType opType = parentType.dyn_cast<ast::OperationType>()) {
     if (name == ast::AllResultsMemberAccessExpr::getMemberName())
       return valueRangeTy;
 
     // Verify member access based on the operation type.
-    if (const ods::Operation *odsOp = lookupODSOperation(opType.getName())) {
+    if (const ods::Operation *odsOp = opType.getODSOperation()) {
       auto results = odsOp->getResults();
 
       // Handle indexed results.
       unsigned index = 0;
       if (llvm::isDigit(name[0]) && !name.getAsInteger(/*Radix=*/10, index) &&
           index < results.size()) {
         return results[index].isVariadic() ? valueRangeTy : valueTy;
       }
FailureOr<ast::OperationExpr *> Parser::createOperationExpr(
   } else if (resultTypeContext == OpResultTypeContext::InterfaceInferrence) {
     assert(results.empty() && "we currently don't check inferrence when "
                               "results were explicitly provided");
     assert(opNameRef &&
            "expected valid operation name when inferring operation results");
     checkOperationResultTypeInferrence(loc, *opNameRef, odsOp);
   }
 
-  return ast::OperationExpr::create(ctx, loc, name, operands, results,
+  return ast::OperationExpr::create(ctx, loc, odsOp, name, operands, results,
                                     attributes);
 }
 
 LogicalResult
 Parser::validateOperationOperands(SMRange loc, Optional<StringRef> name,
                                   const ods::Operation *odsOp,
                                   MutableArrayRef<ast::Expr *> operands) {
   return validateOperationOperandsOrResults(

mlir/lib/Tools/mlir-pdll-lsp-server/PDLLServer.cpp

for (unsigned i = 0, e = tupleType.size(); i < e; ++i) {
         item.filterText = item.label;
         item.insertText = elementNames[i].str();
         completionList.items.emplace_back(item);
       }
     }
   }
 
   void codeCompleteOperationMemberAccess(ast::OperationType opType) final {
-    Optional<StringRef> opName = opType.getName();
-    const ods::Operation *odsOp =
-        opName ? odsContext.lookupOperation(*opName) : nullptr;
+    const ods::Operation *odsOp = opType.getODSOperation();
     if (!odsOp)
       return;
 
     ArrayRef<ods::OperandOrResult> results = odsOp->getResults();
     for (const auto &it : llvm::enumerate(results)) {
       const ods::OperandOrResult &result = it.value();
       const ods::TypeConstraint &constraint = result.getConstraint();
 

mlir/test/mlir-pdll/CodeGen/CPP/general.pdll

 Pattern GeneratedPDLLPattern1 => erase op<>;
 Pattern => erase op<test.op3>;
 
 // -----
 
 // Check the generation of native constraints and rewrites.
 
 // CHECK:      static ::mlir::LogicalResult TestCstPDLFn(::mlir::PatternRewriter &rewriter,
-// CHECK-SAME:                                           ::llvm::ArrayRef<::mlir::PDLValue> values) {
-// CHECK:   ::mlir::Attribute attr = {};
-// CHECK:   if (values[0])
-// CHECK:     attr = values[0].cast<::mlir::Attribute>();
-// CHECK:   ::mlir::Operation * op = {};
-// CHECK:   if (values[1])
-// CHECK:     op = values[1].cast<::mlir::Operation *>();
-// CHECK:   ::mlir::Type type = {};
-// CHECK:   if (values[2])
-// CHECK:     type = values[2].cast<::mlir::Type>();
-// CHECK:   ::mlir::Value value = {};
-// CHECK:   if (values[3])
-// CHECK:     value = values[3].cast<::mlir::Value>();
-// CHECK:   ::mlir::TypeRange typeRange = {};
-// CHECK:   if (values[4])
-// CHECK:     typeRange = values[4].cast<::mlir::TypeRange>();
-// CHECK:   ::mlir::ValueRange valueRange = {};
-// CHECK:   if (values[5])
-// CHECK:     valueRange = values[5].cast<::mlir::ValueRange>();
-
-// CHECK:   return success();
+// CHECK-SAME:     ::mlir::Attribute attr, ::mlir::Operation * op, ::mlir::Type type,
+// CHECK-SAME:     ::mlir::Value value, ::mlir::TypeRange typeRange, ::mlir::ValueRange valueRange) {
+// CHECK-NEXT:   return success();
 // CHECK: }
 
 // CHECK-NOT: TestUnusedCst
 
-// CHECK: static void TestRewritePDLFn(::mlir::PatternRewriter &rewriter, ::mlir::PDLResultList &results,
-// CHECK-SAME:                         ::llvm::ArrayRef<::mlir::PDLValue> values) {
-// CHECK:   ::mlir::Attribute attr = {};
-// CHECK:   ::mlir::Operation * op = {};
-// CHECK:   ::mlir::Type type = {};
-// CHECK:   ::mlir::Value value = {};
-// CHECK:   ::mlir::TypeRange typeRange = {};
-// CHECK:   ::mlir::ValueRange valueRange = {};
-
+// CHECK: static void TestRewritePDLFn(::mlir::PatternRewriter &rewriter,
+// CHECK-SAME:     ::mlir::Attribute attr, ::mlir::Operation * op, ::mlir::Type type,
+// CHECK-SAME:     ::mlir::Value value, ::mlir::TypeRange typeRange, ::mlir::ValueRange valueRange) {
 // CHECK: foo;
 // CHECK: }
 
+// CHECK: static ::mlir::Attribute TestRewriteSinglePDLFn(::mlir::PatternRewriter &rewriter) {
+// CHECK: std::tuple<::mlir::Attribute, ::mlir::Type> TestRewriteTuplePDLFn(::mlir::PatternRewriter &rewriter) {
+
 // CHECK-NOT: TestUnusedRewrite
 
 // CHECK: struct TestCstAndRewrite : ::mlir::PDLPatternModule {
 // CHECK:   registerConstraintFunction("TestCst", TestCstPDLFn);
 // CHECK:   registerRewriteFunction("TestRewrite", TestRewritePDLFn);
 
 Constraint TestCst(attr: Attr, op: Op, type: Type, value: Value, typeRange: TypeRange, valueRange: ValueRange) [{
   return success();
 }];
 Constraint TestUnusedCst() [{ return success(); }];
 
 Rewrite TestRewrite(attr: Attr, op: Op, type: Type, value: Value, typeRange: TypeRange, valueRange: ValueRange) [{ foo; }];
+Rewrite TestRewriteSingle() -> Attr [{}];
+Rewrite TestRewriteTuple() -> (Attr, Type) [{}];
 Rewrite TestUnusedRewrite(op: Op) [{}];
 
 Pattern TestCstAndRewrite {
   let root = op<>(operand: Value, operands: ValueRange) -> (type: Type, types: TypeRange);
   TestCst(attr<"true">, root, type, operand, types, operands);
   rewrite root with {
     TestRewrite(attr<"true">, root, type, operand, types, operands);
+    TestRewriteSingle();
+    TestRewriteTuple();
     erase root;
   };
 }