Defines new PDLInterp operations needed for multi-root matching in PDL.

[sfuniak]

Summary

This is commit 1 of 4 for the multi-root matching in PDL, discussed in https://llvm.discourse.group/t/rfc-multi-root-pdl-patterns-for-kernel-matching/4148 (topic flagged for review).

These operations are:

• pdl.get_accepting_ops: Returns a list of operations accepting the given value or a range of values at the specified position. Thus if there are two operations %op1 = "foo"(%val) and %op2 = "bar"(%val) accepting a value at position 0, %ops = pdl_interp.get_accepting_ops of %val : !pdl.value at 0 will return both of them. This allows us to traverse upwards from a value to operations accepting the value.
• pdl.choose_op: Iteratively chooses one operation from a range of operations. Therefore, writing %op = pdl_interp.choose_op from %ops in the example above will select either %op1or %op2.

Testing: Added the corresponding test cases to mlir/test/Dialect/PDLInterp/ops.mlir.

[rriddle]

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

378

Is there a longer more descriptive name we can go for here? Possibly something that also capture the iterative nature of this?

Also, I'm not entirely sure we should mark this NoSideEffect. Or more specifically, I feel like transformations to this operation have a much greater cost than other "no side effect" operations (and I'm not sure if something will negatively take advantage of that). Say if we hoisted this operation higher up, wouldn't we end up executing much more bitcode than before?

560–564

I mentioned in the RFC, but isn't this just the users of a value? I would stick with the existing IR terminology whenever possible.

[sfuniak]

Thank you for your feedback @rriddle. I will update the diff. Are you okay with moving the "body" of the loop to a region to promote better readability?

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

378

The name choose_op is a leftover from an earlier interpretation of our approach. In that interpretation, we viewed this operation as making a nondeterministic choice (in the sense of NP/nondeterministic polynomial) that the bytecode executor makes. Shall we change this to pdl_interp.iterate_ops as recommended on the RFC, or even pdl_interp.iterate to allow for other types being iterated over?

I can remove the NoSideEffect attribute. Due to its iterative nature, executing this operation higher up incurs bigger cost. We are leaving it up to a later PR to optimize that.

560–564

That's correct, it's just users. I will make the change as recommended by you here and in the RFC.

[sfuniak]

Implemented a pdl_interp.foreach operation.

This change incorporates the community feedback. Instead of listing the operations all in the region of the matcher, we introduce a foreach construct that iterates over the provided operations/values, executing its nested region once for every value of the loop variable. The pdl_interp.foreach is a terminator of the block and specifies the target block that we jump to after the loop gets completed. The terminator of the loop region is pdl_interp.continue, which completes the current iteration and moves on to the next one.

[rriddle]

Really nice!

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

378

I like pdl_interp.iterate (or is it foreach now? either one is alright with me), I already have some ideas about what we can do with that using other types of values.

Also, did you intend to drop the NoSideEffect? It still shows in this diff.

383

Is this foreach now?

826–827

For ranges should we do all users of all of the values? Why just the first value?

mlir/lib/Dialect/PDLInterp/IR/PDLInterp.cpp

121–128

Could we use the declarative format for this? and use a custom directive for the region parsing (which is the part I think needs custom C++ handling)

[rriddle]

In D108543#3024057, @sfuniak wrote:

Are you okay with moving the "body" of the loop to a region to promote better readability?

Yeah that seems fine.

[sfuniak]

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

378

I believe the feedback about NoSideEffect applied to pdl_interp.choose_op (its reincarnation pdl_interp.foreach does not have NoSideEffect). It seems that pdl_interp.continue should have NoSideEffect, much like pdl_interp.finalize. Or am I wrong?

383

Thanks, will update.

826–827

I was not sure what's the right behavior; I do not have a good use case for the range of values. Will make the change you suggested.

mlir/lib/Dialect/PDLInterp/IR/PDLInterp.cpp

121–128

I tried and failed to use declarative format. The issue is the parsing of the loop variable, which is an entry block argument. Do you have ideas how this can be done?

[rriddle]

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

378

Ah, you're right. Sorry about that (my eyes just ignored the actual op)

mlir/lib/Dialect/PDLInterp/IR/PDLInterp.cpp

121–128

Gah, I missed that thanks. This is fine to keep in C++ for now.

[sfuniak]

Cleaned up the semantics of pdl_interp.get_users for range of values.

[sfuniak]

Replaced ForEachOp::create with a custom builder.

