Set the layout rule explicitly for raw buffer operations (microsoft#6701)

The first first fix in microsoft#5392 was not correct. It relied on the layout
rule for the address to be the correct layout rule, but that is not
always the case. The address is just an integer that could exist in any
storage class. The correct solution is to explicitly set the layout rule
for the BitCast operation when expanding the RawBuffer* functions. We
know that the result of the BitCast is a pointer to the physical storage
buffer storage class, so we know the layout need to be the storage
buffer layout.

Fixes microsoft#6554

Author: s-perron

tools/clang/lib/SPIRV/SpirvEmitter.cpp

@@ -14263,6 +14263,7 @@ SpirvEmitter::loadDataFromRawAddress(SpirvInstruction *addressInUInt64,
   SpirvUnaryOp *address = spvBuilder.createUnaryOp(
       spv::Op::OpBitcast, bufferPtrType, addressInUInt64, loc);
   address->setStorageClass(spv::StorageClass::PhysicalStorageBuffer);
+  address->setLayoutRule(spirvOptions.sBufferLayoutRule);
 
   SpirvLoad *loadInst =
       dyn_cast<SpirvLoad>(spvBuilder.createLoad(bufferType, address, loc));
@@ -14291,6 +14292,7 @@ SpirvEmitter::storeDataToRawAddress(SpirvInstruction *addressInUInt64,
   if (!address)
     return nullptr;
   address->setStorageClass(spv::StorageClass::PhysicalStorageBuffer);
+  address->setLayoutRule(spirvOptions.sBufferLayoutRule);
 
   // If the source value has a different layout, it is not safe to directly
   // store it. It needs to be component-wise reconstructed to the new layout.

tools/clang/test/CodeGenSPIRV/intrinsics.vkrawbufferload.hlsl

@@ -4,6 +4,14 @@
 // CHECK: OpExtension "SPV_KHR_physical_storage_buffer"
 // CHECK: OpMemoryModel PhysicalStorageBuffer64 GLSL450
 
+// CHECK:                 OpMemberDecorate %BufferData_0 0 Offset 0
+// CHECK:                 OpMemberDecorate %BufferData_0 1 Offset 4
+// CHECK: %BufferData_0 = OpTypeStruct %float %v3float
+struct BufferData {
+  float  f;
+  float3 v;
+};
+
 uint64_t Address;
 float4 main() : SV_Target0 {
   // CHECK:      [[addr:%[0-9]+]] = OpLoad %ulong
@@ -34,5 +42,13 @@ float4 main() : SV_Target0 {
   // CHECK-NEXT: OpStore %v [[load_3]]
   uint v = vk::RawBufferLoad(Address);
 
+  // CHECK: [[buf:%[0-9]+]] = OpBitcast %_ptr_PhysicalStorageBuffer_BufferData_0
+  // CHECK-NEXT: [[load:%[0-9]+]] = OpLoad %BufferData_0 [[buf]] Aligned 4
+  BufferData d = vk::RawBufferLoad<BufferData>(Address);
+
+  // CHECK: [[buf:%[0-9]+]] = OpBitcast %_ptr_PhysicalStorageBuffer_BufferData_0 %ulong_0
+  // CHECK-NEXT: [[load:%[0-9]+]] = OpLoad %BufferData_0 [[buf]] Aligned 4
+  d = vk::RawBufferLoad<BufferData>(0);
+
   return float4(w.x, x, y, z);
 }

tools/clang/test/CodeGenSPIRV/intrinsics.vkrawbufferstore.hlsl

@@ -75,4 +75,14 @@ void main(uint3 tid : SV_DispatchThreadID) {
   xyzw.z = 84;
   xyzw.w = 69;
   vk::RawBufferStore<XYZW>(Address, xyzw);
+
+  // CHECK: [[xyzwval:%[0-9]+]] = OpLoad %XYZW %xyzw
+  // CHECK-NEXT: [[buf_3:%[0-9]+]] = OpBitcast %_ptr_PhysicalStorageBuffer_XYZW_0 %ulong_0
+  // CHECK-NEXT: [[member1:%[0-9]+]] = OpCompositeExtract %int [[xyzwval]] 0
+  // CHECK-NEXT: [[member2:%[0-9]+]] = OpCompositeExtract %int [[xyzwval]] 1
+  // CHECK-NEXT: [[member3:%[0-9]+]] = OpCompositeExtract %int [[xyzwval]] 2
+  // CHECK-NEXT: [[member4:%[0-9]+]] = OpCompositeExtract %int [[xyzwval]] 3
+  // CHECK-NEXT: [[p_xyzwval:%[0-9]+]] = OpCompositeConstruct %XYZW_0 [[member1]] [[member2]] [[member3]] [[member4]]
+  // CHECK-NEXT: OpStore [[buf_3]] [[p_xyzwval]] Aligned 4
+  vk::RawBufferStore<XYZW>(0, xyzw);
 }