[WGSL] Enable arbitrary arrays in uniform buffers. (shader-slang#5497)

* [WGSL] Enable arbitrary arrays in uniform buffers.

* format code

* Undo irrelevant change and fixups.

* Update expected failure list.

* Fix.

* Rename.

---------

Co-authored-by: slangbot <186143334+slangbot@users.noreply.github.com>

Author: csyonghe

source/slang/slang-emit.cpp

@@ -1455,7 +1455,10 @@ Result linkAndOptimizeIR(
     if (requiredLoweringPassSet.meshOutput)
         legalizeMeshOutputTypes(irModule);
 
-    lowerBufferElementTypeToStorageType(targetProgram, irModule);
+    BufferElementTypeLoweringOptions bufferElementTypeLoweringOptions;
+    bufferElementTypeLoweringOptions.use16ByteArrayElementForConstantBuffer =
+        isWGPUTarget(targetRequest);
+    lowerBufferElementTypeToStorageType(targetProgram, irModule, bufferElementTypeLoweringOptions);
 
     // Rewrite functions that return arrays to return them via `out` parameter,
     // since our target languages doesn't allow returning arrays.

source/slang/slang-ir-insts.h

@@ -3566,6 +3566,7 @@ struct IRBuilder
     IRInst* replaceOperand(IRUse* use, IRInst* newValue);
 
     IRInst* getBoolValue(bool value);
+    IRInst* getIntValue(IRIntegerValue value);
     IRInst* getIntValue(IRType* type, IRIntegerValue value);
     IRInst* getFloatValue(IRType* type, IRFloatingPointValue value);
     IRStringLit* getStringValue(const UnownedStringSlice& slice);

source/slang/slang-ir-lower-buffer-element-type.cpp

@@ -8,6 +8,14 @@ class TargetProgram;
 struct IRTypeLayoutRules;
 struct IRType;
 
+struct BufferElementTypeLoweringOptions
+{
+    bool lowerBufferPointer = false;
+
+    // For WGSL, we can only create arrays that has a stride of 16 bytes for constant buffers.
+    bool use16ByteArrayElementForConstantBuffer = false;
+};
+
 // For each struct type S used as element type of a ConstantBuffer, ParameterBlock or
 // [RW]StructuredBuffer, we create a lowered type L, where matrix types are lowered to arrays of
 // vectors based on major-ness, and loads from the buffer are converted to L_to_S(load(buffer)), and
@@ -18,7 +26,7 @@ struct IRType;
 void lowerBufferElementTypeToStorageType(
     TargetProgram* target,
     IRModule* module,
-    bool lowerBufferPointer = false);
+    BufferElementTypeLoweringOptions options = BufferElementTypeLoweringOptions());
 
 
 // Returns the type layout rules should be used for a buffer resource type.

source/slang/slang-ir-lower-buffer-element-type.h

@@ -2205,7 +2205,12 @@ struct SPIRVLegalizationContext : public SourceEmitterBase
         // pointers in this pass. In the future we should consider separate out IRAddress as
         // the type for IRVar, and use IRPtrType to dedicate pointers in user code, so we can
         // safely lower the pointer load stores early together with other buffer types.
-        lowerBufferElementTypeToStorageType(m_sharedContext->m_targetProgram, m_module, true);
+        BufferElementTypeLoweringOptions bufferElementTypeLoweringOptions;
+        bufferElementTypeLoweringOptions.lowerBufferPointer = true;
+        lowerBufferElementTypeToStorageType(
+            m_sharedContext->m_targetProgram,
+            m_module,
+            bufferElementTypeLoweringOptions);
 
         // The above step may produce empty struct types, so we need to lower them out of existence.
         legalizeEmptyTypes(m_sharedContext->m_targetProgram, m_module, m_sink);

source/slang/slang-ir-spirv-legalize.cpp

@@ -2264,6 +2264,11 @@ IRInst* IRBuilder::getBoolValue(bool inValue)
     return _findOrEmitConstant(keyInst);
 }
 
