Fix depth texture sampling on Metal. (#6168)

8 files changed, 157 insertions, 92 deletions

diff --git a/source/slang/hlsl.meta.slang b/source/slang/hlsl.meta.slang
index 97e2c7be5..ba5c95a0c 100644
--- a/source/slang/hlsl.meta.slang
+++ b/source/slang/hlsl.meta.slang
@@ -855,6 +855,10 @@ float __glsl_texture_level_offset_1d_shadow<TTexture, TCoord, TOffset>(TTexture
     }
 }
 
+__intrinsic_op($(kIROp_MetalCastToDepthTexture))
+__generic<T:ITexelElement, Shape: __ITextureShape, let isArray:int, let isMS:int, let sampleCount:int, let access:int, let isShadow:int, let isCombined:int, let format:int>
+_Texture<T,Shape,isArray,isMS,sampleCount,access,1,isCombined,format> __metal_asDepthTexture(_Texture<T,Shape,isArray,isMS,sampleCount,access,isShadow,isCombined,format> tex);
+
 //@public:
 
 __generic<T:ITexelElement, Shape: __ITextureShape, let isArray:int, let isMS:int, let sampleCount:int, let isShadow:int, let format:int>
@@ -1676,6 +1680,96 @@ extension _Texture<T,Shape,isArray,isMS,sampleCount,access,isShadow,0,format>
         }
     }
 }
+__generic<T:ITexelElement, Shape: __ITextureShape, let isArray:int, let isMS:int, let sampleCount:int, let format:int>
+float __metal_SampleCmp(_Texture<T,Shape,isArray,isMS,sampleCount,0,1,0,format> t, SamplerComparisonState s, vector<float, Shape.dimensions+isArray> location, float compareValue)
+{
+    if (isArray == 1)
+    {
+        switch (Shape.flavor)
+        {
+        case $(SLANG_TEXTURE_2D):
+            __intrinsic_asm "$0.sample_compare($1, ($2).xy, uint(($2).z), $3)";
+        case $(SLANG_TEXTURE_CUBE):
+            __intrinsic_asm "$0.sample_compare($1, ($2).xyz, uint(($2).w), $3)";
+        }
+    }
+    else
+    {
+        switch (Shape.flavor)
+        {
+        case $(SLANG_TEXTURE_2D):
+        case $(SLANG_TEXTURE_CUBE):
+            __intrinsic_asm ".sample_compare";
+        }
+    }
+    __intrinsic_asm "<invalid intrinsic>";
+}
+__generic<T:ITexelElement, Shape: __ITextureShape, let isArray:int, let isMS:int, let sampleCount:int, let format:int>
+float __metal_SampleCmp(_Texture<T,Shape,isArray,isMS,sampleCount,0,1,0,format> t, SamplerComparisonState s, vector<float, Shape.dimensions+isArray> location, float compareValue, constexpr vector<int, Shape.planeDimensions> offset)
+{
+    if (isArray == 1)
+    {
+        switch (Shape.flavor)
+        {
+        case $(SLANG_TEXTURE_2D):
+            __intrinsic_asm "$0.sample_compare($1, ($2).xy, uint(($2).z), $3, $4)";
+        }
+    }
+    else
+    {
+        switch (Shape.flavor)
+        {
+        case $(SLANG_TEXTURE_2D):
+            __intrinsic_asm "$0.sample_compare($1, $2, $3, $4)";
+        }
+    }
+    __intrinsic_asm "<invalid intrinsic>";
+}
+
+__generic<T:ITexelElement, Shape: __ITextureShape, let isArray:int, let isMS:int, let sampleCount:int, let format:int>
+float __metal_SampleCmpLevel(_Texture<T,Shape,isArray,isMS,sampleCount,0,1,0,format> t, SamplerComparisonState s, vector<float, Shape.dimensions+isArray> location, float compareValue, float level)
+{
+    if (isArray == 1)
+    {
+        switch (Shape.flavor)
+        {
+        case $(SLANG_TEXTURE_2D):
+            __intrinsic_asm "$0.sample_compare($1, ($2).xy, uint(($2).z), $3, level($4))";
+        case $(SLANG_TEXTURE_CUBE):
+            __intrinsic_asm "$0.sample_compare($1, ($2).xyz, uint(($2).w), $3, level($4))";
+        }
+    }
+    else
+    {
+        switch (Shape.flavor)
+        {
+        case $(SLANG_TEXTURE_2D):
+        case $(SLANG_TEXTURE_CUBE):
+            __intrinsic_asm "$0.sample_compare($1, $2, $3, level($4))";
+        }
+    }
+    __intrinsic_asm "<invalid intrinsic>";
+}
+__generic<T:ITexelElement, Shape: __ITextureShape, let isArray:int, let isMS:int, let sampleCount:int, let format:int>
+float __metal_SampleCmpLevel(_Texture<T,Shape,isArray,isMS,sampleCount,0,1,0,format> t, SamplerComparisonState s, vector<float, Shape.dimensions+isArray> location, float compareValue, float level, constexpr vector<int, Shape.planeDimensions> offset)
+{
+    switch (Shape.flavor)
+    {
+    case $(SLANG_TEXTURE_2D):
+        if (isArray == 1)
+        {
+            // T sample_compare(sampler s, float2 coord, uint array, float compare_value, lod_options options, int2 offset = int2(0)) const
+            __intrinsic_asm "$0.sample_compare($1, ($2).xy, uint(($2).z), $3, level($4), $5)";
+        }
+        else
+        {
+            // T sample_compare(sampler s, float2 coord, float compare_value, lod_options options, int2 offset = int2(0)) const
+            __intrinsic_asm "$0.sample_compare($1, $2, $3, level($4), $5)";
+        }
+        break;
+    }
+    __intrinsic_asm "<invalid intrinsic>";
+}
 
 __generic<T:ITexelElement, Shape: __ITextureShape, let isArray:int, let isMS:int, let sampleCount:int, let isShadow:int, let format:int>
 extension _Texture<T,Shape,isArray,isMS,sampleCount,0,isShadow,0,format>
