Finalizing operand/result ordering in Python bindings

MLIR
1

Hi folks - I’m forking a discussion on https://reviews.llvm.org/D94810 to this thread, as it produced different opinions on the path forward.

My goal: I’m working e2e examples of using the python API and am trying to make final big corrections to it prior to letting it out of the lab, so to speak. Right now, the way we order operands and results in various contexts is inconsistent and needs to be made consistent in some way across all of the APIs, and I have usability concerns about the ordering as inherited from the C++ side, as I think that language differences may be a key point here. I’d like to fix this once one way or the other and be done with it.

Current state:

  • Docs are mostly up to date: MLIR Python Bindings - MLIR
  • Operation.create() takes arguments (operands: Sequence, results: Sequence, ...)
  • OpView._ods_build_default() takes arguments (operands: Sequence, results: Sequence, ...)
  • FunctionType.get() takes arguments (inputs: Sequence, results: Sequence)
  • For ops with a default builder, the __init__ method is generated with a signature (result..., operand_or_attribute..., *, loc: mlir.ir.Location = None, ip: mlir.ir.InsertionPoint = None)
  • I have previously argued for the Python API to order arguments in order of increasing optionality where that is not in conflict with some other requirement.
  • While porting a lot of Python-API using code, I have tripped on the backwards results-before-operands ordering multiple times, despite being very familiar with the C++ API. Something about it just feels wrong in Python.
  • While result type inference is not yet exposed in the Python API, it is planned, and it would be nice if generated builder functions worked with the way the language deals with optionality for this (expected to be common) case where results are optional.

Note that in Python, we lack a very effective way to perform method overloading, with the pythonic way to typically involve some usage of optional keyword arguments and/or differently named methods to achieve similar aims. There is therefore some more pressure than in C++ to make the default builder be as good as it can be, and we intend to collapse all permutations that differ only by optionality or behavior that can be triggered based on the presence of a keyword argument into a single default builder where at all possible. This should eliminate the need for most of the C++ builder permutations that exist largely for sugaring (and they would not map cleanly anyway (given that sugar in C++ is unlikely to produce a simpler/easier result in Python anyway)).

I am in the process of teaching the OpView._ods_build_default() method to just do the right thing with respect to all native traits that we support, allowing custom builder code to leverage it for easy implementation of user entry-points. It does this by looking at class level attributes that encode details from the ODS for use at runtime and gets us out of the business of generating convoluted Python code to do the same for each op in tablegen. We merely need to decide on the parameter ordering for user-level default and custom builders and map it to this method.

Proposal:

I would like to diverge from the C++ builder ordering, putting operands/attributes before results in generated builder methods. Once result type inference comes in to play, this would allow us to say this:

# With result type inference
op = std.AddFOp(lhs, rhs)
# Explicit result type
op = std.AddFOp(lhs, rhs, result_type)
# Linalg without out args.
op = linalg.MatmulOp(lhs, rhs)

If we follow the C++ convention, this would be:

# With result type inference
op = std.AddFOp(None, lhs, rhs)
# Explicit result type
op = std.AddFOp(result_type, lhs, rhs)
# Linalg without out args.
op = linalg.MatmulOp(None, lhs, rhs)

Alternative considered:

  • Leave it as-is and make everything consistent with the generated C++ builders.
  • Don’t use positional arguments at all, only keyword args.
  • Generate different named builder class methods like MyOp.build_inferred(lhs, rhs) which does not accept the result types at all for the inferred result type case.

Opinions?
@ftynse, @jpienaar, @nicolasvasilache who had comments

2

Which one? There are multiple generated C++ ones, for the inferred type case you have no result type specification ones. That then means it isn’t possible to specify result type when result type can be inferred (well one can always use the low level op creation there). So the AddFOps example would only allow the first form in your example at the end.

I think if we support only one builder then why not the most specialized one ods would have generated in C++? (That does include ones with default args). And we can always change the most specialized one as that is one generated for convenience. So I think I’m saying we could have python and C++ be consistent by choosing the most suitable C++ build method, rather than choosing one or the other.

3

Good question - and you are highlighting the limits of what I know about this. I would greatly appreciate your advice. I’ve been treating “default builder” as a single function (which is how Alex coded it on the python side, if I understand correctly).

That is not a bad idea. Can you point me at the tablegen code that generates these permutations for C++? I assume there is some way to determine the “most specialized”…

4

I would have suggested this as well!

We’d want to generate only std.AddFOp(lhs, rhs) (likely with an optional extra dictionary of attribute).
Even in C++ I’d like to remove the generic builder from the API (or isolate it with a different name), we discussed it recently here: Bug prone ODS builder - #7 by clattner

In C++ we generate the builders with the fields in the order of the ODS definition, there is just a trick around the attributes with a default value I believe.

5

This makes sense to me.

6

Sounds good. I can give this a shot when I have some coding time for this. We will need to extend the CAPI for result type inference.

7

The one down-side is that this will be the first feature that ties the generated python code to only being functional to build dialects that are linked in to LLVM and loaded (i.e. the result type inference is an interface). This was probably unavoidable, but I’ve found link-free use of unregistered dialects to be really useful: the python bindings become an easy, ad-hoc tool for doing arbitrary things to the IR in a few lines of code. It’d be nice to preserve this swiss-army knife aspect of it, imo… especially as real fixes for our linkage story are a ways out/haven’t been designed yet.

What if we just made the currently private _ods_build_default class method be “public” (renamed to build_default) and preserved it as a supported way to always get at the default building functionality for whatever reason? I believe that giving it a different name avoids the class of things in the above bug and will require someone to choose to use it/should be fine. So then, I would always have the option of:

op = std.AddFOp.build_default([result_type], [lhs, rhs])

We could go a step further and generate a default builder with the keyword operands and results spelled out, but I feel that might be encouraging this as a normal thing more than we intend to.

8

Reading the above bug again, I like the proposed name there. Let’s rename _ods_build_default → build_generic.

9

LGTM: this is exactly what I’d like to see on the C++ side as well (segregating the “generic” builder with a different name).