FXML.1923: PDLL support for native constraints with attribute results

[martin-luecke]

This PR adds support for PDLL native constraints to return results.
This is useful for situations where a pattern checks for certain constraints of multiple interdependent attributes and computes a new attribute value based on them. Currently we do this check natively in C++ during matching and after a successful match have to escape to native C++ again to do the computation during the rewriting part of the pattern. With this PR we can do the computation in C++ during matching and use the result in the rewriting part of the pattern. Effectively this enables a choice in the trade-off of memory consumption during matching vs recomputation of values.

This is a simplified example of a situation where this is useful: We have two operations with certain attributes that have interdependent constraints. For instance attr_foo: one_of [0, 2, 4, 8], attr_bar: one_of [0, 2, 4, 8] and attr_foo == attr_bar. The pattern should only match if all conditions are true. The new operation should be created with a new attribute which is computed from the two matched attributes e.g. attr_baz = attr_foo * attr_bar. For the check we already escape to native C++ and have all values at hand so it makes sense to directly compute the new attribute value as well:

Constraint checkAndCompute(attr0: Attr, attr1 : Attr) -> Attr;

Pattern example with benefit(1) { 
    let foo = op<dialect.foo>() {attr = attr_foo : Attr}; 
    let bar = op<dialect.bar>(foo) {attr = attr_bar : Attr}; 
    let attr_baz = checkAndCompute(attr_foo, attr_bar); 
    rewrite bar with { 
        let baz = op<dialect.baz>; 
        setAttribute(baz, attr<"\"attr\"">, attr_baz); 
        replace bar with baz; 
    }; 
} 

To achieve this the following changes were necessary:

• Remove simple check in PDLL parser to allow native constraints to return results

• Change PDL definition of pdl.apply_native_constraint to allow variadic results

• Change PDL_interp definition of pdl_interp.apply_constraint to allow variadic results

• Adjust PDLToPDLInterp Pass:
  The input to the pass is an arbitrary number of PDL patterns. The pass collects all predicates
  that are required to match all of the pdl patterns and tries to establish an ordering that allows
  creation of a single efficient matcher function to match all of them. Hence, the pass does not support
  the creation of a value on the matching side of a pattern out of the box. Simply referring to
  the pdl.apply_native_constraint operation is also not possible because the pdl patterns are deleted
  before the pdl_interp dialect operations are created.
  To solve this we record the type of the results in the predicate for the constraint and
  create positions (how the pass refers to values that are not materialized yet) for them.
  When a position is evaluated (i.e. used by an op on the rhs of the pattern) we emit a placeholder value
  to enable the creation of the pdl_interp dialect operation. The placeholder value is replaced later
  when the actual pdl_interp.apply_constraint operation is created.
  This is required in all scenarios where a native constraint returns a result that is used on the rhs
  because said result has to be an input the the pdl_interp rewriter function that is created for the pattern.
  However, the call to this function is always generated before the pdl_interp.apply_constraint operation is created.

• Modify Bytecode generator and interpreter:
  Constraint functions which return results have a different type compared to existing constraint functions.
  They have the same type as native rewrite functions. For this reason they are for now registered as rewrite functions.
  Other options:

  • change the type for all native constraints to include results
  • Note:
    Now that constraints may produce results it is hard to say what is the exact difference between native constraints
    and native rewrites. The only difference is that native constraints are evaluated during matching and native rewrites
    are evaluated during rewriting. So it might even make sense to unify the two concepts and have a single type for both
    and only a single function to register them, i.e. registerNativeFunction instead of registerConstraintFunction and registerRewriteFunction.

  ByteCode generation:

  • constraint lookup in ConstrainFunctions or RewriteFunctions depending on whether the constraint produces results
  • allocate memory for results in byte code

  ByteCode execution:

  • execute constraint as ConstraintFunction or RewriteFunction depending on whether it produces results
  • save results in byte code

For a follow-up PR:

• add support for constraints which create and return new operations
• add support for returning ranges from constraints

[ehsan-toosi]

First of all, thank you very much for your explanation in the PR description. I appreciate that a lot. That's a nice feature. I have a few questions/suggestions:

• Not for this work perhaps: How hard is it to give these third-kind of native a different name (i.e. utility native) and not mixing what we want with constraints?
• I don't see a useful case at the moment to have declarations of natives for PDLL. When we add new natives, we have to define them here as well which wasn't the case for PDL.
• While I was working with PDLL and sometimes converted it to pdl_interp, I realized that if you have a constraint that is not dependent on the operation or values or attribute you are capturing, after lowering it to pdl_interp, the order of the natives are not matched with the one in PDLL. In fact, it moves it to the beginning of the matching section and that was problematic for some cases. I can help you to recreate that case and see if it would be problematic here as well.
• Important: What if we enable rewrite failure instead of adding this kind of constraint? This is not supported by PDL now. So, the idea is, we match, we go the rewrite section and there we can fail and stop rewriting and we go to the next matching.

[martin-luecke]

Not for this work perhaps: How hard is it to give these third-kind of native a different name (i.e. utility native) and not mixing what we want with constraints?

It really depends on how we want to handle this operation:
In the PDLL frontend the operation is not named anyway so there would be only be changes to the parser as to which operation it emits.
If the aim of your proposal is to not touch how the existing operations are represented and handled we would need to introduce a lot of duplicated code in the lowering pass and in the byte code interpreter as we would handle his new operations almost completely similar to the existing ones.
However, I think this is a quite natural extension to the existing constraints (in the existing code there was even a reference that constraints with results is a TODO). Also note that the semantics of this does not yield a "produce a value on the lhs of the Pattern" behaviour but really still "check a native constraint and possibly return a value". i.e. the pdl.apply_native_constraint still always requires at least one operand to check a constraint on.