@@ -2082,26 +2176,7 @@ extension _Texture<T,Shape,isArray,isMS,sampleCount,0,isShadow,0,format>
                 , "HLSL supports only float and half type textures");
             __intrinsic_asm ".SampleCmp";
         case metal:
-            if (isArray == 1)
-            {
-                switch (Shape.flavor)
-                {
-                case $(SLANG_TEXTURE_2D):
-                    __intrinsic_asm "$0.sample_compare($1, ($2).xy, uint(($2).z), $3)";
-                case $(SLANG_TEXTURE_CUBE):
-                    __intrinsic_asm "$0.sample_compare($1, ($2).xyz, uint(($2).w), $3)";
-                }
-            }
-            else
-            {
-                switch (Shape.flavor)
-                {
-                case $(SLANG_TEXTURE_2D):
-                case $(SLANG_TEXTURE_CUBE):
-                    __intrinsic_asm ".sample_compare";
-                }
-            }
-            __intrinsic_asm "<invalid intrinsic>";
+            return __metal_SampleCmp(__metal_asDepthTexture(this), s, location, compareValue);
         case spirv:
             return spirv_asm
             {
@@ -2183,23 +2258,7 @@ extension _Texture<T,Shape,isArray,isMS,sampleCount,0,isShadow,0,format>
                 , "HLSL supports only float and half type textures");
             __intrinsic_asm ".SampleCmp";
         case metal:
-            if (isArray == 1)
-            {
-                switch (Shape.flavor)
-                {
-                case $(SLANG_TEXTURE_2D):
-                    __intrinsic_asm "$0.sample_compare($1, ($2).xy, uint(($2).z), $3, $4)";
-                }
-            }
-            else
-            {
-                switch (Shape.flavor)
-                {
-                case $(SLANG_TEXTURE_2D):
-                    __intrinsic_asm ".sample_compare";
-                }
-            }
-            __intrinsic_asm "<invalid intrinsic>";
+            return __metal_SampleCmp(__metal_asDepthTexture(this), s, location, compareValue, offset);
         case spirv:
             return spirv_asm
             {
@@ -2311,26 +2370,7 @@ extension _Texture<T,Shape,isArray,isMS,sampleCount,0,isShadow,0,format>
                 , "HLSL supports only float and half type textures");
             __intrinsic_asm ".SampleCmpLevel";
         case metal:
