[Python bindings] Inferring Operation result type based on block builder function

I’ve got an operation which contains a block. With the Python bindings, I’m looking to infer the result type of said operation based on code (the ‘block builder’ python function callback) which builds the block it contains and returns an MLIR Value upon which the block terminator gets built.

The problem I’m running into is I need to create the operation (and specific the result type which I would like to infer), then create the block, then run the block builder (which returns a Value who’s type I would like to use).

In C++, I think I could build a disconnected (un-owned) block first (since I know the block argument types) then create the containing operation and attach the block. I can’t find any way to do this with the existing Python bindings.

Any suggestions? Am I being stupid and not seeing an obvious way?

We are intentionally not allowing detached blocks or regions, this simplifies lifetime management (and the associated overhead) significantly. Can you factor out the result type computation into a separate function that doesn’t build IR?

No. This is a user API though which we expose very thin wrappers around ops. It would be a significant change to build those wrappers first, compute the return type, then do the actual IR building.

I’m perfectly happy to build a dummy operation, add a block, fill it, then move/clone that block to the actual op with the computed result type. I couldn’t find a way to move a block from one operation to another, though. Does this sound like a reasonable way to go?

There has been skepticism expressed in the past about how much these low level IR manipulation APIs are suitable as an actual user (not just “compiler writer”) API. These kinds of impedance mismatches were a part of the intuition.

I’m still not completely tracking what the problem is, as I know I’ve done something similar in the past. Is this because you don’t know the result arity of the operation before trying to construct the IR?

AFK so can’t look in detail, but some kind of a Block.append_to(region) seems not too far to go (with a dummy operation, transferring ownership).

I concur with Alex that it would be good to avoid the accounting that comes with detached ownership of Region and Block unless if there is a strong reason to backtrack on that design decision – getting that right for Operation was not trivial (or free), and it required strict accounting at every point in the API which deals in the given type (so is a pretty invasive thing to do).

To be clear, we’re not directly exposing all IR manipulations to users, but for Operation building we are automatically creating wrappers. Operations which would benefit from nicer syntax (e.g. ones with blocks) and/or we care about, we write custom wrappers. We also have an “escape hatch” for miscellaneous operations which we don’t care enough about using the automatic wrapper. This is works well enough for simple ops like addition, concatenations, etc. Note that we do this at a layer abovethe low-level CIRCT python bindings, so we rely on Python code in the CIRCT bindings to take care of things like result type computation.

This is the relevant patch for anyone interested in the specifics: [PyCDE] Add SystolicArray construct (#2831) · llvm/circt@4120709 · GitHub

For now, we assume one result but in the future we’d like to not be bound to that. For now, it’s just the single result type. IIRC there are issues with changing the result type after IR creation…?

That’s exactly what I was thinking.

I totally understand, which I why I asked before hacking up IRCore.cpp :wink:

I believe that changing the result type is fully in bounds but changing the arity is not.

2 Likes

I’ve prototyped this, but it required adding a new C API function. Is there an existing one which should handle this move? I couldn’t find one and am hesitant to increase the C API surface area.

I haven’t followed closely here, but isn’t mlirRegionAppendOwnedBlock matching what you want here?

/// Takes a block owned by the caller and appends it to the given region.
MLIR_CAPI_EXPORTED void mlirRegionAppendOwnedBlock(MlirRegion region,
                                                   MlirBlock block);

No, it needs to transfer ownership from another Operation. A “remove block from region without erasing” call in conjunction with that one would work…

Given that this isn’t a common use-case, I’m happy to add this functionality in my code rather than MLIR core…

Oh right, that’s not great actually. We should support moving a block from a region to another one in the C API. That said it may come with another fairly large (?) API surface because one may need to remap some values as well.

In my case, I don’t have to.

Considering the C API is intended to be very stable, I don’t want to be adding a single function which isn’t fully baked to it just to support my use-case. I’ll just keep it in my code for now.

It seems unlikely that we could (or would want to pay the price for) making such an API completely safe from an ownership perspective, but I’m supportive of adding such a thing if it is useful.

Dumb question, though: aside from function name and documentation, doesn’t mlirRegionAppendOwnedBlock(MlirRegion region, MlirBlock block) already do what you want in fact? The blocks are just an intrusive list and inserting elsewhere effectively just transfers ownership. We could just define this as the intended behavior.

1 Like

As of now, no. Asserts since the block is already owned. It could be made to work, however, pretty easily.

I haven’t thought through all of the implications, but I think I’d be ok with such an extension (weighed against a lot more duplicated API surface area).

As Mehdi says, a “full” API for cloning/moving/inlining would be a pretty big footprint (needing apis for block/value mapping, etc). This simple “I know what I’m doing, just move the block” API could either be done implicitly when adding (ie. Instead of asserting) or we could add an API to detach a block. The latter is more consistent with the C++ data structures. Either would be fine with me.

1 Like

Same here, both are fine with me. Although I personally have a taste for simpler APIs, i.e. have a “append” only append and have a separate “detach”.

https://reviews.llvm.org/D123165

Should I split this up into two patches?

There is a tradeoff in to just add more steps for clients to realize “something”, so if there is a common sequence of “simpler APIs” calls that is repeated it can be worth looking into introducing a higher level API (in particular to match the kind of thing we do in C++) to avoid “reimplementing the wheel” over and over.

Thanks all: https://github.com/llvm/llvm-project/commit/8d8738f6fed389e63efd6ad4f8c198828acfc604

1 Like