[gfx] Support StructuredBuffer<IInterface>. (#1851)

Co-authored-by: T. Foley <tfoleyNV@users.noreply.github.com>

1 files changed, 107 insertions, 225 deletions

diff --git a/tools/gfx/cpu/render-cpu.cpp b/tools/gfx/cpu/render-cpu.cpp
index 0bdc06ad6..0c5f119b8 100644
--- a/tools/gfx/cpu/render-cpu.cpp
+++ b/tools/gfx/cpu/render-cpu.cpp
@@ -357,7 +357,7 @@ public:
     void*           m_data = nullptr;
 };
 
-class CPUResourceView : public IResourceView, public ComObject
+class CPUResourceView : public ResourceViewBase
 {
 public:
     enum class Kind
@@ -365,15 +365,6 @@ public:
         Buffer,
         Texture,
     };
-
-    SLANG_COM_OBJECT_IUNKNOWN_ALL
-    IResourceView* getInterface(const Guid& guid)
-    {
-        if (guid == GfxGUID::IID_ISlangUnknown || guid == GfxGUID::IID_IResourceView)
-            return static_cast<IResourceView*>(this);
-        return nullptr;
-    }
-
     Kind getViewKind() const { return m_kind; }
     Desc const& getDesc() const { return m_desc; }
 
@@ -576,6 +567,7 @@ public:
         slang::BindingType bindingType;
         Index count;
         Index baseIndex; // Flat index for sub-ojects
+        Index subObjectIndex;
 
         // TODO: The `uniformOffset` field should be removed,
         // since it cannot be supported by the Slang reflection
@@ -610,10 +602,10 @@ public:
     {
         initBase(renderer, layout);
 
-        Index subObjectCount = 0;
-        Index resourceCount = 0;
+        m_subObjectCount = 0;
+        m_resourceCount = 0;
 
-        m_elementTypeLayout = _unwrapParameterGroups(layout);
+        m_elementTypeLayout = _unwrapParameterGroups(layout, m_containerType);
         m_size = m_elementTypeLayout->getSize();
 
         // Compute the binding ranges that are used to store
@@ -645,18 +637,31 @@ public:
                 descriptorSetIndex, rangeIndexInDescriptorSet);
 
             Index baseIndex = 0;
+            Index subObjectIndex = 0;
             switch (slangBindingType)
             {
             case slang::BindingType::ConstantBuffer:
             case slang::BindingType::ParameterBlock:
             case slang::BindingType::ExistentialValue:
-                baseIndex = subObjectCount;
-                subObjectCount += count;
+                baseIndex = m_subObjectCount;
+                subObjectIndex = baseIndex;
+                m_subObjectCount += count;
+                break;
+            case slang::BindingType::RawBuffer:
+            case slang::BindingType::MutableRawBuffer:
+                if (slangLeafTypeLayout->getType()->getElementType() != nullptr)
+                {
+                    // A structured buffer occupies both a resource slot and
+                    // a sub-object slot.
+                    subObjectIndex = m_subObjectCount;
+                    m_subObjectCount += count;
+                }
+                baseIndex = m_resourceCount;
+                m_resourceCount += count;
                 break;
-
             default:
-                baseIndex = resourceCount;
-                resourceCount += count;
+                baseIndex = m_resourceCount;
+                m_resourceCount += count;
                 break;
             }
 
@@ -665,12 +670,10 @@ public:
             bindingRangeInfo.count = count;
             bindingRangeInfo.baseIndex = baseIndex;
             bindingRangeInfo.uniformOffset = uniformOffset;
+            bindingRangeInfo.subObjectIndex = subObjectIndex;
             m_bindingRanges.add(bindingRangeInfo);
         }
 
