[RFC] Handle Argument/Result Attributes While Inlining

MLIR
1

This RFC addresses the lack of a mechanism to handle attributes on input arguments and results while inlining using the DialectInlinerInterface. This is relevant, for example, for correctness when inlining arguments with the byval attribute in LLVM::LLVMFunctionOp in the LLVM dialect, which requires copying the data on the input pointer before passing the value to the inlined operations.

Current state

The DialectInlinerInterface (defined in InliningUtils.cpp) currently exposes the following transformation hooks:

virtual void handleTerminator(Operation *op, Block *newDest) const;

virtual void handleTerminator(Operation *op,
                              ArrayRef<Value> valuesToReplace) const;

virtual Operation *materializeCallConversion(OpBuilder &builder, Value input,
                                             Type resultType,
                                             Location conversionLoc) const;

virtual void processInlinedCallBlocks(
    Operation *call, iterator_range<Region::iterator> inlinedBlocks) const;

The gap to handling argument/result attributes while inlining is twofold:

  • None of the transformation hooks above expose the CallableOpInterface operation, which is also the operation that implements FunctionOpInterface (and thereby exposes argument/result attributes) in practice.
  • The DialectInterfaceInliner does not currently know about argument/result attributes, as they are implemented on FunctionOpInterface.

Proposed change

A possible solution is to add hooks that handle each argument and result attributes pair:

/// For an argument with input value `input` and argument `attributes`, allow adding
/// zero or more operations returned in `newOperations` that handle these attributes, 
/// then return the new or existing value that should be used by inlined operations.
virtual Value handleArgumentAttributes(OpBuilder &builder,
                                       DictionaryAttr attributes,
                                       Value input,
                                       SmallVectorImpl<Operation *> &newOperations) {
  return input;
}

virtual Value handleResultAttributes(OpBuilder &builder,
                                     DictionaryAttr attributes,
                                     Value output,
                                     SmallVectorImpl<Operation *> &newOperations) {
  return input;
}

Returning the newOperations is necessary because of the cleanupState mechanism in inlineCall that requires us to know about newly added operations in case of failure.

Additionally, CallableOpInterface would need to be extended to know about argument/result attributes by adding (for example) the following APIs:

virtual ArrayAttr getCallableArgAttrs() { return {}; }
virtual ArrayAttr getCallableResAttrs() { return {}; }

These two methods would need to be implemented by Func::FuncOp and LLVM::LLVMFuncOp and redirected to FunctionOpInterface::getArgAttrsAttr and FunctionOpInterface::getResAttrsAttr, respectively.

Alternatives

We thought of two alternative directions that could achieve handling of argument/result attributes, which both seem somewhat inelegant to us:

  • Generalize materializeCallConversion into some form of monolithic functions that prepare/handle all inputs and outputs in whatever way necessary, passing both the CallOpInterface and the CallableOpInterface.
  • Extend processInlinedBlocks to a more general function that can do any additional processing necessary, by passing CallableOpInterface, as well as the region being inlined into.

Any suggestions for alternative solutions or modifications to the proposed solution are very welcome!

(Tagging @River707 as the author of the DialectInlinerInterface)

4 Likes
2

Thank you for working on this, I think the dialect inliner interface definitely needs to be more a lot more powerful to handle cases such as yours!

That said, for what reason do you think it is more elegant to have the inliner interface be aware of argument and result attributes, and add hooks handling them, rather than having a more generalized mechanism? On first thought, the alternatives you proposed sound more useful to me, including for other use cases such as eg. inlining a llvm.invoke operation.

Additionally, I’d suggest instead of having your proposed hooks require users to return any newly created operations via the newOperations out parameter, that you could instead install a OpBuilder::Listener instance on the builder passed as the first argument to get any operations created.

3

A solution to this (and any other potential missing hook) could definitely be to just generalize to something along the lines of:

void prepareInline(OpBuilder &builder, CallOpInterface call,
                   CallableOpInterface callee, Region *region);

// All operations in `callee` have already been moved into `call`,
// so it's a bit of an empty shell at this point.
void finalizeInline(OpBuilder &builder, CallOpInterface call,
                    CallableOpInterface callee, Region *region,
                    iterator_range<Region::iterator> inlinedBlocks);

The downside is that we will have little to no reuse of common code, such as injecting appropriate casts (materializeCallConversion), replacing terminators in handleTerminator, and so on. This somewhat undermines the motivation for having an interface in the first place.

If we keep both the existing hooks, and add more generalized hooks like the above, we run into the issue of the responsibility and ordering between hooks being fairly arbitrary/unclear.

That being said, I would be in favor of at least generalizing processInlinedBlocks to the finalizeInline above. Then the implementer of the interface has complete freedom, even if it won’t always be elegant.

4

⚙ D145582 [mlir] Argument and result attribute handling during inlining. implements two new handlers for argument and result handling:

  virtual Value handleArgument(OpBuilder &builder, Operation *call,
                               Operation *callable, Value argument,
                               Type targetType,
                               DictionaryAttr argumentAttrs) const;
  virtual Value handleResult(OpBuilder &builder, Operation *call,
                             Operation *callable, Value result, Type targetType,
                             DictionaryAttr resultAttrs) const;

It also extends the CallableOpInterface as @definelicht suggested.

These new hooks should be sufficient to cover our use cases in the LLVM inliner that include argument marked with byval, zeroext, or signext. The new handlers run right before and after inlining and cannot fail. Unlike materializeCallConversion, which has to rollback if inlining fails at a later stage.

At the moment, the new handlers cannot be used for type conversions since materializeCallConversion enforces that there are no type mismatches after its execution. A follow up revision may either change this or replace materializeCallConversion by the new handlers.