[mlir:PDL] Allow non-bound pdl.attribute/pdl.type operations that create constants

[rriddle]

Summary

This allows for passing in these attributes/types to constraints/rewrites as arguments.

Depends On D114816

[jpienaar]

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

519

Is RangeType marked as buildable? It would seem like if it was buildable then we'd automatically generate this (well if the return type interface was included here too)

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

let description = [{
     ```mlir
     pdl_interp.create_types [i64, i64]
     ```
   }];
 
   let arguments = (ins TypeArrayAttr:$value);
   let results = (outs PDL_RangeOf<PDL_Type>:$result);
   let assemblyFormat = "$value attr-dict";
+
+  let builders = [
+    OpBuilder<(ins "ArrayAttr":$type), [{
+      build($_builder, $_state,
+            pdl::RangeType::get($_builder.getType<pdl::TypeType>()), type);
+    }]>
+  ];
 }
 
 //===----------------------------------------------------------------------===//
 // pdl_interp::EraseOp
 //===----------------------------------------------------------------------===//
 
 def PDLInterp_EraseOp : PDLInterp_Op<"erase"> {
   let summary = "Mark an operation as `erased`";

mlir/lib/Conversion/PDLToPDLInterp/PDLToPDLInterp.cpp

Block *PatternLowering::generateMatcher(MatcherNode &node, Region &region) {
   return block;
 }
 
 Value PatternLowering::getValueAt(Block *&currentBlock, Position *pos) {
   if (Value val = values.lookup(pos))
     return val;
 
   // Get the value for the parent position.
-  Value parentVal = getValueAt(currentBlock, pos->getParent());
+  Value parentVal;
+  if (Position *parent = pos->getParent())
+    parentVal = getValueAt(currentBlock, pos->getParent());
 
   // TODO: Use a location from the position.
-  Location loc = parentVal.getLoc();
+  Location loc = parentVal ? parentVal.getLoc() : builder.getUnknownLoc();
   builder.setInsertionPointToEnd(currentBlock);
   Value value;
   switch (pos->getKind()) {
   case Predicates::OperationPos: {
     auto *operationPos = cast<OperationPosition>(pos);
     if (!operationPos->isUpward()) {
       // Standard (downward) traversal which directly follows the defining op.
       value = builder.create<pdl_interp::GetDefiningOpOp>(
Value PatternLowering::getValueAt(Block *&currentBlock, Position *pos) {
   case Predicates::ResultGroupPos: {
     auto *resPos = cast<ResultGroupPosition>(pos);
     Type valueTy = builder.getType<pdl::ValueType>();
     value = builder.create<pdl_interp::GetResultsOp>(
         loc, resPos->isVariadic() ? pdl::RangeType::get(valueTy) : valueTy,
         parentVal, resPos->getResultGroupNumber());
     break;
   }
+  case Predicates::AttributeLiteralPos: {
+    auto *attrPos = cast<AttributeLiteralPosition>(pos);
+    value =
+        builder.create<pdl_interp::CreateAttributeOp>(loc, attrPos->getValue());
+    break;
+  }
+  case Predicates::TypeLiteralPos: {
+    auto *typePos = cast<TypeLiteralPosition>(pos);
+    Attribute rawTypeAttr = typePos->getValue();
+    if (TypeAttr typeAttr = rawTypeAttr.dyn_cast<TypeAttr>())
+      value = builder.create<pdl_interp::CreateTypeOp>(loc, typeAttr);
+    else
+      value = builder.create<pdl_interp::CreateTypesOp>(
+          loc, rawTypeAttr.cast<ArrayAttr>());
+    break;
+  }
   default:
     llvm_unreachable("Generating unknown Position getter");
     break;
   }
 
   values.insert(pos, value);
   return value;
 }
 
 void PatternLowering::generate(BoolNode *boolNode, Block *&currentBlock,
                                Value val) {
   Location loc = val.getLoc();
   Qualifier *question = boolNode->getQuestion();
   Qualifier *answer = boolNode->getAnswer();
   Region *region = currentBlock->getParent();
 
   // Execute the getValue queries first, so that we create success
   // matcher in the correct (possibly nested) region.
   SmallVector<Value> args;
   if (auto *equalToQuestion = dyn_cast<EqualToQuestion>(question)) {
     args = {getValueAt(currentBlock, equalToQuestion->getValue())};
   } else if (auto *cstQuestion = dyn_cast<ConstraintQuestion>(question)) {
-    for (Position *position : std::get<1>(cstQuestion->getValue()))
+    for (Position *position : cstQuestion->getArgs())
       args.push_back(getValueAt(currentBlock, position));
   }
 
   // Generate the matcher in the current (potentially nested) region
   // and get the failure successor.
   Block *success = generateMatcher(*boolNode->getSuccessNode(), *region);
   Block *failure = failureBlockStack.back();
 
void PatternLowering::generate(BoolNode *boolNode, Block *&currentBlock,
   case Predicates::EqualToQuestion: {
     bool trueAnswer = isa<TrueAnswer>(answer);
     builder.create<pdl_interp::AreEqualOp>(loc, val, args.front(),
                                            trueAnswer ? success : failure,
                                            trueAnswer ? failure : success);
     break;
   }
   case Predicates::ConstraintQuestion: {
-    auto value = cast<ConstraintQuestion>(question)->getValue();
+    auto *cstQuestion = cast<ConstraintQuestion>(question);
     builder.create<pdl_interp::ApplyConstraintOp>(
-        loc, std::get<0>(value), args, std::get<2>(value).cast<ArrayAttr>(),
-        success, failure);
+        loc, cstQuestion->getName(), args, cstQuestion->getParams(), success,
+        failure);
     break;
   }
   default:
     llvm_unreachable("Generating unknown Predicate operation");
   }
 }
 
 template <typename OpT, typename PredT, typename ValT = typename PredT::KeyTy>

mlir/lib/Conversion/PDLToPDLInterp/Predicate.h

enum Kind : unsigned {
   /// Positions, ordered by decreasing priority.
   OperationPos,
   OperandPos,
   OperandGroupPos,
   AttributePos,
   ResultPos,
   ResultGroupPos,
   TypePos,
+  AttributeLiteralPos,
+  TypeLiteralPos,
 
   // Questions, ordered by dependency and decreasing priority.
   IsNotNullQuestion,
   OperationNameQuestion,
   TypeQuestion,
   AttributeQuestion,
   OperandCountAtLeastQuestion,
   OperandCountQuestion,
: public PredicateBase<AttributePosition, Position,
                            Predicates::AttributePos> {
   explicit AttributePosition(const KeyTy &key);
 
   /// Returns the attribute name of this position.
   StringAttr getName() const { return key.second; }
 };
 
 //===----------------------------------------------------------------------===//
+// AttributeLiteralPosition
+
+/// A position describing a literal attribute.
+struct AttributeLiteralPosition
+    : public PredicateBase<AttributeLiteralPosition, Position, Attribute,
+                           Predicates::AttributeLiteralPos> {
+  using PredicateBase::PredicateBase;
+};
+
+//===----------------------------------------------------------------------===//
 // OperandPosition
 
 /// A position describing an operand of an operation.
 struct OperandPosition
     : public PredicateBase<OperandPosition, Position,
                            std::pair<OperationPosition *, unsigned>,
                            Predicates::OperandPos> {
   explicit OperandPosition(const KeyTy &key);
explicit TypePosition(const KeyTy &key) : Base(key) {
     assert((isa<AttributePosition, OperandPosition, OperandGroupPosition,
                 ResultPosition, ResultGroupPosition>(key)) &&
            "expected parent to be an attribute, operand, or result");
     parent = key;
   }
 };
 
 //===----------------------------------------------------------------------===//
+// TypeLiteralPosition
+
+/// A position describing a literal type or type range. The value is stored as
+/// either a TypeAttr, or an ArrayAttr of TypeAttr.
+struct TypeLiteralPosition
+    : public PredicateBase<TypeLiteralPosition, Position, Attribute,
+                           Predicates::TypeLiteralPos> {
+  using PredicateBase::PredicateBase;
+};
+
+//===----------------------------------------------------------------------===//
 // Qualifiers
 //===----------------------------------------------------------------------===//
 
 /// An ordinal predicate consists of a "Question" and a set of acceptable
 /// "Answers" (later converted to ordinal values). A predicate will query some
 /// property of a positional value and decide what to do based on the result.
 ///
 /// This makes top-level predicate representations ordinal (SwitchOp). Later,

 /// Apply a parameterized constraint to multiple position values.
 struct ConstraintQuestion
     : public PredicateBase<
           ConstraintQuestion, Qualifier,
           std::tuple<StringRef, ArrayRef<Position *>, Attribute>,
           Predicates::ConstraintQuestion> {
   using Base::Base;
 
+  /// Return the name of the constraint.
+  StringRef getName() const { return std::get<0>(key); }
+
+  /// Return the arguments of the constraint.
+  ArrayRef<Position *> getArgs() const { return std::get<1>(key); }
+
+  /// Return the constant parameters of the constraint.
+  ArrayAttr getParams() const {
+    return std::get<2>(key).dyn_cast_or_null<ArrayAttr>();
+  }
+
   /// Construct an instance with the given storage allocator.
   static ConstraintQuestion *construct(StorageUniquer::StorageAllocator &alloc,
                                        KeyTy key) {
     return Base::construct(alloc, KeyTy{alloc.copyInto(std::get<0>(key)),
                                         alloc.copyInto(std::get<1>(key)),
                                         std::get<2>(key)});
   }
 };

 
 /// This class provides a storage uniquer that is used to allocate predicate
 /// instances.
 class PredicateUniquer : public StorageUniquer {
 public:
   PredicateUniquer() {
     // Register the types of Positions with the uniquer.
     registerParametricStorageType<AttributePosition>();
+    registerParametricStorageType<AttributeLiteralPosition>();
     registerParametricStorageType<OperandPosition>();
     registerParametricStorageType<OperandGroupPosition>();
     registerParametricStorageType<OperationPosition>();
     registerParametricStorageType<ResultPosition>();
     registerParametricStorageType<ResultGroupPosition>();
     registerParametricStorageType<TypePosition>();
+    registerParametricStorageType<TypeLiteralPosition>();
 
     // Register the types of Questions with the uniquer.
     registerParametricStorageType<AttributeAnswer>();
     registerParametricStorageType<OperationNameAnswer>();
     registerParametricStorageType<TypeAnswer>();
     registerParametricStorageType<UnsignedAnswer>();
     registerSingletonStorageType<FalseAnswer>();
     registerSingletonStorageType<TrueAnswer>();
OperationPosition *getUsersOp(Position *p, Optional<unsigned> operand) {
     return OperationPosition::get(uniquer, p, operand);
   }
 
   /// Returns an attribute position for an attribute of the given operation.
   Position *getAttribute(OperationPosition *p, StringRef name) {
     return AttributePosition::get(uniquer, p, StringAttr::get(ctx, name));
   }
 
+  /// Returns an attribute position for the given attribute.
+  Position *getAttributeLiteral(Attribute attr) {
+    return AttributeLiteralPosition::get(uniquer, attr);
+  }
+
   /// Returns an operand position for an operand of the given operation.
   Position *getOperand(OperationPosition *p, unsigned operand) {
     return OperandPosition::get(uniquer, p, operand);
   }
 
   /// Returns a position for a group of operands of the given operation.
   Position *getOperandGroup(OperationPosition *p, Optional<unsigned> group,
                             bool isVariadic) {
public:
   }
   Position *getAllResults(OperationPosition *p) {
     return getResultGroup(p, /*group=*/llvm::None, /*isVariadic=*/true);
   }
 
   /// Returns a type position for the given entity.
   Position *getType(Position *p) { return TypePosition::get(uniquer, p); }
 
+  /// Returns a type position for the given type value. The value is stored
+  /// as either a TypeAttr, or an ArrayAttr of TypeAttr.
+  Position *getTypeLiteral(Attribute attr) {
+    return TypeLiteralPosition::get(uniquer, attr);
+  }
+
   //===--------------------------------------------------------------------===//
   // Qualifiers
   //===--------------------------------------------------------------------===//
 
   /// An ordinal predicate consists of a "Question" and a set of acceptable
   /// "Answers" (later converted to ordinal values). A predicate will query some
   /// property of a positional value and decide what to do based on the result.
   using Predicate = std::pair<Qualifier *, Qualifier *>;

mlir/lib/Conversion/PDLToPDLInterp/Predicate.cpp

 //===----------------------------------------------------------------------===//
 
 Position::~Position() {}
 
 /// Returns the depth of the first ancestor operation position.
 unsigned Position::getOperationDepth() const {
   if (const auto *operationPos = dyn_cast<OperationPosition>(this))
     return operationPos->getDepth();
-  return parent->getOperationDepth();
+  return parent ? parent->getOperationDepth() : 0;
 }
 
 //===----------------------------------------------------------------------===//
 // AttributePosition
 
 AttributePosition::AttributePosition(const KeyTy &key) : Base(key) {
   parent = key.first;
 }

mlir/lib/Conversion/PDLToPDLInterp/PredicateTree.cpp

TypeSwitch<Position *>(pos)
         getTreePredicates(predList, val, builder, inputs, pos);
       })
       .Case<OperandPosition, OperandGroupPosition>([&](auto *pos) {
         getOperandTreePredicates(predList, val, builder, inputs, pos);
       })
       .Default([](auto *) { llvm_unreachable("unexpected position kind"); });
 }
 
-/// Collect all of the predicates related to constraints within the given
-/// pattern operation.
+static void getAttributePredicates(pdl::AttributeOp op,
+                                   std::vector<PositionalPredicate> &predList,
+                                   PredicateBuilder &builder,
+                                   DenseMap<Value, Position *> &inputs) {
+  Position *&attrPos = inputs[op];
+  if (attrPos)
+    return;
+  Attribute value = op.valueAttr();
+  assert(value && "expected non-tree `pdl.attribute` to contain a value");
+  attrPos = builder.getAttributeLiteral(value);
+}
+
 static void getConstraintPredicates(pdl::ApplyNativeConstraintOp op,
                                     std::vector<PositionalPredicate> &predList,
                                     PredicateBuilder &builder,
                                     DenseMap<Value, Position *> &inputs) {
   OperandRange arguments = op.args();
   ArrayAttr parameters = op.constParamsAttr();
 
   std::vector<Position *> allPositions;
static void getResultPredicates(pdl::ResultsOp op,
   auto *parentPos = cast<OperationPosition>(inputs.lookup(op.parent()));
   bool isVariadic = op.getType().isa<pdl::RangeType>();
   Optional<unsigned> index = op.index();
   resultPos = builder.getResultGroup(parentPos, index, isVariadic);
   if (index)
     predList.emplace_back(resultPos, builder.getIsNotNull());
 }
 
+static void getTypePredicates(Value typeValue,
+                              function_ref<Attribute()> typeAttrFn,
+                              PredicateBuilder &builder,
+                              DenseMap<Value, Position *> &inputs) {
+  Position *&typePos = inputs[typeValue];
+  if (typePos)
+    return;
+  Attribute typeAttr = typeAttrFn();
+  assert(typeAttr &&
+         "expected non-tree `pdl.type`/`pdl.types` to contain a value");
+  typePos = builder.getTypeLiteral(typeAttr);
+}
+
 /// Collect all of the predicates that cannot be determined via walking the
 /// tree.
 static void getNonTreePredicates(pdl::PatternOp pattern,
                                  std::vector<PositionalPredicate> &predList,
                                  PredicateBuilder &builder,
                                  DenseMap<Value, Position *> &inputs) {
   for (Operation &op : pattern.body().getOps()) {
     TypeSwitch<Operation *>(&op)
+        .Case([&](pdl::AttributeOp attrOp) {
+          getAttributePredicates(attrOp, predList, builder, inputs);
+        })
         .Case<pdl::ApplyNativeConstraintOp>([&](auto constraintOp) {
           getConstraintPredicates(constraintOp, predList, builder, inputs);
         })
         .Case<pdl::ResultOp, pdl::ResultsOp>([&](auto resultOp) {
           getResultPredicates(resultOp, predList, builder, inputs);
+        })
+        .Case([&](pdl::TypeOp typeOp) {
+          getTypePredicates(
+              typeOp, [&] { return typeOp.typeAttr(); }, builder, inputs);
+        })
+        .Case([&](pdl::TypesOp typeOp) {
+          getTypePredicates(
+              typeOp, [&] { return typeOp.typesAttr(); }, builder, inputs);
         });
   }
 }
 
 namespace {
 
 /// An op accepting a value at an optional index.
 struct OpIndex {

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

 //===----------------------------------------------------------------------===//
 // pdl::AttributeOp
 //===----------------------------------------------------------------------===//
 
 static LogicalResult verify(AttributeOp op) {
   Value attrType = op.type();
   Optional<Attribute> attrValue = op.value();
 
-  if (!attrValue && isa<RewriteOp>(op->getParentOp()))
+  if (!attrValue) {
+    if (isa<RewriteOp>(op->getParentOp()))
       return op.emitOpError("expected constant value when specified within a "
                             "`pdl.rewrite`");
-  if (attrValue && attrType)
-    return op.emitOpError("expected only one of [`type`, `value`] to be set");
     return verifyHasBindingUse(op);
   }
+  if (attrType)
+    return op.emitOpError("expected only one of [`type`, `value`] to be set");
+  return success();
+}
 
 //===----------------------------------------------------------------------===//
 // pdl::OperandOp
 //===----------------------------------------------------------------------===//
 
 static LogicalResult verify(OperandOp op) { return verifyHasBindingUse(op); }
 
 //===----------------------------------------------------------------------===//
static LogicalResult verify(RewriteOp op) {
 
   return success();
 }
 
 //===----------------------------------------------------------------------===//
 // pdl::TypeOp
 //===----------------------------------------------------------------------===//
 
-static LogicalResult verify(TypeOp op) { return verifyHasBindingUse(op); }
+static LogicalResult verify(TypeOp op) {
+  if (!op.typeAttr())
+    return verifyHasBindingUse(op);
+  return success();
+}
 
 //===----------------------------------------------------------------------===//
 // pdl::TypesOp
 //===----------------------------------------------------------------------===//
 
-static LogicalResult verify(TypesOp op) { return verifyHasBindingUse(op); }
+static LogicalResult verify(TypesOp op) {
+  if (!op.typesAttr())
+    return verifyHasBindingUse(op);
+  return success();
+}
 
 //===----------------------------------------------------------------------===//
 // TableGen'd op method definitions
 //===----------------------------------------------------------------------===//
 
 #define GET_OP_CLASSES
 #include "mlir/Dialect/PDL/IR/PDLOps.cpp.inc"

mlir/test/Conversion/PDLToPDLInterp/pdl-to-pdl-interp-matcher.mlir

pdl.pattern @variadic_results_at : benefit(1) {
     %val = pdl.operand
     %op = pdl.operation -> (%types, %type : !pdl.range<type>, !pdl.type)
     %vals = pdl.results 0 of %op -> !pdl.range<value>
     %root1 = pdl.operation(%vals, %val : !pdl.range<value>, !pdl.value)
     %root2 = pdl.operation(%val, %vals : !pdl.value, !pdl.range<value>)
     pdl.rewrite with "rewriter"(%root1, %root2 : !pdl.operation, !pdl.operation)
   }
 }
+
+// -----
+
+// CHECK-LABEL: module @attribute_literal
+module @attribute_literal {
+  // CHECK: func @matcher(%{{.*}}: !pdl.operation)
+  // CHECK: %[[ATTR:.*]] = pdl_interp.create_attribute 10 : i64
+  // CHECK: pdl_interp.apply_constraint "constraint"(%[[ATTR]] : !pdl.attribute)
+
+  // Check the correct lowering of an attribute that hasn't been bound.
+  pdl.pattern : benefit(1) {
+    %attr = pdl.attribute 10
+    pdl.apply_native_constraint "constraint"(%attr: !pdl.attribute)
+
+    %root = pdl.operation
+    pdl.rewrite %root with "rewriter"
+  }
+}
+
+// -----
+
+// CHECK-LABEL: module @type_literal
+module @type_literal {
+  // CHECK: func @matcher(%{{.*}}: !pdl.operation)
+  // CHECK: %[[TYPE:.*]] = pdl_interp.create_type i32
+  // CHECK: %[[TYPES:.*]] = pdl_interp.create_types [i32, i64]
+  // CHECK: pdl_interp.apply_constraint "constraint"(%[[TYPE]], %[[TYPES]] : !pdl.type, !pdl.range<type>)
+
+  // Check the correct lowering of a type that hasn't been bound.
+  pdl.pattern : benefit(1) {
+    %type = pdl.type : i32
+    %types = pdl.types : [i32, i64]
+    pdl.apply_native_constraint "constraint"(%type, %types: !pdl.type, !pdl.range<type>)
+
+    %root = pdl.operation
+    pdl.rewrite %root with "rewriter"
+  }
+}
+