While I was working with PDLL and sometimes converted it to pdl_interp, I realized that if you have a constraint that is not dependent on the operation or values or attribute you are capturing, after lowering it to pdl_interp, the order of the natives are not matched with the one in PDLL. In fact, it moves it to the beginning of the matching section and that was problematic for some cases. I can help you to recreate that case and see if it would be problematic here as well.

A constraint is meant to always check something on the IR that you are matching (which is why it requires at least one operand) and will by the pass be ordered to a location where all of its operands have been checked and evaluated, i.e. when there are no operands to the beginning of the pattern (Note: This is also influenced by other patterns in the same file, especially when they call the same native constraint). I think what you are referring to are the constraints that should check for state in the pass it is applied in. They are quite special and out of scope for this extension. Nothing with regard to that has been changed so I expect the issues will be similar.

Important: What if we enable rewrite failure instead of adding this kind of constraint? This is not supported by PDL now. So, the idea is, we match, we go the rewrite section and there we can fail and stop rewriting and we go to the next matching.

The semantics of a PDL pattern is: if we have a successful match of the lhs then the rhs should be able to completely perform the rewrite. This clear separation of matching and rewriting enables the creation of an efficient matcher for a set of patterns and avoiding error handling and aborting rewriting in the bytecode interpreter (which would introduce a runtime penalty). I think we should avoid mixing these concepts more than absolutely required as to not lose these properties.

[martin-luecke]

One possible change that I would like your opinion on would be the registering of constraint functions with results.
With this proposal they have to be registered using PDLPatternModule::registerRewriteFunction. This seems very unnatural from the user perspective as they want to register a constraint, not a rewrite.
We could introduce a new PDLPatternModule::registerConstraintFunctionWithResults so the user does not have to think about how constraints with results are handled under the hood.

[ehsan-toosi]

I agree. That would be a better design IMO too. Separating them from the user's perspective reduces the confusion.

[maxbartel]

Really impressive! My only questions are about signaling a failure in a constraint with results.

[maxbartel]

How do we signal a failure with this interface? Don't we need something similar like FailureOr<ResultT>?

[martin-luecke]

Good catch!
I think it makes sense to remove the template magic and under the hood conversion that is enabled by this. The problem is that not all valid native rewrite functions are also valid constraint with results functions. We would have to adjust how the conversion from PDL structures to built-in structure is done for this.
So for the scope of this PR we will only support constraints with results in form of
LogicalResult (PatternRewriter &, PDLResultList &, ArrayRef<PDLValue>)

[maxbartel]

I don't know enough about PDLInterp, could you give me an example when this is relevant?

[martin-luecke]

From my description on top:

This is required in all scenarios where a native constraint returns a result that is used on the rhs
because said result has to be an input the the pdl_interp rewriter function that is created for the pattern.
However, the call to this function is always generated before the pdl_interp.apply_constraint operation is created.

An example for the output of this pass is the following:
PDL input:

module {
  pdl.pattern @simple : benefit(1) {
    %0 = operands
    %1 = operation "test.op"(%0 : !pdl.range<value>) 
    %attr = apply_native_constraint "NativeConstraint"(%1 : !pdl.operation) : !pdl.attribute 
    rewrite %1 {
      %3 = operation "test.success"  {"someAttr" = %attr}
      replace %1 with %3
    }
  }
}

The PDL_interp output is the following:

module {
  pdl_interp.func @matcher(%arg0: !pdl.operation) {
    pdl_interp.check_operation_name of %arg0 is "test.op" -> ^bb2, ^bb1
  ^bb1:  // 4 preds: ^bb0, ^bb2, ^bb3, ^bb4
    pdl_interp.finalize
  ^bb2:  // pred: ^bb0
    pdl_interp.check_result_count of %arg0 is 0 -> ^bb3, ^bb1
  ^bb3:  // pred: ^bb2
    %0 = pdl_interp.apply_constraint "NativeConstraint"(%arg0 : !pdl.operation) : !pdl.attribute -> ^bb4, ^bb1
  ^bb4:  // pred: ^bb3
    pdl_interp.record_match @rewriters::@pdl_generated_rewriter(%0, %arg0 : !pdl.attribute, !pdl.operation) : benefit(1), generatedOps(["test.success"]), loc([%arg0]), root("test.op") -> ^bb1
  }
  module @rewriters {
    pdl_interp.func @pdl_generated_rewriter(%arg0: !pdl.attribute, %arg1: !pdl.operation) {
      %0 = pdl_interp.get_results of %arg1 : !pdl.range<value>
      %1 = pdl_interp.get_value_type of %0 : !pdl.range<type>
      %2 = pdl_interp.create_operation "test.success" {"someAttr" = %arg0}  -> (%1 : !pdl.range<type>)
      pdl_interp.erase %arg1
      pdl_interp.finalize
    }
  }
}

The pdl_interp.record_match operation and the block it is in is not the last thing this pass generates for a matcher. This is done so the operation that is the last check before a successful match can have this block (^bb4) as successor (e.g. %0 = pdl_interp.apply_constraint here). But the pdl_interp.record_match needs the result of %0 = pdl_interp.apply_constraint as operand.
So there is a cyclic dependency here, which we solve using the placeholder.

In short:

• %0 = pdl_interp.apply_constraint needs the success branch as successor.
• pdl_interp.record_match is created with the success branch and needs the result %0 as operand