From 6dca7e39292e6c5672440f6f1dbfb204a79b90d2 Mon Sep 17 00:00:00 2001
From: Yong He <yonghe@outlook.com>
Date: Sat, 3 Feb 2024 12:29:12 -0800
Subject: Fix spirv emit that leads to pathological downstream time. (#3546)

---
 source/slang/slang-emit-spirv.cpp              |  1 +
 source/slang/slang-ir-specialize-resources.cpp | 14 ++++++++---
 source/slang/slang-ir-specialize-resources.h   |  2 +-
 tests/spirv/array-param-gep.slang              | 34 ++++++++++++++++++++++++++
 4 files changed, 47 insertions(+), 4 deletions(-)
 create mode 100644 tests/spirv/array-param-gep.slang

diff --git a/source/slang/slang-emit-spirv.cpp b/source/slang/slang-emit-spirv.cpp
index 7ff6ac7d6..6e1e58755 100644
--- a/source/slang/slang-emit-spirv.cpp
+++ b/source/slang/slang-emit-spirv.cpp
@@ -2262,6 +2262,7 @@ struct SPIRVEmitContext
                 SLANG_UNIMPLEMENTED_X(e.getBuffer());
             }
         case kIROp_Specialize:
+        case kIROp_MissingReturn:
             return nullptr;
         case kIROp_Var:
             return emitVar(parent, inst);
diff --git a/source/slang/slang-ir-specialize-resources.cpp b/source/slang/slang-ir-specialize-resources.cpp
index cfab3fb13..4fd87e173 100644
--- a/source/slang/slang-ir-specialize-resources.cpp
+++ b/source/slang/slang-ir-specialize-resources.cpp
@@ -38,6 +38,7 @@ struct ResourceParameterSpecializationCondition : FunctionCallSpecializeConditio
         // our decision.
         //
         type = unwrapArray(type);
+        bool isArray = type != param->getDataType();
 
         // On all of our (current) targets, a function that
         // takes a `ConstantBuffer<T>` parameter requires
@@ -63,7 +64,7 @@ struct ResourceParameterSpecializationCondition : FunctionCallSpecializeConditio
         if( isKhronosTarget(targetRequest) )
         {
             if (targetRequest->shouldEmitSPIRVDirectly())
-                return isIllegalSPIRVParameterType(type);
+                return isIllegalSPIRVParameterType(type, isArray);
             else
                 return isIllegalGLSLParameterType(type);
         }
@@ -1212,7 +1213,7 @@ bool specializeResourceUsage(
 
 bool isIllegalGLSLParameterType(IRType* type)
 {
-    if (as<IRUniformParameterGroupType>(type))
+    if (as<IRParameterGroupType>(type))
         return true;
     if (as<IRHLSLStructuredBufferTypeBase>(type))
         return true;
@@ -1236,7 +1237,7 @@ bool isIllegalGLSLParameterType(IRType* type)
     return false;
 }
 
-bool isIllegalSPIRVParameterType(IRType* type)
+bool isIllegalSPIRVParameterType(IRType* type, bool isArray)
 {
     if (isIllegalGLSLParameterType(type))
         return true;
@@ -1245,6 +1246,13 @@ bool isIllegalSPIRVParameterType(IRType* type)
     // all Texture types.
     if (as<IRTextureType>(type))
         return true;
+    if (isArray)
+    {
+        if (as<IRSamplerStateTypeBase>(type))
+        {
+            return true;
+        }
+    }
     return false;
 }
 } // namespace Slang
diff --git a/source/slang/slang-ir-specialize-resources.h b/source/slang/slang-ir-specialize-resources.h
index d5665155f..d28ac4175 100644
--- a/source/slang/slang-ir-specialize-resources.h
+++ b/source/slang/slang-ir-specialize-resources.h
@@ -31,6 +31,6 @@ namespace Slang
         IRModule*       irModule);
 
     bool isIllegalGLSLParameterType(IRType* type);
-    bool isIllegalSPIRVParameterType(IRType* type);
+    bool isIllegalSPIRVParameterType(IRType* type, bool isArray);
 
 }
diff --git a/tests/spirv/array-param-gep.slang b/tests/spirv/array-param-gep.slang
new file mode 100644
index 000000000..1a0de0c39
--- /dev/null
+++ b/tests/spirv/array-param-gep.slang
@@ -0,0 +1,34 @@
+//TEST:SIMPLE(filecheck=CHECK): -target spirv -entry computeMain -stage compute -emit-spirv-directly
+
+// Check that we are not generating spirv that loads a global resource array into a SSA register,
+// instead, these arrays should always be accessed via direct OpAccessChain operations to avoid
+// creating a lot of local load/stores in the driver compiler.
+
+struct Scene
+{
+    SamplerState samplers[256];
+    Texture2D textures[100];
+}
+
+ParameterBlock<Scene> scene;
+struct Material
+{
+    int sampler;
+    int texture;
+}
+
+RWStructuredBuffer<float4> result;
+
+float4 shade(Scene scene, Material mat)
+{
+    return scene.textures[mat.texture].SampleLevel(scene.samplers[mat.sampler], float2(0,0), 0);
+}
+
+[numthreads(1,1,1)]
+void computeMain(uniform Material mat)
+{
+    // CHECK: OpEntryPoint
+    // CHECK-NOT: OpLoad {{.*}} %scene{{.*}}samplers
+    // CHECK-NOT: OpLoad {{.*}} %scene{{.*}}textures
+    result[0] = shade(scene, mat);
+}
\ No newline at end of file
-- 
cgit v1.2.3