From 6a9933519dd94e48b59d07c7333dac37186734d0 Mon Sep 17 00:00:00 2001
From: Mehdi Amini <joker.eph@gmail.com>
Date: Tue, 17 Jan 2023 08:28:08 +0000
Subject: [PATCH] Fix crash in LLVM Dialect inliner interface: add support for
 llvm.return

The LLVM inliner was missing the `handleTerminator` method in the
Dialect interface implementation.

Fixes #60093

Differential Revision: https://reviews.llvm.org/D141901
---
 mlir/lib/Dialect/LLVMIR/IR/LLVMDialect.cpp | 15 +++++++++++++++
 mlir/test/Dialect/LLVMIR/inlining.mlir     | 14 ++++++++++++++
 2 files changed, 29 insertions(+)

diff --git a/mlir/lib/Dialect/LLVMIR/IR/LLVMDialect.cpp b/mlir/lib/Dialect/LLVMIR/IR/LLVMDialect.cpp
index 015bff8..a4c6568 100644
--- a/mlir/lib/Dialect/LLVMIR/IR/LLVMDialect.cpp
+++ b/mlir/lib/Dialect/LLVMIR/IR/LLVMDialect.cpp
@@ -2871,6 +2871,21 @@ struct LLVMInlinerInterface : public DialectInlinerInterface {
         })
         .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<Value> valuesToRepl) const final {
+
+    // Only handle "llvm.return" here.
+    auto returnOp = dyn_cast<ReturnOp>(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
index 0ba983f..ce2bf69 100644
--- a/mlir/test/Dialect/LLVMIR/inlining.mlir
+++ b/mlir/test/Dialect/LLVMIR/inlining.mlir
@@ -56,3 +56,17 @@ func.func @test_not_inline(%ptr : !llvm.ptr) -> () {
   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
+}
-- 
2.7.4