-        m_subObjectCount = subObjectCount;
-        m_resourceCount = resourceCount;
-
         SlangInt subObjectRangeCount = m_elementTypeLayout->getSubObjectRangeCount();
         for (SlangInt r = 0; r < subObjectRangeCount; ++r)
         {
@@ -703,6 +706,9 @@ public:
     size_t getSize() { return m_size; }
     Index getResourceCount() const { return m_resourceCount; }
     Index getSubObjectCount() const { return m_subObjectCount; }
+    List<SubObjectRangeInfo>& getSubObjectRanges() { return subObjectRanges; }
+    BindingRangeInfo getBindingRange(Index index) { return m_bindingRanges[index]; }
+    Index getBindingRangeCount() const { return m_bindingRanges.getCount(); }
 };
 
 class CPUEntryPointLayout : public CPUShaderObjectLayout
@@ -763,32 +769,64 @@ public:
     CPUEntryPointLayout* getEntryPoint(Index index) { return m_entryPointLayouts[index]; }
 };
 
-class CPUShaderObject : public ShaderObjectBase
+class CPUShaderObjectData
 {
 public:
-    void* m_data = nullptr;
+    Slang::List<char> m_ordinaryData;
+    // Any "ordinary" / uniform data for this object
+    Slang::RefPtr<CPUBufferResource> m_bufferResource;
+    Slang::RefPtr<CPUBufferView> m_bufferView;
 
-    ~CPUShaderObject()
+    Index getCount() { return m_ordinaryData.getCount(); }
+    void setCount(Index count) { m_ordinaryData.setCount(count); }
+    char* getBuffer() { return m_ordinaryData.getBuffer(); }
+
+    ~CPUShaderObjectData()
     {
-        free(m_data);
+        // m_bufferResource's data is managed by m_ordinaryData so we
+        // set it to null to prevent m_bufferResource from freeing it.
+        if (m_bufferResource)
+            m_bufferResource->m_data = nullptr;
+    }
+
+    /// Returns a StructuredBuffer resource view for GPU access into the buffer content.
+    /// Creates a StructuredBuffer resource if it has not been created.
+    ResourceViewBase* getResourceView(
+        RendererBase* device,
+        slang::TypeLayoutReflection* elementLayout,
+        slang::BindingType bindingType)
+    {
+        SLANG_UNUSED(device);
+        if (!m_bufferResource)
+        {
+            IBufferResource::Desc desc = {};
+            desc.type = IResource::Type::Buffer;
+            desc.elementSize = (int)elementLayout->getSize();
+            m_bufferResource = new CPUBufferResource(desc);
+
+            IResourceView::Desc viewDesc = {};
+            viewDesc.type = IResourceView::Type::UnorderedAccess;
+            viewDesc.format = Format::Unknown;
+            m_bufferView = new CPUBufferView(viewDesc, m_bufferResource);
+            
+        }
+        m_bufferResource->getDesc()->sizeInBytes = m_ordinaryData.getCount();
+        m_bufferResource->m_data = m_ordinaryData.getBuffer();
+        return m_bufferView.Ptr();
     }
+};
+
+class CPUShaderObject
+    : public ShaderObjectBaseImpl<CPUShaderObject, CPUShaderObjectLayout, CPUShaderObjectData>
+{
+    typedef ShaderObjectBaseImpl<CPUShaderObject, CPUShaderObjectLayout, CPUShaderObjectData> Super;
 
-    List<RefPtr<CPUShaderObject>> m_objects;
+public:
     List<RefPtr<CPUResourceView>> m_resources;
 
     virtual SLANG_NO_THROW Result SLANG_MCALL
         init(IDevice* device, CPUShaderObjectLayout* typeLayout);
 
-    CPUShaderObjectLayout* getLayout()
-    {
-        return static_cast<CPUShaderObjectLayout*>(m_layout.Ptr());
-    }
-
-    virtual SLANG_NO_THROW slang::TypeLayoutReflection* SLANG_MCALL getElementTypeLayout() override
-    {
-        return getLayout()->getElementTypeLayout();
-    }
-
     virtual SLANG_NO_THROW UInt SLANG_MCALL getEntryPointCount() override { return 0; }
     virtual SLANG_NO_THROW Result SLANG_MCALL
         getEntryPoint(UInt index, IShaderObject** outEntryPoint) override
@@ -799,143 +837,8 @@ public:
     virtual SLANG_NO_THROW Result SLANG_MCALL
         setData(ShaderOffset const& offset, void const* data, size_t size) override
     {
-        size = Math::Min(size, getLayout()->getSize() - offset.uniformOffset);
-        memcpy((char*)m_data + offset.uniformOffset, data, size);
-        return SLANG_OK;
-    }
-    virtual SLANG_NO_THROW Result SLANG_MCALL getObject(
-        ShaderOffset const& offset,
-        IShaderObject**     outObject) override
-    {
-        auto layout = getLayout();
-
-        auto bindingRangeIndex = offset.bindingRangeIndex;
-        SLANG_ASSERT(bindingRangeIndex >= 0);
-        SLANG_ASSERT(bindingRangeIndex < layout->m_bindingRanges.getCount());
-
-        auto& bindingRange = layout->m_bindingRanges[bindingRangeIndex];
-        auto subObjectIndex = bindingRange.baseIndex + offset.bindingArrayIndex;
-        auto& subObject = m_objects[subObjectIndex];
-
-        returnComPtr(outObject, subObject);
-
-        return SLANG_OK;
-    }
-    virtual SLANG_NO_THROW Result SLANG_MCALL setObject(
-        ShaderOffset const& offset,
-        IShaderObject*      object) override
-    {
-        auto layout = getLayout();
-
-        auto bindingRangeIndex = offset.bindingRangeIndex;
-        SLANG_ASSERT(bindingRangeIndex >= 0);
-        SLANG_ASSERT(bindingRangeIndex < layout->m_bindingRanges.getCount());
-
-        auto& bindingRange = layout->m_bindingRanges[bindingRangeIndex];
-        auto subObjectIndex = bindingRange.baseIndex + offset.bindingArrayIndex;
-
-        CPUShaderObject* subObject = static_cast<CPUShaderObject*>(object);
-        m_objects[subObjectIndex] = subObject;
-
-        switch( bindingRange.bindingType )
-        {
-        default:
-            SLANG_RETURN_ON_FAIL(setData(offset, &subObject->m_data, sizeof(void*)));
-            break;
-
-        // If the range being assigned into represents an interface/existential-type leaf field,
-        // then we need to consider how the `object` being assigned here affects specialization.
-        // We may also need to assign some data from the sub-object into the ordinary data
-        // buffer for the parent object.
-        //
-        case slang::BindingType::ExistentialValue:
-            {
-                auto renderer = getRenderer();
-
-                ComPtr<slang::ISession> slangSession;
-                SLANG_RETURN_ON_FAIL(renderer->getSlangSession(slangSession.writeRef()));
-
-                // A leaf field of interface type is laid out inside of the parent object
-                // as a tuple of `(RTTI, WitnessTable, Payload)`. The layout of these fields
-                // is a contract between the compiler and any runtime system, so we will
-                // need to rely on details of the binary layout.
-
-                // We start by querying the layout/type of the concrete value that the application
-                // is trying to store into the field, and also the layout/type of the leaf
-                // existential-type field itself.
-                //
-                auto concreteTypeLayout = subObject->getElementTypeLayout();
-                auto concreteType = concreteTypeLayout->getType();
-                //
-                auto existentialTypeLayout = layout->getElementTypeLayout()->getBindingRangeLeafTypeLayout(bindingRangeIndex);
-                auto existentialType = existentialTypeLayout->getType();
-
-                // The first field of the tuple (offset zero) is the run-time type information (RTTI)
-                // ID for the concrete type being stored into the field.
-                //
-                // TODO: We need to be able to gather the RTTI type ID from `object` and then
-                // use `setData(offset, &TypeID, sizeof(TypeID))`.
-
-                // The second field of the tuple (offset 8) is the ID of the "witness" for the
-                // conformance of the concrete type to the interface used by this field.
-                //
-                auto witnessTableOffset = offset;
-                witnessTableOffset.uniformOffset += 8;
-                //
-                // Conformances of a type to an interface are computed and then stored by the
-                // Slang runtime, so we can look up the ID for this particular conformance (which
-                // will create it on demand).
-                //
-                // Note: If the type doesn't actually conform to the required interface for
-                // this sub-object range, then this is the point where we will detect that
-                // fact and error out.
-                //
-                uint32_t conformanceID = 0xFFFFFFFF;
-                SLANG_RETURN_ON_FAIL(slangSession->getTypeConformanceWitnessSequentialID(
-                    concreteType, existentialType, &conformanceID));
-                //
-                // Once we have the conformance ID, then we can write it into the object
-                // at the required offset.
-                //
-                SLANG_RETURN_ON_FAIL(setData(witnessTableOffset, &conformanceID, sizeof(conformanceID)));
-
-                // The third field of the tuple (offset 16) is the "payload" that is supposed to
-                // hold the data for a value of the given concrete type.
-                //
-                auto payloadOffset = offset;
-                payloadOffset.uniformOffset += 16;
-
-                // There are two cases we need to consider here for how the payload might be used:
-                //
-                // * If the concrete type of the value being bound is one that can "fit" into the
-                //   available payload space,  then it should be stored in the payload.
-                //
-                // * If the concrete type of the value cannot fit in the payload space, then it
-                //   will need to be stored somewhere else.
-                //
-                if(_doesValueFitInExistentialPayload(concreteTypeLayout, existentialTypeLayout))
-                {
-                    // If the value can fit in the payload area, then we will go ahead and copy
-                    // its bytes into that area.
-                    //
-                    auto valueSize = concreteTypeLayout->getSize();
-                    SLANG_RETURN_ON_FAIL(setData(payloadOffset, subObject->m_data, valueSize));
-                }
-                else
-                {
-                    // If the value cannot fit in the payload area, then we will pass a pointer
-                    // to the sub-object instead.
-                    //
-                    // Note: The Slang compiler does not currently emit code that handles the
-                    // pointer case, but that is the expected implementation for values
-                    // that do not fit into the fixed-size payload.
-                    //
-                    SLANG_RETURN_ON_FAIL(setData(payloadOffset, &subObject->m_data, sizeof(void*)));
-                }
-            }
-            break;
-        }
-
+        size = Math::Min(size, (size_t)m_data.getCount() - offset.uniformOffset);
+        memcpy((char*)m_data.getBuffer() + offset.uniformOffset, data, size);
         return SLANG_OK;
     }
     virtual SLANG_NO_THROW Result SLANG_MCALL
@@ -989,6 +892,31 @@ public:
         return SLANG_OK;
     }
     virtual SLANG_NO_THROW Result SLANG_MCALL
