[llvm] [HLSL] Fix Root signature test error (PR #127261)

[via llvm-commits]

https://github.com/joaosaffran updated https://github.com/llvm/llvm-project/pull/127261

>From 4c7d133d6cc3e49aa4de77150089ce6903fc8b5e Mon Sep 17 00:00:00 2001
From: joaosaffran <joao.saffran at microsoft.com>
Date: Fri, 14 Feb 2025 20:55:00 +0000
Subject: [PATCH 1/3] fix root signature test error

---
 .../RootSignature-MultipleEntryFunctions.ll              | 9 ++++-----
 1 file changed, 4 insertions(+), 5 deletions(-)

diff --git a/llvm/test/CodeGen/DirectX/ContainerData/RootSignature-MultipleEntryFunctions.ll b/llvm/test/CodeGen/DirectX/ContainerData/RootSignature-MultipleEntryFunctions.ll
index 652f8092b7a69..0547b0bae7a7e 100644
--- a/llvm/test/CodeGen/DirectX/ContainerData/RootSignature-MultipleEntryFunctions.ll
+++ b/llvm/test/CodeGen/DirectX/ContainerData/RootSignature-MultipleEntryFunctions.ll
@@ -23,17 +23,16 @@ attributes #0 = { "hlsl.numthreads"="1,1,1" "hlsl.shader"="compute" }
 !6 = !{ !7 } ; list of root signature elements
 !7 = !{ !"RootFlags", i32 2 } ; 1 = allow_input_assembler_input_layout
 
-
-; CHECK-LABEL: Definition for 'main':
-; CHECK-NEXT:   Flags: 0x000001
+; CHECK-LABEL: Definition for 'anotherMain':
+; CHECK-NEXT:   Flags: 0x000002
 ; CHECK-NEXT:   Version: 2
 ; CHECK-NEXT:   NumParameters: 0
 ; CHECK-NEXT:   RootParametersOffset: 0
 ; CHECK-NEXT:   NumStaticSamplers: 0
 ; CHECK-NEXT:   StaticSamplersOffset: 0
 
-; CHECK-LABEL: Definition for 'anotherMain':
-; CHECK-NEXT:   Flags: 0x000002
+; CHECK-LABEL: Definition for 'main':
+; CHECK-NEXT:   Flags: 0x000001
 ; CHECK-NEXT:   Version: 2
 ; CHECK-NEXT:   NumParameters: 0
 ; CHECK-NEXT:   RootParametersOffset: 0

>From e92d788ce7ce8bbc5f14beeea391592edb888e9c Mon Sep 17 00:00:00 2001
From: joaosaffran <joao.saffran at microsoft.com>
Date: Fri, 14 Feb 2025 21:35:38 +0000
Subject: [PATCH 2/3] fix other functions are checked

---
 llvm/lib/Target/DirectX/DXILRootSignature.cpp             | 6 +++---
 .../ContainerData/RootSignature-MultipleEntryFunctions.ll | 8 ++++----
 2 files changed, 7 insertions(+), 7 deletions(-)

diff --git a/llvm/lib/Target/DirectX/DXILRootSignature.cpp b/llvm/lib/Target/DirectX/DXILRootSignature.cpp
index 49fc892eade5d..deadab3d334dd 100644
--- a/llvm/lib/Target/DirectX/DXILRootSignature.cpp
+++ b/llvm/lib/Target/DirectX/DXILRootSignature.cpp
@@ -186,9 +186,9 @@ PreservedAnalyses RootSignatureAnalysisPrinter::run(Module &M,
   OS << "Root Signature Definitions"
      << "\n";
   uint8_t Space = 0;
-  for (const auto &P : RSDMap) {
-    const auto &[Function, RSD] = P;
-    OS << "Definition for '" << Function->getName() << "':\n";
+  for (const auto &F : M) {
+    const auto RSD = RSDMap.at(&F);
+    OS << "Definition for '" << F.getName() << "':\n";
 
     // start root signature header
     Space++;
diff --git a/llvm/test/CodeGen/DirectX/ContainerData/RootSignature-MultipleEntryFunctions.ll b/llvm/test/CodeGen/DirectX/ContainerData/RootSignature-MultipleEntryFunctions.ll
index 0547b0bae7a7e..7adb17d0b022f 100644
--- a/llvm/test/CodeGen/DirectX/ContainerData/RootSignature-MultipleEntryFunctions.ll
+++ b/llvm/test/CodeGen/DirectX/ContainerData/RootSignature-MultipleEntryFunctions.ll
@@ -23,16 +23,16 @@ attributes #0 = { "hlsl.numthreads"="1,1,1" "hlsl.shader"="compute" }
 !6 = !{ !7 } ; list of root signature elements
 !7 = !{ !"RootFlags", i32 2 } ; 1 = allow_input_assembler_input_layout
 
-; CHECK-LABEL: Definition for 'anotherMain':
-; CHECK-NEXT:   Flags: 0x000002
+; CHECK-LABEL: Definition for 'main':
+; CHECK-NEXT:   Flags: 0x000001
 ; CHECK-NEXT:   Version: 2
 ; CHECK-NEXT:   NumParameters: 0
 ; CHECK-NEXT:   RootParametersOffset: 0
 ; CHECK-NEXT:   NumStaticSamplers: 0
 ; CHECK-NEXT:   StaticSamplersOffset: 0
 
-; CHECK-LABEL: Definition for 'main':
-; CHECK-NEXT:   Flags: 0x000001
+; CHECK-LABEL: Definition for 'anotherMain':
+; CHECK-NEXT:   Flags: 0x000002
 ; CHECK-NEXT:   Version: 2
 ; CHECK-NEXT:   NumParameters: 0
 ; CHECK-NEXT:   RootParametersOffset: 0

>From fb2aeeef0e2748e01fc7bd6eacb63ec856fd4a3a Mon Sep 17 00:00:00 2001
From: joaosaffran <joao.saffran at microsoft.com>
Date: Fri, 14 Feb 2025 22:44:08 +0000
Subject: [PATCH 3/3] addressing comments

---
 llvm/lib/Target/DirectX/DXILRootSignature.cpp | 17 ++++++++++-------
 1 file changed, 10 insertions(+), 7 deletions(-)

diff --git a/llvm/lib/Target/DirectX/DXILRootSignature.cpp b/llvm/lib/Target/DirectX/DXILRootSignature.cpp
index deadab3d334dd..ffd9d5bd620a9 100644
--- a/llvm/lib/Target/DirectX/DXILRootSignature.cpp
+++ b/llvm/lib/Target/DirectX/DXILRootSignature.cpp
@@ -187,19 +187,22 @@ PreservedAnalyses RootSignatureAnalysisPrinter::run(Module &M,
      << "\n";
   uint8_t Space = 0;
   for (const auto &F : M) {
-    const auto RSD = RSDMap.at(&F);
+    const auto &RSD = RSDMap.find(&F);
+    if (RSD == RSDMap.end())
+      continue;
+    const auto &RS = RSD->second;
     OS << "Definition for '" << F.getName() << "':\n";
 
     // start root signature header
     Space++;
-    OS << indent(Space) << "Flags: " << format_hex(RSD.Flags, 8) << ":\n";
-    OS << indent(Space) << "Version: " << RSD.Version << ":\n";
-    OS << indent(Space) << "NumParameters: " << RSD.NumParameters << ":\n";
-    OS << indent(Space) << "RootParametersOffset: " << RSD.RootParametersOffset
+    OS << indent(Space) << "Flags: " << format_hex(RS.Flags, 8) << ":\n";
+    OS << indent(Space) << "Version: " << RS.Version << ":\n";
+    OS << indent(Space) << "NumParameters: " << RS.NumParameters << ":\n";
+    OS << indent(Space) << "RootParametersOffset: " << RS.RootParametersOffset
        << ":\n";
-    OS << indent(Space) << "NumStaticSamplers: " << RSD.NumStaticSamplers
+    OS << indent(Space) << "NumStaticSamplers: " << RS.NumStaticSamplers
        << ":\n";
-    OS << indent(Space) << "StaticSamplersOffset: " << RSD.StaticSamplersOffset
+    OS << indent(Space) << "StaticSamplersOffset: " << RS.StaticSamplersOffset
        << ":\n";
     Space--;
     // end root signature header