Add per-buffer data layout control. (#3551)

* Add per-buffer data layout control.

Fixes #3534.

* Fixes.

* Robustness.

* Update test.

* Fix.

12 files changed, 156 insertions, 27 deletions

diff --git a/source/slang/hlsl.meta.slang b/source/slang/hlsl.meta.slang
index 1e0b95c81..a485c6514 100644
--- a/source/slang/hlsl.meta.slang
+++ b/source/slang/hlsl.meta.slang
@@ -2,24 +2,49 @@
 
 typedef uint UINT;
 
-__generic<T>
+[sealed]
+interface IBufferDataLayout
+{
+}
+
+__intrinsic_type($(kIROp_DefaultBufferLayoutType))
+struct DefaultDataLayout : IBufferDataLayout
+{};
+
+__intrinsic_type($(kIROp_Std140BufferLayoutType))
+[require(spirv)]
+[require(glsl)]
+struct Std140DataLayout : IBufferDataLayout
+{};
+
+__intrinsic_type($(kIROp_Std430BufferLayoutType))
+[require(spirv)]
+[require(glsl)]
+struct Std430DataLayout : IBufferDataLayout
+{};
+
+__intrinsic_type($(kIROp_ScalarBufferLayoutType))
+struct ScalarDataLayout : IBufferDataLayout
+{};
+
+__generic<T,L:IBufferDataLayout>
 __intrinsic_op($(kIROp_StructuredBufferGetDimensions))
-uint2 __structuredBufferGetDimensions(AppendStructuredBuffer<T> buffer);
+uint2 __structuredBufferGetDimensions(AppendStructuredBuffer<T,L> buffer);
 
-__generic<T>
+__generic<T,L:IBufferDataLayout>
 __intrinsic_op($(kIROp_StructuredBufferGetDimensions))
-uint2 __structuredBufferGetDimensions(ConsumeStructuredBuffer<T> buffer);
+uint2 __structuredBufferGetDimensions(ConsumeStructuredBuffer<T,L> buffer);
 
 __intrinsic_op($(kIROp_StructuredBufferGetDimensions))
-uint2 __structuredBufferGetDimensions<T>(StructuredBuffer<T> buffer);
+uint2 __structuredBufferGetDimensions<T,L:IBufferDataLayout>(StructuredBuffer<T,L> buffer);
 
 __intrinsic_op($(kIROp_StructuredBufferGetDimensions))
-uint2 __structuredBufferGetDimensions<T>(RWStructuredBuffer<T> buffer);
+uint2 __structuredBufferGetDimensions<T,L:IBufferDataLayout>(RWStructuredBuffer<T,L> buffer);
 
 __intrinsic_op($(kIROp_StructuredBufferGetDimensions))
-uint2 __structuredBufferGetDimensions<T>(RasterizerOrderedStructuredBuffer<T> buffer);
+uint2 __structuredBufferGetDimensions<T,L:IBufferDataLayout>(RasterizerOrderedStructuredBuffer<T,L> buffer);
 
-__generic<T>
+__generic<T, L:IBufferDataLayout=DefaultDataLayout>
 __magic_type(HLSLAppendStructuredBufferType)
 __intrinsic_type($(kIROp_HLSLAppendStructuredBufferType))
 struct AppendStructuredBuffer
@@ -2342,7 +2367,7 @@ void __byteAddressBufferStore<T>(RWByteAddressBuffer buffer, int offset, T value
 __intrinsic_op($(kIROp_ByteAddressBufferStore))
 void __byteAddressBufferStore<T>(RasterizerOrderedByteAddressBuffer buffer, int offset, T value);
 
-__generic<T>
+__generic<T, L:IBufferDataLayout=DefaultDataLayout>
 __magic_type(HLSLStructuredBufferType)
 __intrinsic_type($(kIROp_HLSLStructuredBufferType))
 struct StructuredBuffer
@@ -2375,7 +2400,7 @@ struct StructuredBuffer
     };
 };
 
