Check whether array element is fully specialized (#6000)

* Check whether array element is fully specialized

close #5776

When we start specialize a "specialize" IR, we should
make sure all the elements are fully specialized, but
we miss checking the elements of an array. This change
will check the it.

* add test

* add all wrapper types into the check

* add utility function to check if the type is wrapper type

---------

Co-authored-by: zhangkai <zhangkai@zhangkais-MacBook-Pro.local>
Co-authored-by: Yong He <yonghe@outlook.com>

2 files changed, 93 insertions, 0 deletions

diff --git a/tests/bugs/gh-5776.slang b/tests/bugs/gh-5776.slang
new file mode 100644
index 000000000..625a7b5cc
--- /dev/null
+++ b/tests/bugs/gh-5776.slang
@@ -0,0 +1,86 @@
+//TEST(compute, vulkan):COMPARE_COMPUTE_EX:-vk -compute -shaderobj -output-using-type
+//TEST(compute):COMPARE_COMPUTE_EX:-slang -compute -dx12 -profile sm_6_0 -use-dxil -output-using-type
+//TEST(compute):COMPARE_COMPUTE_EX:-slang -compute -cuda -shaderobj -output-using-type
+//TEST(compute):COMPARE_COMPUTE_EX:-slang -compute -cpu -shaderobj -output-using-type
+//TEST(compute):COMPARE_COMPUTE_EX:-slang -compute -wgpu -output-using-type
+
+
+interface IFoo
+{
+    associatedtype FooType : IFoo;
+}
+
+extension float : IFoo
+{
+    typedef float FooType;
+}
+
+__generic<T:IFoo, let N:int>
+extension Array<T, N> : IFoo
+{
+    typedef Array<T.FooType, N> FooType;
+}
+
+__generic<T:IFoo, let N:int>
+extension vector<T, N> : IFoo
+{
+    typedef vector<T.FooType, N> FooType;
+}
+
+__generic<T:IFoo, let N:int, let M:int>
+extension matrix<T, N, M> : IFoo
+{
+    typedef matrix<T.FooType, N, M> FooType;
+}
+
+struct WrappedBuffer<T : IFoo>
+{
+    StructuredBuffer<T> buffer;
+    int shape;
+
+    T get(int idx) { return buffer[idx]; }
+}
+
+
+struct GradInBuffer<T : IFoo>
+{
+    WrappedBuffer<T.FooType> wrapBuffer;
+}
+
+struct CallData
+{
+    GradInBuffer<float[2]> grad_in1;
+    GradInBuffer<vector<float, 2>> grad_in2;
+    GradInBuffer<float2x2> grad_in3;
+}
+
+
+//TEST_INPUT: set call_data.grad_in1.wrapBuffer.buffer = ubuffer(data=[1.0 2.0 3.0 4.0], stride=4);
+//TEST_INPUT: set call_data.grad_in2.wrapBuffer.buffer = ubuffer(data=[5.0 6.0 7.0 8.0], stride=4);
+//TEST_INPUT: set call_data.grad_in3.wrapBuffer.buffer = ubuffer(data=[1.0 2.0 3.0 4.0 5.0 6.0 7.0 8.0], stride=4);
+ParameterBlock<CallData> call_data;
+
+
+//TEST_INPUT:ubuffer(data=[0.0 0.0 0.0 0.0 0.0 0.0 0.0], stride=4):out, name outputBuffer
+RWStructuredBuffer<float> outputBuffer;
+
+
+[shader("compute")]
+[numthreads(1, 1, 1)]
+void computeMain()
+{
+    float[2] data1 = call_data.grad_in1.wrapBuffer.buffer[0];
+    float[2] data2 = call_data.grad_in1.wrapBuffer.get(1);
+    outputBuffer[0] = data1[0];
+    outputBuffer[1] = data2[0];
+
+    vector<float, 2> data3 = call_data.grad_in2.wrapBuffer.buffer[0];
+    vector<float, 2> data4 = call_data.grad_in2.wrapBuffer.get(1);
+    outputBuffer[2] = data3[0];
+    outputBuffer[3] = data4[0];
+
+    float2x2 data5 = call_data.grad_in3.wrapBuffer.buffer[0];
+    float2x2 data6 = call_data.grad_in3.wrapBuffer.get(1);
+    outputBuffer[4] = data5[0][0];
+    outputBuffer[5] = data6[0][0];
+}
diff --git a/tests/bugs/gh-5776.slang.expected.txt b/tests/bugs/gh-5776.slang.expected.txt
new file mode 100644
index 000000000..ffde6889e
--- /dev/null
+++ b/tests/bugs/gh-5776.slang.expected.txt
@@ -0,0 +1,7 @@
+type: float
+1.000000
+3.000000
+5.000000
+7.000000
+1.000000
+5.000000