Fix spirv intrinsics for partial derivatives. (#3355)

Co-authored-by: Yong He <yhe@nvidia.com>

2 files changed, 135 insertions, 95 deletions

diff --git a/source/slang/hlsl.meta.slang b/source/slang/hlsl.meta.slang
index 2bac47fd6..ac8c59ac6 100644
--- a/source/slang/hlsl.meta.slang
+++ b/source/slang/hlsl.meta.slang
@@ -3811,150 +3811,124 @@ int4 D3DCOLORtoUBYTE4(float4 color)
 }
 
 // Partial-difference derivatives
+${{{{
+const char* diffDimensions[2] = {"x", "y"};
+for (auto xOrY : diffDimensions) {
+}}}}
 __generic<T : __BuiltinFloatingPointType>
-__target_intrinsic(glsl, dFdx)
-[__readNone]
-T ddx(T x);
-
-__generic<T : __BuiltinFloatingPointType, let N : int>
-__target_intrinsic(hlsl)
-__target_intrinsic(glsl, dFdx)
-[__readNone]
-vector<T, N> ddx(vector<T, N> x)
-{
-    VECTOR_MAP_UNARY(T, N, ddx, x);
-}
-
-__generic<T : __BuiltinFloatingPointType, let N : int, let M : int>
-__target_intrinsic(hlsl)
 [__readNone]
-matrix<T, N, M> ddx(matrix<T, N, M> x)
+T dd$(xOrY)(T x)
 {
-    MATRIX_MAP_UNARY(T, N, M, ddx, x);
+    __target_switch
+    {
+    case hlsl:
+    case cpp:
+    case cuda:
+        __intrinsic_asm "dd$(xOrY)";
+    case glsl:
+        __intrinsic_asm "dFd$(xOrY)";
+    case spirv:
+        return spirv_asm {OpDPd$(xOrY) $$T result $x};
+    }
 }
 
-__generic<T : __BuiltinFloatingPointType>
-__target_intrinsic(hlsl)
-__glsl_extension(GL_ARB_derivative_control)
-__target_intrinsic(glsl, dFdxCoarse)
-[__readNone]
-T ddx_coarse(T x);
-
 __generic<T : __BuiltinFloatingPointType, let N : int>
-__target_intrinsic(hlsl)
-__glsl_extension(GL_ARB_derivative_control)
-__target_intrinsic(glsl, dFdxCoarse)
 [__readNone]
-vector<T, N> ddx_coarse(vector<T, N> x)
+vector<T, N> dd$(xOrY)(vector<T, N> x)
 {
-    VECTOR_MAP_UNARY(T, N, ddx_coarse, x);
+    __target_switch
+    {
+    case hlsl:
+    case cpp:
+    case cuda:
+        __intrinsic_asm "dd$(xOrY)";
+    case glsl:
+        __intrinsic_asm "dFd$(xOrY)";
+    case spirv:
+        return spirv_asm {OpDPd$(xOrY) $$vector<T, N> result $x};
+    }
 }
 
 __generic<T : __BuiltinFloatingPointType, let N : int, let M : int>
 __target_intrinsic(hlsl)
 [__readNone]
-matrix<T, N, M> ddx_coarse(matrix<T, N, M> x)
+matrix<T, N, M> dd$(xOrY)(matrix<T, N, M> x)
 {
-    MATRIX_MAP_UNARY(T, N, M, ddx_coarse, x);
+    MATRIX_MAP_UNARY(T, N, M, dd$(xOrY), x);
 }
 
 __generic<T : __BuiltinFloatingPointType>
-__target_intrinsic(hlsl)
-__glsl_extension(GL_ARB_derivative_control)
-__target_intrinsic(glsl, dFdxFine)
-[__readNone]
-T ddx_fine(T x);
-
-__generic<T : __BuiltinFloatingPointType, let N : int>
-__target_intrinsic(hlsl)
 __glsl_extension(GL_ARB_derivative_control)
-__target_intrinsic(glsl, dFdxFine)
 [__readNone]
-vector<T, N> ddx_fine(vector<T, N> x)
+T dd$(xOrY)_coarse(T x)
 {
-    VECTOR_MAP_UNARY(T, N, ddx_fine, x);
-}
-
-__generic<T : __BuiltinFloatingPointType, let N : int, let M : int>
-__target_intrinsic(hlsl)
-[__readNone]
-matrix<T, N, M> ddx_fine(matrix<T, N, M> x)
-{
-    MATRIX_MAP_UNARY(T, N, M, ddx_fine, x);
+    __target_switch
+    {
+    case hlsl: __intrinsic_asm "dd$(xOrY)_coarse";
+    case glsl: __intrinsic_asm "dFd$(xOrY)Coarse";
+    case spirv: return spirv_asm {OpCapability DerivativeControl; result:$$T = OpDPd$(xOrY)Coarse $x};
+    }
 }
 
-__generic<T : __BuiltinFloatingPointType>
-__target_intrinsic(hlsl)
-__target_intrinsic(glsl, dFdy)
-[__readNone]
-T ddy(T x);
-
 __generic<T : __BuiltinFloatingPointType, let N : int>
-__target_intrinsic(hlsl)
-__target_intrinsic(glsl, dFdy)
+__glsl_extension(GL_ARB_derivative_control)
 [__readNone]
-vector<T, N> ddy(vector<T, N> x)
+vector<T, N> dd$(xOrY)_coarse(vector<T, N> x)
 {
-    VECTOR_MAP_UNARY(T, N, ddy, x);
+    __target_switch
+    {
+    case hlsl: __intrinsic_asm "dd$(xOrY)_coarse";
+    case glsl: __intrinsic_asm "dFd$(xOrY)Coarse";
+    case spirv: return spirv_asm {OpCapability DerivativeControl; result:$$vector<T,N> = OpDPd$(xOrY)Coarse $x};
+    }
 }
 
 __generic<T : __BuiltinFloatingPointType, let N : int, let M : int>
 __target_intrinsic(hlsl)
 [__readNone]
-matrix<T, N, M> ddy(matrix<T, N, M> x)
+matrix<T, N, M> dd$(xOrY)_coarse(matrix<T, N, M> x)
 {
-    MATRIX_MAP_UNARY(T, N, M, ddy, x);
+    MATRIX_MAP_UNARY(T, N, M, dd$(xOrY)_coarse, x);
 }
 
 __generic<T : __BuiltinFloatingPointType>
 __glsl_extension(GL_ARB_derivative_control)
-__target_intrinsic(glsl, dFdyCoarse)
-[__readNone]
-T ddy_coarse(T x);
-
-__generic<T : __BuiltinFloatingPointType, let N : int>
-__target_intrinsic(hlsl)
-__glsl_extension(GL_ARB_derivative_control)
-__target_intrinsic(glsl, dFdyCoarse)
 [__readNone]
-vector<T, N> ddy_coarse(vector<T, N> x)
+T dd$(xOrY)_fine(T x)
 {
-    VECTOR_MAP_UNARY(T, N, ddy_coarse, x);
-}
-
-__generic<T : __BuiltinFloatingPointType, let N : int, let M : int>
-__target_intrinsic(hlsl)
-[__readNone]
-matrix<T, N, M> ddy_coarse(matrix<T, N, M> x)
-{
-    MATRIX_MAP_UNARY(T, N, M, ddy_coarse, x);
+    __target_switch
+    {
+    case hlsl: __intrinsic_asm "dd$(xOrY)_fine";
+    case glsl: __intrinsic_asm "dFd$(xOrY)Fine";
+    case spirv: return spirv_asm {OpCapability DerivativeControl; result:$$T = OpDPd$(xOrY)Fine $x};
+    }
 }
 
-__generic<T : __BuiltinFloatingPointType>
-__target_intrinsic(hlsl)
-__glsl_extension(GL_ARB_derivative_control)
-__target_intrinsic(glsl, dFdyFine)
-[__readNone]
-T ddy_fine(T x);
-
 __generic<T : __BuiltinFloatingPointType, let N : int>
-__target_intrinsic(hlsl)
 __glsl_extension(GL_ARB_derivative_control)
-__target_intrinsic(glsl, dFdyFine)
 [__readNone]
-vector<T, N> ddy_fine(vector<T, N> x)
+vector<T, N> dd$(xOrY)_fine(vector<T, N> x)
 {
-    VECTOR_MAP_UNARY(T, N, ddy_fine, x);
+    __target_switch
+    {
+    case hlsl: __intrinsic_asm "dd$(xOrY)_fine";
+    case glsl: __intrinsic_asm "dFd$(xOrY)Fine";
+    case spirv: return spirv_asm {OpCapability DerivativeControl; result:$$vector<T,N> = OpDPd$(xOrY)Fine $x};
+    }
 }
 
 __generic<T : __BuiltinFloatingPointType, let N : int, let M : int>
 __target_intrinsic(hlsl)
 [__readNone]
-matrix<T, N, M> ddy_fine(matrix<T, N, M> x)
+matrix<T, N, M> dd$(xOrY)_fine(matrix<T, N, M> x)
 {
-    MATRIX_MAP_UNARY(T, N, M, ddy_fine, x);
+    MATRIX_MAP_UNARY(T, N, M, dd$(xOrY)_fine, x);
 }
 
+${{{{
+} // for (xOrY)
+}}}}
+
 
 // Radians to degrees
 
