[spirv] Fix issues of RWByteAddressBuffer in the InitList (#3145)

RWByteAddressBuffer object in the initList should be handled like
OpaqueType RWBuffer or Texture2D.

Author: jiaolu

tools/clang/lib/SPIRV/InitListHandler.cpp

@@ -130,7 +130,9 @@ bool InitListHandler::tryToSplitStruct() {
   const QualType initType = init->getAstResultType();
   if (!initType->isStructureType() ||
       // Sampler types will pass the above check but we cannot split it.
-      isSampler(initType))
+      isSampler(initType) ||
+      // Can not split structuredOrByteBuffer
+      isAKindOfStructuredOrByteBuffer(initType))
     return false;
 
   // We are certain the current intializer will be replaced by now.
@@ -214,12 +216,11 @@ SpirvInstruction *InitListHandler::createInitForType(QualType type,
   // Samplers, (RW)Buffers, (RW)Textures
   // It is important that this happens before checking of structure types.
   if (isOpaqueType(type))
-    return createInitForSamplerImageType(type, srcLoc);
+    return createInitForBufferOrImageType(type, srcLoc);
 
   // This should happen before the check for normal struct types
   if (isAKindOfStructuredOrByteBuffer(type)) {
-    emitError("cannot handle structured/byte buffer as initializer", srcLoc);
-    return nullptr;
+    return createInitForBufferOrImageType(type, srcLoc);
   }
 
   if (type->isStructureType())
@@ -427,9 +428,9 @@ InitListHandler::createInitForConstantArrayType(QualType type,
 }
 
 SpirvInstruction *
-InitListHandler::createInitForSamplerImageType(QualType type,
-                                               SourceLocation srcLoc) {
-  assert(isOpaqueType(type));
+InitListHandler::createInitForBufferOrImageType(QualType type,
+                                                SourceLocation srcLoc) {
+  assert(isOpaqueType(type) || isAKindOfStructuredOrByteBuffer(type));
 
   // Samplers, (RW)Buffers, and (RW)Textures are translated into OpTypeSampler
   // and OpTypeImage. They should be treated similar as builtin types.

tools/clang/lib/SPIRV/InitListHandler.h

@@ -133,8 +133,8 @@ class InitListHandler {
   SpirvInstruction *createInitForStructType(QualType type, SourceLocation);
   SpirvInstruction *createInitForConstantArrayType(QualType type,
                                                    SourceLocation);
-  SpirvInstruction *createInitForSamplerImageType(QualType type,
-                                                  SourceLocation);
+  SpirvInstruction *createInitForBufferOrImageType(QualType type,
+                                                   SourceLocation);
 
 private:
   const ASTContext &astContext;

tools/clang/test/CodeGenSPIRV/initializelist.rwbyteaddressbuffer.hlsl

@@ -0,0 +1,19 @@
+// Run: %dxc -T cs_6_0 -E main
+RWBuffer<int> buffer1 : register(u1);
+RWByteAddressBuffer buffer2 : register(u2);
+
+// CHECK: [[Resource:%\w+]] = OpTypeStruct %type_RWByteAddressBuffer %uint
+struct Resource {
+  RWByteAddressBuffer rwbuffer;
+  uint offset;
+};
+
+[numthreads(8, 1, 1)]
+void main(uint globalId : SV_DispatchThreadID,
+          uint localId  : SV_GroupThreadID,
+          uint groupId  : SV_GroupID) {
+// CHECK: [[buffer2:%\w+]] = OpVariable %_ptr_Uniform_type_RWByteAddressBuffer Uniform
+  Resource resourceInfo2 = {buffer2, 2};
+// CHECK: OpCompositeConstruct [[Resource]] [[buffer2]] %uint_2
+  buffer1[0] = resourceInfo2.rwbuffer.Load(0);
+}

tools/clang/unittests/SPIRV/CodeGenSpirvTest.cpp

@@ -983,6 +983,11 @@ TEST_F(FileTest, RWByteAddressBufferAtomicMethods) {
   runFileTest("method.rw-byte-address-buffer.atomic.hlsl");
 }
 
+TEST_F(FileTest, InitializeListRWByteAddressBuffer) {
+  runFileTest("initializelist.rwbyteaddressbuffer.hlsl", Expect::Success,
+              /* runValidation */ false);
+}
+
 // For Buffer/RWBuffer methods
 TEST_F(FileTest, BufferLoad) { runFileTest("method.buffer.load.hlsl"); }
 TEST_F(FileTest, BufferGetDimensions) {