Support for unbounded array of arrays (#1078)

* WIP: Unsized arrays on CPU.

* unbounded-array-of-array working on CPU.

* Remove some left over comments.

10 files changed, 262 insertions, 20 deletions

diff --git a/prelude/slang-cpp-types.h b/prelude/slang-cpp-types.h
index 2ba3fe19e..c8f6357a2 100644
--- a/prelude/slang-cpp-types.h
+++ b/prelude/slang-cpp-types.h
@@ -24,12 +24,20 @@ struct FixedArray
     T m_data[SIZE];
 };
 
-
-// Hmm... I guess a constant buffer should be unwrapped to be just a struct passed in
-/* template <typename T>
-struct ConstantBuffer
+// An array that has no specified size, becomes a 'Array'. This stores the size so it can potentially 
+// do bounds checking.  
+template <typename T>
+struct Array
 {
-}; */
+    const T& operator[](size_t index) const { SLANG_PRELUDE_ASSERT(index < count); return data[index]; }
+    T& operator[](size_t index) { SLANG_PRELUDE_ASSERT(index < count); return data[index]; }
+    
+    T* data;
+    size_t count;
+};
+
+/* Constant buffers become a pointer to the contained type, so ConstantBuffer<T> becomes T* in C++ code.
+*/
 
 template <typename T, int COUNT>
 struct Vector;
diff --git a/source/slang/slang-emit-cpp.cpp b/source/slang/slang-emit-cpp.cpp
index 2aea0cae9..479359d6b 100644
--- a/source/slang/slang-emit-cpp.cpp
+++ b/source/slang/slang-emit-cpp.cpp
@@ -544,6 +544,16 @@ SlangResult CPPSourceEmitter::_calcTypeName(IRType* type, CodeGenTarget target,
             out << ", " << elementCount << ">";
             return SLANG_OK;
         }
+        case kIROp_UnsizedArrayType:
+        {
+            auto arrayType = static_cast<IRUnsizedArrayType*>(type);
+            auto elementType = arrayType->getElementType();
+
+            out << "Array<";
+            SLANG_RETURN_ON_FAIL(_calcTypeName(elementType, target, out));
+            out << ">";
+            return SLANG_OK;
+        }
         default:
         {
             if (isNominalOp(type->op))
diff --git a/source/slang/slang-type-layout.cpp b/source/slang/slang-type-layout.cpp
index 9551cfe4e..ff1f2c147 100644
--- a/source/slang/slang-type-layout.cpp
+++ b/source/slang/slang-type-layout.cpp
@@ -366,6 +366,25 @@ struct CPULayoutRulesImpl : DefaultLayoutRulesImpl
         }
     }
 