[rriddle]

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

828

of the user operation?

830–832

Why are we requiring that the values match the operands of the user? Should we split this op into two? One that just gets the users, and one that does filtering on if the values match the specific operands of the users? This operation right now seems to be merging multiple things into one (which could be represented with components already available): getting the users of an operation, comparing a set of values to the operands of that user.

[sfuniak]

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

830–832

We need to match the operands of the user, because we want to perform upward traversal from value to a specific operand group. But you are right, we could break these up into "standard" get_users + existing ops. Part of the reason why we are struggling in the byte code interpreter is that these are grouped together.

[sfuniak]

Simplified pdl_interp.get_users.

[sfuniak]

Fixed English mistake.

[rriddle]

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

826–827

Didn't we switch to return all users of ranges?

[sfuniak]

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

826–827

See my comment on the other diff. We never returned all users -- rather than grabbing all users of a range, we want users that match all of the range. This is achieved by taking the users of any representative value in a range (in practice, the first one) and filtering it down to those that use the range in its entirety. The newly simplified definition of pdl_interp.get_users captures the former part (gathering the users), whereas the latter (filtering) is performed by pdl_interp.get_operands and pdl_interp.are_equal.

[rriddle]

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

826–827

That is the issue I have with this though. We are creating a disconnect between what you get from the C++ API, a hypothetical ValueRange::getUsers wouldn't return the users of a single value in the range (it would likely act more similarly to ResultRange::getUsers). I'm a bit concerned that this op, which represents a general concept, is being designed for a very specific use case. If we want to only get the users of a specific value of the range, we could compose in other ways; e.g. extracting the first value of the range and then getting the users from that.

[sfuniak]

@rriddle I am happy to make pdl_interp.get_users return all users for a range to match the ResultRange behavior (and updating the PDL-to-PDLInterp lowering). But aren't we equally worried about the inconsistency of pdl_interp.get_defining_op and pdl_interp.get_users? And does this mean I need to introduce another operation to extract the value from a list?

[sfuniak]

Also, as much as I understand that we are defining this new op based on a limited use case, under what circumstances would we want to return all the users of a range? In PDL, the range of values will always show up together as operands, so returning the users of a representative seems just fine.

[sfuniak]

On the second thought, the two operations (pdl_interp.get_defining_op and pdl_interp.get_users) are not symmetric. And I do see value in mimicking c++ API. I will make the change requested.

However, I do believe that the current semantics of pdl_interp.get_defining_op for value range is broken. Consider the following example:

%f#2 = "foo"(%x)
%g = "bar"(%y)
%h = "baz"(%f#0, %g)

and a pattern

pdl.pattern {
  %in = pdl.operand
  %type = pdl.type
  %op0 = pdl.operation (%in : !pdl.value) -> (%type, %type : !pdl.type, !pdl.type)
  %res0 = pdl.results of %op0
  %op1 = pdl.operation (%res0 : !pdl.range<value>)
  pdl.rewrite %op1 with "rewriter"
}

With the current definition of pdl_interp.get_defining_op, the %op1 will match %h (and it should not). We should be checking the defining op of all the values in the range and return null if these ops are not all the same.

[sfuniak]

Updated semantics of pdl_interp.get_users to return all users for a value range.
Added pdl_interp.get_value operation for extracting value with given index from a range.

[sfuniak]

Fixed the comments.

[sfuniak]

@rriddle I hope that my latest update addresses your concerns.

Regarding my (supposed) counter-example for pdl_interp.get_defining_op, it's not quite accurate, because we compare the results with the values whose defining operations we queried, so the matching in this case actually works as expected. But I am wondering if we should do something similar to what you suggested for pdl_interp.get_users - extracting the first value in the range and checking its defining op.

[rriddle]

In D108543#3145839, @sfuniak wrote:

@rriddle I hope that my latest update addresses your concerns.

Regarding my (supposed) counter-example for pdl_interp.get_defining_op, it's not quite accurate, because we compare the results with the values whose defining operations we queried, so the matching in this case actually works as expected. But I am wondering if we should do something similar to what you suggested for pdl_interp.get_users - extracting the first value in the range and checking its defining op.

Thanks for bringing up get_defining_op, I agree that the semantics of that aren't great. For that one, I think we should just detect if the value range is a result range, and if it isn't we just return null.

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

826–827

Another comment on your use of get_users; IIUC you are getting all users of the first value and then comparing against the full range? What happens if other values in the range have more users than the first value?

%a = ...
%b = ...
foo.use1 (%a, %b)
foo.use2 (%b)

?

854–882

Can we have this be a general extract operation that works on any range? We could specialize it further at the bytecode level.

861–862

Can you add a note on the behavior if the index is out-of-bounds? I assume we would return null as a failure (as we do in other cases).

[sfuniak]

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

826–827

Ultimately, we are looking for users that accept the entirety of the value range as its operands. So if other values in the range have more users than the first value does, those users are not going to be relevant (they do not meet the aforementioned criterion of accepting all the values). We could have chosen any value, not necessarily the first one, and we could have also taken the intersection of users of all the values, rather than the users of just one value. But this does not affect correctness, because we still have the equality check between the operands of the user and the value range.

In a sense, the upward and downward traversals are completely symmetric: we extract the users (for upward traversal) or the defining op (for downward traversal) of the first value and then compare against the operands / results. You have to do this comparison because both users and the defining op are just a hint where the possible match might be found. How exactly we do this in the pdl_interp.get_defining_op is up to you, but the current implementation of that op is correct (though it may not generalize well).

854–882

Sure, will do.

861–862

Correct, the results for out-of-bounds should be null. Will update the docs.

[rriddle]

LGTM after resolving comments. Really appreciate the discussion on different use cases.

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

826–827

I see. If we only care about the situations in which the operation accepts all operands, indeed the strategy of using a single value works.

[sfuniak]

Generalized pdl_interp.get_value to pdl_interp.extract

[sfuniak]

Minor cleanup.

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

def PDLInterp_CheckTypesOp
   }];
 
   let arguments = (ins PDL_RangeOf<PDL_Type>:$value,
                        TypeArrayAttr:$types);
   let assemblyFormat = "$value `are` $types attr-dict `->` successors";
 }
 
 //===----------------------------------------------------------------------===//