-            if (isArray == 1)
-            {
-                switch (Shape.flavor)
-                {
-                case $(SLANG_TEXTURE_2D):
-                    __intrinsic_asm "$0.sample_compare($1, ($2).xy, uint(($2).z), $3, level($4))";
-                case $(SLANG_TEXTURE_CUBE):
-                    __intrinsic_asm "$0.sample_compare($1, ($2).xyz, uint(($2).w), $3, level($4))";
-                }
-            }
-            else
-            {
-                switch (Shape.flavor)
-                {
-                case $(SLANG_TEXTURE_2D):
-                case $(SLANG_TEXTURE_CUBE):
-                    __intrinsic_asm "$0.sample_compare($1, $2, $3, level($4))";
-                }
-            }
-            __intrinsic_asm "<invalid intrinsic>";
+            return __metal_SampleCmpLevel(__metal_asDepthTexture(this), s, location, compareValue, level);
         case spirv:
             return spirv_asm
             {
@@ -2362,25 +2402,7 @@ extension _Texture<T,Shape,isArray,isMS,sampleCount,0,isShadow,0,format>
                 , "HLSL supports only float and half type textures");
             __intrinsic_asm ".SampleCmpLevel";
         case metal:
-            if (isShadow == 1)
-            {
-                switch (Shape.flavor)
-                {
-                case $(SLANG_TEXTURE_2D):
-                    if (isArray == 1)
-                    {
-                        // T sample_compare(sampler s, float2 coord, uint array, float compare_value, lod_options options, int2 offset = int2(0)) const
-                        __intrinsic_asm "$0.sample_compare($1, ($2).xy, uint(($2).z), $3, level($4), $5)";
-                    }
-                    else
-                    {
-                        // T sample_compare(sampler s, float2 coord, float compare_value, lod_options options, int2 offset = int2(0)) const
-                        __intrinsic_asm "$0.sample_compare($1, $2, $3, level($4), $5)";
-                    }
-                    break;
-                }
-            }
-            __intrinsic_asm "<invalid intrinsic>";
+            return __metal_SampleCmpLevel(__metal_asDepthTexture(this), s, location, compareValue, level, offset);
         case spirv:
             return spirv_asm
             {
diff --git a/source/slang/slang-emit-c-like.cpp b/source/slang/slang-emit-c-like.cpp
index ac548e354..98bddad9b 100644
--- a/source/slang/slang-emit-c-like.cpp
+++ b/source/slang/slang-emit-c-like.cpp
@@ -1413,6 +1413,7 @@ bool CLikeSourceEmitter::shouldFoldInstIntoUseSites(IRInst* inst)
     // or statement).
     case kIROp_UpdateElement:
     case kIROp_DefaultConstruct:
+    case kIROp_MetalCastToDepthTexture:
         return false;
 
     // Always fold these in, because they are trivial
@@ -3129,6 +3130,7 @@ void CLikeSourceEmitter::_emitInst(IRInst* inst)
     case kIROp_AtomicCompareExchange:
     case kIROp_StructuredBufferGetDimensions:
     case kIROp_MetalAtomicCast:
+    case kIROp_MetalCastToDepthTexture:
         emitInstStmt(inst);
         break;
 
diff --git a/source/slang/slang-emit-metal.cpp b/source/slang/slang-emit-metal.cpp
index 5e4381d21..166f02535 100644
--- a/source/slang/slang-emit-metal.cpp
+++ b/source/slang/slang-emit-metal.cpp
@@ -636,6 +636,22 @@ bool MetalSourceEmitter::tryEmitInstStmtImpl(IRInst* inst)
                 m_writer->emit(");\n");
             return true;
         }
+    case kIROp_MetalCastToDepthTexture:
+        {
+            emitType(inst->getDataType(), getName(inst));
+            m_writer->emit(";\n{\n");
+            m_writer->indent();
+            m_writer->emit("auto _slang_ordinary_texture = ");
+            emitOperand(inst->getOperand(0), getInfo(EmitOp::General));
+            m_writer->emit(";\n");
+            m_writer->emit(getName(inst));
+            m_writer->emit(" = *(");
+            emitType(inst->getDataType());
+            m_writer->emit(" thread*)(&_slang_ordinary_texture);\n");
+            m_writer->dedent();
+            m_writer->emit("}\n");
+            return true;
+        }
     }
     return false;
 }
