[mlir][llvm] Handle invoke op branching to block with its result as an argument

[victor-eds]

Summary

In LLVM, having an invoke instruction branching to a block with a phi
node receiving the invoke instruction result as an argument is
perfectly legal. However, the translation of this construct to MLIR
would result in an invoke with its result being used as a block
argument to a successor, i.e., the operation result would be used in
its definition.

In order to fix this issue due to different IR structures (phi nodes
vs block arguments), this construct is interpreted with an llvm.invoke
operation branching to a dummy block having a single llvm.br operation
passing the required block arguments (including the result of the
llvm.invoke operation) to the actual successor block.

[zero9178]

Thank you! I think this approach makes the most sense.

For reference for others, this was previously discussed here: https://discourse.llvm.org/t/improvements-to-llvm-dialect-exception-related-operations/69503

mlir/lib/Target/LLVMIR/ModuleImport.cpp

1299

Please spell out the autos here and below. MLIR generally only uses auto when the type already appears on the right hand side or for difficult to spell out names (iterators, lambdas etc)

mlir/test/Target/LLVMIR/Import/basic.ll

97

It seems this has sadly not been done consistently in this file :(, but could we use FileCheck variable for the basic block names here?

97

[victor-eds]

Address comments

[victor-eds]

Thanks for the quick review! I've addressed your comments already.

[zero9178]

LGTM!
Please wait a day or so in case other reviewers have anything to add

[gysit]

LGTM, feel free to move the exception related tests into a separate test file, e.g. exception.ll, now that you are extending test coverage and functionality!

[victor-eds]

In D148313#4267806, @gysit wrote:

LGTM, feel free to move the exception related tests into a separate test file, e.g. exception.ll, now that you are extending test coverage and functionality!

Good point! I'll do it in a further PR. Thanks!

mlir/lib/Target/LLVMIR/ModuleImport.cpp

LogicalResult ModuleImport::convertInstruction(llvm::Instruction *inst) {
   if (inst->getOpcode() == llvm::Instruction::Invoke) {
     auto *invokeInst = cast<llvm::InvokeInst>(inst);
 
     SmallVector<Type> types;
     SmallVector<Value> operands;
     if (failed(convertCallTypeAndOperands(invokeInst, types, operands)))
       return failure();
 
-    SmallVector<Value> normalArgs, unwindArgs;
-    (void)convertBranchArgs(invokeInst, invokeInst->getNormalDest(),
-                            normalArgs);
-    (void)convertBranchArgs(invokeInst, invokeInst->getUnwindDest(),
-                            unwindArgs);
+    // Check whether the invoke result is an argument to the normal destination
+    // block.
+    bool invokeResultUsedInPhi = llvm::any_of(
+        invokeInst->getNormalDest()->phis(), [&](const llvm::PHINode &phi) {
+          return phi.getIncomingValueForBlock(invokeInst->getParent()) ==
+                 invokeInst;
+        });
 
+    Block *normalDest = lookupBlock(invokeInst->getNormalDest());
+    Block *directNormalDest = normalDest;
+    if (invokeResultUsedInPhi) {
+      // The invoke result cannot be an argument to the normal destination
+      // block, as that would imply using the invoke operation result in its
+      // definition, so we need to create a dummy block to serve as an
+      // intermediate destination.
+      OpBuilder::InsertionGuard g(builder);
+      directNormalDest = builder.createBlock(normalDest);
+    }
+
+    SmallVector<Value> unwindArgs;
+    if (failed(convertBranchArgs(invokeInst, invokeInst->getUnwindDest(),
+                                 unwindArgs)))
+      return failure();
+
+    // Create the invoke operation. Normal destination block arguments will be
+    // added later on to handle the case in which the operation result is
+    // included in this list.
     InvokeOp invokeOp;
     if (llvm::Function *callee = invokeInst->getCalledFunction()) {
       invokeOp = builder.create<InvokeOp>(
           loc, types,
           SymbolRefAttr::get(builder.getContext(), callee->getName()), operands,
-          lookupBlock(invokeInst->getNormalDest()), normalArgs,
+          directNormalDest, ValueRange(),
           lookupBlock(invokeInst->getUnwindDest()), unwindArgs);
     } else {
       invokeOp = builder.create<InvokeOp>(
-          loc, types, operands, lookupBlock(invokeInst->getNormalDest()),
-          normalArgs, lookupBlock(invokeInst->getUnwindDest()), unwindArgs);
+          loc, types, operands, directNormalDest, ValueRange(),
+          lookupBlock(invokeInst->getUnwindDest()), unwindArgs);
     }
     if (!invokeInst->getType()->isVoidTy())
       mapValue(inst, invokeOp.getResults().front());
     else
       mapNoResultOp(inst, invokeOp);
+
+    SmallVector<Value> normalArgs;
+    if (failed(convertBranchArgs(invokeInst, invokeInst->getNormalDest(),
+                                 normalArgs)))
+      return failure();
+
+    if (invokeResultUsedInPhi) {
+      // The dummy normal dest block will just host an unconditional branch
+      // instruction to the normal destination block passing the required block
+      // arguments (including the invoke operation's result).
+      OpBuilder::InsertionGuard g(builder);
+      builder.setInsertionPointToStart(directNormalDest);
+      builder.create<LLVM::BrOp>(loc, normalArgs, normalDest);
+    } else {
+      // If the invoke operation's result is not a block argument to the normal
+      // destination block, just add the block arguments as usual.
+      assert(llvm::none_of(
+                 normalArgs,
+                 [&](Value val) { return val.getDefiningOp() == invokeOp; }) &&
+             "An llvm.invoke operation cannot pass its result as a block "
+             "argument.");
+      invokeOp.getNormalDestOperandsMutable().append(normalArgs);
+    }
+
     return success();
   }
   if (inst->getOpcode() == llvm::Instruction::GetElementPtr) {
     auto *gepInst = cast<llvm::GetElementPtrInst>(inst);
     Type sourceElementType = convertType(gepInst->getSourceElementType());
     FailureOr<Value> basePtr = convertValue(gepInst->getOperand(0));
     if (failed(basePtr))
       return failure();

mlir/test/Target/LLVMIR/Import/basic.ll

; CHECK: ^bb3:
   ; CHECK: %{{[0-9]+}} = llvm.invoke @bar(%[[a3]]) to ^bb2 unwind ^bb1 : (!llvm.ptr) -> !llvm.ptr
   %6 = invoke ptr @bar(ptr %1) to label %4 unwind label %2
 
 ; CHECK: ^bb4:
   ; CHECK: llvm.return %{{[0-9]+}} : i32
   ret i32 0
 }
 
+declare i32 @foo2()
+
+; CHECK-LABEL: @invokePhi
+; CHECK-SAME:            (%[[cond:.*]]: i1) -> i32
+define i32 @invokePhi(i1 %cond) personality ptr @__gxx_personality_v0 {
+entry:
+  ; CHECK: %[[c0:.*]] = llvm.mlir.constant(0 : i32) : i32
+  ; CHECK: llvm.cond_br %[[cond]], ^[[bb1:.*]], ^[[bb2:.*]]
+  br i1 %cond, label %call, label %nocall
+; CHECK: ^[[bb1]]:
+call:
+  ; CHECK: %[[invoke:.*]] = llvm.invoke @foo2() to ^[[bb3:.*]] unwind ^[[bb5:.*]] : () -> i32
+  %invoke = invoke i32 @foo2() to label %bb0 unwind label %bb1
+; CHECK: ^[[bb2]]:
+nocall:
+  ; CHECK: llvm.br ^[[bb4:.*]](%[[c0]] : i32)
+  br label %bb0
+; CHECK: ^[[bb3]]:
+  ; CHECK: llvm.br ^[[bb4]](%[[invoke]] : i32)
+; CHECK: ^[[bb4]](%[[barg:.*]]: i32):
+bb0:
+  %ret = phi i32 [ 0, %nocall ], [ %invoke, %call ]
+  ; CHECK: llvm.return %[[barg]] : i32
+  ret i32 %ret
+; CHECK: ^[[bb5]]:
+bb1:
+  ; CHECK: %[[lp:.*]] = llvm.landingpad cleanup : i32
+  %resume = landingpad i32 cleanup
+  ; CHECK: llvm.resume %[[lp]] : i32
+  resume i32 %resume
+}
+
+; CHECK-LABEL: @invokePhiComplex
+; CHECK-SAME:                   (%[[cond:.*]]: i1) -> i32
+define i32 @invokePhiComplex(i1 %cond) personality ptr @__gxx_personality_v0 {
+entry:
+  ; CHECK: %[[c0:.*]] = llvm.mlir.constant(0 : i32) : i32
+  ; CHECK: %[[c1:.*]] = llvm.mlir.constant(1 : i32) : i32
+  ; CHECK: %[[c2:.*]] = llvm.mlir.constant(2 : i32) : i32
+  ; CHECK: %[[c20:.*]] = llvm.mlir.constant(20 : i32) : i32  
+  ; CHECK: llvm.cond_br %[[cond]], ^[[bb1:.*]], ^[[bb2:.*]]
+  br i1 %cond, label %call, label %nocall
+; CHECK: ^[[bb1]]:
+call:
+  ; CHECK: %[[invoke:.*]] = llvm.invoke @foo2() to ^[[bb3:.*]] unwind ^[[bb5:.*]] : () -> i32
+  %invoke = invoke i32 @foo2() to label %bb0 unwind label %bb1
+; CHECK: ^[[bb2]]:
+nocall:
+  ; CHECK: llvm.br ^[[bb4:.*]](%[[c0]], %[[c1]], %[[c2]] : i32, i32, i32)
+  br label %bb0
+; CHECK: ^[[bb3]]:
+  ; CHECK: llvm.br ^[[bb4]](%[[invoke]], %[[c20]], %[[invoke]] : i32, i32, i32)
+; CHECK: ^[[bb4]](%[[barg0:.*]]: i32, %[[barg1:.*]]: i32, %[[barg2:.*]]: i32):
+bb0:
+  %a = phi i32 [ 0, %nocall ], [ %invoke, %call ]
+  %b = phi i32 [ 1, %nocall ], [ 20, %call ]
+  %c = phi i32 [ 2, %nocall ], [ %invoke, %call ]
+  ; CHECK: %[[add0:.*]] = llvm.add %[[barg0]], %[[barg1]] : i32
+  ; CHECK: %[[add1:.*]] = llvm.add %[[barg2]], %[[add0]] : i32
+  %d = add i32 %a, %b
+  %e = add i32 %c, %d
+  ; CHECK: llvm.return %[[add1]] : i32
+  ret i32 %e
+; CHECK: ^[[bb5]]:
+bb1:
+  ; CHECK: %[[lp:.*]] = llvm.landingpad cleanup : i32
+  %resume = landingpad i32 cleanup
+  ; CHECK: llvm.resume %[[lp]] : i32
+  resume i32 %resume
+}
+
 ; CHECK-LABEL: @hasGCFunction
 ; CHECK-SAME: garbageCollector = "statepoint-example"
 define void @hasGCFunction() gc "statepoint-example" {
     ret void
 }
 
 ;CHECK-LABEL: @useFreezeOp
 define i32 @useFreezeOp(i32 %x) {