support specialization constant sized array (#6871)

Close #6859

Goal of this PR
We want to support an array whose size can be specialization constant for shared/global variable e.g.

layout (constant_id = 0) const uint BLOCK_SIZE = 64;
shared float buf_a[(BLOCK_SIZE + 5) * 4];
Overview of the solution:

During IndexExpr check, we will loose the restriction to allow SpecConst passing, but the size parameter will not be a constant value because it cannot be folded into a constant, so we will make it follow the same logic as generic parameter value, and the size will be represented by FuncCallIntVal/PolynomialIntVal/DeclRefIntVal.

During IR lowering, we will detect whether there is spec constant in the IntVal, and wrap the IRInst with a SpecConstRateType, and propagate the type though the lowering logic, such that the IntVal representing the array size will have SpecConstRateType.

During spirv emit stage, if we detect that a IRInst has SpecConstRateType, we will emit it as SpecConstantOp.

We have to implement new logic to emit OpSpecConstantOp, the existing emit logic doesn't support emitting OpSpecConstantOp, especially this op can embed arithmetic operation at global scope, where we can only emit arithmetic instruct at local. But there are only few instructs we need to support.

Overview of the solution:
This PR doesn't support generic, and we will create a separate PR to extend that, tracked in #6840.

4 files changed, 186 insertions, 0 deletions

diff --git a/tests/spirv/spec-constant-sized-array-1.slang b/tests/spirv/spec-constant-sized-array-1.slang
new file mode 100644
index 000000000..e974fc37e
--- /dev/null
+++ b/tests/spirv/spec-constant-sized-array-1.slang
@@ -0,0 +1,42 @@
+//TEST:SIMPLE(filecheck=CHECK): -target spirv
+//TEST(compute, vulkan):COMPARE_COMPUTE(filecheck-buffer=BUF):-vk -output-using-type
+
+// CHECK: %[[C0:[0-9A-Za-z_]+]] = OpSpecConstant %int 5
+// CHECK: %[[I3:[0-9A-Za-z_]+]] = OpConstant %int 3
+// CHECK: %[[COP0:[0-9A-Za-z_]+]] = OpSpecConstantOp %int SRem %[[C0]] %[[I3]]
+// CHECK: %[[I2:[0-9A-Za-z_]+]] = OpConstant %int 2
+// CHECK: %[[COP1:[0-9A-Za-z_]+]] = OpSpecConstantOp %int IMul %[[I2]] %[[C0]]
+// CHECK: %[[COP2:[0-9A-Za-z_]+]] = OpSpecConstantOp %int IAdd %[[COP0]] %[[COP1]]
+// CHECK: %[[ARR_TYPE:[0-9A-Za-z_]+]] = OpTypeArray %float %[[COP2]]
+// CHECK: %[[PT_TYPE:[0-9A-Za-z_]+]] = OpTypePointer Workgroup %[[ARR_TYPE]]
+// CHECK: OpVariable %[[PT_TYPE]] Workgroup
+
+
+[SpecializationConstant]
+const int constValue0 = 5;
+
+groupshared float buffer[constValue0 * 2 + constValue0 % 3];
+
+//TEST_INPUT:ubuffer(data=[0], stride=4):out,name=outputBuffer
+RWStructuredBuffer<float> outputBuffer;
+
+[shader("compute")]
+[numthreads(12, 1, 1)]
+void computeMain(uint3 id : SV_GroupThreadID)
+{
+    buffer[id.x] = id.x;
+
+    GroupMemoryBarrier();
+    if (id.x == 0)
+    {
+        static const int size = constValue0 * 2 + constValue0 % 3;
+        float temp = 0;
+        for (uint i = 0; i < size; i++)
+        {
+            temp = temp + buffer[i];
+        }
+        outputBuffer[0] = temp;
+    }
+    // Result will be (0 + size-1) * size / 2
+    // BUF: 66
+}
diff --git a/tests/spirv/spec-constant-sized-array-2.slang b/tests/spirv/spec-constant-sized-array-2.slang
new file mode 100644
index 000000000..95d2f4aee
--- /dev/null
+++ b/tests/spirv/spec-constant-sized-array-2.slang
@@ -0,0 +1,48 @@
+//TEST:SIMPLE(filecheck=CHECK): -target spirv
+//TEST(compute, vulkan):COMPARE_COMPUTE(filecheck-buffer=BUF):-vk -output-using-type
+
+// CHECK: %[[C0:[0-9A-Za-z_]+]] = OpSpecConstant %int 32
+// CHECK: %[[I2:[0-9A-Za-z_]+]] = OpConstant %int 2
+// CHECK: %[[COP0:[0-9A-Za-z_]+]] = OpSpecConstantOp %int SDiv %[[C0]] %[[I2]]
+// CHECK: %[[COP1:[0-9A-Za-z_]+]] = OpSpecConstantOp %int ShiftRightArithmetic %[[C0]] %[[I2]]
+// CHECK: %[[COP2:[0-9A-Za-z_]+]] = OpSpecConstantOp %int IAdd %[[COP0]] %[[COP1]]
+// CHECK: %[[ARR_TYPE:[0-9A-Za-z_]+]] = OpTypeArray %float %[[COP2]]
+// CHECK: %[[PT_TYPE:[0-9A-Za-z_]+]] = OpTypePointer Function %[[ARR_TYPE]]
+
+[SpecializationConstant]
+const int constValue0 = 32;
+
+//TEST_INPUT:ubuffer(data=[0 0], stride=4):out,name=outputBuffer
+RWStructuredBuffer<float> outputBuffer;
+
+static float g_buffer[constValue0 / 2 + (constValue0 >> 2)];
+
+[shader("compute")]
+[numthreads(1, 1, 1)]
+void computeMain()
+{
+    float buffer[constValue0 / 2 + (constValue0 >> 2)];
+    // CHECK: OpVariable %[[PT_TYPE]] Function
+
+    static const int size = constValue0 / 2 + (constValue0 >> 2);   // 16 + 8 = 24
+    for (uint i = 0; i < size; i++)
+    {
+        buffer[i] = i;
+        g_buffer[i] = i * 2;
+    }
+
+    float temp1 = 0.0f;
+    float temp2 = 0.0f;
+    for (uint i = 0; i < size; i++)
+    {
+        temp1 += buffer[i] * 2;
+        temp2 += g_buffer[i];
+    }
+
+    // Result will be (0 + size-1) * size = (0 + 23) * 24 = 552
+    outputBuffer[0] = temp1;
+    // BUF: 552
+
+    outputBuffer[1] = temp2;
+    // BUF-NEXT: 552
+}
diff --git a/tests/spirv/spec-constant-sized-array-3.slang b/tests/spirv/spec-constant-sized-array-3.slang
new file mode 100644
index 000000000..39106aa5b
--- /dev/null
+++ b/tests/spirv/spec-constant-sized-array-3.slang
@@ -0,0 +1,48 @@
+//TEST:SIMPLE(filecheck=CHECK): -target spirv
+//TEST(compute, vulkan):COMPARE_COMPUTE(filecheck-buffer=BUF):-vk -output-using-type
+
+// CHECK: %[[C0:[0-9A-Za-z_]+]] = OpSpecConstant %int 4
+// CHECK: %[[I1:[0-9A-Za-z_]+]] = OpConstant %int 1
+// CHECK: %[[COP0:[0-9A-Za-z_]+]] = OpSpecConstantOp %int BitwiseAnd %[[C0]] %[[I1]]
+// CHECK: %[[I2:[0-9A-Za-z_]+]] = OpConstant %int 2
+// CHECK: %[[COP1:[0-9A-Za-z_]+]] = OpSpecConstantOp %int BitwiseOr %[[COP0]] %[[I2]]
+// CHECK: %[[COP2:[0-9A-Za-z_]+]] = OpSpecConstantOp %int IAdd %[[I1]] %[[COP1]]
+// CHECK: %[[COP3:[0-9A-Za-z_]+]] = OpSpecConstantOp %int ShiftLeftLogical %[[C0]] %[[COP2]]
+// CHECK: %[[ARR_TYPE:[0-9A-Za-z_]+]] = OpTypeArray %float %[[COP3]]
+// CHECK: %[[PT_TYPE:[0-9A-Za-z_]+]] = OpTypePointer Function %[[ARR_TYPE]]
+
+[SpecializationConstant]
+const int constValue0 = 4;
+
+//TEST_INPUT:ubuffer(data=[0], stride=4):out,name=outputBuffer
+RWStructuredBuffer<float> outputBuffer;
+
+static const int size = constValue0 << (1 + (constValue0 & 0x01 | 0x02));   // 4 << 3 = 32
+
+// This test is to verify that if two array has the same spec constant size, they are the same type.
+void func(out float buffer[constValue0 << (1 + (constValue0 & 0x01 | 0x02))], int idx)
+{
+    for (uint i = 0; i < size; i++)
+    {
+        buffer[i] = i;
+    }
+}
+
+[shader("compute")]
+[numthreads(1, 1, 1)]
+void computeMain(uint tid:SV_DispatchThreadID)
+{
+    float buffer[constValue0 << (1 + (constValue0 & 0x01 | 0x02))];
+    // CHECK: OpVariable %[[PT_TYPE]] Function
+
+    func(buffer, tid);
+
+    float temp = buffer[0];
+    for (uint i = 0; i < size; i++)
+    {
+        temp += buffer[i] * 2;
+    }
+    // Result will be (0 + size-1) * size = (0 + 31) * 32 * 2 = 992
+    outputBuffer[0] = temp;
+    // BUF: 992
+}
diff --git a/tests/spirv/spec-constant-sized-array-4.slang b/tests/spirv/spec-constant-sized-array-4.slang
new file mode 100644
index 000000000..0c511bc01
--- /dev/null
+++ b/tests/spirv/spec-constant-sized-array-4.slang
@@ -0,0 +1,48 @@
+//TEST:SIMPLE(filecheck=CHECK): -target spirv
+//TEST(compute, vulkan):COMPARE_COMPUTE(filecheck-buffer=BUF):-vk -output-using-type
+
+// CHECK: %[[C0:[0-9A-Za-z_]+]] = OpSpecConstant %int 32
+// CHECK: %[[C1:[0-9A-Za-z_]+]] = OpSpecConstant %int 2
+// CHECK: %[[COP0:[0-9A-Za-z_]+]] = OpSpecConstantOp %int SDiv %[[C0]] %[[C1]]
+// CHECK: %[[ARR_TYPE:[0-9A-Za-z_]+]] = OpTypeArray %float %[[COP0]]
+// CHECK: %[[PT_TYPE:[0-9A-Za-z_]+]] = OpTypePointer Function %[[ARR_TYPE]]
+
+[SpecializationConstant]
+const int constValue0 = 32;
+
+[SpecializationConstant]
+const int constValue1 = 2;
+
+//TEST_INPUT:ubuffer(data=[0], stride=4):out,name=outputBuffer
+RWStructuredBuffer<float> outputBuffer;
+
+void func(out float buffer[constValue0 / constValue1])
+{
+    for (uint i = 0; i < constValue0 / constValue1; i++)
+    {
+        buffer[i] = i;
+    }
+}
+
+[shader("compute")]
+[numthreads(1, 1, 1)]
+void computeMain()
+{
+    // This test checks that when we are using local constant that are computed from spec constant,
+    // the size can be defined in global scope, and type check can pass the when we call `func`.
+    const int localConst = constValue0 / constValue1;
+    float buffer[localConst];
+    // CHECK: OpVariable %[[PT_TYPE]] Function
+
+    func(buffer);
+
+    float temp = 0.0f;
+    for (uint i = 0; i < localConst; i++)
+    {
+        temp += buffer[i] * 2;
+    }
+
+    // Result will be (0 + localConst-1) * localConst = 15 * 16 = 240
+    outputBuffer[0] = temp;
+    // BUF: 240
+}