Fix GLSL output for byte-address loads of vectors (#1558)

While working on #1557, it became clear that something was going wrong when using `*ByteAddressBuffer.Load<T>` to load a vector type on GLSL/SPIR-V targets.

The root problem was that the IR-level layout logic (which computes the "natural" layout of a type) had not yet been extended to handle vectors. The fix is simple enough, but it highlights the fact that we probably need to go ahead and "complete" that layout logic sooner or later.

This change includes a test case that covers the behavior added here, as well as the case that #1557 fixes. Unfortunately, due to CI system limitations, the HLSL/dxc part of the test is not yet enabled.

3 files changed, 96 insertions, 19 deletions

diff --git a/source/slang/slang-ir-layout.cpp b/source/slang/slang-ir-layout.cpp
index 0003d279a..6ea1093f5 100644
--- a/source/slang/slang-ir-layout.cpp
+++ b/source/slang/slang-ir-layout.cpp
@@ -49,6 +49,30 @@
 namespace Slang
 {
 
+static Result _calcNaturalArraySizeAndAlignment(IRType* elementType, IRInst* elementCountInst, IRSizeAndAlignment* outSizeAndAlignment)
+{
+    auto elementCountLit = as<IRIntLit>(elementCountInst);
+    if(!elementCountLit)
+        return SLANG_FAIL;
+    auto elementCount = elementCountLit->getValue();
+
+    if( elementCount == 0 )
+    {
+        *outSizeAndAlignment = IRSizeAndAlignment(0, 1);
+        return SLANG_OK;
+    }
+
+    IRSizeAndAlignment elementTypeLayout;
+    SLANG_RETURN_ON_FAIL(getNaturalSizeAndAlignment(elementType, &elementTypeLayout));
+
+    auto elementStride = elementTypeLayout.getStride();
+
+    *outSizeAndAlignment = IRSizeAndAlignment(
+        elementStride * (elementCount - 1) + elementTypeLayout.size,
+        elementTypeLayout.alignment);
+    return SLANG_OK;
+}
+
 static Result _calcNaturalSizeAndAlignment(IRType* type, IRSizeAndAlignment* outSizeAndAlignment)
 {
     switch( type->op )
@@ -141,27 +165,21 @@ static Result _calcNaturalSizeAndAlignment(IRType* type, IRSizeAndAlignment* out
         {
             auto arrayType = cast<IRArrayType>(type);
 
-            auto elementCountLit = as<IRIntLit>(arrayType->getElementCount());
-            if(!elementCountLit)
-                return SLANG_FAIL;
-            auto elementCount = elementCountLit->getValue();
-
-            if( elementCount == 0 )
-            {
-                *outSizeAndAlignment = IRSizeAndAlignment(0, 1);
-                return SLANG_OK;
-            }
-
-            auto elementType = arrayType->getElementType();
-            IRSizeAndAlignment elementTypeLayout;
-            SLANG_RETURN_ON_FAIL(getNaturalSizeAndAlignment(elementType, &elementTypeLayout));
+            return _calcNaturalArraySizeAndAlignment(
+                arrayType->getElementType(),
+                arrayType->getElementCount(),
+                outSizeAndAlignment);
+        }
+        break;
 
-            auto elementStride = elementTypeLayout.getStride();
+    case kIROp_VectorType:
+        {
+            auto vecType = cast<IRVectorType>(type);
 
-            *outSizeAndAlignment = IRSizeAndAlignment(
-                elementStride * (elementCount - 1) + elementTypeLayout.size,
-                elementTypeLayout.alignment);
-            return SLANG_OK;
+            return _calcNaturalArraySizeAndAlignment(
+                vecType->getElementType(),
+                vecType->getElementCount(),
+                outSizeAndAlignment);
         }
         break;
 
diff --git a/tests/hlsl-intrinsic/byte-address-buffer/byte-address-16bit-vector.slang b/tests/hlsl-intrinsic/byte-address-buffer/byte-address-16bit-vector.slang
new file mode 100644
index 000000000..403be325e
--- /dev/null
+++ b/tests/hlsl-intrinsic/byte-address-buffer/byte-address-16bit-vector.slang
@@ -0,0 +1,55 @@
+// byte-address-16bit-vector.slang
+
+// Test that loading values using 16-bit vector types from
+// byte-address buffers works as expected, on targets
+// that support it.
+
+//TEST(compute):COMPARE_COMPUTE_EX:-slang -compute -cpu
+
+// Note: disabled on D3D12 because some of the CI services don't have
+// a recent enough build of dxc to support generic load/store.
+//
+//DISABLE_TEST(compute):COMPARE_COMPUTE_EX:-slang -compute -dx12 -use-dxil -profile cs_6_2 -render-features half
+
+//TEST(compute, vulkan):COMPARE_COMPUTE_EX:-slang -vk -compute -profile cs_6_2 -render-features half
+//TEST(compute):COMPARE_COMPUTE_EX:-slang -cuda -compute
+
+// Unavailable on D3D and fxc-based targets, because fxc doesn't
+// support loading anything besides `uint` from a byte-address
+// buffer.
+//
+//TEST_DISABLED(compute):COMPARE_COMPUTE_EX:-slang -compute
+
+
+//TEST_INPUT:ubuffer(data=[0x00010002 0x00030004 0x00050006 0x00070008]):name=inputBuffer
+RWByteAddressBuffer inputBuffer;
+
+//TEST_INPUT:ubuffer(data=[0xdeadbeef 0xdeadbeef 0xdeadbeef 0xdeadbeef]):name=tmpBuffer
+RWByteAddressBuffer tmpBuffer;
+
+//TEST_INPUT:ubuffer(data=[0 0 0 0], stride=4):out,name=outputBuffer
+RWStructuredBuffer<int> outputBuffer;
+
+struct Data
+{
+    uint16_t2 v;
+}
+
+int test(int val)
+{
+    int offsetX = val*4;
+    int offsetY = offsetX & 0xFF;
+
+    Data x = inputBuffer.Load<Data>(offsetX);
+    tmpBuffer.Store<Data>(offsetY, x);
+
+    uint16_t2 y = tmpBuffer.Load<uint16_t2>(offsetX);
+    return int(y.x)*256*16 + int(y.y);
+}
+
+[numthreads(4, 1, 1)]
+void computeMain(uint3 dispatchThreadID : SV_DispatchThreadID)
+{
+    int tid = dispatchThreadID.x;
+    outputBuffer[tid] = test(tid);
+}
\ No newline at end of file
diff --git a/tests/hlsl-intrinsic/byte-address-buffer/byte-address-16bit-vector.slang.expected.txt b/tests/hlsl-intrinsic/byte-address-buffer/byte-address-16bit-vector.slang.expected.txt
new file mode 100644
index 000000000..eefd23c0a
--- /dev/null
+++ b/tests/hlsl-intrinsic/byte-address-buffer/byte-address-16bit-vector.slang.expected.txt
@@ -0,0 +1,4 @@
+2001
+4003
+6005
+8007