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

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.

Differential Revision: https://reviews.llvm.org/D148313

diff --git a/mlir/lib/Target/LLVMIR/ModuleImport.cpp b/mlir/lib/Target/LLVMIR/ModuleImport.cpp
index c5f9d0d..4338f03 100644 (file)
--- a/mlir/lib/Target/LLVMIR/ModuleImport.cpp
+++ b/mlir/lib/Target/LLVMIR/ModuleImport.cpp
@@ -1293,28 +1293,73 @@ LogicalResult ModuleImport::convertInstruction(llvm::Instruction *inst) {
     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) {

diff --git a/mlir/test/Target/LLVMIR/Import/basic.ll b/mlir/test/Target/LLVMIR/Import/basic.ll
index 1ce3ded..d3eebc0 100644 (file)
--- a/mlir/test/Target/LLVMIR/Import/basic.ll
+++ b/mlir/test/Target/LLVMIR/Import/basic.ll
@@ -87,6 +87,77 @@ define i32 @invokeLandingpad() personality ptr @__gxx_personality_v0 {
   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" {