Add support for UAV Counters (#80)

* save work

* Add support for UAV Counters

* Address Comments

* Address comments

* format

* hack: Skip metal in CI

Author: V-FEXrt

include/Support/Pipeline.h

@@ -79,6 +79,7 @@ struct Buffer {
   std::unique_ptr<char[]> Data;
   size_t Size;
   OutputProperties OutputProps;
+  uint32_t Counter;
 
   uint32_t size() const { return Size; }
 
@@ -130,6 +131,7 @@ struct Resource {
   DirectXBinding DXBinding;
   std::optional<VulkanBinding> VKBinding;
   Buffer *BufferPtr = nullptr;
+  bool HasCounter;
 
   bool isRaw() const {
     switch (Kind) {

lib/API/DX/Device.cpp

@@ -89,6 +89,19 @@ static DXGI_FORMAT getRawDXFormat(Resource &R) {
   return DXGI_FORMAT_UNKNOWN;
 }
 
+static uint32_t getUAVBufferSize(Resource &R) {
+  return R.HasCounter
+             ? llvm::alignTo(R.size(), D3D12_UAV_COUNTER_PLACEMENT_ALIGNMENT) +
+                   sizeof(uint32_t)
+             : R.size();
+}
+
+static uint32_t getUAVBufferCounterOffset(Resource &R) {
+  return R.HasCounter
+             ? llvm::alignTo(R.size(), D3D12_UAV_COUNTER_PLACEMENT_ALIGNMENT)
+             : 0;
+}
+
 namespace {
 
 enum DXResourceKind { UAV, SRV, CBV };
@@ -449,9 +462,10 @@ class DXDevice : public offloadtest::Device {
   }
 
   llvm::Expected<ResourceSet> createUAV(Resource &R, InvocationState &IS) {
-    llvm::outs() << "Creating UAV: { Size = " << R.size() << ", Register = u"
+    const uint32_t BufferSize = getUAVBufferSize(R);
+    llvm::outs() << "Creating UAV: { Size = " << BufferSize << ", Register = u"
                  << R.DXBinding.Register << ", Space = " << R.DXBinding.Space
-                 << " }\n";
+                 << ", HasCounter = " << R.HasCounter << " }\n";
     ComPtr<ID3D12Resource> Buffer;
     ComPtr<ID3D12Resource> UploadBuffer;
     ComPtr<ID3D12Resource> ReadBackBuffer;
@@ -461,7 +475,7 @@ class DXDevice : public offloadtest::Device {
     const D3D12_RESOURCE_DESC ResDesc = {
         D3D12_RESOURCE_DIMENSION_BUFFER,
         0,
-        R.size(),
+        BufferSize,
         1,
         1,
         1,
@@ -480,7 +494,7 @@ class DXDevice : public offloadtest::Device {
     const D3D12_HEAP_PROPERTIES UploadHeapProp =
         CD3DX12_HEAP_PROPERTIES(D3D12_HEAP_TYPE_UPLOAD);
     const D3D12_RESOURCE_DESC UploadResDesc =
-        CD3DX12_RESOURCE_DESC::Buffer(R.size());
+        CD3DX12_RESOURCE_DESC::Buffer(BufferSize);
 
     if (auto Err =
             HR::toError(Device->CreateCommittedResource(
@@ -495,7 +509,7 @@ class DXDevice : public offloadtest::Device {
     const D3D12_RESOURCE_DESC ReadBackResDesc = {
         D3D12_RESOURCE_DIMENSION_BUFFER,
         0,
-        R.size(),
+        BufferSize,
         1,
         1,
         1,
@@ -529,24 +543,27 @@ class DXDevice : public offloadtest::Device {
                ComPtr<ID3D12Resource> Buffer) {
     const uint32_t EltSize = R.getElementSize();
     const uint32_t NumElts = R.size() / EltSize;
+    ID3D12Resource *CounterBuffer = R.HasCounter ? Buffer.Get() : nullptr;
+    const uint32_t CounterOffset = getUAVBufferCounterOffset(R);
     DXGI_FORMAT const EltFormat =
         R.isRaw() ? getRawDXFormat(R)
                   : getDXFormat(R.BufferPtr->Format, R.BufferPtr->Channels);
     const D3D12_UNORDERED_ACCESS_VIEW_DESC UAVDesc = {
         EltFormat,
         D3D12_UAV_DIMENSION_BUFFER,
-        {D3D12_BUFFER_UAV{0, NumElts, R.isStructuredBuffer() ? EltSize : 0, 0,
-                          R.isByteAddressBuffer()
-                              ? D3D12_BUFFER_UAV_FLAG_RAW
-                              : D3D12_BUFFER_UAV_FLAG_NONE}}};
+        {D3D12_BUFFER_UAV{
+            0, NumElts, R.isStructuredBuffer() ? EltSize : 0, CounterOffset,
+            R.isByteAddressBuffer() ? D3D12_BUFFER_UAV_FLAG_RAW
+                                    : D3D12_BUFFER_UAV_FLAG_NONE}}};
 
     llvm::outs() << "UAV: HeapIdx = " << HeapIdx << " EltSize = " << EltSize
-                 << " NumElts = " << NumElts << "\n";
+                 << " NumElts = " << NumElts << " HasCounter = " << R.HasCounter
+                 << "\n";
     D3D12_CPU_DESCRIPTOR_HANDLE UAVHandle =
         IS.DescHeap->GetCPUDescriptorHandleForHeapStart();
     UAVHandle.ptr += HeapIdx * Device->GetDescriptorHandleIncrementSize(
                                    D3D12_DESCRIPTOR_HEAP_TYPE_CBV_SRV_UAV);
-    Device->CreateUnorderedAccessView(Buffer.Get(), nullptr, &UAVDesc,
+    Device->CreateUnorderedAccessView(Buffer.Get(), CounterBuffer, &UAVDesc,
                                       UAVHandle);
   }
 
@@ -895,6 +912,11 @@ class DXDevice : public offloadtest::Device {
                                  "Failed to map result."))
         return Err;
       memcpy(R.first->BufferPtr->Data.get(), DataPtr, R.first->size());
+      if (R.first->HasCounter)
+        memcpy(&R.first->BufferPtr->Counter,
+               static_cast<char *>(DataPtr) +
+                   getUAVBufferCounterOffset(*R.first),
+               sizeof(uint32_t));
       R.second.Readback->Unmap(0, nullptr);
       return llvm::Error::success();
     };

lib/Support/Pipeline.cpp

@@ -98,6 +98,7 @@ void MappingTraits<offloadtest::Buffer>::mapping(IO &I,
   I.mapRequired("Format", B.Format);
   I.mapOptional("Channels", B.Channels, 1);
   I.mapOptional("Stride", B.Stride, 0);
+  I.mapOptional("Counter", B.Counter, 0);
   if (!I.outputting() && B.Stride != 0 && B.Channels != 1)
     I.setError("Cannot set a structure stride and more than one channel.");
   switch (B.Format) {
@@ -147,6 +148,7 @@ void MappingTraits<offloadtest::Resource>::mapping(IO &I,
                                                    offloadtest::Resource &R) {
   I.mapRequired("Name", R.Name);
   I.mapRequired("Kind", R.Kind);
+  I.mapOptional("HasCounter", R.HasCounter, 0);
   I.mapRequired("DirectXBinding", R.DXBinding);
   I.mapOptional("VulkanBinding", R.VKBinding);
 }

test/Feature/StructuredBuffer/dec_counter.test

@@ -0,0 +1,49 @@
+#--- source.hlsl
+RWStructuredBuffer<uint> Out : register(u0);
+
+[numthreads(1,1,1)]
+void main(uint GI : SV_GroupIndex) {
+  Out.DecrementCounter();
+  Out.DecrementCounter();
+  Out.DecrementCounter();
+  Out[GI] = Out.DecrementCounter();
+}
+
+//--- pipeline.yaml
+---
+Shaders:
+  - Stage: Compute
+    Entry: main
+    DispatchSize: [1, 1, 1]
+Buffers:
+  - Name: Out
+    Format: Hex32
+    Stride: 4
+    ZeroInitSize: 4
+DescriptorSets:
+  - Resources:
+    - Name: Out
+      Kind: RWStructuredBuffer
+      HasCounter: true
+      DirectXBinding:
+        Register: 0
+        Space: 0
+...
+#--- end
+
+# UNSUPPORTED: Vulkan
+# UNSUPPORTED: Metal
+# UNSUPPORTED: Clang
+
+# RUN: split-file %s %t
+# RUN: %dxc_target -T cs_6_0 -Fo %t.o %t/source.hlsl
+# RUN: %offloader %t/pipeline.yaml %t.o | FileCheck %s
+
+# CHECK: Creating UAV: { Size = 4100, Register = u0, Space = 0, HasCounter = 1 }
+# CHECK: UAV: HeapIdx = 0 EltSize = 4 NumElts = 1 HasCounter = 1
+
+# CHECK: Name: Out
+# CHECK: Counter: 4294967292
+# CHECK: Data: [
+# CHECK: 0xFFFFFFFC
+# CHECK: ]

test/Feature/StructuredBuffer/inc_counter.test

@@ -0,0 +1,49 @@
+#--- source.hlsl
+RWStructuredBuffer<int> Out : register(u0);
+
+[numthreads(1,1,1)]
+void main(uint GI : SV_GroupIndex) {
+  Out.IncrementCounter();
+  Out.IncrementCounter();
+  Out.IncrementCounter();
+  Out[GI] = Out.IncrementCounter();
+}
+
+//--- pipeline.yaml
+---
+Shaders:
+  - Stage: Compute
+    Entry: main
+    DispatchSize: [1, 1, 1]
+Buffers:
+  - Name: Out
+    Format: Hex32
+    Stride: 4
+    ZeroInitSize: 4
+DescriptorSets:
+  - Resources:
+    - Name: Out
+      Kind: RWStructuredBuffer
+      HasCounter: true
+      DirectXBinding:
+        Register: 0
+        Space: 0
+...
+#--- end
+
+# UNSUPPORTED: Vulkan
+# UNSUPPORTED: Metal
+# UNSUPPORTED: Clang
+
+# RUN: split-file %s %t
+# RUN: %dxc_target -T cs_6_0 -Fo %t.o %t/source.hlsl
+# RUN: %offloader %t/pipeline.yaml %t.o | FileCheck %s
+
+# CHECK: Creating UAV: { Size = 4100, Register = u0, Space = 0, HasCounter = 1 }
+# CHECK: UAV: HeapIdx = 0 EltSize = 4 NumElts = 1 HasCounter = 1
+
+# CHECK: Name: Out
+# CHECK: Counter: 4
+# CHECK: Data: [
+# CHECK: 0x3
+# CHECK: ]