+// pdl_interp::ContinueOp
+//===----------------------------------------------------------------------===//
+
+def PDLInterp_ContinueOp
+    : PDLInterp_Op<"continue", [NoSideEffect, HasParent<"ForEachOp">,
+                               Terminator]> {
+  let summary = "Breaks the current iteration";
+  let description = [{
+    `pdl_interp.continue` operation breaks the current iteration within the
+    `pdl_interp.iterate` region and continues with the next iteration from
+    the beginning of the region.
+
+    Example:
+
+    ```mlir
+    pdl_interp.continue
+    ```
+  }];
+
+  let assemblyFormat = "attr-dict";
+}
+
+//===----------------------------------------------------------------------===//
 // pdl_interp::CreateAttributeOp
 //===----------------------------------------------------------------------===//
 
 def PDLInterp_CreateAttributeOp
     : PDLInterp_Op<"create_attribute", [NoSideEffect]> {
   let summary = "Create an interpreter handle to a constant `Attribute`";
   let description = [{
     `pdl_interp.create_attribute` operations generate a handle within the
let description = [{
     ```mlir
     pdl_interp.finalize
     ```
   }];
   let assemblyFormat = "attr-dict";
 }
 
 //===----------------------------------------------------------------------===//
+// pdl_interp::ForEachOp
+//===----------------------------------------------------------------------===//
+
+def PDLInterp_ForEachOp
+    : PDLInterp_Op<"foreach", [Terminator]> {
+  let summary = "Iterates over a range of values or ranges";
+  let description = [{
+    `pdl_interp.foreach` iteratively selects an element from a range of values
+    and executes the region until pdl.continue is reached.
+
+    In the bytecode interpreter, this operation is implemented by looping over
+    the values and, for each selection, running the bytecode until we reach
+    pdl.continue. This may result in multiple matches being reported. Note
+    that the input range is mutated (popped from).
+
+    Example:
+
+    ```mlir
+    pdl_interp.foreach %op : !pdl.operation in %ops {
+      pdl_interp.continue
+    } -> ^next
+    ```
+  }];
+
+  let arguments = (ins PDL_RangeOf<PDL_AnyType>:$values);
+  let regions = (region AnyRegion:$region);
+  let successors = (successor AnySuccessor:$successor);
+
+  let extraClassDeclaration = [{
+    /// Constructs a ForEachOp from the given location, range, and successor.
+    static ForEachOp create(OpBuilder &builder, Location loc, Value range,
+                            Block *successor);
+
+    /// Returns the loop variable.
+    BlockArgument getLoopVariable() { return region().getArgument(0); }
+  }];
+  let parser = [{ return ::parseForEachOp(parser, result); }];
+  let printer = [{ return ::print(p, *this); }];
+  let verifier = [{ return ::verify(*this); }];
+}
+
+//===----------------------------------------------------------------------===//
 // pdl_interp::GetAttributeOp
 //===----------------------------------------------------------------------===//
 
 def PDLInterp_GetAttributeOp : PDLInterp_Op<"get_attribute", [NoSideEffect]> {
   let summary = "Get a specified attribute value from an `Operation`";
   let description = [{
     `pdl_interp.get_attribute` operations try to get a specific attribute from
     an operation. If the operation does not have that attribute, a null value is
OpBuilder<(ins "Value":$operation), [{
       build($_builder, $_state,
             pdl::RangeType::get($_builder.getType<pdl::ValueType>()), operation,
             IntegerAttr());
     }]>,
   ];
 }
 
 //===----------------------------------------------------------------------===//
+// pdl_interp::GetUsersOp
+//===----------------------------------------------------------------------===//
+
+def PDLInterp_GetUsersOp
+    : PDLInterp_Op<"get_users", [NoSideEffect]> {
+  let summary = "Get the users of a `Value`";
+  let description = [{
+    `pdl_interp.get_users` returns the operations using a value (or first value
+    from a range of values) at the specified operand position.
+
+    Example:
+
+    ```mlir
+    %ops = pdl_interp.get_users of %value : !pdl.value at 2
+    ```
+  }];
+
+  let arguments = (ins PDL_InstOrRangeOf<PDL_Value>:$value,
+                       Confined<I32Attr, [IntNonNegative]>:$index);
+  let results = (outs PDL_RangeOf<PDL_Operation>:$operations);
+  let assemblyFormat = "`of` $value `:` type($value) `at` $index attr-dict";
+}
+
+//===----------------------------------------------------------------------===//
 // pdl_interp::GetValueTypeOp
 //===----------------------------------------------------------------------===//
 
 def PDLInterp_GetValueTypeOp : PDLInterp_Op<"get_value_type", [NoSideEffect,
      TypesMatchWith<"`value` type matches arity of `result`",
                     "result", "value", "getGetValueTypeOpValueType($_self)">]> {
   let summary = "Get the result type of a specified `Value`";
   let description = [{
     `pdl_interp.get_value_type` operations get the resulting type of a specific
     value or range thereof.
 
     Example:
 
     ```mlir
     // Get the type of a single value.
     %type = pdl_interp.get_value_type of %value : !pdl.type
 
     // Get the types of a value range.
     %type = pdl_interp.get_value_type of %values : !pdl.range<type>
     ```
   }];
 
   let arguments = (ins PDL_InstOrRangeOf<PDL_Value>:$value);
   let results = (outs PDL_InstOrRangeOf<PDL_Type>:$result);
   let assemblyFormat = "`of` $value `:` type($result) attr-dict";
 
   let builders = [
     OpBuilder<(ins "Value":$value), [{
       Type valType = value.getType();
       Type typeType = $_builder.getType<pdl::TypeType>();
       build($_builder, $_state,
             valType.isa<pdl::RangeType>() ? pdl::RangeType::get(typeType)
                                           : typeType,
             value);
     }]>
   ];
 }
 
 //===----------------------------------------------------------------------===//
 // pdl_interp::InferredTypesOp
 //===----------------------------------------------------------------------===//
 
 def PDLInterp_InferredTypesOp : PDLInterp_Op<"inferred_types"> {
   let summary = "Generate a handle to a range of Types that are \"inferred\"";
   let description = [{
     `pdl_interp.inferred_types` operations generate handles to ranges of types
     that should be inferred. This signals to other operations, such as
     `pdl_interp.create_operation`, that these types should be inferred.

mlir/lib/Dialect/PDLInterp/IR/PDLInterp.cpp

if (attrNames.empty())
     return;
   p << " {";
   interleaveComma(llvm::seq<int>(0, attrNames.size()), p,
                   [&](int i) { p << attrNames[i] << " = " << attrArgs[i]; });
   p << '}';
 }
 
 //===----------------------------------------------------------------------===//
+// pdl_interp::ForEachOp
+//===----------------------------------------------------------------------===//
+
+ForEachOp ForEachOp::create(OpBuilder &builder, Location loc, Value range,
+                            Block *successor) {
+  ForEachOp op = builder.create<pdl_interp::ForEachOp>(loc, range, successor);
+  auto range_type = range.getType().cast<pdl::RangeType>();
+  op.region().emplaceBlock();
+  op.region().addArgument(range_type.getElementType());
+  return op;
+}
+
+static ParseResult parseForEachOp(OpAsmParser &parser, OperationState &result) {
+  // Parse the loop variable followed by type.
+  OpAsmParser::OperandType loopVariable;
+  Type loopVariableType;
+  if (parser.parseRegionArgument(loopVariable) ||
+      parser.parseColonType(loopVariableType))
+    return failure();
+
+  // Parse the "in" keyword.
+  if (parser.parseKeyword("in", " after loop variable"))
+    return failure();
+
+  // Parse the operand (value range).
+  OpAsmParser::OperandType operandInfo;
+  if (parser.parseOperand(operandInfo))
+    return failure();
+
+  // Resolve the operand.
+  Type rangeType = pdl::RangeType::get(loopVariableType);
+  if (parser.resolveOperand(operandInfo, rangeType, result.operands))
+    return failure();
+
+  // Parse the body region.
+  Region *body = result.addRegion();
+  if (parser.parseRegion(*body, {loopVariable}, {loopVariableType}))
+    return failure();
+
+  // Parse the attribute dictionary.
+  if (parser.parseOptionalAttrDict(result.attributes))
+    return failure();
+
+  // Parse the successor.
+  Block *successor;
+  if (parser.parseArrow() || parser.parseSuccessor(successor))
+    return failure();
+  result.addSuccessors(successor);
+
+  return success();
+}
+
+static void print(OpAsmPrinter &p, ForEachOp op) {
+  BlockArgument arg = op.getLoopVariable();
+  p << ' ' << arg << " : " << arg.getType() << " in " << op.values();
+  p.printRegion(op.region(), /*printEntryBlockArgs=*/false);
+  p.printOptionalAttrDict(op->getAttrs());
+  p << " -> ";
+  p.printSuccessor(op.successor());
+}
+
+static LogicalResult verify(ForEachOp op) {
+  // Verify that the operation has exactly one argument.
+  if (op.region().getNumArguments() != 1)
+    return op.emitOpError("requires exactly one argument");
+
+  // Verify that the loop variable and the operand (value range)
+  // have compatible types.
+  BlockArgument arg = op.getLoopVariable();
+  Type rangeType = pdl::RangeType::get(arg.getType());
+  if (rangeType != op.values().getType())
+    return op.emitOpError("operand must be a range of loop variable type");
+
+  return success();
+}
+
+//===----------------------------------------------------------------------===//
 // pdl_interp::GetValueTypeOp
 //===----------------------------------------------------------------------===//
 
 /// Given the result type of a `GetValueTypeOp`, return the expected input type.
 static Type getGetValueTypeOpValueType(Type type) {
   Type valueTy = pdl::ValueType::get(type.getContext());
   return type.isa<pdl::RangeType>() ? pdl::RangeType::get(valueTy) : valueTy;
 }
 
 //===----------------------------------------------------------------------===//
 // TableGen Auto-Generated Op and Interface Definitions
 //===----------------------------------------------------------------------===//
 
 #define GET_OP_CLASSES
 #include "mlir/Dialect/PDLInterp/IR/PDLInterpOps.cpp.inc"

mlir/test/Dialect/PDLInterp/ops.mlir

func @operations(%attribute: !pdl.attribute,
   // attributes
   %op2 = pdl_interp.create_operation "foo.op" {"attr" = %attribute, "attr1" = %attribute}
 
   // operands, and results
   %op3 = pdl_interp.create_operation "foo.op"(%input : !pdl.value) -> (%type : !pdl.type)
 
   pdl_interp.finalize
 }
+
+// -----
+
+func @foreach(%ops: !pdl.range<operation>) {
+  // iterate over a range of operations
+  pdl_interp.foreach %op : !pdl.operation in %ops {
+    %val = pdl_interp.get_result 0 of %op
+    pdl_interp.continue
+  } -> ^end
+
+  ^end:
+    pdl_interp.finalize
+}
+
+// -----
+
+func @users(%input: !pdl.value, %inputs: !pdl.range<value>) {
+  // single value
+  %ops1 = pdl_interp.get_users of %input : !pdl.value at 2
+
+  // a range of values
+  %ops2 = pdl_interp.get_users of %inputs : !pdl.range<value> at 3
+
+  pdl_interp.finalize
+}