diff --git a/source/slang/slang-ir-inst-defs.h b/source/slang/slang-ir-inst-defs.h
index 179e3a5f4..591ffabf0 100644
--- a/source/slang/slang-ir-inst-defs.h
+++ b/source/slang/slang-ir-inst-defs.h
@@ -555,6 +555,8 @@ INST(MetalSetVertex, metalSetVertex, 2, 0)
 INST(MetalSetPrimitive, metalSetPrimitive, 2, 0)
 INST(MetalSetIndices, metalSetIndices, 2, 0)
 
+INST(MetalCastToDepthTexture, MetalCastToDepthTexture, 1, 0)
+
 // Construct a vector from a scalar
 //
 // %dst = MakeVectorFromScalar %T %N %val
diff --git a/source/slang/slang-ir-metal-legalize.cpp b/source/slang/slang-ir-metal-legalize.cpp
index 0d58bdd14..e9f693622 100644
--- a/source/slang/slang-ir-metal-legalize.cpp
+++ b/source/slang/slang-ir-metal-legalize.cpp
@@ -206,7 +206,17 @@ static void processInst(IRInst* inst)
     case kIROp_Leq:
         legalizeBinaryOp(inst);
         break;
-
+    case kIROp_MetalCastToDepthTexture:
+        {
+            // If the operand is already a depth texture, don't do anything.
+            auto textureType = as<IRTextureTypeBase>(inst->getOperand(0)->getDataType());
+            if (textureType && getIntVal(textureType->getIsShadowInst()) == 1)
+            {
+                inst->replaceUsesWith(inst->getOperand(0));
+                inst->removeAndDeallocate();
+            }
+            break;
+        }
     default:
         for (auto child : inst->getModifiableChildren())
         {
diff --git a/source/slang/slang-ir.cpp b/source/slang/slang-ir.cpp
index ac51ae451..a328057ac 100644
--- a/source/slang/slang-ir.cpp
+++ b/source/slang/slang-ir.cpp
@@ -8248,6 +8248,7 @@ bool IRInst::mightHaveSideEffects(SideEffectAnalysisOptions options)
     case kIROp_GetArrayLength:
     case kIROp_ResolveVaryingInputRef:
     case kIROp_GetPerVertexInputArray:
+    case kIROp_MetalCastToDepthTexture:
         return false;
 
     case kIROp_ForwardDifferentiate:
diff --git a/tests/metal/depth-texture.slang b/tests/metal/depth-texture.slang
new file mode 100644
index 000000000..ac15b3f1a
--- /dev/null
+++ b/tests/metal/depth-texture.slang
@@ -0,0 +1,12 @@
+//TEST:SIMPLE(filecheck=CHECK): -target metallib
+//CHECK: sample_compare_depth_2d
+Texture2D texture;
+SamplerComparisonState sampler;
+
+RWStructuredBuffer<float> output;
+
+[numthreads(1,1,1)]
+void computeMain()
+{
+    output[0] = texture.SampleCmpLevelZero(sampler, float2(0, 0), 0);
+}
\ No newline at end of file
diff --git a/tests/metal/texture.slang b/tests/metal/texture.slang
index 12274edef..90a42e3d6 100644
--- a/tests/metal/texture.slang
+++ b/tests/metal/texture.slang
@@ -566,25 +566,25 @@ bool TEST_texture<T>(
         // float SampleCmp()
         // ==================
 
-        // METAL: d2D{{.*}}.sample_compare(
+        // METAL: {{.*}}.sample_compare(
         // METALLIB: call {{.*}}.sample_compare_depth_2d.f32(
         && float(0) == d2D.SampleCmp(shadowSampler, float2(u, u), 0)
 
-        // METAL: d2DArray{{.*}}.sample_compare(
+        // METAL: {{.*}}.sample_compare(
         // METALLIB: call {{.*}}.sample_compare_depth_2d_array.f32(
         && float(0) == d2DArray.SampleCmp(shadowSampler, float3(u, u, 0), 0)
 
-        // METAL: dCube{{.*}}.sample_compare(
+        // METAL: {{.*}}.sample_compare(
         // METALLIB: call {{.*}}.sample_compare_depth_cube.f32(
         && float(0) == dCube.SampleCmp(shadowSampler, normalize(float3(u, 1 - u, u)), 0)
 
-        // METAL: dCubeArray{{.*}}.sample_compare(
+        // METAL: {{.*}}.sample_compare(
         // METALLIB: call {{.*}}.sample_compare_depth_cube_array.f32(
         && float(0) == dCubeArray.SampleCmp(shadowSampler, float4(normalize(float3(u, 1 - u, u)), 0), 0)
 
         // Offset variant
 
-        // METAL: d2D{{.*}}.sample_compare(
+        // METAL: {{.*}}.sample_compare(
         // METALLIB: call {{.*}}.sample_compare_depth_2d.f32(
         && float(0) == d2D.SampleCmp(shadowSampler, float2(u2, u), 0, int2(0, 0))
 
@@ -592,25 +592,25 @@ bool TEST_texture<T>(
         //  float SampleCmpLevelZero()
         // ===================================
 
-        // METAL: d2D{{.*}}.sample_compare(
+        // METAL: {{.*}}.sample_compare(
         // METALLIB: call {{.*}}.sample_compare_depth_2d.f32(
         && float(0) == d2D.SampleCmpLevelZero(shadowSampler, float2(u, u), 0)
 
-        // METAL: d2DArray{{.*}}.sample_compare(
+        // METAL: {{.*}}.sample_compare(
         // METALLIB: call {{.*}}.sample_compare_depth_2d_array.f32(
         && float(0) == d2DArray.SampleCmpLevelZero(shadowSampler, float3(u, u, 0), 0)
 
-        // METAL: dCube{{.*}}.sample_compare(
+        // METAL: {{.*}}.sample_compare(
         // METALLIB: call {{.*}}.sample_compare_depth_cube.f32(
         && float(0) == dCube.SampleCmpLevelZero(shadowSampler, normalize(float3(u, 1 - u, u)), 0)
 
-        // METAL: dCubeArray{{.*}}.sample_compare(
+        // METAL: {{.*}}.sample_compare(
         // METALLIB: call {{.*}}.sample_compare_depth_cube_array.f32(
         && float(0) == dCubeArray.SampleCmpLevelZero(shadowSampler, float4(normalize(float3(u, 1-u, u)), 0), 0)
 
         // Offset variant
 
-        // METAL: d2D{{.*}}.sample_compare(
+        // METAL: {{.*}}.sample_compare(
         // METALLIB: call {{.*}}.sample_compare_depth_2d.f32(
         && float(0) == d2D.SampleCmpLevelZero(shadowSampler, float2(u2, u), 0, int2(0, 0))
 
@@ -620,25 +620,25 @@ bool TEST_texture<T>(
 
         // These require SM 6.7 for dx12 but functional compute tests currently do not run with the cs_6_7.
 #if !defined(EXCLUDE_SM_6_7)
-        // METAL: d2D{{.*}}.sample_compare(
+        // METAL: {{.*}}.sample_compare(
         // METALLIB: call {{.*}}.sample_compare_depth_2d.f32(
         && float(0) == d2D.SampleCmpLevel(shadowSampler, float2(u, u), 0, 0.0)
 
-        // METAL: d2DArray{{.*}}.sample_compare(
+        // METAL: {{.*}}.sample_compare(
         // METALLIB: call {{.*}}.sample_compare_depth_2d_array.f32(
         && float(0) == d2DArray.SampleCmpLevel(shadowSampler, float3(u, u, 0), 0, 0.0)
 
-        // METAL: dCube{{.*}}.sample_compare(
+        // METAL: {{.*}}.sample_compare(
         // METALLIB: call {{.*}}.sample_compare_depth_cube.f32(
         && float(0) == dCube.SampleCmpLevel(shadowSampler, normalize(float3(u, 1 - u, u)), 0, 0.0)
 
-        // METAL: dCubeArray{{.*}}.sample_compare(
+        // METAL: {{.*}}.sample_compare(
         // METALLIB: call {{.*}}.sample_compare_depth_cube_array.f32(
         && float(0) == dCubeArray.SampleCmpLevel(shadowSampler, float4(normalize(float3(u, 1-u, u)), 0), 0, 0.0)
 
         // Offset variant
 
-        // METAL: d2D{{.*}}.sample_compare(
+        // METAL: {{.*}}.sample_compare(
         // METALLIB: call {{.*}}.sample_compare_depth_2d.f32(
         && float(0) == d2D.SampleCmpLevel(shadowSampler, float2(u2, u), 0, 0.0, int2(0, 0))
 #endif