[mlir][pdl] Add a new !pdl.range<> type

[rriddle]

Summary

This type represents a range of positional values. It will be used in followup revisions to add support for variadic constructs to PDL, such as operand and result ranges.

Depends On D95716

[jpienaar]

mlir/include/mlir/Dialect/PDL/IR/PDLTypes.td

58

So the range itself is not captured as a parem (or am I missing it?) but this the type corresponding to the runtime range

64

Could we add kwarg style comments for the different bodies?

[rriddle]

Rebase

mlir/include/mlir/Dialect/PDL/IR/PDLTypes.td

58

Reworded. If we look at something like ValueRange in the same way as Value, a PDL_Range<Value> is not conceptually different from a PDL_Value. The PDL_Range is a handle to multiple values, the PDL_Value is a handle to one.

64

Obsolete now that the getChecked body is no longer necessary.

mlir/include/mlir/Dialect/PDL/IR/PDLOps.td

let description = [{
     // Apply `myConstraint` to the entities defined by `input`, `attr`, and
     // `op`. `42`, `"abc"`, and `i32` are constant parameters passed to the
     // constraint.
     pdl.apply_constraint "myConstraint"[42, "abc", i32](%input, %attr, %op : !pdl.value, !pdl.attribute, !pdl.operation)
     ```
   }];
 
   let arguments = (ins StrAttr:$name,
-                       Variadic<PDL_PositionalValue>:$args,
+                       Variadic<PDL_AnyType>:$args,
                        OptionalAttr<ArrayAttr>:$constParams);
   let assemblyFormat = [{
     $name ($constParams^)? `(` $args `:` type($args) `)` attr-dict
   }];
 
   let builders = [
     OpBuilderDAG<(ins "StringRef":$name, CArg<"ValueRange", "{}">:$args,
       CArg<"ArrayRef<Attribute>", "{}">:$params), [{
let description = [{
     Example:
 
     ```mlir
     %ret = pdl.create_native "myNativeFunc"[42, "gt"](%arg0, %arg1) : !pdl.attribute
     ```
   }];
 
   let arguments = (ins StrAttr:$name,
-                       Variadic<PDL_PositionalValue>:$args,
+                       Variadic<PDL_AnyType>:$args,
                        OptionalAttr<ArrayAttr>:$constParams);
-  let results = (outs PDL_PositionalValue:$result);
+  let results = (outs PDL_AnyType:$result);
   let assemblyFormat = [{
     $name ($constParams^)? (`(` $args^ `:` type($args) `)`)? `:` type($result)
     attr-dict
   }];
   let verifier = ?;
 }
 
 //===----------------------------------------------------------------------===//
pdl.rewrite %root {
       %op = pdl.operation "foo.op"(%arg0, %arg1)
       pdl.replace %root with %op
     }
     ```
   }];
 
   let arguments = (ins PDL_Operation:$root,
                        OptionalAttr<StrAttr>:$name,
-                       Variadic<PDL_PositionalValue>:$externalArgs,
+                       Variadic<PDL_AnyType>:$externalArgs,
                        OptionalAttr<ArrayAttr>:$externalConstParams);
   let regions = (region AnyRegion:$body);
   let assemblyFormat = [{
     $root (`with` $name^ ($externalConstParams^)?
                          (`(` $externalArgs^ `:` type($externalArgs) `)`)?)?
           ($body^)?
     attr-dict-with-keyword
   }];

mlir/include/mlir/Dialect/PDL/IR/PDLTypes.h

 #define MLIR_DIALECT_PDL_IR_PDLTYPES_H_
 
 #include "mlir/IR/Types.h"
 
 //===----------------------------------------------------------------------===//
 // PDL Dialect Types
 //===----------------------------------------------------------------------===//
 
+namespace mlir {
+namespace pdl {
+/// This class represents the base class of all PDL types.
+class PDLType : public Type {
+public:
+  using Type::Type;
+
+  static bool classof(Type type);
+};
+} // namespace pdl
+} // namespace mlir
+
 #define GET_TYPEDEF_CLASSES
 #include "mlir/Dialect/PDL/IR/PDLOpsTypes.h.inc"
 
 #endif // MLIR_DIALECT_PDL_IR_PDLTYPES_H_

mlir/include/mlir/Dialect/PDL/IR/PDLTypes.td

 #define MLIR_DIALECT_PDL_IR_PDLTYPES
 
 include "mlir/Dialect/PDL/IR/PDLDialect.td"
 
 //===----------------------------------------------------------------------===//
 // PDL Types
 //===----------------------------------------------------------------------===//
 