+        setObject(ShaderOffset const& offset, IShaderObject* object) override
+    {
+        SLANG_RETURN_ON_FAIL(Super::setObject(offset, object));
+
+        auto bindingRangeIndex = offset.bindingRangeIndex;
+        auto& bindingRange = getLayout()->m_bindingRanges[bindingRangeIndex];
+
+        CPUShaderObject* subObject = static_cast<CPUShaderObject*>(object);
+
+        switch (bindingRange.bindingType)
+        {
+        default:
+            {
+                void* bufferPtr = subObject->m_data.getBuffer();
+                SLANG_RETURN_ON_FAIL(setData(offset, &bufferPtr, sizeof(void*)));
+            }
+            break;
+        case slang::BindingType::ExistentialValue:
+        case slang::BindingType::RawBuffer:
+        case slang::BindingType::MutableRawBuffer:
+            break;
+        }
+        return SLANG_OK;
+    }
+    virtual SLANG_NO_THROW Result SLANG_MCALL
         setSampler(ShaderOffset const& offset, ISamplerState* sampler) override
     {
         SLANG_UNUSED(sampler);
@@ -1003,52 +931,7 @@ public:
         return SLANG_OK;
     }
 
-    // Appends all types that are used to specialize the element type of this shader object in `args` list.
-    virtual Result collectSpecializationArgs(ExtendedShaderObjectTypeList& args) override
-    {
-        // TODO: the logic here is a copy-paste of `GraphicsCommonShaderObject::collectSpecializationArgs`,
-        // consider moving the implementation to `ShaderObjectBase` and share the logic among different implementations.
-
-        auto& subObjectRanges = getLayout()->subObjectRanges;
-        // The following logic is built on the assumption that all fields that involve existential types (and
-        // therefore require specialization) will results in a sub-object range in the type layout.
-        // This allows us to simply scan the sub-object ranges to find out all specialization arguments.
-        for (Index subObjIndex = 0; subObjIndex < subObjectRanges.getCount(); subObjIndex++)
-        {
-            // Retrieve the corresponding binding range of the sub object.
-            auto bindingRange = getLayout()->m_bindingRanges[subObjectRanges[subObjIndex].bindingRangeIndex];
-            switch (bindingRange.bindingType)
-            {
-            case slang::BindingType::ExistentialValue:
-            {
-                // A binding type of `ExistentialValue` means the sub-object represents a interface-typed field.
-                // In this case the specialization argument for this field is the actual specialized type of the bound
-                // shader object. If the shader object's type is an ordinary type without existential fields, then the
-                // type argument will simply be the ordinary type. But if the sub object's type is itself a specialized
-                // type, we need to make sure to use that type as the specialization argument.
-
-                // TODO: need to implement the case where the field is an array of existential values.
-                ExtendedShaderObjectType specializedSubObjType;
-                SLANG_RETURN_ON_FAIL(m_objects[subObjIndex]->getSpecializedShaderObjectType(&specializedSubObjType));
-                args.add(specializedSubObjType);
-                break;
-            }
-            case slang::BindingType::ParameterBlock:
-            case slang::BindingType::ConstantBuffer:
-                // Currently we only handle the case where the field's type is
-                // `ParameterBlock<SomeStruct>` or `ConstantBuffer<SomeStruct>`, where `SomeStruct` is a struct type
-                // (not directly an interface type). In this case, we just recursively collect the specialization arguments
-                // from the bound sub object.
-                SLANG_RETURN_ON_FAIL(m_objects[subObjIndex]->collectSpecializationArgs(args));
-                // TODO: we need to handle the case where the field is of the form `ParameterBlock<IFoo>`. We should treat
-                // this case the same way as the `ExistentialValue` case here, but currently we lack a mechanism to distinguish
-                // the two scenarios.
-                break;
-            }
-            // TODO: need to handle another case where specialization happens on resources fields e.g. `StructuredBuffer<IFoo>`.
-        }
-        return SLANG_OK;
-    }
+    char* getDataBuffer() { return m_data.getBuffer(); }
 };
 
 class CPUEntryPointShaderObject : public CPUShaderObject
