[Mlir-commits] [mlir] [mlir][spirv] Handle a missing case when inlining spirv.ReturnValue (PR #80733)

[Artem Tyurin]

https://github.com/agentcooper created https://github.com/llvm/llvm-project/pull/80733

Fixes https://github.com/llvm/llvm-project/issues/73285.

>From c78ff1308f9117f243c83a7024f9558ecdaceac4 Mon Sep 17 00:00:00 2001
From: Artem Tyurin <artem.tyurin at gmail.com>
Date: Mon, 5 Feb 2024 20:17:48 +0100
Subject: [PATCH] [mlir][spirv] Handle a missing case when inlining
 spirv.ReturnValue

Fixes https://github.com/llvm/llvm-project/issues/73285.
---
 mlir/lib/Dialect/SPIRV/IR/SPIRVDialect.cpp    |  4 +++-
 .../Dialect/SPIRV/Transforms/inlining.mlir    | 22 +++++++++++++++++++
 2 files changed, 25 insertions(+), 1 deletion(-)

diff --git a/mlir/lib/Dialect/SPIRV/IR/SPIRVDialect.cpp b/mlir/lib/Dialect/SPIRV/IR/SPIRVDialect.cpp
index db26fa81790d1..e914f46bdef64 100644
--- a/mlir/lib/Dialect/SPIRV/IR/SPIRVDialect.cpp
+++ b/mlir/lib/Dialect/SPIRV/IR/SPIRVDialect.cpp
@@ -95,7 +95,9 @@ struct SPIRVInlinerInterface : public DialectInlinerInterface {
       OpBuilder(op).create<spirv::BranchOp>(op->getLoc(), newDest);
       op->erase();
     } else if (auto retValOp = dyn_cast<spirv::ReturnValueOp>(op)) {
-      llvm_unreachable("unimplemented spirv.ReturnValue in inliner");
+      OpBuilder(op).create<spirv::BranchOp>(retValOp->getLoc(), newDest,
+                                            retValOp->getOperands());
+      op->erase();
     }
   }
 
diff --git a/mlir/test/Dialect/SPIRV/Transforms/inlining.mlir b/mlir/test/Dialect/SPIRV/Transforms/inlining.mlir
index 8d663b4edc452..3aadb19ec1582 100644
--- a/mlir/test/Dialect/SPIRV/Transforms/inlining.mlir
+++ b/mlir/test/Dialect/SPIRV/Transforms/inlining.mlir
@@ -224,5 +224,27 @@ spirv.module Logical GLSL450 {
   spirv.ExecutionMode @inline_into_selection_region "LocalSize", 32, 1, 1
 }
 
+// -----
+
+spirv.module Logical GLSL450 {
+  // CHECK-LABEL: @foo
+  spirv.func @foo(%arg0: i32) -> i32 "None" {
+    // CHECK-NOT: spirv.FunctionCall
+    // CHECK-NEXT: spirv.Constant 1
+    %res = spirv.FunctionCall @bar(%arg0) : (i32) -> i32
+    spirv.ReturnValue %res : i32
+  }
+
+  spirv.func @bar(%arg1: i32) -> i32 "None" attributes {sym_visibility = "private"} {
+    %cst1_i32 = spirv.Constant 1 : i32
+    %0 = spirv.IEqual %arg1, %cst1_i32 : i32
+    spirv.BranchConditional %0, ^bb1(%cst1_i32 : i32), ^bb2
+  ^bb1(%1: i32):
+    spirv.ReturnValue %1 : i32
+  ^bb2:
+    spirv.ReturnValue %cst1_i32 : i32
+  }
+}
+
 // TODO: Add tests for inlining structured control flow into
 // structured control flow.