From c79e867cd2bbf414f53de169cd4480666303f0dc Mon Sep 17 00:00:00 2001
From: Justin Bogner <mail@justinbogner.com>
Date: Tue, 25 Feb 2025 12:36:29 -0800
Subject: [PATCH] [DirectX] Update CBuffer to refer to a `dx.Layout` type
 (#128697)

This adds support cbuffers based on llvm/wg-hlsl#171 - the type argument
of the CBuffer TargetExtType is either a `dx.Layout` type which reports
its own size, or it's a normal type and we can simply refer to
DataLayout.
---
 llvm/include/llvm/Analysis/DXILResource.h     | 22 +++++++++++++++-
 llvm/lib/Analysis/DXILResource.cpp            |  9 ++++++-
 .../DXILResource/buffer-frombinding.ll        | 26 +++++++++++++++++++
 3 files changed, 55 insertions(+), 2 deletions(-)

diff --git a/llvm/include/llvm/Analysis/DXILResource.h b/llvm/include/llvm/Analysis/DXILResource.h
index d4b1a9e2ca340..d399457e16916 100644
--- a/llvm/include/llvm/Analysis/DXILResource.h
+++ b/llvm/include/llvm/Analysis/DXILResource.h
@@ -166,7 +166,6 @@ class CBufferExtType : public TargetExtType {
   CBufferExtType &operator=(const CBufferExtType &) = delete;
 
   Type *getResourceType() const { return getTypeParameter(0); }
-  uint32_t getCBufferSize() const { return getIntParameter(0); }
 
   static bool classof(const TargetExtType *T) {
     return T->getName() == "dx.CBuffer";
@@ -197,6 +196,27 @@ class SamplerExtType : public TargetExtType {
   }
 };
 
+/// The dx.Layout target extension type
+///
+/// `target("dx.Layout", <Type>, <size>, [offsets...])`
+class LayoutExtType : public TargetExtType {
+public:
+  LayoutExtType() = delete;
+  LayoutExtType(const LayoutExtType &) = delete;
+  LayoutExtType &operator=(const LayoutExtType &) = delete;
+
+  Type *getWrappedType() const { return getTypeParameter(0); }
+  uint32_t getSize() const { return getIntParameter(0); }
+  uint32_t getOffsetOfElement(int I) const { return getIntParameter(I + 1); }
+
+  static bool classof(const TargetExtType *T) {
+    return T->getName() == "dx.Layout";
+  }
+  static bool classof(const Type *T) {
+    return isa<TargetExtType>(T) && classof(cast<TargetExtType>(T));
+  }
+};
+
 //===----------------------------------------------------------------------===//
 
 class ResourceTypeInfo {
diff --git a/llvm/lib/Analysis/DXILResource.cpp b/llvm/lib/Analysis/DXILResource.cpp
index 4ffc9dbebda8d..c0319c9a354a6 100644
--- a/llvm/lib/Analysis/DXILResource.cpp
+++ b/llvm/lib/Analysis/DXILResource.cpp
@@ -382,7 +382,14 @@ ResourceTypeInfo::UAVInfo ResourceTypeInfo::getUAV() const {
 
 uint32_t ResourceTypeInfo::getCBufferSize(const DataLayout &DL) const {
   assert(isCBuffer() && "Not a CBuffer");
-  return cast<CBufferExtType>(HandleTy)->getCBufferSize();
+
+  Type *ElTy = cast<CBufferExtType>(HandleTy)->getResourceType();
+
+  if (auto *LayoutTy = dyn_cast<LayoutExtType>(ElTy))
+    return LayoutTy->getSize();
+
+  // TODO: What should we do with unannotated arrays?
+  return DL.getTypeAllocSize(ElTy);
 }
 
 dxil::SamplerType ResourceTypeInfo::getSamplerType() const {
diff --git a/llvm/test/Analysis/DXILResource/buffer-frombinding.ll b/llvm/test/Analysis/DXILResource/buffer-frombinding.ll
index ab7151c57473f..a416124221dcb 100644
--- a/llvm/test/Analysis/DXILResource/buffer-frombinding.ll
+++ b/llvm/test/Analysis/DXILResource/buffer-frombinding.ll
@@ -106,6 +106,30 @@ define void @test_typedbuffer() {
   ; CHECK:   Element Type: f32
   ; CHECK:   Element Count: 4
 
+  %cb0 = call target("dx.CBuffer", {float})
+      @llvm.dx.resource.handlefrombinding(i32 1, i32 0, i32 1, i32 0, i1 false)
+  ; CHECK: Binding [[CB0:[0-9]+]]:
+  ; CHECK:   Binding:
+  ; CHECK:     Record ID: 0
+  ; CHECK:     Space: 1
+  ; CHECK:     Lower Bound: 0
+  ; CHECK:     Size: 1
+  ; CHECK:   Class: CBuffer
+  ; CHECK:   Kind: CBuffer
+  ; CHECK:   CBuffer size: 4
+
+  %cb1 = 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 [[CB1:[0-9]+]]:
+  ; CHECK:   Binding:
+  ; CHECK:     Record ID: 1
+  ; CHECK:     Space: 1
+  ; CHECK:     Lower Bound: 8
+  ; CHECK:     Size: 1
+  ; CHECK:   Class: CBuffer
+  ; CHECK:   Kind: CBuffer
+  ; CHECK:   CBuffer size: 4
+
   ; CHECK-NOT: Binding {{[0-9]+}}:
 
   ret void
@@ -118,5 +142,7 @@ 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 =
+; CHECK-DAG: Call bound to [[CB1]]: %cb1 =
 
 attributes #0 = { nocallback nofree nosync nounwind willreturn memory(none) }