-class PDL_Type<string name, string typeMnemonic> : TypeDef<PDL_Dialect, name> {
+class PDL_Type<string name, string typeMnemonic>
+    : TypeDef<PDL_Dialect, name, "::mlir::pdl::PDLType"> {
   let mnemonic = typeMnemonic;
 }
 
 //===----------------------------------------------------------------------===//
 // pdl::AttributeType
 //===----------------------------------------------------------------------===//
 
 def PDL_Attribute : PDL_Type<"Attribute", "attribute"> {
def PDL_Operation : PDL_Type<"Operation", "operation"> {
   let summary = "PDL handle to an `mlir::Operation *`";
   let description = [{
     This type represents a handle to an instance of an `mlir::Operation *`,
     bound to a value that is usable within a PDL pattern or rewrite.
   }];
 }
 
 //===----------------------------------------------------------------------===//
+// pdl::RangeType
+//===----------------------------------------------------------------------===//
+
+def PDL_Range : PDL_Type<"Range", "range"> {
+  let summary = "Range of PDL handles for a given sub-type";
+  let description = [{
+    This type represents a range of handles to instances of the given PDL
+    element type, i.e. `Attribute`, `Operation`, `Type`, or `Value`.
+  }];
+  let parameters = (ins "Type":$elementType);
+
+  let builders = [
+    TypeBuilderWithInferredContext<(ins "Type":$elementType), [{
+      return Base::get(elementType.getContext(), elementType);
+    }], [{
+      return Base::getChecked($_loc, elementType);
+    }]>,
+  ];
+  let genVerifyInvariantsDecl = 1;
+  let skipDefaultBuilders = 1;
+}
+
+//===----------------------------------------------------------------------===//
 // pdl::TypeType
 //===----------------------------------------------------------------------===//
 
 def PDL_Type : PDL_Type<"Type", "type"> {
   let summary = "PDL handle to an `mlir::Type`";
   let description = [{
     This type represents a handle to an instance of an `mlir::Type`, bound to a
     value that is usable within a PDL pattern or rewrite.
let description = [{
     value that is usable within a PDL pattern or rewrite.
   }];
 }
 
 //===----------------------------------------------------------------------===//
 // Additional Type Constraints
 //===----------------------------------------------------------------------===//
 
-// A positional value is a location on a pattern DAG, which may be an attribute,
-// operation, or operand/result.
-def PDL_PositionalValue :
-    AnyTypeOf<[PDL_Attribute, PDL_Operation, PDL_Type, PDL_Value],
-              "Positional Value">;
+def PDL_AnyType : Type<
+  CPred<"$_self.isa<::mlir::pdl::PDLType>()">, "pdl type",
+        "::mlir::pdl::PDLType">;
 
 #endif // MLIR_DIALECT_PDL_IR_PDLTYPES

mlir/include/mlir/Dialect/PDLInterp/IR/PDLInterpOps.td

let description = [{
     ```mlir
     // Apply `myConstraint` to the entities defined by `input`, `attr`, and
     // `op`.
     pdl_interp.apply_constraint "myConstraint"[42, "abc", i32](%input, %attr, %op : !pdl.value, !pdl.attribute, !pdl.operation) -> ^matchDest, ^failureDest
     ```
   }];
 
   let arguments = (ins StrAttr:$name,
-                       Variadic<PDL_PositionalValue>:$args,
+                       Variadic<PDL_AnyType>:$args,
                        OptionalAttr<ArrayAttr>:$constParams);
   let assemblyFormat = [{
     $name ($constParams^)? `(` $args `:` type($args) `)` attr-dict `->`
     successors
   }];
 }
 
 //===----------------------------------------------------------------------===//
let description = [{
 
     // Rewriter operating on the root operation along with additional arguments
     // and constant parameters.
     pdl_interp.apply_rewrite "rewriter"[42](%value : !pdl.value) on %root
     ```
   }];
   let arguments = (ins StrAttr:$name,
                        PDL_Operation:$root,
-                       Variadic<PDL_PositionalValue>:$args,
+                       Variadic<PDL_AnyType>:$args,
                        OptionalAttr<ArrayAttr>:$constParams);
   let assemblyFormat = [{
     $name ($constParams^)? (`(` $args^ `:` type($args) `)`)? `on` $root
     attr-dict
   }];
 }
 
 //===----------------------------------------------------------------------===//
let description = [{
 
     Example:
 
     ```mlir
     pdl_interp.are_equal %result1, %result2 : !pdl.value -> ^matchDest, ^failureDest
     ```
   }];
 
-  let arguments = (ins PDL_PositionalValue:$lhs,
-                       PDL_PositionalValue:$rhs);
+  let arguments = (ins PDL_AnyType:$lhs, PDL_AnyType:$rhs);
   let assemblyFormat = "operands `:` type($lhs) attr-dict `->` successors";
 }
 
 //===----------------------------------------------------------------------===//
 // pdl_interp::BranchOp
 //===----------------------------------------------------------------------===//
 
 def PDLInterp_BranchOp : PDLInterp_Op<"branch", [NoSideEffect, Terminator]> {
let description = [{
     Example:
 
     ```mlir
     %ret = pdl_interp.create_native "myNativeFunc"[42, "gt"](%arg0, %arg1 : !pdl.value, !pdl.value) : !pdl.attribute
     ```
   }];
 
   let arguments = (ins StrAttr:$name,
-                       Variadic<PDL_PositionalValue>:$args,
+                       Variadic<PDL_AnyType>:$args,
                        OptionalAttr<ArrayAttr>:$constParams);
-  let results = (outs PDL_PositionalValue:$result);
+  let results = (outs PDL_AnyType:$result);
   let assemblyFormat = [{
     $name ($constParams^)? (`(` $args^ `:` type($args) `)`)? `:` type($result)
     attr-dict
   }];
   let verifier = ?;
 }
 
 //===----------------------------------------------------------------------===//
let description = [{
 
     Example:
 
     ```mlir
     pdl_interp.is_not_null %value : !pdl.value -> ^matchDest, ^failureDest
     ```
   }];
 
-  let arguments = (ins PDL_PositionalValue:$value);
+  let arguments = (ins PDL_AnyType:$value);
   let assemblyFormat = "$value `:` type($value) attr-dict `->` successors";
 }
 
 //===----------------------------------------------------------------------===//
 // pdl_interp::RecordMatchOp
 //===----------------------------------------------------------------------===//
 
 def PDLInterp_RecordMatchOp
     : PDLInterp_Op<"record_match", [AttrSizedOperandSegments, Terminator]> {
   let summary = "Record the metadata for a successful pattern match";
   let description = [{
     `pdl_interp.record_match` operations record a successful pattern match with
     the interpreter and branch to the next part of the matcher. The metadata
     recorded by these operations correspond to a specific `pdl.pattern`, as well
     as what values were used during that match that should be propagated to the
     rewriter.
 
     Example:
 
     ```mlir
     pdl_interp.record_match @rewriters::myRewriter(%root : !pdl.operation) : benefit(1), loc([%root, %op1]), root("foo.op") -> ^nextDest
     ```
   }];
 
-  let arguments = (ins Variadic<PDL_PositionalValue>:$inputs,
+  let arguments = (ins Variadic<PDL_AnyType>:$inputs,
                        Variadic<PDL_Operation>:$matchedOps,
                        SymbolRefAttr:$rewriter,
                        OptionalAttr<StrAttr>:$rootKind,
                        OptionalAttr<StrArrayAttr>:$generatedOps,
                        Confined<I16Attr, [IntNonNegative]>:$benefit);
   let successors = (successor AnySuccessor:$dest);
   let assemblyFormat = [{
     $rewriter (`(` $inputs^ `:` type($inputs) `)`)? `:`

mlir/lib/Dialect/PDL/IR/CMakeLists.txt

 add_mlir_dialect_library(MLIRPDL
   PDL.cpp
+  PDLTypes.cpp
 
   ADDITIONAL_HEADER_DIRS
   ${MLIR_MAIN_INCLUDE_DIR}/mlir/Dialect/PDL
 
   DEPENDS
   MLIRPDLOpsIncGen
 
   LINK_LIBS PUBLIC
   MLIRIR
   MLIRInferTypeOpInterface
   MLIRSideEffectInterfaces
   )

mlir/lib/Dialect/PDL/IR/PDL.cpp

 //===- PDL.cpp - Pattern Descriptor Language Dialect ----------------------===//
 //
 // 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/Dialect/PDL/IR/PDL.h"
 #include "mlir/Dialect/PDL/IR/PDLOps.h"
 #include "mlir/Dialect/PDL/IR/PDLTypes.h"
 #include "mlir/IR/BuiltinTypes.h"
-#include "mlir/IR/DialectImplementation.h"
 #include "mlir/Interfaces/InferTypeOpInterface.h"
 #include "llvm/ADT/StringSwitch.h"
-#include "llvm/ADT/TypeSwitch.h"
 
 using namespace mlir;
 using namespace mlir::pdl;
 
 //===----------------------------------------------------------------------===//
 // PDLDialect
 //===----------------------------------------------------------------------===//
 

 }
 
 //===----------------------------------------------------------------------===//
 // TableGen'd op method definitions
 //===----------------------------------------------------------------------===//
 
 #define GET_OP_CLASSES
 #include "mlir/Dialect/PDL/IR/PDLOps.cpp.inc"
-
-//===----------------------------------------------------------------------===//
-// TableGen'd type method definitions
-//===----------------------------------------------------------------------===//
-
-#define GET_TYPEDEF_CLASSES
-#include "mlir/Dialect/PDL/IR/PDLOpsTypes.cpp.inc"
-
-Type PDLDialect::parseType(DialectAsmParser &parser) const {
-  StringRef keyword;
-  if (parser.parseKeyword(&keyword))
-    return Type();
-  if (Type type = generatedTypeParser(getContext(), parser, keyword))
-    return type;
-
-  parser.emitError(parser.getNameLoc(), "invalid 'pdl' type: `")
-      << keyword << "'";
-  return Type();
-}
-
-void PDLDialect::printType(Type type, DialectAsmPrinter &printer) const {
-  if (failed(generatedTypePrinter(type, printer)))
-    llvm_unreachable("unknown 'pdl' type");
-}

mlir/lib/Dialect/PDL/IR/PDLTypes.cpp

+//===- PDLTypes.cpp - Pattern Descriptor Language Types -------------------===//
+//
+// 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/Dialect/PDL/IR/PDLTypes.h"
+#include "mlir/Dialect/PDL/IR/PDL.h"
+#include "mlir/IR/Builders.h"
+#include "mlir/IR/DialectImplementation.h"
+#include "llvm/ADT/TypeSwitch.h"
+
+using namespace mlir;
+using namespace mlir::pdl;
+
+//===----------------------------------------------------------------------===//
+// TableGen'd type method definitions
+//===----------------------------------------------------------------------===//
+
+#define GET_TYPEDEF_CLASSES
+#include "mlir/Dialect/PDL/IR/PDLOpsTypes.cpp.inc"
+
+//===----------------------------------------------------------------------===//
+// PDLDialect
+//===----------------------------------------------------------------------===//
+
+static Type parsePDLType(DialectAsmParser &parser) {
+  StringRef keyword;
+  if (parser.parseKeyword(&keyword))
+    return Type();
+  if (Type type = generatedTypeParser(parser.getBuilder().getContext(), parser,
+                                      keyword))
+    return type;
+
+  // FIXME: This ends up with a double error being emitted if `RangeType` also
+  // emits an error. We should rework the `generatedTypeParser` to better
+  // support when the keyword is valid but the individual type parser itself
+  // emits an error.
+  parser.emitError(parser.getNameLoc(), "invalid 'pdl' type: `")
+      << keyword << "'";
+  return Type();
+}
+
+Type PDLDialect::parseType(DialectAsmParser &parser) const {
+  return parsePDLType(parser);
+}
+
+void PDLDialect::printType(Type type, DialectAsmPrinter &printer) const {
+  if (failed(generatedTypePrinter(type, printer)))
+    llvm_unreachable("unknown 'pdl' type");
+}
+
+//===----------------------------------------------------------------------===//
+// PDL Types
+//===----------------------------------------------------------------------===//
+
+bool PDLType::classof(Type type) {
+  return llvm::isa<PDLDialect>(type.getDialect());
+}
+
+//===----------------------------------------------------------------------===//
+// RangeType
+//===----------------------------------------------------------------------===//
+
+Type RangeType::parse(MLIRContext *context, DialectAsmParser &parser) {
+  if (parser.parseLess())
+    return Type();
+
+  llvm::SMLoc elementLoc = parser.getCurrentLocation();
+  Type elementType = parsePDLType(parser);
+  if (!elementType || parser.parseGreater())
+    return Type();
+
+  if (elementType.isa<RangeType>()) {
+    parser.emitError(elementLoc)
+        << "element of pdl.range cannot be another range, but got"
+        << elementType;
+    return Type();
+  }
+  return RangeType::get(elementType);
+}
+
+void RangeType::print(DialectAsmPrinter &printer) const {
+  printer << "range<";
+  generatedTypePrinter(getElementType(), printer);
+  printer << ">";
+}
+
+LogicalResult RangeType::verifyConstructionInvariants(Location loc,
+                                                      Type elementType) {
+  if (!elementType.isa<PDLType>() || elementType.isa<RangeType>()) {
+    return emitError(loc)
+           << "expected element of pdl.range to be one of [!pdl.attribute, "
+              "!pdl.operation, !pdl.type, !pdl.value], but got "
+           << elementType;
+  }
+  return success();
+}

mlir/test/Dialect/PDL/invalid-types.mlir

+// RUN: mlir-opt %s -split-input-file -verify-diagnostics
+
+//===----------------------------------------------------------------------===//
+// pdl::RangeType
+//===----------------------------------------------------------------------===//
+
+// expected-error@+2 {{element of pdl.range cannot be another range, but got'!pdl.range<value>'}}
+// expected-error@+1 {{invalid 'pdl' type}}
+#invalid_element = !pdl.range<range<value>>