[PSA] Updating Inliner Interface

MLIR
1

We plan to update the inliner interface beginning of next week. The revision adds support for attribute based argument and result handling which has been discussed in
[RFC] Handle Argument/Result Attributes While Inlining. The change is motivated by the LLVM dialect inliner, which, for example, has to handle the byval, signext, or zeroext attributes. A base revision introduces the new argument and result handlers and a follow up revision replaces the existing type conversion hook.

Current Interface

Up to now, type conversions have been performed using:

Operation *materializeCallConversion(OpBuilder &builder, Value input, Type resultType, Location loc)

The hook can only produce a single cast operation. This limitation is needed to make the hook revertible, which is a requirement since the inlining may fail after calling the hook.

Updated Interface

The updated design splits the type conversion into two steps:

bool isTypeConvertible(Operation *call, Operation *callable, Type sourceType, Type targetType, DictionaryAttr argOrResAttrs, bool isResult)

allows the inliner to check if a type conversion is possible. Only once it knows the inlining is possible it calls two separate hooks to perform the actual argument and result conversion:

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

These hooks run right before and after inlining the callee region and are not revertible.

Once the second revision lands, downstream projects need to replace materializeCallConversion by the newly introduced isTypeConvertible, handleArgument, and handleArgument hooks. The base revision does not require any changes.

The Flang inliner flang/lib/Optimizer/Dialect/FIRDialect.cpp and the Toy dialect inliner mlir/examples/toy/Ch4/mlir/Dialect.cpp are upstream examples for the required changes. In case of Flang, the inlining interface currently implements the following conversion hook:

  mlir::Operation *materializeCallConversion(mlir::OpBuilder &builder,
                                             mlir::Value input,
                                             mlir::Type resultType,
                                             mlir::Location loc) const final {
    return builder.create<fir::ConvertOp>(loc, resultType, input);

After landing the revisions, it will change to:

  /// This hook checks if `sourceType` is convertible to `targetType`.
  bool isTypeConvertible(mlir::Operation *call, mlir::Operation *callable,
                         mlir::Type sourceType, mlir::Type targetType,
                         mlir::DictionaryAttr argOrResAttrs,
                         bool isResult) const final {
    return true;
  }

  /// This hook converts the `argument` to `targetType` if needed.
  mlir::Value handleArgument(mlir::OpBuilder &builder, mlir::Operation *call,
                             mlir::Operation *callable, mlir::Value argument,
                             mlir::Type targetType,
                             mlir::DictionaryAttr argumentAttrs) const final {
    // The convert operation folds if the argument type matches the target type.
    return builder.createOrFold<fir::ConvertOp>(call->getLoc(), targetType,
                                                argument);
  }

  /// This hook converts the `result` to `targetType` if needed.
  mlir::Value handleResult(mlir::OpBuilder &builder, mlir::Operation *call,
                           mlir::Operation *callable, mlir::Value result,
                           mlir::Type targetType,
                           mlir::DictionaryAttr resultAttrs) const final {
    // The convert operation folds if the result type matches the target type.
    return builder.createOrFold<fir::ConvertOp>(call->getLoc(), targetType,
                                                result);```
1 Like
2

Update

I landed the base revision of the inliner interface change (https://reviews.llvm.org/D145582). It adds two new methods to the CallableOpInterface, which provide access to the argument and result attributes. Downstream users thus need to implement the following methods for any operation that implements CallableOpInterface:

    /// Returns the argument attributes for all callable region arguments or
    /// null if there are none.
    ::mlir::ArrayAttr getCallableArgAttrs() {
      return getArgAttrs().value_or(nullptr);
    }

    /// Returns the result attributes for all callable region results or
    /// null if there are none.
    ::mlir::ArrayAttr getCallableResAttrs() {
      return getResAttrs().value_or(nullptr);
    }

The above implementation should work for all callable operations that implement the function op interface. Any other callables may just return nullptr if they do not have argument or result attributes.