+    SimpleArrayLayoutInfo GetArrayLayout( SimpleLayoutInfo elementInfo, LayoutSize elementCount) override
+    {
+        if (elementCount.isInfinite())
+        {
+            // This is an unsized array, get information for element
+            auto info = Super::GetArrayLayout(elementInfo, LayoutSize(1));
+
+            // So it is actually a Array<T> on CPU which is a pointer and a size
+            info.size = sizeof(void*) * 2;
+            info.alignment = sizeof(void*);
+
+            return info;
+        }
+        else
+        {
+            return Super::GetArrayLayout(elementInfo, elementCount);
+        }
+    }
+
     UniformLayoutInfo BeginStructLayout() override
     {
         return Super::BeginStructLayout();
diff --git a/tests/compute/unbounded-array-of-array.slang b/tests/compute/unbounded-array-of-array.slang
new file mode 100644
index 000000000..760740282
--- /dev/null
+++ b/tests/compute/unbounded-array-of-array.slang
@@ -0,0 +1,35 @@
+//DISABLE_TEST:CPU_REFLECTION: -profile cs_5_0 -entry computeMain -target cpp 
+//TEST(compute):COMPARE_COMPUTE_EX:-cpu -compute 
+
+struct IntAoa { RWStructuredBuffer<int> array[]; }
+
+//TEST_INPUT:ubuffer(data=[0 0 0 0 0 0 0 0], stride=4):dxbinding(0),glbinding(0),out,name outputBuffer
+RWStructuredBuffer<int> outputBuffer;
+
+//TEST_INPUT:array(size=2):name g_aoa.array
+ParameterBlock<IntAoa> g_aoa;
+
+//TEST_INPUT:ubuffer(data=[1 2 3 4], stride=4):name=g_aoa.array[0]
+//TEST_INPUT:ubuffer(data=[8 17 34], stride=4):name=g_aoa.array[1]
+
+[numthreads(8, 1, 1)]
+void computeMain(uint3 dispatchThreadID : SV_DispatchThreadID)
+{
+    int index = int(dispatchThreadID.x);
+    
+    int baseIndex = index >> 2;
+    int innerIndex = index & 3;
+
+    RWStructuredBuffer<int> buffer = g_aoa.array[baseIndex];
+
+    // Get the size 
+    uint bufferCount, bufferStride;
+    buffer.GetDimensions(bufferCount, bufferStride);
+
+    if (innerIndex >= bufferCount)
+    {
+        innerIndex = bufferCount - 1;
+    }
+
+	outputBuffer[index] = buffer[innerIndex]; 
+}
\ No newline at end of file
diff --git a/tests/compute/unbounded-array-of-array.slang.expected.txt b/tests/compute/unbounded-array-of-array.slang.expected.txt
new file mode 100644
index 000000000..aeec8496c
--- /dev/null
+++ b/tests/compute/unbounded-array-of-array.slang.expected.txt
@@ -0,0 +1,8 @@
+1
+2
+3
+4
+8
+11
+22
+22
diff --git a/tools/render-test/cpu-compute-util.cpp b/tools/render-test/cpu-compute-util.cpp
index 688e6b254..ef5755275 100644
--- a/tools/render-test/cpu-compute-util.cpp
+++ b/tools/render-test/cpu-compute-util.cpp
@@ -191,9 +191,24 @@ static CPUComputeUtil::Resource* _newOneTexture2D(int elemCount)
             const auto kind = typeLayout->getKind();
             switch (kind)
             {
-                case slang::TypeReflection::Kind::Vector:
-                case slang::TypeReflection::Kind::Matrix:
                 case slang::TypeReflection::Kind::Array:
+                {
+                    auto elementCount = int(typeLayout->getElementCount());
+                    if (elementCount == 0)
+                    {
+                        if (srcEntry.type == ShaderInputType::Array)
+                        {
+                            // We need to set the size
+                            CPUMemoryBinding::Buffer buffer;
+                            SLANG_RETURN_ON_FAIL(binding.setArrayCount(location, srcEntry.arrayDesc.size, buffer));
+                        }
+                        break;
+                    }
+                    SLANG_RETURN_ON_FAIL(binding.setInplace(location, srcEntry.bufferData.getBuffer(), srcEntry.bufferData.getCount() * sizeof(unsigned int)));
+                    break;
+                }
+                case slang::TypeReflection::Kind::Vector:
+                case slang::TypeReflection::Kind::Matrix:                
                 case slang::TypeReflection::Kind::Scalar:
                 case slang::TypeReflection::Kind::Struct:
                 {
diff --git a/tools/render-test/cpu-memory-binding.cpp b/tools/render-test/cpu-memory-binding.cpp
index 3bf1f1bbd..24d16756d 100644
--- a/tools/render-test/cpu-memory-binding.cpp
+++ b/tools/render-test/cpu-memory-binding.cpp
@@ -156,13 +156,25 @@ SlangResult CPUMemoryBinding::_add(slang::VariableLayoutReflection* varLayout, s
         {
             auto elementTypeLayout = typeLayout->getElementTypeLayout();
             auto elementCount = int(typeLayout->getElementCount());
-            const size_t stride = elementTypeLayout->getSize();
-            uint8_t* cur = (uint8_t*)dst;
-            for (int i = 0; i < elementCount; ++i)
+
+            if (elementCount == 0)
+            {
+                // We don't currently know the size that this array will be. So let's initially size it to 0.
+
+                CPPPrelude::Array<uint8_t>& dstBuf = *(CPPPrelude::Array<uint8_t>*)dst;
+                dstBuf.data = nullptr;
+                dstBuf.count = 0;
+            }
+            else
             {
-                Buffer elementBuffer;
-                _add(nullptr, elementTypeLayout, cur, elementBuffer);
-                cur += stride;
+                const size_t stride = elementTypeLayout->getSize();
+                uint8_t* cur = (uint8_t*)dst;
+                for (int i = 0; i < elementCount; ++i)
+                {
+                    Buffer elementBuffer;
+                    _add(nullptr, elementTypeLayout, cur, elementBuffer);
+                    cur += stride;
+                }
             }
             break;
         }
@@ -305,6 +317,11 @@ CPUMemoryBinding::Location CPUMemoryBinding::Location::toIndex(int index) const
     {
         return *this;
     }
+    SLANG_ASSERT(index >= 0);
+    if (index < 0)
+    {
+        return Location();
+    }
 
     auto typeLayout = m_typeLayout;
     uint8_t* cur = m_cur;
@@ -318,13 +335,24 @@ CPUMemoryBinding::Location CPUMemoryBinding::Location::toIndex(int index) const
             const auto elementCount = int(typeLayout->getElementCount());
             const auto elementStride = typeLayout->getElementStride(SLANG_PARAMETER_CATEGORY_UNIFORM);
 
-            if (index < 0 || index >= elementCount)
+            if (elementCount == 0)
             {
-                SLANG_ASSERT(index < elementCount);
+                CPPPrelude::Array<uint8_t>& array = *(CPPPrelude::Array<uint8_t>*)cur;
+                if (index < array.count)
+                {
+                    return Location(elementTypeLayout, array.data + elementStride * index);
+                }
                 return Location();
             }
-
-            return Location(elementTypeLayout, cur + elementStride * index);
+            else
+            {
+                if (index >= elementCount)
+                {
+                    SLANG_ASSERT(index < elementCount);
+                    return Location();
+                }
+                return Location(elementTypeLayout, cur + elementStride * index);
+            }
         }
         default: break;
     }
@@ -332,6 +360,13 @@ CPUMemoryBinding::Location CPUMemoryBinding::Location::toIndex(int index) const
     return Location();
 }
 
+SlangResult CPUMemoryBinding::initValue(slang::TypeLayoutReflection* typeLayout, void* dst)
+{
+    auto size = typeLayout->getSize();
+    // Zeroing works for built in types, and object types
+    memset(dst, 0, size);
+    return SLANG_OK;
+}
 
 SlangResult CPUMemoryBinding::setBufferContents(const Location& location, const void* initialData, size_t sizeInBytes)
 {
@@ -520,4 +555,87 @@ SlangResult CPUMemoryBinding::setInplace(const Location& location, const void* d
     return SLANG_OK;
 }
 
+Index CPUMemoryBinding::findBufferIndex(const void* ptr) const
+{
+    const Index count = m_allBuffers.getCount();
+
+    for (Index i = 0; i < count; ++i)
+    {
+        if (m_allBuffers[i].m_data == ptr)
+        {
+            return i;
+        }
+    }
+    return -1;
+}
+
+SlangResult CPUMemoryBinding::setArrayCount(const Location& location, int count, Buffer& outBuffer)
+{
+    if (!location.isValid())
+    {
+        return SLANG_FAIL;
+    }
+
+    auto typeLayout = location.getTypeLayout();
+    uint8_t* cur = location.getPtr();
+
+    const auto kind = typeLayout->getKind();
+
+    if (!(typeLayout->getKind() == slang::TypeReflection::Kind::Array && typeLayout->getElementCount() == 0))
+    {
+        return SLANG_FAIL;
+    }
+
+    CPPPrelude::Array<uint8_t>& array = *(CPPPrelude::Array<uint8_t>*)cur;
+    uint8_t* elements = array.data;
+
+    // Making smaller, just reduce the count.
+    // NOTE! Nothing is done here about deallocating resources which are perhaps no longer reachable.
+    // This isn't a leakage problem tho, as all buffers are released automatically when scope is left.
+    if (count <= array.count)
+    {
+        array.count = count;
+        return SLANG_OK;
+    }
+
+    const size_t elementStride = typeLayout->getElementStride(SLANG_PARAMETER_CATEGORY_UNIFORM);
+
+    const Index bufferIndex = elements ? findBufferIndex(elements) : -1;
+    if (bufferIndex > 0)
+    {
+        int maxCount = int(m_allBuffers[bufferIndex].m_sizeInBytes / elementStride);
+        if (count <= maxCount)
+        {
+            // Just initialize the space
+            memset(elements + elementStride * array.count, 0, (count - array.count) * elementStride);
+            array.count = count;
+            return SLANG_OK;
+        }
+    }
+
+    // Ok allocate a buffer that can hold all the elements
+
+    const size_t newBufferSize = count * elementStride;
+    Buffer newBuffer = _allocateBuffer(newBufferSize);
+
+    // Copy over the data from the old buffer if there is any
+    if (elements && array.count)
+    {
+        memcpy(newBuffer.m_data, elements, array.count * elementStride);
+    }
+
+    // Remove the old buffer as no longer needed
+    if (bufferIndex >= 0)
+    {
+        m_allBuffers.removeAt(bufferIndex);
+    }
+
+    // Set data 
+    array.count = count;
+    array.data = newBuffer.m_data;
+
+    outBuffer = newBuffer;
+    return SLANG_OK;
+}
+
 } // renderer_test
diff --git a/tools/render-test/cpu-memory-binding.h b/tools/render-test/cpu-memory-binding.h
index 36bc15c6b..8ca8e0681 100644
--- a/tools/render-test/cpu-memory-binding.h
+++ b/tools/render-test/cpu-memory-binding.h
@@ -43,6 +43,8 @@ struct CPUMemoryBinding
     slang::VariableLayoutReflection* getParameterByName(const char* name);
     slang::VariableLayoutReflection* getEntryPointParameterByName(const char* name);
 
+        /// Finds which buffer starts at the ptr index
+    Slang::Index findBufferIndex(const void* ptr) const;
 
     Location find(const char* name);
 
@@ -50,6 +52,15 @@ struct CPUMemoryBinding
     SlangResult setNewBuffer(const Location& location, const void* initialData, size_t sizeInBytes, Buffer& outBuffer);
     SlangResult setObject(const Location& location, void* object);
     SlangResult setInplace(const Location& location, const void* data, size_t sizeInBytes);
+        /// Initialize memory with a 'sensible' value based on type. Pointer types become null.
+    SlangResult initValue(slang::TypeLayoutReflection* typeLayout, void* dst);
+    SlangResult initValue(const Location& location) { return initValue(location.getTypeLayout(), location.getPtr()); }
+
+        /// Set the size of a *non fixed size* array at location.
+        /// A non fixed size array is reflected as having a count of 0 elements.
+        /// Only returns a buffer in outBuffer if a new buffer is created. 
+    SlangResult setArrayCount(const Location& location, int count, Buffer& outBuffer);
+    
     SlangResult init(slang::ShaderReflection* reflection, int entryPointIndex);
     CPUMemoryBinding();
 
diff --git a/tools/render-test/shader-input-layout.cpp b/tools/render-test/shader-input-layout.cpp
index b5b65f52d..37aa5dab6 100644
--- a/tools/render-test/shader-input-layout.cpp
+++ b/tools/render-test/shader-input-layout.cpp
@@ -138,7 +138,11 @@ namespace renderer_test
                     {
                         ShaderInputLayoutEntry entry;
 
-                        if (parser.LookAhead("cbuffer"))
+                        if (parser.LookAhead("array"))
+                        {
+                            entry.type = ShaderInputType::Array;
+                        }
+                        else if (parser.LookAhead("cbuffer"))
                         {
                             entry.type = ShaderInputType::Buffer;
                             entry.bufferDesc.type = InputBufferType::ConstantBuffer;
@@ -253,7 +257,15 @@ namespace renderer_test
                                 else if (word == "size")
                                 {
                                     parser.Read("=");
-                                    entry.textureDesc.size = parser.ReadInt();
+                                    auto size = parser.ReadInt();
+                                    if (entry.type == ShaderInputType::Array)
+                                    {
+                                        entry.arrayDesc.size = size;
+                                    }
+                                    else
+                                    {
+                                        entry.textureDesc.size = size;
+                                    }
                                 }
                                 else if (word == "random")
                                 {
diff --git a/tools/render-test/shader-input-layout.h b/tools/render-test/shader-input-layout.h
index 009be7514..8a00980d9 100644
--- a/tools/render-test/shader-input-layout.h
+++ b/tools/render-test/shader-input-layout.h
@@ -12,7 +12,7 @@ using namespace gfx;
 
 enum class ShaderInputType
 {
-    Buffer, Texture, Sampler, CombinedTextureSampler
+    Buffer, Texture, Sampler, CombinedTextureSampler, Array
 };
 
 enum class InputTextureContent
@@ -51,6 +51,11 @@ struct InputSamplerDesc
     bool isCompareSampler = false;
 };
 
+struct ArrayDesc
+{
+    int size = 0;
+};
+
 class ShaderInputLayoutEntry
 {
 public:
@@ -59,6 +64,7 @@ public:
     InputTextureDesc textureDesc;
     InputBufferDesc bufferDesc;
     InputSamplerDesc samplerDesc;
+    ArrayDesc arrayDesc;
     bool isOutput = false;
     bool isCPUOnly = false;
     int hlslBinding = -1;