From 499bf9b073ae57559a156f912a7bd4c9fee7c7fd Mon Sep 17 00:00:00 2001
From: Justin Bogner <mail@justinbogner.com>
Date: Sat, 22 Feb 2025 12:38:53 -0800
Subject: [PATCH] [DXIL][Analysis] Make sure resource accessors are contiguous

When some resource types were present, but not all of them, we were
ending up in a situation where we would fail to initialize the `FirstX`
variables and get incorrect iterators.

Fixes #128560.
---
 llvm/lib/Analysis/DXILResource.cpp            |  5 +++++
 .../DXILResource/buffer-frombinding.ll        | 12 ++++++++++++
 .../CodeGen/DirectX/Metadata/cbuffer-only.ll  | 19 +++++++++++++++++++
 3 files changed, 36 insertions(+)
 create mode 100644 llvm/test/CodeGen/DirectX/Metadata/cbuffer-only.ll

diff --git a/llvm/lib/Analysis/DXILResource.cpp b/llvm/lib/Analysis/DXILResource.cpp
index 4ffc9dbebda8d..2f75b285356a3 100644
--- a/llvm/lib/Analysis/DXILResource.cpp
+++ b/llvm/lib/Analysis/DXILResource.cpp
@@ -749,6 +749,11 @@ void DXILBindingMap::populate(Module &M, DXILResourceTypeMap &DRTM) {
       NextID = 0;
     }
 
+    // We need to make sure the types of resource are ordered even if some are
+    // missing.
+    FirstCBuffer = std::min({FirstCBuffer, FirstSampler});
+    FirstUAV = std::min({FirstUAV, FirstCBuffer});
+
     // Adjust the resource binding to use the next ID.
     RBI.setBindingID(NextID++);
   }
diff --git a/llvm/test/Analysis/DXILResource/buffer-frombinding.ll b/llvm/test/Analysis/DXILResource/buffer-frombinding.ll
index ab7151c57473f..d6afa52bbad1a 100644
--- a/llvm/test/Analysis/DXILResource/buffer-frombinding.ll
+++ b/llvm/test/Analysis/DXILResource/buffer-frombinding.ll
@@ -106,6 +106,17 @@ define void @test_typedbuffer() {
   ; CHECK:   Element Type: f32
   ; CHECK:   Element Count: 4
 
+  %cb0 = call target("dx.CBuffer", target("dx.Layout", {float}, 4, 0))
+      @llvm.dx.resource.handlefrombinding(i32 1, i32 8, i32 1, i32 0, i1 false)
+  ; CHECK: Binding [[CB0:[0-9]+]]:
+  ; CHECK:   Binding:
+  ; CHECK:     Record ID: 0
+  ; CHECK:     Space: 1
+  ; CHECK:     Lower Bound: 8
+  ; CHECK:     Size: 1
+  ; CHECK:   Class: CBuffer
+  ; CHECK:   Kind: CBuffer
+
   ; CHECK-NOT: Binding {{[0-9]+}}:
 
   ret void
@@ -118,5 +129,6 @@ define void @test_typedbuffer() {
 ; CHECK-DAG: Call bound to [[UAV1]]: %uav1 =
 ; CHECK-DAG: Call bound to [[UAV2]]: %uav2_1 =
 ; CHECK-DAG: Call bound to [[UAV2]]: %uav2_2 =
+; CHECK-DAG: Call bound to [[CB0]]: %cb0 =
 
 attributes #0 = { nocallback nofree nosync nounwind willreturn memory(none) }
diff --git a/llvm/test/CodeGen/DirectX/Metadata/cbuffer-only.ll b/llvm/test/CodeGen/DirectX/Metadata/cbuffer-only.ll
new file mode 100644
index 0000000000000..c07e8d9a2b591
--- /dev/null
+++ b/llvm/test/CodeGen/DirectX/Metadata/cbuffer-only.ll
@@ -0,0 +1,19 @@
+; Regression test for https://github.com/llvm/llvm-project/issues/128560 -
+; check that cbuffers are populated correctly when there aren't any other kinds
+; of resource.
+
+; RUN: opt -S -passes=dxil-translate-metadata %s | FileCheck %s
+
+target triple = "dxil-pc-shadermodel6.6-compute"
+
+define void @cbuffer_is_only_binding() {
+  %cbuf = call target("dx.CBuffer", target("dx.Layout", {float}, 4, 0))
+      @llvm.dx.resource.handlefrombinding(i32 1, i32 8, i32 1, i32 0, i1 false)
+  ; CHECK: %cbuffer = type
+
+  ret void
+}
+
+; CHECK:      @[[CB0:.*]] = external constant %cbuffer
+
+; CHECK: !{i32 0, ptr @[[CB0]], !""