[mlir][LLVM] Fix mapping of result values of `llvm.invoke` during export

The result value of a llvm.invoke operation is currently not mapped to the corresponding llvm::Value* when exporting to LLVM IR. This leads to any later operations using the result to crash as it receives a nullptr.

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

2 files changed, 48 insertions, 9 deletions

diff --git a/mlir/lib/Target/LLVMIR/Dialect/LLVMIR/LLVMToLLVMIRTranslation.cpp b/mlir/lib/Target/LLVMIR/Dialect/LLVMIR/LLVMToLLVMIRTranslation.cpp
index 8d19f466fe5b..4f5e636c0a8e 100644
--- a/mlir/lib/Target/LLVMIR/Dialect/LLVMIR/LLVMToLLVMIRTranslation.cpp
+++ b/mlir/lib/Target/LLVMIR/Dialect/LLVMIR/LLVMToLLVMIRTranslation.cpp
@@ -342,22 +342,29 @@ convertOperationImpl(Operation &opInst, llvm::IRBuilderBase &builder,
   if (auto invOp = dyn_cast<LLVM::InvokeOp>(opInst)) {
     auto operands = moduleTranslation.lookupValues(opInst.getOperands());
     ArrayRef<llvm::Value *> operandsRef(operands);
+    llvm::Value *result;
     if (auto attr = opInst.getAttrOfType<FlatSymbolRefAttr>("callee")) {
-      builder.CreateInvoke(moduleTranslation.lookupFunction(attr.getValue()),
-                           moduleTranslation.lookupBlock(invOp.getSuccessor(0)),
-                           moduleTranslation.lookupBlock(invOp.getSuccessor(1)),
-                           operandsRef);
+      result = builder.CreateInvoke(
+          moduleTranslation.lookupFunction(attr.getValue()),
+          moduleTranslation.lookupBlock(invOp.getSuccessor(0)),
+          moduleTranslation.lookupBlock(invOp.getSuccessor(1)), operandsRef);
     } else {
       auto *calleePtrType =
           cast<llvm::PointerType>(operandsRef.front()->getType());
       auto *calleeType =
           cast<llvm::FunctionType>(calleePtrType->getElementType());
-      builder.CreateInvoke(calleeType, operandsRef.front(),
-                           moduleTranslation.lookupBlock(invOp.getSuccessor(0)),
-                           moduleTranslation.lookupBlock(invOp.getSuccessor(1)),
-                           operandsRef.drop_front());
+      result = builder.CreateInvoke(
+          calleeType, operandsRef.front(),
+          moduleTranslation.lookupBlock(invOp.getSuccessor(0)),
+          moduleTranslation.lookupBlock(invOp.getSuccessor(1)),
+          operandsRef.drop_front());
     }
-    return success();
+    // InvokeOp can only have 0 or 1 result
+    if (invOp->getNumResults() != 0) {
+      moduleTranslation.mapValue(opInst.getResult(0), result);
+      return success();
+    }
+    return success(result->getType()->isVoidTy());
   }
 
   if (auto lpOp = dyn_cast<LLVM::LandingpadOp>(opInst)) {
diff --git a/mlir/test/Target/LLVMIR/llvmir.mlir b/mlir/test/Target/LLVMIR/llvmir.mlir
index f5b6d60662ad..e87eebaca515 100644
--- a/mlir/test/Target/LLVMIR/llvmir.mlir
+++ b/mlir/test/Target/LLVMIR/llvmir.mlir
@@ -1296,6 +1296,38 @@ llvm.func @invokeLandingpad() -> i32 attributes { personality = @__gxx_personali
   %8 = llvm.invoke @bar(%6) to ^bb2 unwind ^bb1 : (!llvm.ptr<i8>) -> !llvm.ptr<i8>
 }
 
+// -----
+
+llvm.mlir.global external constant @_ZTIi() : !llvm.ptr<i8>
+llvm.func @foo() -> i8
+llvm.func @__gxx_personality_v0(...) -> i32
+
+// CHECK-LABEL: @invoke_result
+// CHECK-SAME: %[[a0:[0-9]+]]
+llvm.func @invoke_result(%arg0 : !llvm.ptr<i8>) attributes { personality = @__gxx_personality_v0 } {
+    %0 = llvm.mlir.addressof @_ZTIi : !llvm.ptr<ptr<i8>>
+// CHECK: %[[a1:[0-9]+]] = invoke i8 @foo()
+// CHECK-NEXT: to label %[[normal:[0-9]+]] unwind label %[[unwind:[0-9]+]]
+    %1 = llvm.invoke @foo() to ^bb1 unwind ^bb2 : () -> i8
+
+// CHECK: [[normal]]:
+// CHECK-NEXT: store i8 %[[a1]], i8* %[[a0]]
+// CHECK-NEXT: ret void
+^bb1:
+    llvm.store %1, %arg0 : !llvm.ptr<i8>
+    llvm.return
+
+// CHECK: [[unwind]]:
+// CHECK-NEXT: landingpad { i8*, i32 }
+// CHECK-NEXT: catch i8** @_ZTIi
+// CHECK-NEXT: ret void
+^bb2:
+    %7 = llvm.landingpad (catch %0 : !llvm.ptr<ptr<i8>>) : !llvm.struct<(ptr<i8>, i32)>
+    llvm.return
+}
+
+// -----
+
 // CHECK-LABEL: @callFreezeOp
 llvm.func @callFreezeOp(%x : i32) {
   // CHECK: freeze i32 %{{[0-9]+}}