-__generic<T>
+__generic<T, L:IBufferDataLayout=DefaultDataLayout>
 __magic_type(HLSLConsumeStructuredBufferType)
 __intrinsic_type($(kIROp_HLSLConsumeStructuredBufferType))
 struct ConsumeStructuredBuffer
@@ -3093,7 +3118,7 @@ static const struct {
 for(auto item : kMutableStructuredBufferCases) {
 }}}}
 
-__generic<T>
+__generic<T, L:IBufferDataLayout=DefaultDataLayout>
 __magic_type(HLSL$(item.name)Type)
 __intrinsic_type($(item.op))
 struct $(item.name)
diff --git a/source/slang/slang-emit-cpp.cpp b/source/slang/slang-emit-cpp.cpp
index 95a6ea4ff..58c635d90 100644
--- a/source/slang/slang-emit-cpp.cpp
+++ b/source/slang/slang-emit-cpp.cpp
@@ -352,6 +352,9 @@ SlangResult CPPSourceEmitter::calcTypeName(IRType* type, CodeGenTarget target, S
                     // Assumes ordering of types matches ordering of operands.
 
                     UInt operandCount = type->getOperandCount();
+                    if (as<IRHLSLStructuredBufferTypeBase>(type))
+                        operandCount = 1;
+
                     if (operandCount)
                     {
                         m_writer->emit("<");
diff --git a/source/slang/slang-emit-glsl.cpp b/source/slang/slang-emit-glsl.cpp
index 248b803f4..bae34b6fa 100644
--- a/source/slang/slang-emit-glsl.cpp
+++ b/source/slang/slang-emit-glsl.cpp
@@ -122,7 +122,28 @@ void GLSLSourceEmitter::_emitGLSLStructuredBuffer(IRGlobalParam* varDecl, IRHLSL
     _requireGLSLVersion(430);
 
     m_writer->emit("layout(");
-    m_writer->emit(getTargetReq()->getForceGLSLScalarBufferLayout() ? "scalar" : "std430");
+    auto layoutTypeOp = structuredBufferType->getDataLayout()
+        ? structuredBufferType->getDataLayout()->getOp()
+        : kIROp_DefaultBufferLayoutType;
+    switch (layoutTypeOp)
+    {
+    case kIROp_DefaultBufferLayoutType:
+        m_writer->emit(getTargetReq()->getForceGLSLScalarBufferLayout() ? "scalar" : "std430");
+        break;
+    case kIROp_Std430BufferLayoutType:
+        m_writer->emit("std430");
+        break;
+    case kIROp_Std140BufferLayoutType:
+        m_writer->emit("std140");
+        break;
+    case kIROp_ScalarBufferLayoutType:
+        _requireGLSLExtension(toSlice("GL_EXT_scalar_block_layout"));
+        m_writer->emit("scalar");
+        break;
+    default:
+        m_writer->emit("std430");
+        break;
+    }
 
     bool isReadOnly = (as<IRHLSLStructuredBufferType>(structuredBufferType) != nullptr);
     auto layout = getVarLayout(varDecl);
diff --git a/source/slang/slang-ir-byte-address-legalize.cpp b/source/slang/slang-ir-byte-address-legalize.cpp
index eec620133..fd2b599f8 100644
--- a/source/slang/slang-ir-byte-address-legalize.cpp
+++ b/source/slang/slang-ir-byte-address-legalize.cpp
@@ -775,13 +775,13 @@ struct ByteAddressBufferLegalizationContext
         // Our task in this function is to compute the type for
         // a structure buffer that is equivalent to `byteAddressBufferType`,
         // but with the given `elementType`.
-
+        IRInst* typeOperands[2] = { elementType, m_builder.getType(kIROp_DefaultBufferLayoutType) };
         switch( byteAddressBufferType->getOp() )
         {
             // The basic `*ByteAddressBuffer` types map directly to the `*StructuredBuffer<elementType>` cases.
-        case kIROp_HLSLByteAddressBufferType:                   return m_builder.getType(kIROp_HLSLStructuredBufferType, elementType);
-        case kIROp_HLSLRWByteAddressBufferType:                 return m_builder.getType(kIROp_HLSLRWStructuredBufferType, elementType);
-        case kIROp_HLSLRasterizerOrderedByteAddressBufferType:  return m_builder.getType(kIROp_HLSLRasterizerOrderedStructuredBufferType, elementType);
+        case kIROp_HLSLByteAddressBufferType:                   return m_builder.getType(kIROp_HLSLStructuredBufferType, 2, typeOperands);
+        case kIROp_HLSLRWByteAddressBufferType:                 return m_builder.getType(kIROp_HLSLRWStructuredBufferType, 2, typeOperands);
+        case kIROp_HLSLRasterizerOrderedByteAddressBufferType:  return m_builder.getType(kIROp_HLSLRasterizerOrderedStructuredBufferType, 2, typeOperands);
 
         case kIROp_ArrayType:
         case kIROp_UnsizedArrayType:
diff --git a/source/slang/slang-ir-inst-defs.h b/source/slang/slang-ir-inst-defs.h
index 23fcb3a92..e183058ac 100644
--- a/source/slang/slang-ir-inst-defs.h
+++ b/source/slang/slang-ir-inst-defs.h
@@ -133,6 +133,11 @@ INST(Nop, nop, 0, 0)
         INST(SamplerComparisonStateType, SamplerComparisonState, 0, HOISTABLE)
     INST_RANGE(SamplerStateTypeBase, SamplerStateType, SamplerComparisonStateType)
 
+    INST(DefaultBufferLayoutType, DefaultLayout, 0, HOISTABLE)
+    INST(Std140BufferLayoutType, Std140Layout, 0, HOISTABLE)
+    INST(Std430BufferLayoutType, Std430Layout, 0, HOISTABLE)
+    INST(ScalarBufferLayoutType, ScalarLayout, 0, HOISTABLE)
+
     INST(TextureFootprintType, TextureFootprintType, 1, HOISTABLE)
 
     INST(TextureShape1DType, TextureShape1DType, 0, HOISTABLE)
diff --git a/source/slang/slang-ir-layout.cpp b/source/slang/slang-ir-layout.cpp
index 9bbdb1aef..7c34cea8d 100644
--- a/source/slang/slang-ir-layout.cpp
+++ b/source/slang/slang-ir-layout.cpp
@@ -316,6 +316,12 @@ case kIROp_##TYPE##Type:                                        \
         return SLANG_OK;
     }
     break;
+    case kIROp_ScalarBufferLayoutType:
+    case kIROp_Std140BufferLayoutType:
+    case kIROp_Std430BufferLayoutType:
+    case kIROp_DefaultBufferLayoutType:
+        *outSizeAndAlignment = IRSizeAndAlignment(0, 4);
+        return SLANG_OK;
     default:
         break;
     }
@@ -324,6 +330,7 @@ case kIROp_##TYPE##Type:                                        \
         *outSizeAndAlignment = IRSizeAndAlignment(8, 8);
         return SLANG_OK;
     }
+
     return SLANG_FAIL;
 }
 
