Using (DRR) on variadic operand

I would like to do an extremely simple rewrite on an Op with a single variadic operand.

Say I have an Op called MyMaxOp with operand Variadic<AnyTensor>:$input

I’d like to do something like this

MaxPattern : Pat< 
(MyMaxOp $input__1, $input__2, $input__3_to__n),
(MyMaxOp (MyMaxOp $input__1, $input__2) $input__3_to__n)
>;

In other words when more then two operands are provided to MaxOp group the first two operands into a single MaxOp operand. This would then recursively rewrite the variadic MyMaxOp to nested MyMaxOp, such that MaxOp would only ever be used on two operands at a time.

Unfortunately the DRR documentation is a little bit of a tease. Right after it defines the terminology for working with Variadic inputs and outputs, it has a good old-fashioned TODO.

If there are any examples for matching variadics in DRR, or any hints on syntax, that would be greatly appreciated

Variadic matching support is not complete/some welcome enhancements possible there. mlir/test/lib/Dialect/Test/TestOps.td has some examples. But I can explain the syntax, So in the s-exp’esque form of the DAG matchers, each internal node has 1 value. So we introduce a convention for multi-result ops to designate which value should be matched on. In the above you’d have (Foo:$input__0 $x, $y, $z) and here it names the results of Foo (prefix $input) and says the first result is used in matching (I think you actually just have 1 result for max so this convention doesn’t help).

In what you want to do, you’ll have MyMaxOp has 1 operand that is variadic, so in DRR form there is only 1 operand too (it matches 1:1 with ODS def of the operands and attributes). So you’d have to do something like

MaxPattern : Pat<(MyMaxOp $input),
  (MyMaxOp (ConcatOperands (MyMaxOp (PickFirstTwo $input)),
                           (DropFrontTwo $input))),
  [(NumOperandsGreaterThanTwo $input)]>; 

DRR side (where I just made up all the helper operators and one would implement using custom operatives, look at BindMultipleNativeCodeCallResult as example).

We’ve delayed on these as we were focussing on PDL/PDLL, so unfortunately you may run into some TODOs.

1 Like

(Concat one may even require some changes in DRR backend, shouldnt be too big changes but could be missing still)

1 Like

I was hoping I would be lucky and that there might be a concise way of handling this in DRR. A sml, Haskell like way of head and tail matching arrays would be super neat for Variadics.

Something like using (arg1:arg2:arg_left_overs) To match the first two variadic operands. Or even using a recursive datatype for the matching. I’m not familiar enough with TableGen to know the constraints you guys are under, in terms of what features and syntax would actually be viable.

I’ll try out your suggestions. Thanks for pointing me to where I could find examples, and the quick primer on syntax. That was exactly what I was looking for!

Created a small example ⚙ D121046 [mlir] Add test case using variadic operand pattern (need to hoist some of this out into helpers and update docs, but wanted to make sure things work as expected).

Yeah we are a bit constrained syntax wise (some even dislike __ short hand). We could do something like $arg1___arg2___arg_remainder syntax wise (so have __ represent fixed indexing, ___ as some separator). But it may get quite magical. Or something like (Cons $arg1, (Cons $arg2, $arg_left_overs)) could be made to work.

Thanks @jpienaar, you are the man!