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, 42 insertions, 13 deletions

diff --git a/source/slang/slang-emit.cpp b/source/slang/slang-emit.cpp
index 09c2efea9..804a44b81 100644
--- a/source/slang/slang-emit.cpp
+++ b/source/slang/slang-emit.cpp
@@ -120,6 +120,7 @@
 #include "slang-ir-variable-scope-correction.h"
 #include "slang-ir-vk-invert-y.h"
 #include "slang-ir-wgsl-legalize.h"
+#include "slang-ir-wrap-cbuffer-element.h"
 #include "slang-ir-wrap-structured-buffers.h"
 #include "slang-legalize-types.h"
 #include "slang-lower-to-ir.h"
@@ -1276,6 +1277,23 @@ Result linkAndOptimizeIR(
     // We don't need the legalize pass for C/C++ based types
     if (options.shouldLegalizeExistentialAndResourceTypes)
     {
+        if (isMetalTarget(targetRequest))
+        {
+            // Metal is a special target in that we want to legalize constant buffer
+            // types as if it is a bindful target, and skip legalizing parameter block
+            // types as if it is a bindless target.
+            // To achieve this, we want to ensure that all resource typed fields in parameter blocks
+            // are translated into descriptor handles first before running the resource type
+            // legalization pass for metal, so that type legalization pass won't mess around with
+            // them.
+            BufferElementTypeLoweringOptions bufferElementTypeLoweringOptions = {};
+            bufferElementTypeLoweringOptions.loweringPolicyKind =
+                BufferElementTypeLoweringPolicyKind::MetalParameterBlock;
+            lowerBufferElementTypeToStorageType(
+                targetProgram,
+                irModule,
+                bufferElementTypeLoweringOptions);
+        }
 
         // The Slang language allows interfaces to be used like
         // ordinary types (including placing them in constant
@@ -1403,17 +1421,17 @@ Result linkAndOptimizeIR(
     // Some information for `static_assert` is available only after the specialization.
     checkStaticAssert(irModule->getModuleInst(), sink);
 
-    // For HLSL (and fxc/dxc) only, we need to "wrap" any
-    // structured buffers defined over matrix types so
-    // that they instead use an intermediate `struct`.
-    // This is required to get those targets to respect
-    // the options for matrix layout set via `#pragma`
-    // or command-line options.
-    //
     switch (target)
     {
     case CodeGenTarget::HLSL:
         {
+            // For HLSL(fxc) only, we need to "wrap" any
+            // structured buffers defined over matrix types so
+            // that they instead use an intermediate `struct`.
+            // This is required to get those targets to respect
+            // the options for matrix layout set via `#pragma`
+            // or command-line options.
+            //
             wrapStructuredBuffersOfMatrices(irModule);
 #if 0
             dumpIRIfEnabled(codeGenContext, irModule, "STRUCTURED BUFFERS WRAPPED");
@@ -1421,7 +1439,12 @@ Result linkAndOptimizeIR(
             validateIRModuleIfEnabled(codeGenContext, irModule);
         }
         break;
-
+    case CodeGenTarget::Metal:
+    case CodeGenTarget::MetalLib:
+        // Metal does not allow `ConstantBuffer<StructuredBuffer<T>>`, so we need to create
+        // a wrapper struct for the `StructuredBuffer<T>`.
+        wrapCBufferElementsForMetal(irModule);
+        break;
     default:
         break;
     }
@@ -1828,11 +1851,16 @@ Result linkAndOptimizeIR(
             rcpWOfPositionInput(irModule);
     }
 
-    bool emitSpirvDirectly = targetProgram->shouldEmitSPIRVDirectly();
-
-    BufferElementTypeLoweringOptions bufferElementTypeLoweringOptions;
-    bufferElementTypeLoweringOptions.use16ByteArrayElementForConstantBuffer =
-        isWGPUTarget(targetRequest);
+    BufferElementTypeLoweringOptions bufferElementTypeLoweringOptions = {};
+    if (isWGPUTarget(targetRequest))
+        bufferElementTypeLoweringOptions.loweringPolicyKind =
+            BufferElementTypeLoweringPolicyKind::WGSL;
+    else if (isKhronosTarget(targetRequest))
+        bufferElementTypeLoweringOptions.loweringPolicyKind =
+            BufferElementTypeLoweringPolicyKind::KhronosTarget;
+    else
+        bufferElementTypeLoweringOptions.loweringPolicyKind =
+            BufferElementTypeLoweringPolicyKind::Default;
     lowerBufferElementTypeToStorageType(targetProgram, irModule, bufferElementTypeLoweringOptions);
 
     // If we are generating code for glsl or metal, perform address space propagation now.
@@ -1850,6 +1878,7 @@ Result linkAndOptimizeIR(
 
     performForceInlining(irModule);
 
+    bool emitSpirvDirectly = targetProgram->shouldEmitSPIRVDirectly();
     if (emitSpirvDirectly)
     {
         performIntrinsicFunctionInlining(irModule);