GLSL Passthrough support for SSBO types (#3446)

* GLSL Passthrough support for SSBO types

* GLSL Passthrough support for SSBO types

* Correctly apply glsl local size layout to entry points during lowering

* Test for glsl layout correctness

* typo

* Reflect GLSL SSBO as raw buffers

* Functional test for glsl ssbo

* Allow allow glsl for render tests

* Functional test for ssbo passthrough

* Functional test for ssbo passthrough with spirv-direct

* fix windows build error

---------

Co-authored-by: Yong He <yonghe@outlook.com>

12 files changed, 137 insertions, 25 deletions

diff --git a/source/slang/slang-ast-type.h b/source/slang/slang-ast-type.h
index 713446e93..1d2ebc566 100644
--- a/source/slang/slang-ast-type.h
+++ b/source/slang/slang-ast-type.h
@@ -383,7 +383,7 @@ class GLSLOutputParameterGroupType : public VaryingParameterGroupType
 };
 
 
-// type for GLLSL `buffer` blocks
+// type for GLSL `buffer` blocks
 class GLSLShaderStorageBufferType : public PointerLikeType
 {
     SLANG_AST_CLASS(GLSLShaderStorageBufferType)
diff --git a/source/slang/slang-emit-c-like.cpp b/source/slang/slang-emit-c-like.cpp
index 8aa550a03..05c525965 100644
--- a/source/slang/slang-emit-c-like.cpp
+++ b/source/slang/slang-emit-c-like.cpp
@@ -29,6 +29,7 @@
 
 #include "slang-emit-source-writer.h"
 #include "slang-mangled-lexer.h"
+#include "slang/slang-ir.h"
 #include <assert.h>
 
 namespace Slang {
@@ -3396,7 +3397,6 @@ void CLikeSourceEmitter::emitFuncDecorationsImpl(IRFunc* func)
     }
 }
 
