[Mlir-commits] [mlir] [mlir][spirv] Remove destroyed values from ValueIDMap (PR #164098)

[Igor Wodiany]

https://github.com/IgWod updated https://github.com/llvm/llvm-project/pull/164098

>From fa1556c8f26ccd3914f4b9aa1d1091c2a1dbd81f Mon Sep 17 00:00:00 2001
From: Igor Wodiany <9111077+IgWod at users.noreply.github.com>
Date: Sat, 18 Oct 2025 16:30:25 +0000
Subject: [PATCH 1/2] [mlir][spirv] Remove destroyed values from ValueIDMap

When serializing SPIR-V MLIR containing externally linked function
with debug enabled, the serialization crashes as `printValueIDMap`
tries to print a key value that has been already destroyed. This
happen as for externally linked function the body of the function
is erased, that causes arguments to be destroyed as well, but the
valueIDMap was never updated.
---
 mlir/lib/Target/SPIRV/Serialization/SerializeOps.cpp | 9 +++++++++
 mlir/test/Target/SPIRV/function-decorations.mlir     | 3 ++-
 2 files changed, 11 insertions(+), 1 deletion(-)

diff --git a/mlir/lib/Target/SPIRV/Serialization/SerializeOps.cpp b/mlir/lib/Target/SPIRV/Serialization/SerializeOps.cpp
index e9b180a70bb23..3c2c65bf64ed6 100644
--- a/mlir/lib/Target/SPIRV/Serialization/SerializeOps.cpp
+++ b/mlir/lib/Target/SPIRV/Serialization/SerializeOps.cpp
@@ -311,6 +311,15 @@ LogicalResult Serializer::processFuncOp(spirv::FuncOp op) {
     op.addEntryBlock();
     if (failed(processFuncParameter(op)))
       return failure();
+
+    // Erasing the body of the function destroys arguments, so we need to remove
+    // them from the map to avoid problems when processing invalid values used
+    // as keys. We have already serialized function arguments so we probably can
+    // remove them from the map as external function will not have any uses.
+    for (Value arg : op.getArguments())
+      if (valueIDMap.count(arg))
+        valueIDMap.erase(arg);
+
     // Don't need to process the added block, there is nothing to process,
     // the fake body was added just to get the arguments, remove the body,
     // since it's use is done.
diff --git a/mlir/test/Target/SPIRV/function-decorations.mlir b/mlir/test/Target/SPIRV/function-decorations.mlir
index ef627145aac3e..6098e42f063a2 100644
--- a/mlir/test/Target/SPIRV/function-decorations.mlir
+++ b/mlir/test/Target/SPIRV/function-decorations.mlir
@@ -1,4 +1,5 @@
-// RUN: mlir-translate -no-implicit-module -test-spirv-roundtrip -split-input-file -verify-diagnostics %s | FileCheck %s
+// RUN: mlir-translate --no-implicit-module --test-spirv-roundtrip --split-input-file %s | FileCheck %s
+// RUN: mlir-translate --no-implicit-module --test-spirv-roundtrip --split-input-file --debug %s | FileCheck %s
 
 spirv.module Logical GLSL450 requires #spirv.vce<v1.0, [Shader, Linkage], []> {
     spirv.func @linkage_attr_test_kernel()  "DontInline"  attributes {}  {

>From 62cb401b09485ed7321e1b5f78555ff7fbd81c0b Mon Sep 17 00:00:00 2001
From: Igor Wodiany <9111077+IgWod at users.noreply.github.com>
Date: Sat, 18 Oct 2025 19:44:39 +0000
Subject: [PATCH 2/2] Address feedback

---
 mlir/lib/Target/SPIRV/Serialization/SerializeOps.cpp | 3 +--
 1 file changed, 1 insertion(+), 2 deletions(-)

diff --git a/mlir/lib/Target/SPIRV/Serialization/SerializeOps.cpp b/mlir/lib/Target/SPIRV/Serialization/SerializeOps.cpp
index 3c2c65bf64ed6..85e92c7ced394 100644
--- a/mlir/lib/Target/SPIRV/Serialization/SerializeOps.cpp
+++ b/mlir/lib/Target/SPIRV/Serialization/SerializeOps.cpp
@@ -317,8 +317,7 @@ LogicalResult Serializer::processFuncOp(spirv::FuncOp op) {
     // as keys. We have already serialized function arguments so we probably can
     // remove them from the map as external function will not have any uses.
     for (Value arg : op.getArguments())
-      if (valueIDMap.count(arg))
-        valueIDMap.erase(arg);
+      valueIDMap.erase(arg);
 
     // Don't need to process the added block, there is nothing to process,
     // the fake body was added just to get the arguments, remove the body,