Fix a bug in the forward derivative of cross product (#4006)

* Fix a bug in fwd-diff for cross product

* Also add a test for the reverse-mode AD

---------

Co-authored-by: Yong He <yonghe@outlook.com>

3 files changed, 73 insertions, 3 deletions

diff --git a/source/slang/diff.meta.slang b/source/slang/diff.meta.slang
index 5817ac7cf..0026a76f9 100644
--- a/source/slang/diff.meta.slang
+++ b/source/slang/diff.meta.slang
@@ -1018,17 +1018,17 @@ DifferentialPair<vector<T, 3>> __d_cross(DifferentialPair<vector<T, 3>> a, Diffe
     T aybz = a.p.y * b.p.z;
     T azby = a.p.z * b.p.y;
     T px = aybz - azby;
-    T dx = (b.p.z - azby) * a.d.y + (a.p.y - azby) * b.d.z + (aybz - b.p.y) * a.d.z + (aybz - a.p.z) * b.d.y;
+    T dx = a.d.y * b.p.z + a.p.y * b.d.z - a.d.z * b.p.y - a.p.z * b.d.y;
 
     T azbx = a.p.z * b.p.x;
     T axbz = a.p.x * b.p.z;
     T py = azbx - axbz;
-    T dy = (b.p.x - axbz) * a.d.z + (a.p.z - axbz) * b.d.x + (azbx - b.p.z) * a.d.x + (azbx - a.p.x) * b.d.z;
+    T dy = a.d.z * b.p.x + a.p.z * b.d.x - a.d.x * b.p.z - a.p.x * b.d.z;
 
     T axby = a.p.x * b.p.y;
     T aybx = a.p.y * b.p.x;
     T pz = axby - aybx;
-    T dz = (b.p.y - aybx) * a.d.x + (a.p.x - aybx) * b.d.y + (axby - b.p.x) * a.d.y + (axby - a.p.y) * b.d.x;
+    T dz = a.d.x * b.p.y + a.p.x * b.d.y - a.d.y * b.p.x - a.p.y * b.d.x;
     
     return DifferentialPair<vector<T, 3>>(vector<T, 3>(px, py, pz), vector<T, 3>.Differential(dx, dy, dz));
 }
diff --git a/tests/autodiff/geometry/ray-triangle-intersect.slang b/tests/autodiff/geometry/ray-triangle-intersect.slang
new file mode 100644
index 000000000..3e729c2ab
--- /dev/null
+++ b/tests/autodiff/geometry/ray-triangle-intersect.slang
@@ -0,0 +1,63 @@
+//TEST(compute):COMPARE_COMPUTE_EX:-slang -compute -shaderobj -output-using-type
+//TEST(compute, vulkan):COMPARE_COMPUTE_EX:-vk -compute -shaderobj -output-using-type
+
+//TEST_INPUT:ubuffer(data=[0 0 0 0 0 0], stride=4):out,name=outputBuffer
+RWStructuredBuffer<float> outputBuffer;
+
+[Differentiable]
+[PreferRecompute]
+float3 diffRayIntersectTriangle(no_diff float3 rayOrigin, float3 rayDir, no_diff float3 p[3])
+{
+    float3 e1 = p[1] - p[0];
+    float3 e2 = p[2] - p[0];
+    float3 pVec = cross(rayDir, e2);
+    float divisor = dot(pVec, e1);
+    float3 s = rayOrigin - p[0];
+    float u = dot(s, pVec) / divisor;
+    float3 qVec = cross(s, e1);
+    float v = dot(rayDir, qVec) / divisor;
+    float t = dot(e2, qVec) / divisor;
+    return float3(u, v, t);
+}
+
+[Differentiable]
+[PreferRecompute]
+float3 diffRayIntersectTriangle2(no_diff float3 rayOrigin, float3 rayTarget, no_diff float3 p[3])
+{
+    float3 rayDir = normalize(rayTarget - rayOrigin);
+    float3 uvt = diffRayIntersectTriangle(rayOrigin, rayDir, p);
+    float3 result = (1.f - uvt.x - uvt.y) * p[0] + uvt.x * p[1] + uvt.y * p[2];
+    return result;
+}
+
+[numthreads(1, 1, 1)]
+void computeMain(uint3 dispatchThreadID : SV_DispatchThreadID)
+{
+    float3 shadePos = float3(0.674034, 0.0, 0.123171);
+    float3 targetPos = float3(0.5, 0.2, -1.0);
+    float3 triPos[3] = { float3(0.0, 1.0, -1.0), float3(1.0, 1.0, 0.0), float3(0.0, 1.0, 0.0) };
+
+    // Forward-mode
+    DifferentialPair<float3> dpIsectPos = fwd_diff(diffRayIntersectTriangle2)(
+        shadePos,
+        DifferentialPair<float3>(targetPos, float3(1.0, 0.0, 0.0)),
+        triPos
+    );
+
+    outputBuffer[0] = dpIsectPos.d[0]; // Expect: 5.0
+    outputBuffer[1] = dpIsectPos.d[1]; // Expect: 0.0
+    outputBuffer[2] = dpIsectPos.d[2]; // Expect: 0.0
+
+    // Reverse-mode
+    DifferentialPair<float3> dpTargetPos = diffPair(targetPos, float3(0.f));
+    bwd_diff(diffRayIntersectTriangle2)(
+        shadePos,
+        dpTargetPos,
+        triPos,
+        float3(1.f, 1.f, 1.f)
+    );
+
+    outputBuffer[3] = dpTargetPos.d[0]; // Expect: 5.0
+    outputBuffer[4] = dpTargetPos.d[1]; // Expect: 32.4301
+    outputBuffer[5] = dpTargetPos.d[2]; // Expect: 5.0
+}
diff --git a/tests/autodiff/geometry/ray-triangle-intersect.slang.expected.txt b/tests/autodiff/geometry/ray-triangle-intersect.slang.expected.txt
new file mode 100644
index 000000000..318cb3cb1
--- /dev/null
+++ b/tests/autodiff/geometry/ray-triangle-intersect.slang.expected.txt
@@ -0,0 +1,7 @@
+type: float
+5.0
+0.0
+0.0
+5.0
+32.43011856
+5.0