@@ -4541,13 +4515,13 @@ __generic<T : __BuiltinFloatingPointType>
 [__readNone]
 __target_intrinsic(hlsl)
 __target_intrinsic(glsl)
-__target_intrinsic(spirv, "OpFWidth resultType resultId _0")
+__target_intrinsic(spirv, "OpFwidth resultType resultId _0")
 T fwidth(T x);
 
 __generic<T : __BuiltinFloatingPointType, let N : int>
 __target_intrinsic(hlsl)
 __target_intrinsic(glsl)
-__target_intrinsic(spirv, "OpFWidth resultType resultId _0")
+__target_intrinsic(spirv, "OpFwidth resultType resultId _0")
 [__readNone]
 vector<T, N> fwidth(vector<T, N> x)
 {
diff --git a/tests/spirv/fwidth.slang b/tests/spirv/fwidth.slang
new file mode 100644
index 000000000..a55dfa3b0
--- /dev/null
+++ b/tests/spirv/fwidth.slang
@@ -0,0 +1,66 @@
+//TEST:SIMPLE(filecheck=CHECK): -stage fragment -target spirv -emit-spirv-directly -entry main
+
+RWStructuredBuffer<float> output;
+
+void main()
+{
+    {
+        // CHECK: OpFwidth
+        float w = 1.0;
+        float b = fwidth(w);
+        output[0] = b;
+
+        // CHECK: OpDPdx
+        float b1 = ddx(w);
+        output[1] = b1;
+        // CHECK: OpDPdy
+        float b2 = ddy(w);
+        output[2] = b2;
+
+        // CHECK: OpDPdxCoarse
+        float b3 = ddx_coarse(w);
+        output[3] = b3;
+
+        // CHECK: OpDPdyCoarse
+        float b4 = ddy_coarse(w);
+        output[4] = b4;
+
+        // CHECK: OpDPdxFine
+        float b5 = ddx_fine(w);
+        output[5] = b5;
+
+        // CHECK: OpDPdyFine
+        float b6 = ddy_fine(w);
+        output[6] = b6;
+    }
+
+    {
+        // CHECK: OpFwidth
+        float3 w = 1.0;
+        float3 b = fwidth(w);
+        output[7] = b.x;
+
+        // CHECK: OpDPdx
+        float3 b1 = ddx(w);
+        output[8] = b1.x;
+        // CHECK: OpDPdy
+        float3 b2 = ddy(w);
+        output[9] = b2.x;
+
+        // CHECK: OpDPdxCoarse
+        float3 b3 = ddx_coarse(w);
+        output[10] = b3.x;
+
+        // CHECK: OpDPdyCoarse
+        float3 b4 = ddy_coarse(w);
+        output[11] = b4.x;
+
+        // CHECK: OpDPdxFine
+        float3 b5 = ddx_fine(w);
+        output[12] = b5.x;
+
+        // CHECK: OpDPdyFine
+        float3 b6 = ddy_fine(w);
+        output[13] = b6.x;
+    }
+}