Lowering and type conversion of block arguments

MLIR
#1

Hi folks,

I’m struggling to get the right way to convert types of the block arguments when using the lowering infra.

Let’s say I have a simple dialect with three ops (value, branch, and return), e.g.:

module @a_module  {
  func @f() -> !foo.foo_value {
    %0 = "foo.value"() : () -> !foo.foo_value
    "foo.br"(%0)[^bb1] : (!foo.foo_value) -> ()
  ^bb1(%1: !foo.foo_value):  // pred: ^bb0
    %2 = "foo.return"(%1) : (!foo.foo_value) -> !foo.foo_value
  }
}

Then I lower it to std dialect via direct mapping (foo.value -> std.constant i8, foo.branch -> std.br, foo.return -> std.return), e.g.:

module @a_module  {
  func @f() -> !foo.foo_value {
    %c42_i8 = constant 42 : i8
    br ^bb1(%c42_i8 : i8)
  ^bb1(%0: !foo.foo_value):  // pred: ^bb0
    %1 = "foo.return"(%0) : (!foo.foo_value) -> !foo.foo_value
  }
}

It almost works except the bb1 argument and @f return types are not converted.
AIUI, I should be calling convertRegionTypes somewhere in the lowering pipeline.

The issue is that I don’t quite understand when should I call it?

I tried to call it from foo.branch lowering, but it (expectedly) breaks the IR as I modify things outside of the current operation.
I think the right place to call it would be in the context of std.FuncOp, but I want to preserve the FuncOp (i.e. it is a legal operation).

I also tried creating custom std.FuncOp, but then I don’t see how to lower it into the std.FuncOp

Any hints would be highly appreciated as I’m pretty much stuck at this point :sweat_smile:

P.S. here is a sample in case it helps https://github.com/AlexDenisov/mlir-type-conversion

#2

The part I was missing is how to “force” conversion of already legal op.

I overlooked dynamically legal

I ended up with something along these lines:

  mlir::LogicalResult matchAndRewrite(mlir::FuncOp op, llvm::ArrayRef<mlir::Value> operands,
                                      mlir::ConversionPatternRewriter &rewriter) const final {
    auto newType = typeConverter.convertType(op.getType()).cast<mlir::FunctionType>();
    // TODO: Propagate attributes?
    auto newFunc = rewriter.create<mlir::FuncOp>(op.getLoc(), op.sym_name(), newType);
    rewriter.inlineRegionBefore(op.body(), newFunc.addEntryBlock());
    rewriter.eraseBlock(&newFunc.getBlocks().back());
    if (mlir::failed(rewriter.convertRegionTypes(&newFunc.getBody(), typeConverter))) {
      llvm::errs() << "failed to convert region types\n";
      return mlir::failure();
    }
    rewriter.replaceOp(op, { newFunc->getResult(0) });

    return mlir::success();
  }
////

target.addDynamicallyLegalOp<mlir::FuncOp>(
      [&](mlir::FuncOp op) { return typeConverter.isLegal(op.getType()); });

Not sure if this is the right way, but so far looks like it does the trick.

I tried to call it from foo.branch lowering, but it (expectedly) breaks the IR as I modify things outside of the current operation.

Here I’ve seen cases (while printing the function during the conversion) with an <<UNKNOWN SSA VALUE>>, but in the end the IR is valid.
So a another side question: is it expected intermediate behaviour or there is still a problem with my code?

Thank you,
Alex.