Fix legalization crash when processing metal parameter blocks. (#8591)

Closes #7606.

When Slang compile for a bindful target, we will run the resource type
legalization pass to hoist resource typed struct fields outside of the
struct type and define them as global parameters and passing them around
via dedicated function parameters.
When we compile for a bindless target, we don't run this pass.
However, Metal is a hybrid bindful and bindless target. We need to run
type legalization for the constant buffer, but skip type legalization
for parameter block.

The previous attempt to support this behavior is to hack the type
legalization pass to return `LegalVal::simple` when it sees a
`ParameterBlock<T>`. However, whenever the code is accessing
`parameterBlock.someNestedField`, the type of the nested field may get a
`LegalType::tuple`, and now we will run into inconsistent scenarios
where we have a `LegalVal::simple` on the operand val, and but the
legalization logic is expecting that val to be a `LegalType::tuple`.
This breaks a lot of assumptions and invariants in the type legalization
pass, resulting unstable/fragile behavior.

To systematically solve this problem, this change generalizes the
existing legalize buffer element type pass to translate
`ParameterBlock<Texture2D>` (and similar cases) to
`ParameterBlock<Texture2D.Handle>`. So that such parameter block will
always be legalized to `LegalType:::simple` during type legalization,
and we will never run into any inconsistent cases. This allowed us to
get rid of the hacky logic in the type legalization pass to try to
workaround the inconsistencies.

1 files changed, 143 insertions, 0 deletions

diff --git a/tests/metal/parameter-block.slang b/tests/metal/parameter-block.slang
new file mode 100644
index 000000000..75127f460
--- /dev/null
+++ b/tests/metal/parameter-block.slang
@@ -0,0 +1,143 @@
+//TEST:SIMPLE(filecheck=CHECK): -stage compute -entry computeMain -target metal
+//TEST:SIMPLE(filecheck=CHECK-ASM): -stage compute -entry computeMain -target metallib
+
+// This test verifies that resource types within structs and arrays
+// are preserved in Metal. This is done by lowering resource tyeps into
+// descriptor handle types during lower-buffer-element-type pass.
+
+struct Val
+{
+    RWStructuredBuffer<float> buffer;
+}
+
+struct Val_f
+{
+    float f;
+}
+
+struct ResourceStruct {
+    RWStructuredBuffer<float> buffer;
+    Texture2D<float4> texture;
+}
+
+struct NestedStruct {
+    Array<ResourceStruct, 2> resources;
+    ResourceStruct resource;
+    int padding;
+}
+
+struct ComplexStruct {
+    Array<NestedStruct, 3> nested;
+    RWStructuredBuffer<float> directBuffer;
+}
+
+struct ParentStruct {
+    struct ChildStruct {
+        Texture2D<float> texture;
+        int i;
+        float f;
+    } c;
+
+    Array<Texture2D<float>, 2> tex_array;
+}
+
+struct OtherParent {
+      ParentStruct::ChildStruct c2;
+}
+
+struct StructWithParameterBlock {
+    RWStructuredBuffer<float> buffer;
+    ParameterBlock<Array<Val, 2>> arr;
+}
+
+uniform StructWithParameterBlock struct_with_parameter_block;
+
+
+// Test various parameter block configurations
+ParameterBlock<ResourceStruct> simpleBlock;
+ParameterBlock<Array<ResourceStruct, 4>> arrayBlock; 
+ParameterBlock<NestedStruct> nestedBlock;
+ParameterBlock<ComplexStruct> complexBlock;
+ParameterBlock<Array<ComplexStruct, 2>> arrayComplexBlock;
+ParameterBlock<ParentStruct> parentBlock;
+
+ParameterBlock<RWStructuredBuffer<Val_f>> valBlock;
+
+
+RWStructuredBuffer<float> outputBuffer;
+
+void func(ParentStruct::ChildStruct c) {
+    outputBuffer[0] = c.texture.Load(int3(0));
+    Texture2D<float> tex = c.texture;
+    outputBuffer[0] = tex.Load(int3(0));
+
+}
+
+void func2(ParentStruct p) {
+    Array<Texture2D<float>, 2> tex_array_local = p.tex_array;
+    outputBuffer[1] = tex_array_local[0].Load(int3(0));
+    outputBuffer[2] = p.c.texture.Load(int3(3));
+}
+
+void func3(Array<Texture2D<float>, 2> arr) {
+    outputBuffer[3] = arr[0].Load(int3(0));
+}
+
+void func4(ParameterBlock<ParentStruct> block) {
+    ParentStruct p = block;
+    outputBuffer[8] = p.c.texture.Load(int3(0));
+}
+
+// CHECK-ASM: {{.*}} @computeMain
+[numthreads(1, 1, 1)]
+void computeMain(uint3 tid: SV_DispatchThreadID)
+{
+
+    // CHECK: {{.*}}->buffer{{.*}}
+    outputBuffer[0] = simpleBlock.buffer[0];
+    
+    // CHECK: arrayBlock{{.*}}->data{{.*}}[int(1)].buffer
+    ResourceStruct fromArray = arrayBlock[1];
+    outputBuffer[1] = fromArray.buffer[1];
+        
+    // CHECK: complexBlock{{.*}}->nested{{.*}}->data{{.*}}[int(1)])->resources{{.*}})->data{{.*}}[int(1)])->buffer{{.*}}
+    outputBuffer[3] = complexBlock.nested[1].resources[1].buffer[3];
+    
+    // CHECK: {{.*}}arrayComplexBlock{{.*}}->data{{.*}}[int(0)])->nested{{.*}})->data{{.*}}[int(2)])->resources{{.*}})->data{{.*}}[int(0)].buffer
+    ComplexStruct complexFromArray = arrayComplexBlock[0];
+    outputBuffer[4] = complexFromArray.nested[2].resources[0].buffer[4];
+    
+    // CHECK: {{.*}}->directBuffer{{.*}}
+    outputBuffer[5] = complexBlock.directBuffer[5];
+    
+    // Test parameter block accessed via a local variable
+    // CHECK: nestedBlock{{.*}}->resource{{.*}}.buffer_0
+    ResourceStruct resStruct = nestedBlock.resource;
+    outputBuffer[6] = resStruct.buffer[2];
+
+    // CHECK: {{.*}}valBlock{{.*}}
+    outputBuffer[7] = valBlock[0].f;
+
+    // Test parameter block inside a struct
+    // CHECK: {{.*}}struct_with_parameter_block_arr{{.*}}
+    StructWithParameterBlock local = struct_with_parameter_block;
+    outputBuffer[8] = local.arr[0].buffer[0];
+
+    // Test function with parameter block containing a struct with a resource type
+    func(parentBlock.c);
+
+    // Test function with a copy of the parameter block
+    OtherParent other_p;
+    other_p.c2 = parentBlock.c;
+    func(other_p.c2);
+
+    ParentStruct p;
+    p = parentBlock;
+    func2(p);
+
+    // Test function with a field from the parameter block
+    func3(p.tex_array);
+
+    // Test function with a parameter block in the function signature
+    func4(parentBlock);
+}
\ No newline at end of file