Update the GetDimension hlsl builtin for spirv path. In case of sampler, a combined sampled image needs an OpImage to be generated. (#3424)

4 files changed, 86 insertions, 20 deletions

diff --git a/source/slang/hlsl.meta.slang b/source/slang/hlsl.meta.slang
index b88b861ed..144141b79 100644
--- a/source/slang/hlsl.meta.slang
+++ b/source/slang/hlsl.meta.slang
@@ -957,7 +957,7 @@ extension __TextureImpl<T,Shape,isArray,isMS,sampleCount,0,isShadow,0,format>
     }
 }
 
-// Texture.GetDimensions
+// Texture.GetDimensions and Sampler.GetDimensions
 ${{{{
 const char* kTextureShapeTypeNames[] = {
     "__Shape1D", "__Shape2D", "__Shape3D", "__ShapeCube"};
diff --git a/source/slang/slang-stdlib-textures.cpp b/source/slang/slang-stdlib-textures.cpp
index 9ebfc54d6..a6c955b36 100644
--- a/source/slang/slang-stdlib-textures.cpp
+++ b/source/slang/slang-stdlib-textures.cpp
@@ -66,7 +66,8 @@ void TextureTypeInfo::writeFuncBody(
     const char* funcName,
     const String& glsl,
     const String& cuda,
-    const String& spirv)
+    const String& spirvDefault,
+    const String& spirvCombined)
 {
     BraceScope funcScope{i, sb};
     {
@@ -88,12 +89,25 @@ void TextureTypeInfo::writeFuncBody(
             sb << i << "case cuda:\n";
             sb << i << "__intrinsic_asm \"" << cuda << "\";\n";
         }
-        if(spirv.getLength())
+        if(spirvDefault.getLength() && spirvCombined.getLength())
         {
             sb << i << "case spirv:\n";
-            sb << i << "return spirv_asm\n";
-            BraceScope spirvScope{i, sb, ";\n"};
-            sb << spirv << "\n";
+            sb << i << "if (isCombined != 0)\n";
+            sb << i << "{\n";
+            {
+                sb << i << "return spirv_asm\n";
+                BraceScope spirvCombinedScope{i, sb, ";\n"};
+                sb << spirvCombined << "\n";
+            }
+            sb << i << "}\n";
+            sb << i << "else\n";
+            sb << i << "{\n";
+            {
+                sb << i << "return spirv_asm\n";
+                BraceScope spirvDefaultScope{i, sb, ";\n"};
+                sb << spirvDefault << "\n";
+            }
+            sb << i << "}\n";
         }
     }
 }
@@ -102,7 +116,8 @@ void TextureTypeInfo::writeFuncWithSig(
     const char* funcName,
     const String& sig,
     const String& glsl,
-    const String& spirv,
+    const String& spirvDefault,
+    const String& spirvCombined,
     const String& cuda,
     const ReadNoneMode readNoneMode)
 {
@@ -111,7 +126,7 @@ void TextureTypeInfo::writeFuncWithSig(
     sb << i << "[__readNone]\n";
     sb << i << "[ForceInline]\n";
     sb << i << sig << "\n";
-    writeFuncBody(funcName, glsl, cuda, spirv);
+    writeFuncBody(funcName, glsl, cuda, spirvDefault, spirvCombined);
     sb << "\n";
 }
 
@@ -120,7 +135,8 @@ void TextureTypeInfo::writeFunc(
     const char* funcName,
     const String& params,
     const String& glsl,
-    const String& spirv,
+    const String& spirvDefault,
+    const String& spirvCombined,
     const String& cuda,
     const ReadNoneMode readNoneMode)
 {
@@ -128,7 +144,8 @@ void TextureTypeInfo::writeFunc(
         funcName,
         cat(returnType, " ", funcName, "(", params, ")"),
         glsl,
-        spirv,
+        spirvDefault,
+        spirvCombined,
         cuda,
         readNoneMode
     );
@@ -271,16 +288,17 @@ void TextureTypeInfo::writeGetDimensionFunctions()
                 glsl << "\";\n";
             }
 
-            StringBuilder spirv;
+            // SPIRV ASM generation
+            auto generateSpirvAsm = [&](StringBuilder& spirv, UnownedStringSlice imageVar) 
             {
-                spirv << "OpCapability ImageQuery; ";
                 spirv << "%vecSize:$$uint";
                 if (sizeDimCount > 1) spirv << sizeDimCount;
                 spirv << " = ";
                 if (isMultisample)
-                    spirv << "OpImageQuerySize $this;";
+                    spirv << "OpImageQuerySize " << imageVar << ";";
                 else
-                    spirv << "OpImageQuerySizeLod $this $0;";
+                    spirv << "OpImageQuerySizeLod " << imageVar <<" $0;";
+
                 auto convertAndStore = [&](UnownedStringSlice uintSourceVal, const char* destParam)
                     {
                         if (UnownedStringSlice(rawT) == "uint")
@@ -343,17 +361,28 @@ void TextureTypeInfo::writeGetDimensionFunctions()
 
                 if (isMultisample)
                 {
-                    spirv << "%_sampleCount : $$uint = OpImageQuerySamples $this;";
+                    spirv << "%_sampleCount : $$uint = OpImageQuerySamples" << imageVar << ";";
                     convertAndStore(UnownedStringSlice("_sampleCount"), "sampleCount");
                 }
 
                 if (includeMipInfo)
                 {
-                    spirv << "%_levelCount : $$uint = OpImageQueryLevels $this;";
+                    spirv << "%_levelCount : $$uint = OpImageQueryLevels" << imageVar << ";";
                     convertAndStore(UnownedStringSlice("_levelCount"), "numberOfLevels");
                 }
+            };
+            StringBuilder spirvCombined;
+            {
+                spirvCombined << "OpCapability ImageQuery; ";
+                spirvCombined << "%image:__imageType(this) = OpImage $this; ";
+                generateSpirvAsm(spirvCombined, toSlice("%image"));
             }
 
+            StringBuilder spirvDefault;
+            {
+                spirvDefault << "OpCapability ImageQuery; ";
+                generateSpirvAsm(spirvDefault, toSlice("$this"));
+            }
             sb << "    __glsl_version(450)\n";
             sb << "    __glsl_extension(GL_EXT_samplerless_texture_functions)\n";
             writeFunc(
@@ -361,7 +390,8 @@ void TextureTypeInfo::writeGetDimensionFunctions()
                 "GetDimensions",
                 params,
                 glsl,
-                spirv,
+                spirvDefault,
+                spirvCombined,
                 "",
                 ReadNoneMode::Always);
         }
diff --git a/source/slang/slang-stdlib-textures.h b/source/slang/slang-stdlib-textures.h
index 35aed3e27..38ca41862 100644
--- a/source/slang/slang-stdlib-textures.h
+++ b/source/slang/slang-stdlib-textures.h
@@ -66,13 +66,15 @@ public:
         const char* funcName,
         const String& glsl,
         const String& cuda,
-        const String& spirv
+        const String& spirvDefault,
+        const String& spirvCombined
     );
     void writeFuncWithSig(
         const char* funcName,
         const String& sig,
         const String& glsl = String{},
-        const String& spirv = String{},
+        const String& spirvDefault = String{},
+        const String& spirvCombined = String{},
         const String& cuda = String{},
         const ReadNoneMode readNoneMode = ReadNoneMode::Never
     );
@@ -81,7 +83,8 @@ public:
         const char* funcName,
         const String& params,
         const String& glsl = String{},
-        const String& spirv = String{},
+        const String& spirvDefault = String{},
+        const String& spirvCombined = String{},
         const String& cuda = String{},
         const ReadNoneMode readNoneMode = ReadNoneMode::Never
     );
diff --git a/tests/cross-compile/vk-sampler-getdimension.slang b/tests/cross-compile/vk-sampler-getdimension.slang
new file mode 100644
index 000000000..ca0a5ce11
--- /dev/null
+++ b/tests/cross-compile/vk-sampler-getdimension.slang
@@ -0,0 +1,33 @@
+// vk-sampler-getdimension.slang
+
+//TEST:SIMPLE(filecheck=CHECK):-target spirv -entry main -stage fragment -emit-spirv-directly
+
+//  Input to the fragment shader.
+struct PSin
+{
+  float3 position : POSITION;
+};
+
+// Output of the fragment shader
+struct PSout
+{
+  float4 color : SV_Target;
+};
+
+[[vk::binding(1)]] Sampler3D g_Volume;
+
+// Fragment Shader
+[shader("pixel")]
+PSout main(PSin stage, bool isFrontFacing : SV_IsFrontFace)
+{
+  // CHECK: %image{{.*}} = OpImage
+  // CHECK: OpImageQuerySizeLod %{{.*}} %image{{.*}}
+  PSout output;
+  // Find normal at position
+  uint3 dim;
+  uint levels;
+  g_Volume.GetDimensions(0, dim.x, dim.y, dim.z, levels);
+  output.color = g_Volume.Sample(float3(dim));
+  
+  return output;
+}