diff --git a/source/slang/slang-ir-lower-append-consume-structured-buffer.cpp b/source/slang/slang-ir-lower-append-consume-structured-buffer.cpp
index 5a7c7ee1c..1d4f10bb9 100644
--- a/source/slang/slang-ir-lower-append-consume-structured-buffer.cpp
+++ b/source/slang/slang-ir-lower-append-consume-structured-buffer.cpp
@@ -33,8 +33,12 @@ namespace Slang
 
         builder.addDecoration(elementBufferKey, kIROp_CounterBufferDecoration, counterBufferKey);
 
-        auto elementBufferType = builder.getType(kIROp_HLSLRWStructuredBufferType, elementType);
-        auto counterBufferType = builder.getType(kIROp_HLSLRWStructuredBufferType, builder.getIntType());
+        IRInst* operands[2] = { elementType, type->getDataLayout() };
+        auto elementBufferType = builder.getType(kIROp_HLSLRWStructuredBufferType, 2, operands);
+
+        operands[0] = builder.getIntType();
+        operands[1] = builder.getType(kIROp_DefaultBufferLayoutType);
+        auto counterBufferType = builder.getType(kIROp_HLSLRWStructuredBufferType, 2, operands);
 
         builder.createStructField(structType, elementBufferKey, elementBufferType);
         builder.createStructField(structType, counterBufferKey, counterBufferType);