-
 void CLikeSourceEmitter::emitStruct(IRStructType* structType)
 {
     // If the selected `struct` type is actually an intrinsic
diff --git a/source/slang/slang-emit-glsl.cpp b/source/slang/slang-emit-glsl.cpp
index 0fc0c855a..3d3f9687e 100644
--- a/source/slang/slang-emit-glsl.cpp
+++ b/source/slang/slang-emit-glsl.cpp
@@ -9,6 +9,7 @@
 
 #include "slang-legalize-types.h"
 #include "slang-ir-layout.h"
+#include "slang/slang-ir.h"
 
 #include <assert.h>
 
@@ -229,8 +230,7 @@ void GLSLSourceEmitter::_emitGLSLStructuredBuffer(IRGlobalParam* varDecl, IRHLSL
     m_writer->emit(";\n");
 }
 
-
-void GLSLSourceEmitter::_emitGLSLByteAddressBuffer(IRGlobalParam* varDecl, IRByteAddressBufferTypeBase* byteAddressBufferType)
+void GLSLSourceEmitter::emitSSBOHeader(IRGlobalParam* varDecl, IRType* bufferType)
 {
     // TODO: A lot of this logic is copy-pasted from `emitIRStructuredBuffer_GLSL`.
     // It might be worthwhile to share the common code to avoid regressions sneaking
@@ -247,9 +247,9 @@ void GLSLSourceEmitter::_emitGLSLByteAddressBuffer(IRGlobalParam* varDecl, IRByt
     auto layout = getVarLayout(varDecl);
     if (layout)
     {
-        // We can use ShaderResource/DescriptorSlot interchangably here. 
+        // We can use ShaderResource/DescriptorSlot interchangably here.
         // This is possible because vk-shift-*
-        bool isReadOnly = (as<IRHLSLByteAddressBufferType>(byteAddressBufferType) != nullptr);
+        bool isReadOnly = (as<IRHLSLByteAddressBufferType>(bufferType) != nullptr);
 
         const LayoutResourceKindFlags kinds = (isReadOnly ? LayoutResourceKindFlag::ShaderResource : LayoutResourceKindFlag::UnorderedAccess)
             | LayoutResourceKindFlag::DescriptorTableSlot;
@@ -283,12 +283,17 @@ void GLSLSourceEmitter::_emitGLSLByteAddressBuffer(IRGlobalParam* varDecl, IRByt
     HLSLRasterizerOrderedByteAddressBufferType  - Allows read/write access
     */
 
-    if (as<IRHLSLByteAddressBufferType>(byteAddressBufferType))
+    if (as<IRHLSLByteAddressBufferType>(bufferType))
     {
         m_writer->emit("readonly ");
     }
 
     m_writer->emit("buffer ");
+}
+
+void GLSLSourceEmitter::_emitGLSLByteAddressBuffer(IRGlobalParam* varDecl, IRByteAddressBufferTypeBase* byteAddressBufferType)
+{
+    emitSSBOHeader(varDecl, byteAddressBufferType);
 
     // Generate a dummy name for the block
     m_writer->emit("_S");
@@ -307,6 +312,21 @@ void GLSLSourceEmitter::_emitGLSLByteAddressBuffer(IRGlobalParam* varDecl, IRByt
     m_writer->emit(";\n");
 }
 
+void GLSLSourceEmitter::_emitGLSLSSBO(IRGlobalParam* varDecl, IRGLSLShaderStorageBufferType* ssboType)
+{
+    emitSSBOHeader(varDecl, ssboType);
+
+    const auto structType = cast<IRStructType>(ssboType->getOperand(0));
+    m_writer->emit(getName(structType));
+    m_writer->emit("_Block");
+    emitStructDeclarationsBlock(structType, true);
+
+    m_writer->emit(getName(varDecl));
+    emitArrayBrackets(varDecl->getDataType());
+
+    m_writer->emit(";\n");
+}
+
 void GLSLSourceEmitter::_emitGLSLParameterGroup(IRGlobalParam* varDecl, IRUniformParameterGroupType* type)
 {
     auto varLayout = getVarLayout(varDecl);
@@ -1282,6 +1302,11 @@ bool GLSLSourceEmitter::tryEmitGlobalParamImpl(IRGlobalParam* varDecl, IRType* v
         _emitGLSLByteAddressBuffer(varDecl, byteAddressBufferType);
         return true;
     }
+    else if (const auto glslSSBOType = as<IRGLSLShaderStorageBufferType>(unwrapArray(varType)))
+    {
+        _emitGLSLSSBO(varDecl, glslSSBOType);
+        return true;
+    }
 
     // We want to skip the declaration of any system-value variables
     // when outputting GLSL (well, except in the case where they
@@ -1603,6 +1628,25 @@ void GLSLSourceEmitter::emitBufferPointerTypeDefinition(IRInst* ptrType)
     m_writer->emit("};\n");
 }
 
+// Is this type only used by SSBO declarations, if so then we don't need to
+// emit it and it'll be emitted inline there.
+static bool isSSBOInternalStructType(IRInst* inst)
+{
+    if(!as<IRStructType>(inst))
+        return false;
+
+    bool onlySSBOUses = true;
+    for(auto use = inst->firstUse; use; use = use->nextUse)
+    {
+        if(!as<IRGLSLShaderStorageBufferType>(use->user))
+        {
+            onlySSBOUses = false;
+            break;
+        }
+    }
+    return onlySSBOUses;
+}
+
 void GLSLSourceEmitter::emitGlobalInstImpl(IRInst* inst)
 {
     switch (inst->getOp())
@@ -1610,6 +1654,11 @@ void GLSLSourceEmitter::emitGlobalInstImpl(IRInst* inst)
     case kIROp_HLSLConstBufferPointerType:
         emitBufferPointerTypeDefinition(inst);
         break;
+    // No need to use structs which are just taking part in a SSBO declaration
+    case kIROp_StructType:
+        if(isSSBOInternalStructType(inst))
+            break;
+        [[fallthrough]];
     default:
         Super::emitGlobalInstImpl(inst);
         break;
diff --git a/source/slang/slang-emit-glsl.h b/source/slang/slang-emit-glsl.h
index f15ea03b4..1a95cd346 100644
--- a/source/slang/slang-emit-glsl.h
+++ b/source/slang/slang-emit-glsl.h
@@ -61,6 +61,8 @@ protected:
 
     void _emitGLSLByteAddressBuffer(IRGlobalParam* varDecl, IRByteAddressBufferTypeBase* byteAddressBufferType);
     void _emitGLSLParameterGroup(IRGlobalParam* varDecl, IRUniformParameterGroupType* type);
+    void _emitGLSLSSBO(IRGlobalParam* varDecl, IRGLSLShaderStorageBufferType* ssboType);
+    void emitSSBOHeader(IRGlobalParam* varDecl, IRType* bufferType);
 
     void _emitGLSLPerVertexVaryingFragmentInput(IRGlobalParam* param, IRType* type);
 
diff --git a/source/slang/slang-emit.cpp b/source/slang/slang-emit.cpp
index 7ad4ad5d4..311fce83d 100644
--- a/source/slang/slang-emit.cpp
+++ b/source/slang/slang-emit.cpp
@@ -242,7 +242,8 @@ Result linkAndOptimizeIR(
     // un-specialized IR.
     dumpIRIfEnabled(codeGenContext, irModule);
 
-    lowerGLSLShaderStorageBufferObjects(irModule, sink);
+    if(!isKhronosTarget(targetRequest))
+        lowerGLSLShaderStorageBufferObjectsToStructuredBuffers(irModule, sink);
 
     translateGLSLGlobalVar(codeGenContext, irModule);
 
diff --git a/source/slang/slang-ir-layout.cpp b/source/slang/slang-ir-layout.cpp
index a6684c7b1..9bbdb1aef 100644
--- a/source/slang/slang-ir-layout.cpp
+++ b/source/slang/slang-ir-layout.cpp
@@ -152,10 +152,18 @@ case kIROp_##TYPE##Type:                                        \
         auto structType = cast<IRStructType>(type);
         IRSizeAndAlignment structLayout;
         IRIntegerValue offset = 0;
+        bool seenFinalUnsizedArrayField = false;
         for (auto field : structType->getFields())
         {
+            // If we failed to catch an unsized array earlier in the pipeline,
+            // this will pick it up before generating nonsense results for
+            // subsequent offsets
+            SLANG_ASSERT(!seenFinalUnsizedArrayField);
+
             IRSizeAndAlignment fieldTypeLayout;
             SLANG_RETURN_ON_FAIL(getSizeAndAlignment(target, rules, field->getFieldType(), &fieldTypeLayout));
+            seenFinalUnsizedArrayField = fieldTypeLayout.size == IRSizeAndAlignment::kIndeterminateSize;
+
             structLayout.size = align(offset, fieldTypeLayout.alignment);
             structLayout.alignment = std::max(structLayout.alignment, fieldTypeLayout.alignment);
 
@@ -202,6 +210,15 @@ case kIROp_##TYPE##Type:                                        \
     }
     break;
 
+    case kIROp_UnsizedArrayType:
+    {
+        auto unsizedArrayType = cast<IRUnsizedArrayType>(type);
+        getSizeAndAlignment(target, rules, unsizedArrayType->getElementType(), outSizeAndAlignment);
+        outSizeAndAlignment->size = IRSizeAndAlignment::kIndeterminateSize;
+        return SLANG_OK;
+    }
+    break;
+
     case kIROp_VectorType:
     {
         auto vecType = cast<IRVectorType>(type);
diff --git a/source/slang/slang-ir-layout.h b/source/slang/slang-ir-layout.h
index b5eb54497..19360aa20 100644
--- a/source/slang/slang-ir-layout.h
+++ b/source/slang/slang-ir-layout.h
@@ -17,6 +17,7 @@
 //
 
 #include "slang-ir.h"
+#include <limits>
 
 namespace Slang
 {
@@ -40,9 +41,12 @@ struct IRSizeAndAlignment
         , alignment(alignment)
     {}
 
-    IRIntegerValue  size = 0;
+    IRIntegerValue size = 0;
 
-    int             alignment = 1;
+    // Used for an unsized array as a final struct member
+    const static IRIntegerValue kIndeterminateSize = std::numeric_limits<IRIntegerValue>::min();
+
+    int alignment = 1;
 
     inline IRIntegerValue getStride()
     {
diff --git a/source/slang/slang-ir-lower-glsl-ssbo-types.cpp b/source/slang/slang-ir-lower-glsl-ssbo-types.cpp
index 010c898ad..f0bdb7e83 100644
--- a/source/slang/slang-ir-lower-glsl-ssbo-types.cpp
+++ b/source/slang/slang-ir-lower-glsl-ssbo-types.cpp
@@ -188,6 +188,20 @@ namespace Slang
         }
     }
 
+    static void lowerSSBOsToPointers(IRModule* module)
+    {
+        IRBuilder builder{module};
+        overAllInsts(module, [&](IRInst* inst){
+            if(const auto ssbo = as<IRGLSLShaderStorageBufferType>(inst))
+            {
+                builder.setInsertAfter(ssbo);
+                const auto ptr = builder.getPtrType(ssbo->getElementType());
+                ssbo->replaceUsesWith(ptr);
+                ssbo->removeAndDeallocate();
+            }
+        });
+    }
+
     static void diagnoseRemainingSSBOs(IRModule* module, DiagnosticSink* sink)
     {
         overAllInsts(module, [&](IRInst* inst){
@@ -198,10 +212,16 @@ namespace Slang
         });
     }
 
-    void lowerGLSLShaderStorageBufferObjects(IRModule* module, DiagnosticSink* sink)
+    void lowerGLSLShaderStorageBufferObjectsToStructuredBuffers(IRModule* module, DiagnosticSink* sink)
     {
         lowerArrayLikeSSBOs(module);
         lowerStructLikeSSBOs(module);
         diagnoseRemainingSSBOs(module, sink);
     }
+
+    void lowerGLSLShaderStorageBufferObjectsToPointers(IRModule* module, DiagnosticSink* sink)
+    {
+        lowerSSBOsToPointers(module);
+        diagnoseRemainingSSBOs(module, sink);
+    }
 }
diff --git a/source/slang/slang-ir-lower-glsl-ssbo-types.h b/source/slang/slang-ir-lower-glsl-ssbo-types.h
index fb8ea3917..532c9e115 100644
--- a/source/slang/slang-ir-lower-glsl-ssbo-types.h
+++ b/source/slang/slang-ir-lower-glsl-ssbo-types.h
@@ -7,8 +7,13 @@ namespace Slang
     struct IRModule;
     class DiagnosticSink;
 
-    /// Lower tuple types to ordinary `struct`s.
-    void lowerGLSLShaderStorageBufferObjects(
+    void lowerGLSLShaderStorageBufferObjectsToPointers(
         IRModule* module,
-        DiagnosticSink* sink);
+        DiagnosticSink* sink
+    );
+
+    void lowerGLSLShaderStorageBufferObjectsToStructuredBuffers(
+        IRModule* module,
+        DiagnosticSink* sink
+    );
 }
diff --git a/source/slang/slang-ir-spirv-legalize.cpp b/source/slang/slang-ir-spirv-legalize.cpp
index 681fb36b1..1675fe279 100644
--- a/source/slang/slang-ir-spirv-legalize.cpp
+++ b/source/slang/slang-ir-spirv-legalize.cpp
@@ -614,14 +614,18 @@ struct SPIRVLegalizationContext : public SourceEmitterBase
                     varLayoutInst->removeAndDeallocate();
                 }
             }
-            else
+            else if (auto structuredBufferType = as<IRHLSLStructuredBufferTypeBase>(innerType))
             {
-                if (auto structuredBufferType = as<IRHLSLStructuredBufferTypeBase>(innerType))
-                {
-                    innerType = lowerStructuredBufferType(structuredBufferType).structType;
-                    storageClass = SpvStorageClassStorageBuffer;
-                    needLoad = false;
-                }
+                innerType = lowerStructuredBufferType(structuredBufferType).structType;
+                storageClass = SpvStorageClassStorageBuffer;
+                needLoad = false;
+            }
+            else if (auto glslShaderStorageBufferType = as<IRGLSLShaderStorageBufferType>(innerType))
+            {
+                innerType = glslShaderStorageBufferType->getElementType();
+                builder.addDecoration(innerType, kIROp_SPIRVBlockDecoration);
+                storageClass = SpvStorageClassStorageBuffer;
+                needLoad = false;
             }
 
             auto innerElementType = innerType;
diff --git a/source/slang/slang-reflection-api.cpp b/source/slang/slang-reflection-api.cpp
index f0540deb3..28b22b833 100644
--- a/source/slang/slang-reflection-api.cpp
+++ b/source/slang/slang-reflection-api.cpp
@@ -168,6 +168,11 @@ static SlangParameterCategory maybeRemapParameterCategory(
                     return SLANG_PARAMETER_CATEGORY_DESCRIPTOR_TABLE_SLOT;
                 break;
 
+            case SLANG_TYPE_KIND_SHADER_STORAGE_BUFFER:
+                if (category == SLANG_PARAMETER_CATEGORY_UNIFORM)
+                    return SLANG_PARAMETER_CATEGORY_DESCRIPTOR_TABLE_SLOT;
+                break;
+
                 // TODO: implement more helpers here
 
             default:
@@ -382,6 +387,7 @@ SLANG_API SlangTypeKind spReflectionType_GetKind(SlangReflectionType* inType)
     CASE(HLSLRWByteAddressBufferType);
     CASE(HLSLRasterizerOrderedByteAddressBufferType);
     CASE(UntypedBufferResourceType);
+    CASE(GLSLShaderStorageBufferType);
 #undef CASE
 
     else if (const auto arrayType = as<ArrayExpressionType>(type))
@@ -668,6 +674,7 @@ SLANG_API SlangResourceShape spReflectionType_GetResourceShape(SlangReflectionTy
     CASE(HLSLRasterizerOrderedByteAddressBufferType,    SLANG_BYTE_ADDRESS_BUFFER,      SLANG_RESOURCE_ACCESS_RASTER_ORDERED);
     CASE(RaytracingAccelerationStructureType,           SLANG_ACCELERATION_STRUCTURE,   SLANG_RESOURCE_ACCESS_READ);
     CASE(UntypedBufferResourceType,                     SLANG_BYTE_ADDRESS_BUFFER,      SLANG_RESOURCE_ACCESS_READ);
+    CASE(GLSLShaderStorageBufferType,                   SLANG_BYTE_ADDRESS_BUFFER,      SLANG_RESOURCE_ACCESS_READ_WRITE);
 #undef CASE
 
     return SLANG_RESOURCE_NONE;
@@ -703,9 +710,7 @@ SLANG_API SlangResourceAccess spReflectionType_GetResourceAccess(SlangReflection
     CASE(HLSLRWByteAddressBufferType,                   SLANG_BYTE_ADDRESS_BUFFER,  SLANG_RESOURCE_ACCESS_READ_WRITE);
     CASE(HLSLRasterizerOrderedByteAddressBufferType,    SLANG_BYTE_ADDRESS_BUFFER,  SLANG_RESOURCE_ACCESS_RASTER_ORDERED);
     CASE(UntypedBufferResourceType,                     SLANG_BYTE_ADDRESS_BUFFER,  SLANG_RESOURCE_ACCESS_READ);
-
-    // This isn't entirely accurate, but I can live with it for now
-    CASE(GLSLShaderStorageBufferType, SLANG_STRUCTURED_BUFFER, SLANG_RESOURCE_ACCESS_READ_WRITE);
+    CASE(GLSLShaderStorageBufferType,                   SLANG_BYTE_ADDRESS_BUFFER,  SLANG_RESOURCE_ACCESS_READ_WRITE);
 #undef CASE
 
     return SLANG_RESOURCE_ACCESS_NONE;
@@ -1310,6 +1315,11 @@ namespace Slang
                 return SLANG_BINDING_TYPE_MUTABLE_RAW_BUFFER;
             }
         }
+        else if( as<GLSLShaderStorageBufferType>(type) )
+        {
+            // TODO Immutable buffers
+            return SLANG_BINDING_TYPE_MUTABLE_RAW_BUFFER;
+        }
         else if( as<ConstantBufferType>(type) )
         {
             return SLANG_BINDING_TYPE_CONSTANT_BUFFER;
diff --git a/source/slang/slang-type-layout.cpp b/source/slang/slang-type-layout.cpp
index 074aa3350..ccd34a966 100644
--- a/source/slang/slang-type-layout.cpp
+++ b/source/slang/slang-type-layout.cpp
@@ -3905,7 +3905,7 @@ static TypeLayoutResult _createTypeLayout(
     // This case is mostly to allow users to add new resource types...
     CASE(UntypedBufferResourceType,         RawBuffer);
 
-    CASE(GLSLShaderStorageBufferType,       MutableBuffer);
+    CASE(GLSLShaderStorageBufferType,       MutableRawBuffer);
 
 #undef CASE