Handle local arrays of RWStructrued buffers (microsoft#5523)

This commit adds support for local arrays of RWStructuredBuffers.

Author: s-perron

tools/clang/lib/SPIRV/AstTypeProbe.cpp

@@ -1001,6 +1001,9 @@ bool isAKindOfStructuredOrByteBuffer(QualType type) {
 }
 
 bool isOrContainsAKindOfStructuredOrByteBuffer(QualType type) {
+  while (type->isArrayType())
+    type = type->getAsArrayTypeUnsafe()->getElementType();
+
   if (const RecordType *recordType = type->getAs<RecordType>()) {
     StringRef name = recordType->getDecl()->getName();
     if (name == "StructuredBuffer" || name == "RWStructuredBuffer" ||

tools/clang/lib/SPIRV/LowerTypeVisitor.cpp

@@ -530,10 +530,24 @@ const SpirvType *LowerTypeVisitor::lowerType(QualType type,
   // Array type
   if (const auto *arrayType = astContext.getAsArrayType(type)) {
     const auto elemType = arrayType->getElementType();
-    const auto *loweredElemType =
-        lowerType(arrayType->getElementType(), rule, isRowMajor, srcLoc);
-    llvm::Optional<uint32_t> arrayStride = llvm::None;
 
+    // If layout rule is void, it means these resource types are used for
+    // declaring local resources. This should be lowered to a pointer to the
+    // array.
+    //
+    // The pointer points to the Uniform storage class, and the element type
+    // should have the corresponding layout.
+    bool isLocalStructuredOrByteBuffer =
+        isAKindOfStructuredOrByteBuffer(elemType) &&
+        rule == SpirvLayoutRule::Void;
+
+    SpirvLayoutRule elementLayoutRule =
+        (isLocalStructuredOrByteBuffer ? getCodeGenOptions().sBufferLayoutRule
+                                       : rule);
+    const SpirvType *loweredElemType =
+        lowerType(elemType, elementLayoutRule, isRowMajor, srcLoc);
+
+    llvm::Optional<uint32_t> arrayStride = llvm::None;
     if (rule != SpirvLayoutRule::Void &&
         // We won't have stride information for structured/byte buffers since
         // they contain runtime arrays.
@@ -544,13 +558,23 @@ const SpirvType *LowerTypeVisitor::lowerType(QualType type,
       arrayStride = stride;
     }
 
+    const SpirvType *spirvArrayType = nullptr;
     if (const auto *caType = astContext.getAsConstantArrayType(type)) {
       const auto size = static_cast<uint32_t>(caType->getSize().getZExtValue());
-      return spvContext.getArrayType(loweredElemType, size, arrayStride);
+      spirvArrayType =
+          spvContext.getArrayType(loweredElemType, size, arrayStride);
+    } else {
+      assert(type->isIncompleteArrayType());
+      spirvArrayType =
+          spvContext.getRuntimeArrayType(loweredElemType, arrayStride);
+    }
+
+    if (isLocalStructuredOrByteBuffer) {
+      return spvContext.getPointerType(spirvArrayType,
+                                       spv::StorageClass::Uniform);
     }
 
-    assert(type->isIncompleteArrayType());
-    return spvContext.getRuntimeArrayType(loweredElemType, arrayStride);
+    return spirvArrayType;
   }
 
   // Reference types

tools/clang/test/CodeGenSPIRV/type.rwstructured-buffer.array.counter.indirect2.hlsl

@@ -0,0 +1,30 @@
+// RUN: %dxc -T ps_6_6 -E main -O0 -fvk-allow-rwstructuredbuffer-arrays
+
+struct PSInput
+{
+	uint idx : COLOR;
+};
+
+// CHECK: OpDecorate %g_rwbuffer DescriptorSet 2
+// CHECK: OpDecorate %g_rwbuffer Binding 0
+// CHECK: OpDecorate %counter_var_g_rwbuffer DescriptorSet 2
+// CHECK: OpDecorate %counter_var_g_rwbuffer Binding 1
+
+// CHECK: %g_rwbuffer = OpVariable %_ptr_Uniform__arr_type_RWStructuredBuffer_uint_uint_5 Uniform
+// CHECK: %counter_var_g_rwbuffer = OpVariable %_ptr_Uniform__arr_type_ACSBuffer_counter_uint_5 Uniform
+RWStructuredBuffer<uint> g_rwbuffer[5] : register(u0, space2);
+
+float4 main(PSInput input) : SV_TARGET
+{
+    RWStructuredBuffer<uint> l_rwbuffer[5] = g_rwbuffer;
+
+// CHECK: [[ac1:%\d+]] = OpAccessChain %_ptr_Uniform_type_ACSBuffer_counter %counter_var_g_rwbuffer %int_0
+// CHECK: [[ac2:%\d+]] = OpAccessChain %_ptr_Uniform_int [[ac1]] %uint_0
+// CHECK: OpAtomicIAdd %int [[ac2]] %uint_1 %uint_0 %int_1
+    l_rwbuffer[0].IncrementCounter();
+
+// CHECK: [[ac1:%\d+]] = OpAccessChain %_ptr_Uniform_type_RWStructuredBuffer_uint %g_rwbuffer {{%\d+}}
+// CHECK: [[ac2:%\d+]] = OpAccessChain %_ptr_Uniform_uint [[ac1]] %int_0 %uint_0
+// CHECK: OpLoad %uint [[ac2]]
+    return l_rwbuffer[input.idx][0];
+}

tools/clang/unittests/SPIRV/CodeGenSpirvTest.cpp

@@ -162,6 +162,9 @@ TEST_F(FileTest, RWStructuredBufferArrayCounterFlattened) {
 TEST_F(FileTest, RWStructuredBufferArrayCounterIndirect) {
   runFileTest("type.rwstructured-buffer.array.counter.indirect.hlsl");
 }
+TEST_F(FileTest, RWStructuredBufferArrayCounterIndirect2) {
+  runFileTest("type.rwstructured-buffer.array.counter.indirect2.hlsl");
+}
 TEST_F(FileTest, RWStructuredBufferArrayBindAttributes) {
   runFileTest("type.rwstructured-buffer.array.binding.attributes.hlsl");
 }