@@ -1173,8 +1056,8 @@ private:
         varyingInput.endGroupID.y = y;
         varyingInput.endGroupID.z = z;
 
-        auto globalParamsData = m_currentRootObject->m_data;
-        auto entryPointParamsData = entryPointObject->m_data;
+        auto globalParamsData = m_currentRootObject->getDataBuffer();
+        auto entryPointParamsData = entryPointObject->getDataBuffer();
         func(&varyingInput, entryPointParamsData, globalParamsData);
     }
 
@@ -1379,10 +1262,7 @@ SlangResult CPUShaderObject::init(IDevice* device, CPUShaderObjectLayout* typeLa
     //
     auto slangLayout = getLayout()->getElementTypeLayout();
     size_t uniformSize = slangLayout->getSize();
-    if (uniformSize)
-    {
-        m_data = malloc(uniformSize);
-    }
+    m_data.setCount(uniformSize);
 
     // If the layout specifies that we have any resources or sub-objects,
     // then we need to size the appropriate arrays to account for them.
@@ -1405,6 +1285,8 @@ SlangResult CPUShaderObject::init(IDevice* device, CPUShaderObjectLayout* typeLa
         //
         if (!subObjectLayout)
             continue;
+        auto _debugname = subObjectLayout->getElementTypeLayout()->getName();
+        
         //
         // Otherwise, we will allocate a sub-object to fill
         // in each entry in this range, based on the layout