diff --git a/source/slang/slang-ir-lower-buffer-element-type.cpp b/source/slang/slang-ir-lower-buffer-element-type.cpp
index 6cd8495b4..182b101df 100644
--- a/source/slang/slang-ir-lower-buffer-element-type.cpp
+++ b/source/slang/slang-ir-lower-buffer-element-type.cpp
@@ -477,7 +477,11 @@ namespace Slang
             {
                 IRBuilder builder(newElementType);
                 builder.setInsertAfter(newElementType);
-                return builder.getType(originalPtrLikeType->getOp(), newElementType);
+                ShortList<IRInst*> operands;
+                for (UInt i = 0; i < originalPtrLikeType->getOperandCount(); i++)
+                    operands.add(originalPtrLikeType->getOperand(i));
+                operands[0] = newElementType;
+                return builder.getType(originalPtrLikeType->getOp(), (UInt)operands.getCount(), operands.getArrayView().getBuffer());
             }
             SLANG_UNREACHABLE("unhandled ptr like or buffer type");
         }
@@ -538,9 +542,14 @@ namespace Slang
 
                 builder.setInsertBefore(bufferType);
 
+                ShortList<IRInst*> typeOperands;
+                for (UInt i = 0; i < bufferType->getOperandCount(); i++)
+                    typeOperands.add(bufferType->getOperand(i));
+                typeOperands[0] = loweredBufferElementTypeInfo.loweredType;
                 auto loweredBufferType = builder.getType(
                     bufferType->getOp(),
-                    loweredBufferElementTypeInfo.loweredType);
+                    (UInt)typeOperands.getCount(),
+                    typeOperands.getArrayView().getBuffer());
 
                 // We treat a value of a buffer type as a pointer, and use a work list to translate
                 // all loads and stores through the pointer values that needs lowering.
@@ -823,7 +832,24 @@ namespace Slang
         case kIROp_HLSLAppendStructuredBufferType:
         case kIROp_HLSLConsumeStructuredBufferType:
         case kIROp_HLSLRasterizerOrderedStructuredBufferType:
+        {
+            auto structBufferType = as<IRHLSLStructuredBufferTypeBase>(bufferType);
+            auto layoutTypeOp = structBufferType->getDataLayout()
+                ? structBufferType->getDataLayout()->getOp()
+                : kIROp_DefaultBufferLayoutType;
+            switch (layoutTypeOp)
+            {
+            case kIROp_DefaultBufferLayoutType:
+                return IRTypeLayoutRules::getStd430();
+            case kIROp_Std140BufferLayoutType:
+                return IRTypeLayoutRules::getStd140();
+            case kIROp_Std430BufferLayoutType:
+                return IRTypeLayoutRules::getStd430();
+            case kIROp_ScalarBufferLayoutType:
+                return IRTypeLayoutRules::getNatural();
+            }
             return IRTypeLayoutRules::getStd430();
+        }
         case kIROp_ConstantBufferType:
         case kIROp_ParameterBlockType:
             return IRTypeLayoutRules::getStd140();
diff --git a/source/slang/slang-ir-specialize.cpp b/source/slang/slang-ir-specialize.cpp
index 3f75a8b32..b82daa9a4 100644
--- a/source/slang/slang-ir-specialize.cpp
+++ b/source/slang/slang-ir-specialize.cpp
@@ -2080,11 +2080,16 @@ struct SpecializationContext
                 baseElementType,
                 slotOperandCount,
                 type->getExistentialArgs());
