Merge pull request #10 from llvm-beanz/user/cbieneman/reland-root-descriptor-fix

Fix setting root descriptor table

Author: llvm-beanz

lib/API/DX/Device.cpp

@@ -513,16 +513,13 @@ class DXDevice : public hlsltest::Device {
 
     uint32_t Inc = Device->GetDescriptorHandleIncrementSize(
         D3D12_DESCRIPTOR_HEAP_TYPE_CBV_SRV_UAV);
+    CD3DX12_GPU_DESCRIPTOR_HANDLE Handle{IS.DescHeap->GetGPUDescriptorHandleForHeapStart()};
 
-    uint32_t Offset = 0;
-    for (uint32_t Idx = 0; Idx < P.Sets.size(); ++Idx, Offset += Inc) {
-      D3D12_GPU_DESCRIPTOR_HANDLE Handle =
-          IS.DescHeap->GetGPUDescriptorHandleForHeapStart();
-      Handle.ptr += Offset;
+    for (uint32_t Idx = 0; Idx < P.Sets.size(); ++Idx) {
       IS.CmdList->SetComputeRootDescriptorTable(Idx, Handle);
-      // TODO: This probably computes the wrong offsets if I have multiple
-      // descriptor tables in use.
+      Handle.Offset(P.Sets[Idx].Resources.size(), Inc);
     }
+
     IS.CmdList->Dispatch(P.DispatchSize[0], P.DispatchSize[1],
                          P.DispatchSize[2]);
 

test/Basic/DescriptorSets.test

@@ -0,0 +1,61 @@
+#--- DescriptorSets.hlsl
+
+#if defined(__spirv__) || defined(__SPIRV__)
+#define REGISTER(Idx, Space)
+#else
+#define REGISTER(Idx, Space) : register(Idx, Space)
+#endif
+
+[[vk::binding(0)]] RWBuffer<float4> In REGISTER(u0, space0);
+[[vk::binding(1)]] RWBuffer<float4> Out1 REGISTER(u1, space4);
+[[vk::binding(0,1)]] RWBuffer<float4> Out2 REGISTER(u2, space4);
+
+[numthreads(1,1,1)]
+void main(uint GI : SV_GroupIndex) {
+  Out1[GI] = In[GI] * In[GI];
+  Out2[GI] = In[GI] * In[GI] * In[GI];
+}
+//--- DescriptorSets.yaml
+---
+DispatchSize: [1, 1, 1]
+DescriptorSets:
+  - Resources:
+    - Access: ReadWrite
+      Format: Float32
+      Channels: 4
+      Data: [ 2, 4, 6, 8]
+      DirectXBinding:
+        Register: 0
+        Space: 0
+    - Access: ReadWrite
+      Format: Float32
+      Channels: 4
+      ZeroInitSize: 16
+      DirectXBinding:
+        Register: 1
+        Space: 4
+  - Resources:
+    - Access: ReadWrite
+      Format: Float32
+      Channels: 4
+      ZeroInitSize: 16
+      DirectXBinding:
+        Register: 2
+        Space: 4
+...
+#--- end
+
+# UNSUPPORTED: Clang
+# RUN: split-file %s %t
+# RUN: %if DirectX %{ dxc -T cs_6_0 -Fo %t.dxil %t/DescriptorSets.hlsl %}
+# RUN: %if DirectX %{ %offloader %t/DescriptorSets.yaml %t.dxil | FileCheck %s %}
+# RUN: %if Vulkan %{ dxc -T cs_6_0 -spirv -fspv-target-env=vulkan1.1 -Fo %t.spv %t/DescriptorSets.hlsl %}
+# RUN: %if Vulkan %{ %offloader %t/DescriptorSets.yaml %t.spv | FileCheck %s %}
+
+# RUN: %if Metal %{ dxc -T cs_6_0 -Fo %t.dxil %t/DescriptorSets.hlsl %}
+# RUN: %if Metal %{ metal-shaderconverter %t.dxil -o=%t.metallib %}
+# RUN: %if Metal %{ %offloader %t/DescriptorSets.yaml %t.metallib | FileCheck %s %}
+
+# CHECK: Data:
+# CHECK: Data: [ 4, 16, 36, 64 ]
+# CHECK: Data: [ 8, 64, 216, 512 ]