[Mlir-commits] [mlir] [MLIR][LLVM] Avoid importing broken calls and invokes (PR #125041)

[Christian Ulmann]

https://github.com/Dinistro created https://github.com/llvm/llvm-project/pull/125041

This commit adds a check to catch calls/invokes that do not satisfy the return type of their callee. This is not verified in LLVM IR but is considered UB. Importing this into MLIR will lead to verification errors, thus we should avoid this early on.

>From f6a261ab3bee53190d14599c2b341dbf2f8a5569 Mon Sep 17 00:00:00 2001
From: Christian Ulmann <christian.ulmann at nextsilicon.com>
Date: Thu, 30 Jan 2025 09:19:34 +0000
Subject: [PATCH] [MLIR][LLVM] Avoid importing broken calls and invokes

This commit adds a check to catch calls/invokes that do not satisfy the
return type of their callee. This is not verified in LLVM IR but is
considered UB. Importing this into MLIR will lead to verification
errors, thus we should avoid this early on.
---
 mlir/lib/Target/LLVMIR/ModuleImport.cpp       | 42 ++++++++++++-------
 .../Target/LLVMIR/Import/import-failure.ll    | 32 +++++++++++++-
 2 files changed, 58 insertions(+), 16 deletions(-)

diff --git a/mlir/lib/Target/LLVMIR/ModuleImport.cpp b/mlir/lib/Target/LLVMIR/ModuleImport.cpp
index 1d1a985c46fb5b9..d2e5390e397933d 100644
--- a/mlir/lib/Target/LLVMIR/ModuleImport.cpp
+++ b/mlir/lib/Target/LLVMIR/ModuleImport.cpp
@@ -1620,7 +1620,7 @@ LogicalResult ModuleImport::convertInstruction(llvm::Instruction *inst) {
     return success();
   }
   if (inst->getOpcode() == llvm::Instruction::Call) {
-    auto callInst = cast<llvm::CallInst>(inst);
+    auto *callInst = cast<llvm::CallInst>(inst);
     llvm::Value *calledOperand = callInst->getCalledOperand();
 
     FailureOr<SmallVector<Value>> operands =
@@ -1629,10 +1629,10 @@ LogicalResult ModuleImport::convertInstruction(llvm::Instruction *inst) {
       return failure();
 
     auto callOp = [&]() -> FailureOr<Operation *> {
-      if (auto asmI = dyn_cast<llvm::InlineAsm>(calledOperand)) {
-        Type resultTy = convertType(callInst->getType());
-        if (!resultTy)
-          return failure();
+      Type resultTy = convertType(callInst->getType());
+      if (!resultTy)
+        return failure();
+      if (auto *asmI = dyn_cast<llvm::InlineAsm>(calledOperand)) {
         return builder
             .create<InlineAsmOp>(
                 loc, resultTy, *operands,
@@ -1642,17 +1642,21 @@ LogicalResult ModuleImport::convertInstruction(llvm::Instruction *inst) {
                 /*is_align_stack=*/false, /*asm_dialect=*/nullptr,
                 /*operand_attrs=*/nullptr)
             .getOperation();
-      } else {
-        LLVMFunctionType funcTy = convertFunctionType(callInst);
-        if (!funcTy)
-          return failure();
-
-        FlatSymbolRefAttr callee = convertCalleeName(callInst);
-        auto callOp = builder.create<CallOp>(loc, funcTy, callee, *operands);
-        if (failed(convertCallAttributes(callInst, callOp)))
-          return failure();
-        return callOp.getOperation();
       }
+      LLVMFunctionType funcTy = convertFunctionType(callInst);
+      if (!funcTy)
+        return failure();
+
+      if (funcTy.getReturnType() != resultTy)
+        return emitError(loc)
+               << "incompatible call and function return types: " << resultTy
+               << " vs. " << funcTy.getReturnType();
+
+      FlatSymbolRefAttr callee = convertCalleeName(callInst);
+      auto callOp = builder.create<CallOp>(loc, funcTy, callee, *operands);
+      if (failed(convertCallAttributes(callInst, callOp)))
+        return failure();
+      return callOp.getOperation();
     }();
 
     if (failed(callOp))
@@ -1720,6 +1724,14 @@ LogicalResult ModuleImport::convertInstruction(llvm::Instruction *inst) {
     if (!funcTy)
       return failure();
 
+    Type resultTy = convertType(invokeInst->getType());
+    if (!resultTy)
+      return failure();
+    if (funcTy.getReturnType() != resultTy)
+      return emitError(loc)
+             << "incompatible invoke and function return types: " << resultTy
+             << " vs. " << funcTy.getReturnType();
+
     FlatSymbolRefAttr calleeName = convertCalleeName(invokeInst);
 
     // Create the invoke operation. Normal destination block arguments will be
diff --git a/mlir/test/Target/LLVMIR/Import/import-failure.ll b/mlir/test/Target/LLVMIR/Import/import-failure.ll
index b616cb81e0a8a5c..910f60fd81e8a02 100644
--- a/mlir/test/Target/LLVMIR/Import/import-failure.ll
+++ b/mlir/test/Target/LLVMIR/Import/import-failure.ll
@@ -1,4 +1,4 @@
-; RUN: not mlir-translate -import-llvm -emit-expensive-warnings -split-input-file %s 2>&1 | FileCheck %s
+; RUN: not mlir-translate -import-llvm -emit-expensive-warnings -split-input-file %s 2>&1 -o /dev/null | FileCheck %s
 
 ; CHECK:      <unknown>
 ; CHECK-SAME: error: unhandled instruction: indirectbr ptr %dst, [label %bb1, label %bb2]
@@ -353,3 +353,33 @@ declare void @llvm.experimental.noalias.scope.decl(metadata)
 ; CHECK:      import-failure.ll
 ; CHECK-SAME: warning: unhandled data layout token: ni:42
 target datalayout = "e-ni:42-i64:64"
+
+; // -----
+
+; CHECK:      <unknown>
+; CHECK-SAME: incompatible call and function return types: 'i64' vs. '!llvm.ptr'
+define i64 @incompatible_call_and_callee_types() {
+L.entry:
+  %0 = call i64 @callee()
+  ret i64 %0
+}
+
+declare ptr @callee()
+
+; // -----
+
+; CHECK:      <unknown>
+; CHECK-SAME: incompatible invoke and function return types: '!llvm.void' vs. 'i32'
+define void @f() personality ptr @__gxx_personality_v0 {
+entry:
+  invoke void @g() to label %bb1 unwind label %bb2
+bb1:
+  ret void
+bb2:
+  %0 = landingpad i32 cleanup
+  unreachable
+}
+
+declare i32 @g()
+
+declare i32 @__gxx_personality_v0(...)