-
-            auto newPtrLikeType = builder.getType(
+            // Create a new type inst where the first operand is replaced
+            // with wrappedElementType.
+            ShortList<IRInst*> operands;
+            operands.add(wrappedElementType);
+            for (UInt i = 1; i < baseType->getOperandCount(); i++)
+                operands.add(baseType->getOperand(i));
+            IRInst* newPtrLikeType = builder.getType(
                 baseType->getOp(),
-                1,
-                &wrappedElementType);
+                operands.getCount(),
+                operands.getArrayView().getBuffer());
             addUsersToWorkList(type);
             addToWorkList(newPtrLikeType);
             addToWorkList(wrappedElementType);
diff --git a/source/slang/slang-ir.h b/source/slang/slang-ir.h
index 828d50c48..2193e0351 100644
--- a/source/slang/slang-ir.h
+++ b/source/slang/slang-ir.h
@@ -1466,7 +1466,14 @@ struct IRBuiltinGenericType : IRType
 };
 
 SIMPLE_IR_PARENT_TYPE(PointerLikeType, BuiltinGenericType);
-SIMPLE_IR_PARENT_TYPE(HLSLStructuredBufferTypeBase, BuiltinGenericType)
+
+struct IRHLSLStructuredBufferTypeBase : IRBuiltinGenericType
+{
+    IRType* getDataLayout() { return (IRType*)getOperand(1); }
+
+    IR_PARENT_ISA(HLSLStructuredBufferTypeBase)
+};
+
 SIMPLE_IR_TYPE(HLSLStructuredBufferType, HLSLStructuredBufferTypeBase)
 SIMPLE_IR_TYPE(HLSLRWStructuredBufferType, HLSLStructuredBufferTypeBase)
 SIMPLE_IR_TYPE(HLSLRasterizerOrderedStructuredBufferType, HLSLStructuredBufferTypeBase)
diff --git a/tests/doc/doc.slang.expected b/tests/doc/doc.slang.expected
index ed9892b33..1163e6485 100644
--- a/tests/doc/doc.slang.expected
+++ b/tests/doc/doc.slang.expected
@@ -120,7 +120,7 @@ void Hey::doAnotherThing(int a);
 # inputBuffer
 
 ```
-RWStructuredBuffer<int> inputBuffer
+RWStructuredBuffer<int, DefaultDataLayout> inputBuffer
 ```
 
 ## Description
@@ -330,7 +330,7 @@ An enum
 # outputBuffer
 
 ```
-RWStructuredBuffer<int> outputBuffer
+RWStructuredBuffer<int, DefaultDataLayout> outputBuffer
 ```
 
 ## Description
diff --git a/tests/spirv/scalar-structured-buffer.slang b/tests/spirv/scalar-structured-buffer.slang
new file mode 100644
index 000000000..383ead9f5
--- /dev/null
+++ b/tests/spirv/scalar-structured-buffer.slang
@@ -0,0 +1,26 @@
+//TEST(compute, vulkan):COMPARE_COMPUTE_EX(filecheck-buffer=CHECK): -vk -compute
+
+float getData<L:IBufferDataLayout>(StructuredBuffer<float3, L> p, int index)
+{
+    return p[index].x;
+}
+
+// GLSL: layout({{.*}}scalar{{.*}})
+
+//TEST_INPUT:set scalarBuffer = ubuffer(data=[1 2 3 4 5 6], stride=12)
+StructuredBuffer<float3, ScalarDataLayout> scalarBuffer;
+//TEST_INPUT:set outputBuffer = out ubuffer(data=[0 0 0 0 0 0], stride=12)
+RWStructuredBuffer<float3, ScalarDataLayout> outputBuffer;
+
+[numthreads(1,1,1)]
+void computeMain()
+{
+    // CHECK: 1
+    // CHECK: 1
+    // CHECK: 1
+    // CHECK: 4
+    // CHECK: 4
+    // CHECK: 4
+    outputBuffer[0] = float3(getData(scalarBuffer, 0));
+    outputBuffer[1] = float3(getData(scalarBuffer, 1));
+}