+IRInst* IRBuilder::getIntValue(IRIntegerValue value)
+{
+    return getIntValue(getIntType(), value);
+}
+
 IRInst* IRBuilder::getIntValue(IRType* type, IRIntegerValue inValue)
 {
     IRConstant keyInst;

source/slang/slang-ir.cpp

@@ -3,14 +3,13 @@
 // Note! This test doesn't work on CUDA or CPU targets, because both these targets 
 // assume matrices are tightly packed, whereas GPU targets align rows to 16 bytes.
 
+//TEST(compute):COMPARE_COMPUTE_EX(filecheck-buffer=ALIGNED):-slang -compute -wgpu -output-using-type -xslang -matrix-layout-row-major -shaderobj
 //TEST(compute):COMPARE_COMPUTE_EX(filecheck-buffer=PACKED):-cpu -compute -output-using-type -compile-arg -O3 -xslang -matrix-layout-row-major -shaderobj
 //TEST(compute):COMPARE_COMPUTE_EX(filecheck-buffer=ALIGNED):-slang -compute -output-using-type -xslang -matrix-layout-row-major -shaderobj
 //TEST(compute):COMPARE_COMPUTE_EX(filecheck-buffer=ALIGNED):-slang -compute -dx12 -output-using-type -xslang -matrix-layout-row-major -shaderobj
 //TEST(compute, vulkan):COMPARE_COMPUTE_EX(filecheck-buffer=ALIGNED):-vk -compute -output-using-type -xslang -matrix-layout-row-major -shaderobj
 //TEST(compute):COMPARE_COMPUTE_EX(filecheck-buffer=PACKED):-cuda -compute -output-using-type -xslang -matrix-layout-row-major -shaderobj
 //TEST(compute):COMPARE_COMPUTE(filecheck-buffer=PACKED): -slang -output-using-type -shaderobj -mtl
-//DISABLE_TEST(compute):COMPARE_COMPUTE:-wgpu
-
 
 // matrix<R, C>
 //TEST_INPUT:cbuffer(data=[1.0 2.0 3.0 4.0 5.0 6.0 0.0 0.0  10.0 20.0 0.0 0.0 ]):name matrixBuffer

tests/compute/non-square-row-major.slang

@@ -55,7 +55,6 @@ tests/compute/interface-shader-param-in-struct.slang.4 syn (wgpu)
 tests/compute/interface-shader-param.slang.5 syn (wgpu)
 tests/compute/kernel-context-threading.slang.6 syn (wgpu)
 tests/compute/loop-unroll.slang.7 syn (wgpu)
-tests/compute/non-square-row-major.slang.6 syn (wgpu)
 tests/compute/parameter-block (wgpu)
 tests/compute/texture-get-dimensions (wgpu)
 tests/compute/texture-sampling (wgpu)

tests/expected-failure-github.txt

@@ -0,0 +1,23 @@
+//TEST(compute):COMPARE_COMPUTE_EX(filecheck-buffer=CHECK):-slang -compute -wgpu -output-using-type
+
+struct Params{
+    int2 data[2];
+    int scalarData[2];
+}
+
+//TEST_INPUT:set params = cbuffer(data=[0 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15])
+ConstantBuffer<Params> params;
+
+//TEST_INPUT:set outputBuffer = out ubuffer(data=[0 0 0 0], stride=4)
+RWStructuredBuffer<float> outputBuffer;
+
+float helper(float2 v)
+{return v.y;}
+
+[numthreads(1,1,1)]
+void computeMain()
+{
+    // CHECK: 17
+    outputBuffer[0] = helper(params.data[1]) + params.scalarData[1];
+}
+

tests/wgsl/uniform-array-2.slang

@@ -0,0 +1,21 @@
+//TEST(compute):COMPARE_COMPUTE_EX(filecheck-buffer=CHECK):-slang -compute -wgpu -output-using-type
+
+struct MyPoint { int x; int y; int type; }
+
+struct Params{
+    MyPoint data[2];
+}
+
+//TEST_INPUT:set params = cbuffer(data=[0 1 2 3 4 5 6 7 8 9 10 11])
+ConstantBuffer<Params> params;
+
+//TEST_INPUT:set outputBuffer = out ubuffer(data=[0 0 0 0], stride=4)
+RWStructuredBuffer<float> outputBuffer;
+
+[numthreads(1,1,1)]
+void computeMain()
+{
+    // CHECK: 5
+    outputBuffer[0] = params.data[1].y;
+}
+