Added forward-mode autodiff support for more instructions (#2331)

* Merge slang-ir-diff-jvp.cpp

* Added support and tests for other float vector types

* Added swizzle test and code to handle it (tests failing currently)

* Fixed one test, the other is still pending

* Fixed instruction cloning logic to avoid modifying original function

* Fixed an issue with custom 'pow_jvp' and added support for vector contructor

* Minor update to comments

* Fixed support for division

* Fixed an issue with uninitialized diagnostic sink

* Moved derivative processing to after mandatory inlining.

Skip instructions that don't have side-effects and aren't used by anything.

* WIP: Handling unconditional control flow and multi-block functions

* Support for unconditional multi-block functions

* Added a dead code elimination step to the derivative pass

* Changed name of 'hasNoSideEffects()'

13 files changed, 312 insertions, 82 deletions

diff --git a/source/slang/slang-check-expr.cpp b/source/slang/slang-check-expr.cpp
index 67e8bf650..1895da70b 100644
--- a/source/slang/slang-check-expr.cpp
+++ b/source/slang/slang-check-expr.cpp
@@ -1525,15 +1525,14 @@ namespace Slang
 
     Type* primalToJVPParamType(ASTBuilder* builder, Type* primalType)
     {
-        // Only float and float3 types can be differentiated for now.
+        // Only float and vector<float> types can be differentiated for now.
         
         if (primalType->equals(builder->getFloatType()))
             return primalType;
         else if (auto primalVectorType = as<VectorExpressionType>(primalType))
         {
-            // TODO(sai): There's probably a more elegant way to check if a type is a float3?
-            if (getIntVal(primalVectorType->elementCount) == 3 && primalVectorType->elementType->equals(builder->getFloatType()))
-                return primalVectorType;
+            if (auto jvpElementType = primalToJVPParamType(builder, primalVectorType->elementType))
+                return builder->getVectorType(jvpElementType, primalVectorType->elementCount);
         }
         else if (auto primalOutType = as<OutType>(primalType))
         {
diff --git a/source/slang/slang-ir-diff-jvp.cpp b/source/slang/slang-ir-diff-jvp.cpp
index f5afccd0c..fd1d0086d 100644
--- a/source/slang/slang-ir-diff-jvp.cpp
+++ b/source/slang/slang-ir-diff-jvp.cpp
@@ -4,6 +4,7 @@
 #include "slang-ir.h"
 #include "slang-ir-insts.h"
 #include "slang-ir-clone.h"
+#include "slang-ir-dce.h"
 
 namespace Slang
 {
@@ -103,6 +104,11 @@ struct JVPTranscriber
         if (IRType* typeD = differentiateType(builder, paramP->getFullType()))
         {
             IRParam* paramD = builder->emitParam(typeD);
+
+            auto nameHintD = getJVPVarName(paramP);
+            if (nameHintD.getLength() > 0)
+                builder->addNameHintDecoration(paramD, nameHintD.getUnownedSlice());
+
             SLANG_ASSERT(paramD);
             return paramD;
         }
@@ -118,6 +124,7 @@ struct JVPTranscriber
         {   
             auto newParamP = builder->emitParam(inoutTypeP->getValueType());
             cloneEnv.mapOldValToNew.Add(paramP, newParamP);
+            cloneInstDecorationsAndChildren(&cloneEnv, builder->getSharedBuilder(), paramP, newParamP);
 
             return newParamP;
         }
@@ -138,7 +145,7 @@ struct JVPTranscriber
         List<IRParam*> newParamListP;
         for (auto paramP : paramListP)
         {
-            if(requiresPrimalClone(builder, paramP))
+            if(isPurelyFunctional(builder, paramP))
                 newParamListP.add(as<IRParam>(emitInputParam(builder, paramP)));
         }
 
@@ -154,12 +161,30 @@ struct JVPTranscriber
         return newParamListD;
     }
 
+    // Returns "d<var-name>" to use as a name hint for variables and parameters.
+    // If no primal name is available, returns a blank string.
+    // 
+    String getJVPVarName(IRInst* varP)
+    {
+        if (auto namehintDecoration = varP->findDecoration<IRNameHintDecoration>())
+        {
+            return ("d" + String(namehintDecoration->getName()));
+        }
+
+        return String("");
+    }
+
     IRInst* differentiateVar(IRBuilder* builder, IRVar* varP)
     {
         if (IRType* typeD = differentiateType(builder, varP->getDataType()->getValueType()))
         {
             IRVar* varD = builder->emitVar(typeD);
             SLANG_ASSERT(varD);
+
+            auto nameHintD = getJVPVarName(varP);
+            if (nameHintD.getLength() > 0)
+                builder->addNameHintDecoration(varD, nameHintD.getUnownedSlice());
+
             return varD;
         }
         return nullptr;
@@ -270,7 +295,7 @@ struct JVPTranscriber
 
     IRInst* differentiateReturn(IRBuilder* builder, IRReturn* returnP)
     {
-        IRInst* returnVal = findCloneForOperand(&cloneEnv, returnP->getVal());
+        IRInst* returnVal = returnP->getVal();
         if (auto returnValD = getDifferentialInst(returnVal, nullptr))
         {   
             IRReturn* returnD = as<IRReturn>(builder->emitReturn(returnValD));
@@ -366,6 +391,71 @@ struct JVPTranscriber
         return nullptr;
     }
 
+    IRInst* differentiateSwizzle(IRBuilder* builder, IRSwizzle* swizzleP)
+    {
+        if (auto baseD = getDifferentialInst(swizzleP->getBase(), nullptr))
+        {
+            List<IRInst*> swizzleIndices;
+            for (UIndex ii = 0; ii < swizzleP->getElementCount(); ii++)
+                swizzleIndices.add(swizzleP->getElementIndex(ii));
+            
+            return builder->emitSwizzle(differentiateType(builder, swizzleP->getDataType()),
+                                        baseD,
+                                        swizzleP->getElementCount(),
+                                        swizzleIndices.getBuffer());
+        }
+        return nullptr;
+    }
+
+    IRInst* differentiateByPassthrough(IRBuilder* builder, IRInst* origInst)
+    {
+        UCount operandCount = origInst->getOperandCount();
+
+        List<IRInst*> diffOperands;
+        for (UIndex ii = 0; ii < operandCount; ii++)
+        {
+            // If the operand has a differential version, replace the original with the 
+            // differential.
+            // Otherwise, abandon the differentiation attempt and assume that origInst 
+            // cannot (or does not need to) be differentiated.
+            // 
+            if (auto diffInst = getDifferentialInst(origInst->getOperand(ii), nullptr))
+                diffOperands.add(diffInst);
+            else
+                return nullptr;
+        }
+        
+        return builder->emitIntrinsicInst(
+                    differentiateType(builder, origInst->getDataType()),
+                    origInst->getOp(),
+                    operandCount,
+                    diffOperands.getBuffer());
+    }
+
+    IRInst* handleControlFlow(IRBuilder* builder, IRInst* origInst)
+    {
+        switch(origInst->getOp())
+        {
+            case kIROp_unconditionalBranch:
+                auto origBranch = as<IRUnconditionalBranch>(origInst);
+
+                // Branches with extra operands not handled currently.
+                if (origBranch->getOperandCount() > 1)
+                    break;
+
+                if (auto diffBlock = getDifferentialInst(origBranch->getTargetBlock(), nullptr))
+                    return builder->emitBranch(as<IRBlock>(diffBlock));
+                else
+                    return nullptr;
+        }
+
+        getSink()->diagnose(
+            origInst->sourceLoc,
+            Diagnostics::unimplemented,
+            "attempting to differentiate unhandled control flow");
+        return nullptr;
+    }
+
     // In differential computation, the 'default' differential value is always zero.
     // This is a consequence of differential computing being inherently linear. As a 
     // result, it's useful to have a method to generate zero literals of any (arithmetic) type.
@@ -390,11 +480,11 @@ struct JVPTranscriber
 
     // Logic for whether a primal instruction needs to be replicated
     // in the differential function. We detect and avoid replicating 
-    // side-effect instructions.
+    // 'side-effect' instructions.
     // 
-    bool requiresPrimalClone(IRBuilder*, IRInst* instP)
+    bool isPurelyFunctional(IRBuilder*, IRInst* instP)
     {
-        if (as<IRReturn>(instP))
+        if (as<IRTerminatorInst>(instP))
             return false;
         else if (auto paramP = as<IRParam>(instP))
         {
@@ -425,38 +515,37 @@ struct JVPTranscriber
 
     IRInst* transcribe(IRBuilder* builder, IRInst* oldInstP)
     {
-        IRInst* instP = oldInstP;
 
-        // Clone the old instruction, but only if it's safe to do so.
+        // Clone the old instruction into the new differential function.
+        // 
+        IRInst* instP = cloneInst(&cloneEnv, builder, oldInstP);
+
+        SLANG_ASSERT(instP);
+
+        IRInst* instD = differentiateInst(builder, instP);
+        
+        // In case it's not safe to clone the old instruction, 
+        // remove it from the graph.
         // For instance, instructions that handle control flow 
         // (return statements) shouldn't be replicated.
         //
-        if (requiresPrimalClone(builder, oldInstP))
-            instP = cloneInst(&cloneEnv, builder, oldInstP);
+        if (isPurelyFunctional(builder, oldInstP))
+            mapDifferentialInst(instP, instD);
         else
         {
-            // We replace the operands of the old instruction with their clones,
-            // if available.
-            // 
-            for(UInt ii = 0; ii < oldInstP->getOperandCount(); ++ii)
-            {
-                auto oldOperand = oldInstP->getOperand(ii);
-                auto newOperand = findCloneForOperand(&cloneEnv, oldOperand);
+            // This inst should never have been used.
+            SLANG_ASSERT(instP->firstUse == nullptr);
 
-                instP->getOperands()[ii].init(instP, newOperand);
-            }
+            instP->removeAndDeallocate();
+            mapDifferentialInst(oldInstP, instD);
         }
-        SLANG_ASSERT(instP);
-
-        IRInst* instD = differentiateInst(builder, instP);
-
-        mapDifferentialInst(instP, instD);
 
         return instD;
     }
 
     IRInst* differentiateInst(IRBuilder* builder, IRInst* instP)
     {
+        // Handle common operations
         switch (instP->getOp())
         {
         case kIROp_Var:
@@ -474,6 +563,7 @@ struct JVPTranscriber
         case kIROp_Add:
         case kIROp_Mul:
         case kIROp_Sub:
+        case kIROp_Div:
             return differentiateBinaryArith(builder, instP);
 
         case kIROp_Construct:
@@ -481,13 +571,33 @@ struct JVPTranscriber
         
         case kIROp_Call:
             return differentiateCall(builder, as<IRCall>(instP));
+        
+        case kIROp_swizzle:
+            return differentiateSwizzle(builder, as<IRSwizzle>(instP));
+        
+        case kIROp_constructVectorFromScalar:
+            return differentiateByPassthrough(builder, instP);
 
-        default:
-            getSink()->diagnose(instP->sourceLoc,
+        case kIROp_unconditionalBranch:
+        case kIROp_conditionalBranch:
+            return handleControlFlow(builder, instP);
+
+        }
+    
+        // If none of the cases have been hit, check if the instruction is a
+        // type.
+        // For now we don't have logic to differentiate types that appear in blocks.
+        // So, we ignore them.
+        //
+        if (as<IRType>(instP))
+            return nullptr;
+        
+        
+        // If we reach this statement, the instruction type is likely unhandled.
+        getSink()->diagnose(instP->sourceLoc,
                     Diagnostics::unimplemented,
                     "this instruction cannot be differentiated");
-            return nullptr;
-        }
+        return nullptr;
     }
 };
 
@@ -531,27 +641,6 @@ struct IRWorkQueue
 
 struct JVPDerivativeContext
 {
-    // This type passes over the module and generates
-    // forward-mode derivative versions of functions 
-    // that are explicitly marked for it.
-    //
-    IRModule*                       module;
-
-    // Shared builder state for our derivative passes.
-    SharedIRBuilder                 sharedBuilderStorage;
-
-    // A transcriber object that handles the main job of 
-    // processing instructions while maintaining state.
-    //
-    JVPTranscriber                  transcriberStorage;
-    
-    // Diagnostic object from the compile request for
-    // error messages.
-    DiagnosticSink*                 sink;
-
-    // Work queue to hold a stream of instructions that need
-    // to be checked for references to derivative functions.
-    IRWorkQueue                    workQueueStorage;
 
     DiagnosticSink* getSink()
     {
@@ -726,13 +815,22 @@ struct JVPDerivativeContext
             builder->addNameHintDecoration(jvpFn, jvpName);
 
         builder->setInsertInto(jvpFn);
-
-        // Start with _extremely_ basic functions
-        SLANG_ASSERT(primalFn->getFirstBlock() == primalFn->getLastBlock());
         
+        // Emit a block instruction for every block in the function, and map it as the 
+        // corresponding differential.
+        //
+        for (auto block = primalFn->getFirstBlock(); block; block = block->getNextBlock())
+        {
+            auto jvpBlock = builder->emitBlock();
+            transcriberStorage.mapDifferentialInst(block, jvpBlock);
+        }
+
+        // Go back over the blocks, and process the children of each block.
         for (auto block = primalFn->getFirstBlock(); block; block = block->getNextBlock())
         {
-            emitJVPBlock(builder, primalFn->getFirstBlock());
+            auto jvpBlock = as<IRBlock>(transcriberStorage.getDifferentialInst(block, block));
+            SLANG_ASSERT(jvpBlock);
+            emitJVPBlock(builder, block, jvpBlock);
         }
 
         return jvpFn;
@@ -759,7 +857,6 @@ struct JVPDerivativeContext
         return name;
     }
 
-
     IRBlock* emitJVPBlock(IRBuilder*    builder, 
                           IRBlock*      primalBlock,
                           IRBlock*      jvpBlock = nullptr)
@@ -789,6 +886,35 @@ struct JVPDerivativeContext
         return jvpBlock;
     }
 
+    JVPDerivativeContext(IRModule* module, DiagnosticSink* sink) : module(module), sink(sink)
+    {
+        transcriberStorage.sink = sink;
+    }
+
+    protected:
+
+    // This type passes over the module and generates
+    // forward-mode derivative versions of functions 
+    // that are explicitly marked for it.
+    //
+    IRModule*                       module;
+
+    // Shared builder state for our derivative passes.
+    SharedIRBuilder                 sharedBuilderStorage;
+
+    // A transcriber object that handles the main job of 
+    // processing instructions while maintaining state.
+    //
+    JVPTranscriber                  transcriberStorage;
+    
+    // Diagnostic object from the compile request for
+    // error messages.
+    DiagnosticSink*                 sink;
+
+    // Work queue to hold a stream of instructions that need
+    // to be checked for references to derivative functions.
+    IRWorkQueue                    workQueueStorage;
+
 };
 
 // Set up context and call main process method.
@@ -798,9 +924,13 @@ bool processJVPDerivativeMarkers(
         DiagnosticSink*                     sink,
         IRJVPDerivativePassOptions const&)
 {
-    JVPDerivativeContext context;
-    context.module = module;
-    context.sink = sink;
+    JVPDerivativeContext context(module, sink);
+    
+    // Simplify module to remove dead code.
+    IRDeadCodeEliminationOptions options;
+    options.keepExportsAlive = true;
+    options.keepLayoutsAlive = true;
+    eliminateDeadCode(module, options);
 
     return context.processModule();
 }
diff --git a/source/slang/slang-lower-to-ir.cpp b/source/slang/slang-lower-to-ir.cpp
index 628f37c5b..81201f5f8 100644
--- a/source/slang/slang-lower-to-ir.cpp
+++ b/source/slang/slang-lower-to-ir.cpp
@@ -8516,15 +8516,6 @@ RefPtr<IRModule> generateIRForTranslationUnit(
 #endif
 
     validateIRModuleIfEnabled(compileRequest, module);
-    
-    // Process higher-order-function calls before any optimization passes
-    // to allow the optimizations to affect the generated funcitons.
-    // 1. Process JVP derivative functions.
-    processJVPDerivativeMarkers(module, compileRequest->getSink());
-    // 2. Process VJP derivative functions.
-    // processVJPDerivativeMarkers(module); // Disabled currently. No impl yet.
-    // 3. Replace JVP & VJP calls.
-    processDerivativeCalls(module);
 
 
     // We will perform certain "mandatory" optimization passes now.
@@ -8560,6 +8551,15 @@ RefPtr<IRModule> generateIRForTranslationUnit(
     // temporaries whenever possible.
     constructSSA(module);
 
+    // Process higher-order-function calls before any optimization passes
+    // to allow the optimizations to affect the generated funcitons.
+    // 1. Process JVP derivative functions.
+    processJVPDerivativeMarkers(module, compileRequest->getSink());
+    // 2. Process VJP derivative functions.
+    // processVJPDerivativeMarkers(module); // Disabled currently. No impl yet.
+    // 3. Replace JVP & VJP calls.
+    processDerivativeCalls(module);
+
     // Do basic constant folding and dead code elimination
     // using Sparse Conditional Constant Propagation (SCCP)
     //
diff --git a/tests/autodiff/arithmetic-jvp.slang b/tests/autodiff/arithmetic-jvp.slang
index 8632aa332..4e4d200e1 100644
--- a/tests/autodiff/arithmetic-jvp.slang
+++ b/tests/autodiff/arithmetic-jvp.slang
@@ -1,7 +1,7 @@
 //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], stride=4):out,name=outputBuffer
+//TEST_INPUT:ubuffer(data=[0 0 0 0 0], stride=4):out,name=outputBuffer
 RWStructuredBuffer<float> outputBuffer;
 
 __differentiate_jvp float f(float x)
@@ -27,11 +27,16 @@ __differentiate_jvp float h(float x, float y)
     return m * n + 2 * x * y;
 }
 
+__differentiate_jvp float j(float x, float y)
+{
+    float m = x / y;
+    return m * y;
+}
 
 [numthreads(1, 1, 1)]
 void computeMain(uint3 dispatchThreadID : SV_DispatchThreadID)
 {
-     {
+    {
         float a = 2.0;
         float b = 1.5;
         float da = 1.0;
@@ -41,5 +46,6 @@ void computeMain(uint3 dispatchThreadID : SV_DispatchThreadID)
         outputBuffer[1] = __jvp(f)(a, 0.0);       // Expect: 0
         outputBuffer[2] = __jvp(g)(a, da);        // Expect: 2
         outputBuffer[3] = __jvp(h)(a, b, da, db); // Expect: 8
+        outputBuffer[4] = __jvp(j)(a, b, da, db); // Expect: 1
     }
 }
diff --git a/tests/autodiff/arithmetic-jvp.slang.expected.txt b/tests/autodiff/arithmetic-jvp.slang.expected.txt
index 0545c08a1..597fe4035 100644
--- a/tests/autodiff/arithmetic-jvp.slang.expected.txt
+++ b/tests/autodiff/arithmetic-jvp.slang.expected.txt
@@ -2,4 +2,5 @@ type: float
 1.0
 0.0
 2.0
-8.0
\ No newline at end of file
+8.0
+1.0
\ No newline at end of file
diff --git a/tests/autodiff/inout-parameters-jvp.slang b/tests/autodiff/inout-parameters-jvp.slang
index 989e56c02..40e9d30ca 100644
--- a/tests/autodiff/inout-parameters-jvp.slang
+++ b/tests/autodiff/inout-parameters-jvp.slang
@@ -4,6 +4,14 @@
 //TEST_INPUT:ubuffer(data=[0 0 0 0], stride=4):out,name=outputBuffer
 RWStructuredBuffer<float> outputBuffer;
 
+
+__differentiate_jvp void g(float x, float y, inout float z)
+{
+    float m = x + y;
+    float n = x - y;
+    z += m * n + 2 * x * y;
+}
+
 __differentiate_jvp void h(float x, float y, inout float z)
 {
     float m = x + y;
@@ -27,4 +35,9 @@ void computeMain(uint3 dispatchThreadID : SV_DispatchThreadID)
     outputBuffer[0] = dz; // Expect: 12.0
     outputBuffer[1] = z;  // Expect: 1.0
 
+    __jvp(g)(x, y, z, dx, dy, dz);
+
+    outputBuffer[2] = dz; // Expect: 21.5
+    outputBuffer[3] = z;  // Expect: 1.0
+
 }
\ No newline at end of file
diff --git a/tests/autodiff/inout-parameters-jvp.slang.expected.txt b/tests/autodiff/inout-parameters-jvp.slang.expected.txt
index d8a590c0e..c48ef7bf6 100644
--- a/tests/autodiff/inout-parameters-jvp.slang.expected.txt
+++ b/tests/autodiff/inout-parameters-jvp.slang.expected.txt
@@ -1,5 +1,5 @@
 type: float
 12.0
 1.0
-0.0
-0.0
\ No newline at end of file
+21.5
+1.0
\ No newline at end of file
diff --git a/tests/autodiff/nested-jvp.slang b/tests/autodiff/nested-jvp.slang
index c37641c91..222396ec8 100644
--- a/tests/autodiff/nested-jvp.slang
+++ b/tests/autodiff/nested-jvp.slang
@@ -20,7 +20,7 @@ float max_(float x, float y)
 
 float pow_jvp(float x, float n, float dx, float dn)
 {
-    return dx * n * pow(x, n-1) + dn * pow(x, n) * log(x);
+    return dx * n * pow(x, n-1) + ((dn != 0.0) ? (dn * pow(x, n) * log(x)) : 0.0);
 }
 
 
@@ -36,6 +36,10 @@ __differentiate_jvp float3 fresnel(float3 f0, float3 f90, float cosTheta)
     return f0 + (f90 - f0) * pow_(max_(1 - cosTheta, 0.0), 5);
 }
 
+__differentiate_jvp float g(float a, float b, float c)
+{
+    return fresnel(float3(a), float3(b), 2 * c * c).y;
+}
 
 [numthreads(1, 1, 1)]
 void computeMain(uint3 dispatchThreadID : SV_DispatchThreadID)
@@ -50,5 +54,17 @@ void computeMain(uint3 dispatchThreadID : SV_DispatchThreadID)
         float d_cosTheta = 1.0;
 
         outputBuffer[0] = __jvp(fresnel)(f0, f90, cosTheta, d_f0, d_f90, d_cosTheta).y; // Expect: -0.031250
+
+        float a = 1.0;
+        float b = -0.4;
+        float c = 0.5;
+
+        float da = -0.4;
+        float db = -1.0;
+        float dc = 0.2;
+
+        outputBuffer[1] = __jvp(g)(a, b, c, da, db, dc);    // Expect: -0.24375
+        outputBuffer[2] = g(a, b, c);                       // Expect: 0.95625
+        outputBuffer[3] = __jvp(g)(a, b, 3.0, da, db, dc);  // Expect: -0.4;
     }
 }
diff --git a/tests/autodiff/nested-jvp.slang.expected.txt b/tests/autodiff/nested-jvp.slang.expected.txt
index 107153351..f200c0a88 100644
--- a/tests/autodiff/nested-jvp.slang.expected.txt
+++ b/tests/autodiff/nested-jvp.slang.expected.txt
@@ -1,5 +1,5 @@
 type: float
 -0.031250
-0.0
-0.0
-0.0
\ No newline at end of file
+-0.243750
+0.95625
+-0.4
\ No newline at end of file
diff --git a/tests/autodiff/vector-arithmetic-jvp.slang b/tests/autodiff/vector-arithmetic-jvp.slang
index 5852b49a7..2b43f1752 100644
--- a/tests/autodiff/vector-arithmetic-jvp.slang
+++ b/tests/autodiff/vector-arithmetic-jvp.slang
@@ -16,6 +16,20 @@ __differentiate_jvp float3 g(float3 x, float3 y)
     return a * b + 2 * x * y;
 }
 
+__differentiate_jvp float2 h(float2 x, float2 y)
+{
+    float2 a = x + y;
+    float2 b = x - y;
+    return a * b + 2 * x * y;
+}
+
+__differentiate_jvp float4 j(float4 x, float4 y)
+{
+    float4 a = x + y;
+    float4 b = x - y;
+    return a * b + 2 * x * y;
+}
+
 [numthreads(1, 1, 1)]
 void computeMain(uint3 dispatchThreadID : SV_DispatchThreadID)
 {
@@ -24,7 +38,15 @@ void computeMain(uint3 dispatchThreadID : SV_DispatchThreadID)
         float3 b = float3(1.5, 1.5, 1.5);
         float3 da = float3(1.0, 1.0, 1.0);
 
-        outputBuffer[0] = __jvp(f)(a, da).z;                          // Expect: 1
-        outputBuffer[1] = __jvp(g)(a, b, da, float3(2.0, 1.0, 0.0)).y;       // Expect: 8
+        float2 a2 = float2(2.0, 1.0);
+        float2 b2 = float2(1.5, -2.0);
+
+        float4 a4 = float4(2.0, 1.0, 0.0, 2.0);
+        float4 b4 = float4(1.5, -2.0, 1.0, 1.5);
+
+        outputBuffer[0] = __jvp(f)(a, da).z;                                            // Expect: 1
+        outputBuffer[1] = __jvp(g)(a, b, da, float3(2.0, 1.0, 0.0)).y;                  // Expect: 8
+        outputBuffer[2] = __jvp(h)(a2, b2, float2(1.0, 0.0), float2(1.0, 1.0)).x;       // Expect: 8
+        outputBuffer[3] = __jvp(j)(a4, b4, float4(1.0), float4(2.0)).w;                 // Expect: 9
     }
 }
diff --git a/tests/autodiff/vector-arithmetic-jvp.slang.expected.txt b/tests/autodiff/vector-arithmetic-jvp.slang.expected.txt
index adbb9a448..9202deda2 100644
--- a/tests/autodiff/vector-arithmetic-jvp.slang.expected.txt
+++ b/tests/autodiff/vector-arithmetic-jvp.slang.expected.txt
@@ -1,5 +1,5 @@
 type: float
 1.0
 8.0
-0.0
-0.0
\ No newline at end of file
+8.0
+9.0
\ No newline at end of file
diff --git a/tests/autodiff/vector-swizzle-jvp.slang b/tests/autodiff/vector-swizzle-jvp.slang
new file mode 100644
index 000000000..6722b54dc
--- /dev/null
+++ b/tests/autodiff/vector-swizzle-jvp.slang
@@ -0,0 +1,38 @@
+//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], stride=4):out,name=outputBuffer
+RWStructuredBuffer<float> outputBuffer;
+
+__differentiate_jvp float2 f(float3 x)
+{
+    return x.zy;
+}
+
+__differentiate_jvp float2 g(float3 x, float4 y)
+{
+    float3 a = x + y.zyx;
+    float2 b = x.zx - y.yw;
+    return a.yz * b + 2 * x.xy * y.yz;
+}
+
+[numthreads(1, 1, 1)]
+void computeMain(uint3 dispatchThreadID: SV_DispatchThreadID)
+{
+    {
+        float3 a = float3(2.0, 2.0, 2.0);
+        float3 da = float3(1.0, 0.5, 1.0);
+
+        outputBuffer[0] = __jvp(f)(a, da).x;  // Expect: 1
+        outputBuffer[1] = __jvp(f)(a, da).y;  // Expect: 0.5
+        
+        float3 x = float3(0.5, 2.0, 0.5);
+        float4 y = float4(-1.5, 1.0, 4.0, 2.0);
+        float3 dx = float3(1.0, 0.0, -1.0);
+        float4 dy = float4(0.0, 0.5, -0.25, 1.0);
+
+        outputBuffer[2] = __jvp(g)(x, y, dx, dy).x; // Expect: -2.25
+        outputBuffer[3] = __jvp(g)(x, y, dx, dy).y; // Expect: 0.5
+        
+    }
+}
diff --git a/tests/autodiff/vector-swizzle-jvp.slang.expected.txt b/tests/autodiff/vector-swizzle-jvp.slang.expected.txt
new file mode 100644
index 000000000..3a49054c9
--- /dev/null
+++ b/tests/autodiff/vector-swizzle-jvp.slang.expected.txt
@@ -0,0 +1,5 @@
+type: float
+1.0
+0.5
+-2.25
+0.5
\ No newline at end of file