Setting scf::IfOp resultType from within the builder

MLIR
#1

The IfOp builder takes a resultTypes up front. My problem is I have the thenBuilder and elseBuilder which have some user code getting added to blocks, each returns a Value and only after I have executed those functions can I get what value to yield and return out of the IfOp.

Is there a way to modify the resultTypes? or my choice is only to create the IfOp, get the returnTypes and then create a new one, this time setting the resultTypes?

Could this build function below have got resultTypes returned by thenBuider and elseBuilder (also performing some validation on both to ensure they are same number with same elementTypes)?

void IfOp::build(OpBuilder &builder, OperationState &result,
                 TypeRange resultTypes, Value cond,
                 function_ref<void(OpBuilder &, Location)> thenBuilder,
                 function_ref<void(OpBuilder &, Location)> elseBuilder) {
  assert(thenBuilder && "the builder callback for 'then' must be present");

  result.addOperands(cond);
  result.addTypes(resultTypes);

  OpBuilder::InsertionGuard guard(builder);
  Region *thenRegion = result.addRegion();
  builder.createBlock(thenRegion);
  thenBuilder(builder, result.location);

  Region *elseRegion = result.addRegion();
  if (!elseBuilder)
    return;

  builder.createBlock(elseRegion);
  elseBuilder(builder, result.location);
}
Scf.whileOp type inference with dynamic shapes
#2

I’ve run into the same thing before but only when using the C and python builder APIs, which do not have the same restriction. I don’t know if it’s the best way, but it is reasonable to build invalid ops, so long as you correct them before verification. I suspect it is fine if you construct it with an empty range for result types, then when you know more, just setResultTypes on the underlying operation.

The verifier enforces that the terminator operand types match the if op’s result types, and so long as you arrange for that to happen, it is less important how you get there.

There may be a more principled answer…

#3

That won’t work. You can’t change the number of results after an operation is created.

#4

Ah, right.

(More characters to satisfy discourse)

#5

Can you create and populate the region body for the then/else branch and only after that create the IfOp?

#6

Below are the only builders I have access to, it creates its own regions from what I see?

Or are you suggesting somehow I move regions inside the IfOp after the fact? Is there some way to move regions?

  static void build(::mlir::OpBuilder &odsBuilder, ::mlir::OperationState &odsState, Value cond, bool withElseRegion);
  static void build(::mlir::OpBuilder &odsBuilder, ::mlir::OperationState &odsState, TypeRange resultTypes, Value cond, bool withElseRegion);
  static void build(::mlir::OpBuilder &odsBuilder, ::mlir::OperationState &odsState, TypeRange resultTypes, Value cond, function_ref<void(OpBuilder &, Location)> thenBuilder = buildTerminatedBody, function_ref<void(OpBuilder &, Location)> elseBuilder = nullptr);
  static void build(::mlir::OpBuilder &odsBuilder, ::mlir::OperationState &odsState, Value cond, function_ref<void(OpBuilder &, Location)> thenBuilder = buildTerminatedBody, function_ref<void(OpBuilder &, Location)> elseBuilder = nullptr);
#7

Yes, but if you already built the body/Region for each branch, the constructor that takes thenBuilder and elseBuilder can be invoked with lambdas that will splice your bodies in place.

#8

Thanks, let me try to see how to do that and report back.

Below is what I have right now, just for reference.

    auto resultType = mlir::UnrankedTensorType::get(getElementTypeOrSelf(inputValues->at(0).getType()));
    auto trueBlockOp = [&](mlir::OpBuilder &nested, mlir::Location loc) {     
        trueResult = _trueBlock();
        builder->create<YieldOp>(nameLoc,
                                 trueResult.value);
    };
    auto falseBlockOp = [&](mlir::OpBuilder &nested, mlir::Location loc) {  
        falseResult = _falseBlock();
        builder->create<YieldOp>(nameLoc,
                                 falseResult.value);
    };
    IfOp op = builder->create<IfOp>(nameLoc,
                                    resultType,
                                    inputValues->at(0),
                                    trueBlockOp,
                                    falseBlockOp);
#9

This seems to do the trick, thanks, is this what you meant as well?

Although long term do we see value in having an ifOp builder in scf which handles this for the user (as a convenience)?
Instead of having to create the op twice and moving things around? I am sure I will run into the same thing for whileOp too, but I can use the same trick I think.

    auto trueBlockOp = [&](mlir::OpBuilder &nested, mlir::Location loc) {
        trueResult = _trueBlock();  
        builder->create<YieldOp>(nameLoc,
                                 trueResult.value);
    };
    auto falseBlockOp = [&](mlir::OpBuilder &nested, mlir::Location loc) {
        falseResult = _falseBlock();
        builder->create<YieldOp>(nameLoc,
                                 falseResult.value);
    };
    
    IfOp op = builder->create<IfOp>(nameLoc,
                                    inputValues->at(0),
                                    trueBlockOp,
                                    falseBlockOp);
    assert(getElementTypeOrSelf(trueResult.value.getType()) == getElementTypeOrSelf(falseResult.value.getType()));
    auto resultType = mlir::UnrankedTensorType::get(getElementTypeOrSelf(trueResult.value.getType()));
    
    auto trueBlockWithResultOp = [&](mlir::OpBuilder &nested, mlir::Location loc) {
        nested.getBlock()->getParent()->takeBody(op.thenRegion());
    };
    auto falseBlockWithResultOp = [&](mlir::OpBuilder &nested, mlir::Location loc) {
        nested.getBlock()->getParent()->takeBody(op.elseRegion());
    };
    IfOp opNew = builder->create<IfOp>(nameLoc,
                                       resultType,
                                       inputValues->at(0),
                                       trueBlockWithResultOp,
                                       falseBlockWithResultOp);
    op.erase();
#10

It isn’t clear to me why the creation of the first IfOp is necessary here?
I don’t know what is the type of trueResult, but can’t you do just something that looks like this:

    Block trueBlock;
    trueResult = _trueBlock(&trueBlock);  
    OpBuilder::atEnd(&trueBlock).create<YieldOp>(nameLoc, trueResult.value);

    Block falseBlock;
    falseResult = _falseBlock(&falseBlock);
    OpBuilder::atEnd(&falseBlock).create<YieldOp>(nameLoc, falseResult.value);
    assert(trueResult.getType() == falseResult.getType());


    auto trueBlockWithResultOp = [&](mlir::OpBuilder &nested, mlir::Location loc) {
        nested.getBlock()->getOperations().splice(nested.getBlock()->end(), trueBlock.getOperations());
    };
    auto falseBlockWithResultOp = [&](mlir::OpBuilder &nested, mlir::Location loc) {
        nested.getBlock()->getOperations().splice(nested.getBlock()->end(), falseBlock.getOperations());
    };
    IfOp opNew = builder->create<IfOp>(nameLoc,
                                       trueResult.getType(),
                                       inputValues->at(0),
                                       trueBlockWithResultOp,
                                       falseBlockWithResultOp);
1 Like
#11

You are absolutely correct, I am able to use this and then I don’t have to create 2 IfOps, I can splice ops from these blocks.
That resolves my dilemma I think, thanks all!

#12

Alternatively, we could add a builder that derives the result types from the terminators of the regions, but that would require the body-builder callbacks to always create the proper terminator (and fail an assertion otherwise).

#13

That IMO would be useful.

#14

This should be straightforward to implement, would you like to give it a try?

#15

Sure will take a shot.

1 Like