diff --git a/mlir/lib/Dialect/LLVMIR/IR/LLVMDialect.cpp b/mlir/lib/Dialect/LLVMIR/IR/LLVMDialect.cpp --- a/mlir/lib/Dialect/LLVMIR/IR/LLVMDialect.cpp +++ b/mlir/lib/Dialect/LLVMIR/IR/LLVMDialect.cpp @@ -2871,6 +2871,21 @@ }) .Default([](auto) { return false; }); } + /// Handle the given inlined terminator by replacing it with a new operation + /// as necessary. Required when the region has only one block. + void handleTerminator(Operation *op, + ArrayRef valuesToRepl) const final { + + // Only handle "llvm.return" here. + auto returnOp = dyn_cast(op); + if (!returnOp) + return; + + // Replace the values directly with the return operands. + assert(returnOp.getNumOperands() == valuesToRepl.size()); + for (const auto &it : llvm::enumerate(returnOp.getOperands())) + valuesToRepl[it.index()].replaceAllUsesWith(it.value()); + } }; } // end anonymous namespace diff --git a/mlir/test/Dialect/LLVMIR/inlining.mlir b/mlir/test/Dialect/LLVMIR/inlining.mlir --- a/mlir/test/Dialect/LLVMIR/inlining.mlir +++ b/mlir/test/Dialect/LLVMIR/inlining.mlir @@ -56,3 +56,17 @@ call @with_mem_attr(%ptr) : (!llvm.ptr) -> () return } + +// ----- +// Check that llvm.return is correctly handled + +func.func @func(%arg0 : i32) -> i32 { + llvm.return %arg0 : i32 +} +// CHECK-LABEL: @llvm_ret +// CHECK-NOT: call +// CHECK: return %arg0 +func.func @llvm_ret(%arg0 : i32) -> i32 { + %res = call @func(%arg0) : (i32) -> (i32) + return %res : i32 +}