Refactor Auto-diff passes (#2526)

* Initial refactor

* Refactor passes tests

* Removed Differential Bottom references from the IR side

1 files changed, 0 insertions, 3197 deletions

diff --git a/source/slang/slang-ir-diff-jvp.cpp b/source/slang/slang-ir-diff-jvp.cpp
deleted file mode 100644
index c9ca687e4..000000000
--- a/source/slang/slang-ir-diff-jvp.cpp
+++ /dev/null
@@ -1,3197 +0,0 @@
-// slang-ir-diff-jvp.cpp
-#include "slang-ir-diff-jvp.h"
-
-#include "slang-ir-clone.h"
-#include "slang-ir-dce.h"
-#include "slang-ir-eliminate-phis.h"
-#include "slang-ir-util.h"
-#include "slang-ir-inst-pass-base.h"
-
-// origX, primalX, diffX
-// origX -> primalX (cloneEnv)
-// origX -> diffX (instMapD)
-
-namespace Slang
-{
-
-namespace
-{
-
-IRInst* _lookupWitness(IRBuilder* builder, IRInst* witness, IRInst* requirementKey)
-{
-    if (auto witnessTable = as<IRWitnessTable>(witness))
-    {
-        for (auto entry : witnessTable->getEntries())
-        {
-            if (entry->getRequirementKey() == requirementKey)
-                return entry->getSatisfyingVal();
-        }
-    }
-    else if (auto witnessTableParam = as<IRParam>(witness))
-    {
-        return builder->emitLookupInterfaceMethodInst(
-            builder->getTypeKind(),
-            witnessTableParam,
-            requirementKey);
-    }
-    return nullptr;
-}
-
-}
-
-struct DifferentialPairTypeBuilder
-{
-
-    IRStructField* findField(IRInst* type, IRStructKey* key)
-    {
-        if (auto irStructType = as<IRStructType>(type))
-        {
-            for (auto field : irStructType->getFields())
-            {
-                if (field->getKey() == key)
-                {
-                    return field;
-                }
-            }
-        }
-        else if (auto irSpecialize = as<IRSpecialize>(type))
-        {
-            if (auto irGeneric = as<IRGeneric>(irSpecialize->getBase()))
-            {
-                if (auto irGenericStructType = as<IRStructType>(findInnerMostGenericReturnVal(irGeneric)))
-                {
-                    return findField(irGenericStructType, key);
-                }
-            }
-        }
-
-        return nullptr;
-    }
-
-    IRInst* findSpecializationForParam(IRInst* specializeInst, IRInst* genericParam)
-    {
-        // Get base generic that's being specialized.
-        auto genericType = as<IRGeneric>(as<IRSpecialize>(specializeInst)->getBase());
-        SLANG_ASSERT(genericType);
-        
-        // Find the index of genericParam in the base generic.
-        int paramIndex = -1;
-        int currentIndex = 0;
-        for (auto param : genericType->getParams())
-        {
-            if (param == genericParam)
-                paramIndex = currentIndex;
-            currentIndex ++;
-        }
-
-        SLANG_ASSERT(paramIndex >= 0);
-
-        // Return the corresponding operand in the specialization inst.
-        return specializeInst->getOperand(1 + paramIndex);
-    }
-
-    IRInst* emitFieldAccessor(IRBuilder* builder, IRInst* baseInst, IRStructKey* key)
-    {
-        IRInst* pairType = nullptr;
-        if (auto basePtrType = as<IRPtrTypeBase>(baseInst->getDataType()))
-        {
-            auto baseTypeInfo = lowerDiffPairType(builder, basePtrType->getValueType());
-
-            // TODO(sai): Not sure at the moment how to handle diff-bottom pointer types,
-            // especially since we probably don't need diff bottom anymore.
-            // 
-            SLANG_ASSERT(!baseTypeInfo.isTrivial);
-
-            pairType = builder->getPtrType(kIROp_PtrType, (IRType*)baseTypeInfo.loweredType);
-        }
-        else
-        {
-            auto baseTypeInfo = lowerDiffPairType(builder, baseInst->getDataType());
-            if (baseTypeInfo.isTrivial)
-            {
-                if (key == globalPrimalKey)
-                    return baseInst;
-                else
-                    return builder->getDifferentialBottom();
-            }
-
-            pairType = baseTypeInfo.loweredType;
-        }
-
-        if (auto basePairStructType = as<IRStructType>(pairType))
-        {
-            return as<IRFieldExtract>(builder->emitFieldExtract(
-                    findField(basePairStructType, key)->getFieldType(),
-                    baseInst,
-                    key
-                ));
-        }
-        else if (auto ptrType = as<IRPtrTypeBase>(pairType))
-        {
-            if (auto ptrInnerSpecializedType = as<IRSpecialize>(ptrType->getValueType()))
-            {
-                auto genericType = findInnerMostGenericReturnVal(as<IRGeneric>(ptrInnerSpecializedType->getBase()));
-                if (auto genericBasePairStructType = as<IRStructType>(genericType))
-                {
-                    return as<IRFieldAddress>(builder->emitFieldAddress(
-                        builder->getPtrType((IRType*)
-                            findSpecializationForParam(
-                                ptrInnerSpecializedType,
-                                findField(ptrInnerSpecializedType, key)->getFieldType())),
-                        baseInst,
-                        key
-                    ));
-                }
-            }
-            else if (auto ptrBaseStructType = as<IRStructType>(ptrType->getValueType()))
-            {
-                return as<IRFieldAddress>(builder->emitFieldAddress(
-                    builder->getPtrType((IRType*)
-                            findField(ptrBaseStructType, key)->getFieldType()),
-                    baseInst,
-                    key));
-            }
-        }
-        else if (auto specializedType = as<IRSpecialize>(pairType))
-        {
-            // TODO: Stopped here -> The type being emitted is incorrect. don't emit the generic's
-            // type, emit the specialization type.
-            // 
-            auto genericType = findInnerMostGenericReturnVal(as<IRGeneric>(specializedType->getBase()));
-            if (auto genericBasePairStructType = as<IRStructType>(genericType))
-            {
-                return as<IRFieldExtract>(builder->emitFieldExtract(
-                    (IRType*)findSpecializationForParam(
-                        specializedType,
-                        findField(genericBasePairStructType, key)->getFieldType()),
-                    baseInst,
-                    key
-                ));
-            }
-            else if (auto genericPtrType = as<IRPtrTypeBase>(genericType))
-            {
-                if (auto genericPairStructType = as<IRStructType>(genericPtrType->getValueType()))
-                {
-                    return as<IRFieldAddress>(builder->emitFieldAddress(
-                            builder->getPtrType((IRType*)
-                                findSpecializationForParam(
-                                    specializedType,
-                                    findField(genericPairStructType, key)->getFieldType())),
-                            baseInst,
-                            key
-                        ));
-                }
-            }
-        }
-        else
-        {
-            SLANG_UNEXPECTED("Unrecognized field. Cannot emit field accessor");
-        }
-        return nullptr;
-    }
-
-    IRInst* emitPrimalFieldAccess(IRBuilder* builder, IRInst* baseInst)
-    {
-        return emitFieldAccessor(builder, baseInst, this->globalPrimalKey);
-    }
-
-    IRInst* emitDiffFieldAccess(IRBuilder* builder, IRInst* baseInst)
-    {
-        return emitFieldAccessor(builder, baseInst, this->globalDiffKey);
-    }
-
-    IRStructKey* _getOrCreateDiffStructKey()
-    {
-        if (!this->globalDiffKey)
-        {
-            IRBuilder builder(sharedContext->sharedBuilder);
-            // Insert directly at top level (skip any generic scopes etc.)
-            builder.setInsertInto(sharedContext->moduleInst);
-
-            this->globalDiffKey = builder.createStructKey();
-            builder.addNameHintDecoration(this->globalDiffKey , UnownedTerminatedStringSlice("differential"));
-        }
-
-        return this->globalDiffKey;
-    }
-
-    IRStructKey* _getOrCreatePrimalStructKey()
-    {
-        if (!this->globalPrimalKey)
-        {
-            // Insert directly at top level (skip any generic scopes etc.)
-            IRBuilder builder(sharedContext->sharedBuilder);
-            builder.setInsertInto(sharedContext->moduleInst);
-
-            this->globalPrimalKey = builder.createStructKey();
-            builder.addNameHintDecoration(this->globalPrimalKey , UnownedTerminatedStringSlice("primal"));
-        }
-
-        return this->globalPrimalKey;
-    }
-
-    IRInst* _createDiffPairType(IRType* origBaseType, IRType* diffType)
-    {
-        switch (origBaseType->getOp())
-        {
-        case kIROp_lookup_interface_method:
-        case kIROp_Specialize:
-        case kIROp_Param:
-            return nullptr;
-        default:
-            break;
-        }
-        if (diffType->getOp() != kIROp_DifferentialBottomType)
-        {
-            IRBuilder builder(sharedContext->sharedBuilder);
-            builder.setInsertBefore(diffType);
-
-            auto pairStructType = builder.createStructType();
-            builder.createStructField(pairStructType, _getOrCreatePrimalStructKey(), origBaseType);
-            builder.createStructField(pairStructType, _getOrCreateDiffStructKey(), (IRType*)diffType);
-            return pairStructType;
-        }
-        return origBaseType;
-    }
-
-    struct LoweredPairTypeInfo
-    {
-        IRInst* loweredType;
-        bool isTrivial;
-    };
-
-    IRInst* getDiffTypeFromPairType(IRBuilder* builder, IRDifferentialPairType* type)
-    {
-        auto witnessTable = type->getWitness();
-        return _lookupWitness(builder, witnessTable, sharedContext->differentialAssocTypeStructKey);
-    }
-
-    IRInst* getDiffTypeWitnessFromPairType(IRBuilder* builder, IRDifferentialPairType* type)
-    {
-        auto witnessTable = type->getWitness();
-        return _lookupWitness(builder, witnessTable, sharedContext->differentialAssocTypeWitnessStructKey);
-    }
-
-    LoweredPairTypeInfo lowerDiffPairType(IRBuilder* builder, IRType* originalPairType)
-    {
-        LoweredPairTypeInfo result = {};
-        
-        if (pairTypeCache.TryGetValue(originalPairType, result))
-            return result;
-        auto pairType = as<IRDifferentialPairType>(originalPairType);
-        if (!pairType)
-        {
-            result.isTrivial = true;
-            result.loweredType = originalPairType;
-            return result;
-        }
-        auto primalType = pairType->getValueType();
-        if (as<IRParam>(primalType))
-        {
-            result.isTrivial = false;
-            result.loweredType = nullptr;
-            return result;
-        }
-
-        auto diffType = getDiffTypeFromPairType(builder, pairType);
-        if (!diffType)
-            return result;
-        result.loweredType = _createDiffPairType(pairType->getValueType(), (IRType*)diffType);
-        result.isTrivial = (diffType->getOp() == kIROp_DifferentialBottomType);
-        pairTypeCache.Add(originalPairType, result);
-
-        return result;
-    }
-
-    Dictionary<IRInst*, LoweredPairTypeInfo> pairTypeCache;
-
-    IRStructKey* globalPrimalKey = nullptr;
-
-    IRStructKey* globalDiffKey = nullptr;
-
-    IRInst* genericDiffPairType = nullptr;
-
-    List<IRInst*> generatedTypeList;
-
-    AutoDiffSharedContext* sharedContext = nullptr;
-};
-
-struct JVPTranscriber
-{
-
-    // Stores the mapping of arbitrary 'R-value' instructions to instructions that represent
-    // their differential values.
-    Dictionary<IRInst*, IRInst*>            instMapD;
-
-    // Set of insts currently being transcribed. Used to avoid infinite loops.
-    HashSet<IRInst*>                        instsInProgress;
-
-    // Cloning environment to hold mapping from old to new copies for the primal
-    // instructions.
-    IRCloneEnv                              cloneEnv;
-
-    // Diagnostic sink for error messages.
-    DiagnosticSink*                         sink;
-
-    // Type conformance information.
-    AutoDiffSharedContext*                  autoDiffSharedContext;
-
-    // Builder to help with creating and accessing the 'DifferentiablePair<T>' struct
-    DifferentialPairTypeBuilder*            pairBuilder;
-
-    DifferentiableTypeConformanceContext    differentiableTypeConformanceContext;
-
-    List<InstPair>                          followUpFunctionsToTranscribe;
-
-    SharedIRBuilder* sharedBuilder;
-    // Witness table that `DifferentialBottom:IDifferential`.
-    IRWitnessTable* differentialBottomWitness = nullptr;
-    Dictionary<InstPair, IRInst*> differentialPairTypes;
-
-    JVPTranscriber(AutoDiffSharedContext* shared, SharedIRBuilder* inSharedBuilder)
-        : differentiableTypeConformanceContext(shared), sharedBuilder(inSharedBuilder)
-    {
-        
-    }
-
-    DiagnosticSink* getSink()
-    {
-        SLANG_ASSERT(sink);
-        return sink;
-    }
-
-    void mapDifferentialInst(IRInst* origInst, IRInst* diffInst)
-    {
-        if (hasDifferentialInst(origInst))
-        {
-            if (lookupDiffInst(origInst) != diffInst)
-            {
-                SLANG_UNEXPECTED("Inconsistent differential mappings");
-            }
-        }
-        else
-        {
-            instMapD.Add(origInst, diffInst);
-        }
-    }
-
-    void mapPrimalInst(IRInst* origInst, IRInst* primalInst)
-    {
-        if (cloneEnv.mapOldValToNew.ContainsKey(origInst) && cloneEnv.mapOldValToNew[origInst] != primalInst)
-        {
-            getSink()->diagnose(origInst->sourceLoc,
-                Diagnostics::internalCompilerError,
-                "inconsistent primal instruction for original");
-        }
-        else
-        {
-            cloneEnv.mapOldValToNew[origInst] = primalInst;
-        }
-    }
-
-    IRInst* lookupDiffInst(IRInst* origInst)
-    {
-        return instMapD[origInst];
-    }
-
-    IRInst* lookupDiffInst(IRInst* origInst, IRInst* defaultInst)
-    {
-        return (hasDifferentialInst(origInst)) ? instMapD[origInst] : defaultInst;
-    }
-
-    bool hasDifferentialInst(IRInst* origInst)
-    {
-        return instMapD.ContainsKey(origInst);
-    }
-
-    IRInst* lookupPrimalInst(IRInst* origInst)
-    {
-        return cloneEnv.mapOldValToNew[origInst];
-    }
-
-    IRInst* lookupPrimalInst(IRInst* origInst, IRInst* defaultInst)
-    {
-        return (hasPrimalInst(origInst)) ? lookupPrimalInst(origInst) : defaultInst;
-    }
-
-    bool hasPrimalInst(IRInst* origInst)
-    {
-        return cloneEnv.mapOldValToNew.ContainsKey(origInst);
-    }
-    
-    IRInst* findOrTranscribeDiffInst(IRBuilder* builder, IRInst* origInst)
-    {
-        if (!hasDifferentialInst(origInst))
-        {
-            transcribe(builder, origInst);
-            SLANG_ASSERT(hasDifferentialInst(origInst));
-        }
-
-        return lookupDiffInst(origInst);
-    }
-
-    IRInst* findOrTranscribePrimalInst(IRBuilder* builder, IRInst* origInst)
-    {
-        if (!hasPrimalInst(origInst))
-        {
-            transcribe(builder, origInst);
-            SLANG_ASSERT(hasPrimalInst(origInst));
-        }
-
-        return lookupPrimalInst(origInst);
-    }
-
-    IRFuncType* differentiateFunctionType(IRBuilder* builder, IRFuncType* funcType)
-    {
-        List<IRType*> newParameterTypes;
-        IRType* diffReturnType;
-
-        for (UIndex i = 0; i < funcType->getParamCount(); i++)
-        {
-            auto origType = funcType->getParamType(i);
-            origType = (IRType*) lookupPrimalInst(origType, origType);
-            if (auto diffPairType = tryGetDiffPairType(builder, origType))
-                newParameterTypes.add(diffPairType);
-            else
-                newParameterTypes.add(origType);
-        }
-
-        // Transcribe return type to a pair.
-        // This will be void if the primal return type is non-differentiable.
-        //
-        auto origResultType = (IRType*) lookupPrimalInst(funcType->getResultType(), funcType->getResultType());
-        if (auto returnPairType = tryGetDiffPairType(builder, origResultType))
-            diffReturnType = returnPairType;
-        else
-            diffReturnType = builder->getVoidType();
-
-        return builder->getFuncType(newParameterTypes, diffReturnType);
-    }
-
-    IRWitnessTable* getDifferentialBottomWitness()
-    {
-        IRBuilder builder(sharedBuilder);
-        builder.setInsertInto(sharedBuilder->getModule()->getModuleInst());
-        auto result =
-            as<IRWitnessTable>(differentiableTypeConformanceContext.lookUpConformanceForType(
-                builder.getDifferentialBottomType()));
-        SLANG_ASSERT(result);
-        return result;
-    }
-
-    // Get or construct `:IDifferentiable` conformance for a DifferentiablePair.
-    IRWitnessTable* getDifferentialPairWitness(IRInst* inDiffPairType)
-    {
-        IRBuilder builder(sharedBuilder);
-        builder.setInsertInto(inDiffPairType->parent);
-        auto diffPairType = as<IRDifferentialPairType>(inDiffPairType);
-        SLANG_ASSERT(diffPairType);
-
-        auto table = builder.createWitnessTable(autoDiffSharedContext->differentiableInterfaceType, diffPairType);
-
-        // Differentiate the pair type to get it's differential (which is itself a pair)
-        auto diffDiffPairType = differentiateType(&builder, diffPairType);
-
-        // And place it in the synthesized witness table.
-        builder.createWitnessTableEntry(table, autoDiffSharedContext->differentialAssocTypeStructKey, diffDiffPairType);
-        // Omit the method synthesis here, since we can just intercept those directly at `getXXMethodForType`.
-
-        // Record this in the context for future lookups
-        differentiableTypeConformanceContext.differentiableWitnessDictionary[diffPairType] = table;
-
-        return table;
-    }
-
-    IRType* getOrCreateDiffPairType(IRInst* primalType, IRInst* witness)
-    {
-        IRBuilder builder(sharedBuilder);
-        builder.setInsertInto(primalType->parent);
-        return builder.getDifferentialPairType(
-            (IRType*)primalType,
-            witness);
-    }
-
-    IRType* getOrCreateDiffPairType(IRInst* primalType)
-    {
-        IRBuilder builder(sharedBuilder);
-        builder.setInsertInto(primalType->parent);
-        auto witness = as<IRWitnessTable>(
-            differentiableTypeConformanceContext.lookUpConformanceForType((IRType*)primalType));
-
-        if (!witness)
-        {
-            if (auto primalPairType = as<IRDifferentialPairType>(primalType))
-            {
-                witness = getDifferentialPairWitness(primalPairType);
-            }
-            else
-            {
-                witness = getDifferentialBottomWitness();
-            }
-        }
-
-        return builder.getDifferentialPairType(
-            (IRType*)primalType,
-            witness);
-    }
-
-    IRType* differentiateType(IRBuilder* builder, IRType* origType)
-    {
-        IRInst* diffType = nullptr;
-        if (!instMapD.TryGetValue(origType, diffType))
-        {
-            diffType = _differentiateTypeImpl(builder, origType);
-            instMapD[origType] = diffType;
-        }
-        return (IRType*)diffType;
-    }
-
-    IRType* _differentiateTypeImpl(IRBuilder* builder, IRType* origType)
-    {
-        if (auto ptrType = as<IRPtrTypeBase>(origType))
-            return builder->getPtrType(
-                origType->getOp(),
-                differentiateType(builder, ptrType->getValueType()));
-
-        // If there is an explicit primal version of this type in the local scope, load that
-        // otherwise use the original type. 
-        //
-        IRInst* primalType = lookupPrimalInst(origType, origType);
-        
-        // Special case certain compound types (PtrType, FuncType, etc..)
-        // otherwise try to lookup a differential definition for the given type.
-        // If one does not exist, then we assume it's not differentiable.
-        // 
-        switch (primalType->getOp())
-        {
-        case kIROp_Param:
-            if (as<IRTypeType>(primalType->getDataType()))
-                return (IRType*)(differentiableTypeConformanceContext.getDifferentialForType(
-                    builder,
-                    (IRType*)primalType));
-            else if (as<IRWitnessTableType>(primalType->getDataType()))
-                return (IRType*)primalType;
-        
-        case kIROp_ArrayType:
-            {
-                auto primalArrayType = as<IRArrayType>(primalType);
-                if (auto diffElementType = differentiateType(builder, primalArrayType->getElementType()))
-                    return builder->getArrayType(
-                        diffElementType,
-                        primalArrayType->getElementCount());
-                else
-                    return nullptr;
-            }
-
-        case kIROp_DifferentialPairType:
-            {
-                auto primalPairType = as<IRDifferentialPairType>(primalType);
-                return getOrCreateDiffPairType(
-                    pairBuilder->getDiffTypeFromPairType(builder, primalPairType),
-                    pairBuilder->getDiffTypeWitnessFromPairType(builder, primalPairType));
-            }
-        
-        case kIROp_FuncType:
-            return differentiateFunctionType(builder, as<IRFuncType>(primalType));
-
-        case kIROp_OutType:
-            if (auto diffValueType = differentiateType(builder, as<IROutType>(primalType)->getValueType()))
-                return builder->getOutType(diffValueType);
-            else   
-                return nullptr;
-
-        case kIROp_InOutType:
-            if (auto diffValueType = differentiateType(builder, as<IRInOutType>(primalType)->getValueType()))
-                return builder->getInOutType(diffValueType);
-            else
-                return nullptr;
-
-        case kIROp_TupleType:
-            {
-                auto tupleType = as<IRTupleType>(primalType);
-                List<IRType*> diffTypeList;
-                // TODO: what if we have type parameters here?
-                for (UIndex ii = 0; ii < tupleType->getOperandCount(); ii++)
-                    diffTypeList.add(
-                        differentiateType(builder, (IRType*)tupleType->getOperand(ii)));
-
-                return builder->getTupleType(diffTypeList);
-            }
-
-        default:
-            return (IRType*)(differentiableTypeConformanceContext.getDifferentialForType(builder, (IRType*)primalType));
-        }
-    }
-
-    IRType* tryGetDiffPairType(IRBuilder* builder, IRType* primalType)
-    {
-        // If this is a PtrType (out, inout, etc..), then create diff pair from
-        // value type and re-apply the appropropriate PtrType wrapper.
-        // 
-        if (auto origPtrType = as<IRPtrTypeBase>(primalType))
-        {   
-            if (auto diffPairValueType = tryGetDiffPairType(builder, origPtrType->getValueType()))
-                return builder->getPtrType(primalType->getOp(), diffPairValueType);
-            else 
-                return nullptr;
-        }
-        auto diffType = differentiateType(builder, primalType);
-        if (diffType)
-            return (IRType*)getOrCreateDiffPairType(primalType);
-        return nullptr;
-    }
-
-    InstPair transcribeParam(IRBuilder* builder, IRParam* origParam)
-    {
-        auto primalDataType = lookupPrimalInst(origParam->getDataType(), origParam->getDataType());
-        // Do not differentiate generic type (and witness table) parameters
-        if (as<IRTypeType>(primalDataType) || as<IRWitnessTableType>(primalDataType))
-        {
-            return InstPair(
-                cloneInst(&cloneEnv, builder, origParam),
-                nullptr);    
-        }
-
-        // Is this param a phi node or a function parameter?
-        auto func = as<IRGlobalValueWithCode>(origParam->getParent()->getParent());
-        bool isFuncParam = (func && origParam->getParent() == func->getFirstBlock());
-        if (isFuncParam)
-        {
-            if (auto diffPairType = tryGetDiffPairType(builder, (IRType*)primalDataType))
-            {
-                IRInst* diffPairParam = builder->emitParam(diffPairType);
-
-                auto diffPairVarName = makeDiffPairName(origParam);
-                if (diffPairVarName.getLength() > 0)
-                    builder->addNameHintDecoration(diffPairParam, diffPairVarName.getUnownedSlice());
-
-                SLANG_ASSERT(diffPairParam);
-            
-                if (auto pairType = as<IRDifferentialPairType>(diffPairType))
-                {
-                    return InstPair(
-                        builder->emitDifferentialPairGetPrimal(diffPairParam),
-                        builder->emitDifferentialPairGetDifferential(
-                            (IRType*)pairBuilder->getDiffTypeFromPairType(builder, pairType),
-                            diffPairParam));
-                }
-                else if (auto pairPtrType = as<IRPtrTypeBase>(diffPairType))
-                {
-                    auto ptrInnerPairType = as<IRDifferentialPairType>(pairPtrType->getValueType());
-
-                    return InstPair(
-                        builder->emitDifferentialPairAddressPrimal(diffPairParam),
-                        builder->emitDifferentialPairAddressDifferential(
-                            builder->getPtrType(
-                                kIROp_PtrType,
-                                (IRType*)pairBuilder->getDiffTypeFromPairType(builder, ptrInnerPairType)),
-                            diffPairParam));
-                }
-            }
-
-            return InstPair(
-                    cloneInst(&cloneEnv, builder, origParam),
-                    nullptr);
-        }
-        else
-        {
-            auto primal = cloneInst(&cloneEnv, builder, origParam);
-            IRInst* diff = nullptr;
-            if (IRType* diffType = differentiateType(builder, (IRType*)primalDataType))
-            {
-                diff = builder->emitParam(diffType);
-            }
-            return InstPair(primal, diff);
-        }
-        
-    }
-
-    // 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* origVar)
-    {
-        if (auto namehintDecoration = origVar->findDecoration<IRNameHintDecoration>())
-        {
-            return ("d" + String(namehintDecoration->getName()));
-        }
-
-        return String("");
-    }
-
-    // Returns "dp<var-name>" to use as a name hint for parameters.
-    // If no primal name is available, returns a blank string.
-    // 
-    String makeDiffPairName(IRInst* origVar)
-    {
-        if (auto namehintDecoration = origVar->findDecoration<IRNameHintDecoration>())
-        {
-            return ("dp" + String(namehintDecoration->getName()));
-        }
-
-        return String("");
-    }
-
-    InstPair transcribeVar(IRBuilder* builder, IRVar* origVar)
-    {
-        if (IRType* diffType = differentiateType(builder, origVar->getDataType()->getValueType()))
-        {
-            IRVar* diffVar = builder->emitVar(diffType);
-            SLANG_ASSERT(diffVar);
-
-            auto diffNameHint = getJVPVarName(origVar);
-            if (diffNameHint.getLength() > 0)
-                builder->addNameHintDecoration(diffVar, diffNameHint.getUnownedSlice());
-
-            return InstPair(cloneInst(&cloneEnv, builder, origVar), diffVar);
-        }
-
-        return InstPair(cloneInst(&cloneEnv, builder, origVar), nullptr);
-    }
-
-    InstPair transcribeBinaryArith(IRBuilder* builder, IRInst* origArith)
-    {
-        SLANG_ASSERT(origArith->getOperandCount() == 2);
-
-        IRInst* primalArith = cloneInst(&cloneEnv, builder, origArith);
-        
-        auto origLeft = origArith->getOperand(0);
-        auto origRight = origArith->getOperand(1);
-
-        auto primalLeft = findOrTranscribePrimalInst(builder, origLeft);
-        auto primalRight = findOrTranscribePrimalInst(builder, origRight);
-    
-        auto diffLeft = findOrTranscribeDiffInst(builder, origLeft);
-        auto diffRight = findOrTranscribeDiffInst(builder, origRight);
-
-
-        if (diffLeft || diffRight)
-        {
-            diffLeft = diffLeft ? diffLeft : getDifferentialZeroOfType(builder, primalLeft->getDataType());
-            diffRight = diffRight ? diffRight : getDifferentialZeroOfType(builder, primalRight->getDataType());
-
-            auto resultType = primalArith->getDataType();
-            switch(origArith->getOp())
-            {
-            case kIROp_Add:
-                return InstPair(primalArith, builder->emitAdd(resultType, diffLeft, diffRight));
-            case kIROp_Mul:
-                return InstPair(primalArith, builder->emitAdd(resultType,
-                    builder->emitMul(resultType, diffLeft, primalRight),
-                    builder->emitMul(resultType, primalLeft, diffRight)));
-            case kIROp_Sub:
-                return InstPair(primalArith, builder->emitSub(resultType, diffLeft, diffRight));
-            case kIROp_Div:
-                return InstPair(primalArith, builder->emitDiv(resultType, 
-                    builder->emitSub(
-                        resultType,
-                        builder->emitMul(resultType, diffLeft, primalRight),
-                        builder->emitMul(resultType, primalLeft, diffRight)),
-                    builder->emitMul(
-                        primalRight->getDataType(), primalRight, primalRight
-                    )));
-            default:
-                getSink()->diagnose(origArith->sourceLoc,
-                    Diagnostics::unimplemented,
-                    "this arithmetic instruction cannot be differentiated");
-            }
-        }
-
-        return InstPair(primalArith, nullptr);
-    }
-
-    
-    InstPair transcribeBinaryLogic(IRBuilder* builder, IRInst* origLogic)
-    {
-        SLANG_ASSERT(origLogic->getOperandCount() == 2);
-
-        // TODO: Check other boolean cases.
-        if (as<IRBoolType>(origLogic->getDataType()))
-        {
-            // Boolean operations are not differentiable. For the linearization
-            // pass, we do not need to do anything but copy them over to the ne
-            // function.
-            auto primalLogic = cloneInst(&cloneEnv, builder, origLogic);
-            return InstPair(primalLogic, nullptr);
-        }
-        
-        SLANG_UNEXPECTED("Logical operation with non-boolean result");
-    }
-
-    InstPair transcribeLoad(IRBuilder* builder, IRLoad* origLoad)
-    {
-        auto origPtr = origLoad->getPtr();
-        auto primalPtr = lookupPrimalInst(origPtr, nullptr);
-        auto primalPtrValueType = as<IRPtrTypeBase>(primalPtr->getFullType())->getValueType();
-
-        if (auto diffPairType = as<IRDifferentialPairType>(primalPtrValueType))
-        {
-            // Special case load from an `out` param, which will not have corresponding `diff` and
-            // `primal` insts yet.
-            
-            auto load = builder->emitLoad(primalPtr);
-            auto primalElement = builder->emitDifferentialPairGetPrimal(load);
-            auto diffElement = builder->emitDifferentialPairGetDifferential(
-                (IRType*)pairBuilder->getDiffTypeFromPairType(builder, diffPairType), load);
-            return InstPair(primalElement, diffElement);
-        }
-
-        auto primalLoad = cloneInst(&cloneEnv, builder, origLoad);
-        IRInst* diffLoad = nullptr;
-        if (auto diffPtr = lookupDiffInst(origPtr, nullptr))
-        {
-            // Default case, we're loading from a known differential inst.
-            diffLoad = as<IRLoad>(builder->emitLoad(diffPtr));
-        } 
-        return InstPair(primalLoad, diffLoad);
-    }
-
-    InstPair transcribeStore(IRBuilder* builder, IRStore* origStore)
-    {
-        IRInst* origStoreLocation = origStore->getPtr();
-        IRInst* origStoreVal = origStore->getVal();
-        auto primalStoreLocation = lookupPrimalInst(origStoreLocation, nullptr);
-        auto diffStoreLocation = lookupDiffInst(origStoreLocation, nullptr);
-        auto primalStoreVal = lookupPrimalInst(origStoreVal, nullptr);
-        auto diffStoreVal = lookupDiffInst(origStoreVal, nullptr);
-
-        if (!diffStoreLocation)
-        {
-            auto primalLocationPtrType = as<IRPtrTypeBase>(primalStoreLocation->getDataType());
-            if (auto diffPairType = as<IRDifferentialPairType>(primalLocationPtrType->getValueType()))
-            {
-                auto valToStore = builder->emitMakeDifferentialPair(diffPairType, primalStoreVal, diffStoreVal);
-                auto store = builder->emitStore(primalStoreLocation, valToStore);
-                return InstPair(store, nullptr);
-            }
-        }
-
-        auto primalStore = cloneInst(&cloneEnv, builder, origStore);
-
-        IRInst* diffStore = nullptr;
-
-        // If the stored value has a differential version, 
-        // emit a store instruction for the differential parameter.
-        // Otherwise, emit nothing since there's nothing to load.
-        // 
-        if (diffStoreLocation && diffStoreVal)
-        {
-            // Default case, storing the entire type (and not a member)
-            diffStore = as<IRStore>(
-                builder->emitStore(diffStoreLocation, diffStoreVal));
-            
-            return InstPair(primalStore, diffStore);
-        }
-
-        return InstPair(primalStore, nullptr);
-    }
-
-    InstPair transcribeReturn(IRBuilder* builder, IRReturn* origReturn)
-    {
-        IRInst* origReturnVal = origReturn->getVal();
-
-        auto returnDataType = (IRType*) lookupPrimalInst(origReturnVal->getDataType(), origReturnVal->getDataType());
-        if (as<IRFunc>(origReturnVal) || as<IRGeneric>(origReturnVal) || as<IRStructType>(origReturnVal) || as<IRFuncType>(origReturnVal))
-        {
-            // If the return value is itself a function, generic or a struct then this
-            // is likely to be a generic scope. In this case, we lookup the differential
-            // and return that.
-            IRInst* primalReturnVal = findOrTranscribePrimalInst(builder, origReturnVal);
-            IRInst* diffReturnVal = findOrTranscribeDiffInst(builder, origReturnVal);
-            
-            // Neither of these should be nullptr.
-            SLANG_RELEASE_ASSERT(primalReturnVal && diffReturnVal);
-            IRReturn* diffReturn = as<IRReturn>(builder->emitReturn(diffReturnVal));
-
-            return InstPair(diffReturn, diffReturn);
-        }
-        else if (auto pairType = tryGetDiffPairType(builder, returnDataType))
-        {   
-            IRInst* primalReturnVal = findOrTranscribePrimalInst(builder, origReturnVal);
-            IRInst* diffReturnVal = findOrTranscribeDiffInst(builder, origReturnVal);
-            if(!diffReturnVal)
-                diffReturnVal = getDifferentialZeroOfType(builder, returnDataType);
-
-            // If the pair type can be formed, this must be non-null.
-            SLANG_RELEASE_ASSERT(diffReturnVal);
-
-            auto diffPair = builder->emitMakeDifferentialPair(pairType, primalReturnVal, diffReturnVal);
-            IRReturn* pairReturn = as<IRReturn>(builder->emitReturn(diffPair));
-            return InstPair(pairReturn, pairReturn);
-        }
-        else
-        {
-            // If the return type is not differentiable, emit the primal value only.
-            IRInst* primalReturnVal = findOrTranscribePrimalInst(builder, origReturnVal);
-
-            IRInst* primalReturn = builder->emitReturn(primalReturnVal);
-            return InstPair(primalReturn, nullptr);
-            
-        }
-    }
-
-    // Since int/float literals are sometimes nested inside an IRConstructor
-    // instruction, we check to make sure that the nested instr is a constant
-    // and then return nullptr. Literals do not need to be differentiated.
-    //
-    InstPair transcribeConstruct(IRBuilder* builder, IRInst* origConstruct)
-    {   
-        IRInst* primalConstruct = cloneInst(&cloneEnv, builder, origConstruct);
-        
-        // Check if the output type can be differentiated. If it cannot be 
-        // differentiated, don't differentiate the inst
-        // 
-        auto primalConstructType = (IRType*) lookupPrimalInst(origConstruct->getDataType(), origConstruct->getDataType());
-        if (auto diffConstructType = differentiateType(builder, primalConstructType))
-        {
-            UCount operandCount = origConstruct->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, use a zero.
-                // 
-                if (auto diffInst = lookupDiffInst(origConstruct->getOperand(ii), nullptr))
-                    diffOperands.add(diffInst);
-                else 
-                {
-                    auto operandDataType = origConstruct->getOperand(ii)->getDataType();
-                    operandDataType = (IRType*) lookupPrimalInst(operandDataType, operandDataType);
-                    diffOperands.add(getDifferentialZeroOfType(builder, operandDataType));
-                }
-            }
-            
-            return InstPair(
-                primalConstruct, 
-                builder->emitIntrinsicInst(
-                    diffConstructType,
-                    origConstruct->getOp(),
-                    operandCount,
-                    diffOperands.getBuffer()));
-        }
-        else
-        {
-            return InstPair(primalConstruct, nullptr);
-        }
-    }
-
-    // Differentiating a call instruction here is primarily about generating
-    // an appropriate call list based on whichever parameters have differentials 
-    // in the current transcription context.
-    // 
-    InstPair transcribeCall(IRBuilder* builder, IRCall* origCall)
-    {   
-        
-        IRInst* origCallee = origCall->getCallee();
-
-        if (!origCallee)
-        {
-            // Note that this can only happen if the callee is a result
-            // of a higher-order operation. For now, we assume that we cannot
-            // differentiate such calls safely.
-            // TODO(sai): Should probably get checked in the front-end.
-            //
-            getSink()->diagnose(origCall->sourceLoc,
-                Diagnostics::internalCompilerError,
-                "attempting to differentiate unresolved callee");
-            
-            return InstPair(nullptr, nullptr);
-        }
-
-        // Since concrete functions are globals, the primal callee is the same
-        // as the original callee.
-        //
-        auto primalCallee = origCallee;
-
-        IRInst* diffCallee = nullptr;
-
-        if (auto derivativeReferenceDecor = primalCallee->findDecoration<IRForwardDerivativeDecoration>())
-        {
-            // If the user has already provided an differentiated implementation, use that.
-            diffCallee = derivativeReferenceDecor->getForwardDerivativeFunc();
-        }
-        else if (primalCallee->findDecoration<IRForwardDifferentiableDecoration>())
-        {
-            // If the function is marked for auto-diff, push a `differentiate` inst for a follow up pass
-            // to generate the implementation.
-            diffCallee = builder->emitForwardDifferentiateInst(
-                differentiateFunctionType(builder, as<IRFuncType>(primalCallee->getFullType())),
-                primalCallee);
-        }
-        else
-        {
-            // The callee is non differentiable, just return primal value with null diff value.
-            IRInst* primalCall = cloneInst(&cloneEnv, builder, origCall);
-            return InstPair(primalCall, nullptr);
-        }
-
-        List<IRInst*> args;
-        // Go over the parameter list and create pairs for each input (if required)
-        for (UIndex ii = 0; ii < origCall->getArgCount(); ii++)
-        {
-            auto origArg = origCall->getArg(ii);
-            auto primalArg = findOrTranscribePrimalInst(builder, origArg);
-            SLANG_ASSERT(primalArg);
-
-            auto primalType = primalArg->getDataType();
-            auto diffArg = findOrTranscribeDiffInst(builder, origArg);
-
-            if (!diffArg)
-                diffArg = getDifferentialZeroOfType(builder, primalType);
-
-            if (auto pairType = tryGetDiffPairType(builder, primalType))
-            {
-                // If a pair type can be formed, this must be non-null.
-                SLANG_RELEASE_ASSERT(diffArg);
-                auto diffPair = builder->emitMakeDifferentialPair(pairType, primalArg, diffArg);
-                args.add(diffPair);
-            }
-            else
-            {
-                // Add original/primal argument.
-                args.add(primalArg);
-            }
-        }
-        
-        IRType* diffReturnType = nullptr;
-        diffReturnType = tryGetDiffPairType(builder, origCall->getFullType());
-
-        if (!diffReturnType)
-        {
-            SLANG_RELEASE_ASSERT(origCall->getFullType()->getOp() == kIROp_VoidType);
-            diffReturnType = builder->getVoidType();
-        }
-
-        auto callInst = builder->emitCallInst(
-            diffReturnType,
-            diffCallee,
-            args);
-
-        if (diffReturnType->getOp() != kIROp_VoidType)
-        {
-            IRInst* primalResultValue = builder->emitDifferentialPairGetPrimal(callInst);
-            auto diffType = differentiateType(builder, origCall->getFullType());
-            IRInst* diffResultValue = builder->emitDifferentialPairGetDifferential(diffType, callInst);
-            return InstPair(primalResultValue, diffResultValue);
-        }
-        else
-        {
-            // Return the inst itself if the return value is void.
-            // This is fine since these values should never actually be used anywhere.
-            // 
-            return InstPair(callInst, callInst);
-        }
-    }
-
-    InstPair transcribeSwizzle(IRBuilder* builder, IRSwizzle* origSwizzle)
-    {
-        IRInst* primalSwizzle = cloneInst(&cloneEnv, builder, origSwizzle);
-
-        if (auto diffBase = lookupDiffInst(origSwizzle->getBase(), nullptr))
-        {
-            List<IRInst*> swizzleIndices;
-            for (UIndex ii = 0; ii < origSwizzle->getElementCount(); ii++)
-                swizzleIndices.add(origSwizzle->getElementIndex(ii));
-            
-            return InstPair(
-                primalSwizzle,
-                builder->emitSwizzle(
-                    differentiateType(builder, primalSwizzle->getDataType()),
-                    diffBase,
-                    origSwizzle->getElementCount(),
-                    swizzleIndices.getBuffer()));
-        }
-
-        return InstPair(primalSwizzle, nullptr);
-    }
-
-    InstPair transcribeByPassthrough(IRBuilder* builder, IRInst* origInst)
-    {
-        IRInst* primalInst = cloneInst(&cloneEnv, builder, 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 = lookupDiffInst(origInst->getOperand(ii), nullptr))
-                diffOperands.add(diffInst);
-            else
-                return InstPair(primalInst, nullptr);
-        }
-        
-        return InstPair(
-            primalInst, 
-            builder->emitIntrinsicInst(
-                differentiateType(builder, primalInst->getDataType()),
-                origInst->getOp(),
-                operandCount,
-                diffOperands.getBuffer()));
-    }
-
-    InstPair transcribeControlFlow(IRBuilder* builder, IRInst* origInst)
-    {
-        switch(origInst->getOp())
-        {
-            case kIROp_unconditionalBranch:
-            case kIROp_loop:
-                auto origBranch = as<IRUnconditionalBranch>(origInst);
-
-                // Grab the differentials for any phi nodes.
-                List<IRInst*> newArgs;
-                for (UIndex ii = 0; ii < origBranch->getArgCount(); ii++)
-                {
-                    auto origArg = origBranch->getArg(ii);
-                    auto primalArg = lookupPrimalInst(origArg);
-                    newArgs.add(primalArg);
-
-                    if (differentiateType(builder, primalArg->getDataType()))
-                    {
-                        auto diffArg = lookupDiffInst(origArg, nullptr);
-                        if (diffArg)
-                            newArgs.add(diffArg);
-                    }
-                }
-
-                IRInst* diffBranch = nullptr;
-                if (auto diffBlock = findOrTranscribeDiffInst(builder, origBranch->getTargetBlock()))
-                {
-                    if (auto origLoop = as<IRLoop>(origInst))
-                    {
-                        auto breakBlock = findOrTranscribeDiffInst(builder, origLoop->getBreakBlock());
-                        auto continueBlock = findOrTranscribeDiffInst(builder, origLoop->getContinueBlock());
-                        List<IRInst*> operands;
-                        operands.add(breakBlock);
-                        operands.add(continueBlock);
-                        operands.addRange(newArgs);
-                        diffBranch = builder->emitIntrinsicInst(
-                            nullptr,
-                            kIROp_loop,
-                            operands.getCount(),
-                            operands.getBuffer());
-                    }
-                    else
-                    {
-                        diffBranch = builder->emitBranch(
-                            as<IRBlock>(diffBlock),
-                            newArgs.getCount(),
-                            newArgs.getBuffer());
-                    }
-                }
-
-                // For now, every block in the original fn must have a corresponding
-                // block to compute *both* primals and derivatives (i.e linearized block)
-                SLANG_ASSERT(diffBranch);
-
-                return InstPair(diffBranch, diffBranch);
-        }
-
-        getSink()->diagnose(
-            origInst->sourceLoc,
-            Diagnostics::unimplemented,
-            "attempting to differentiate unhandled control flow");
-
-        return InstPair(nullptr, nullptr);
-    }
-
-    InstPair transcribeConst(IRBuilder* builder, IRInst* origInst)
-    {
-        switch(origInst->getOp())
-        {
-            case kIROp_FloatLit:
-                return InstPair(origInst, builder->getFloatValue(origInst->getDataType(), 0.0f));
-            case kIROp_VoidLit:
-                return InstPair(origInst, origInst);
-            case kIROp_IntLit:
-                return InstPair(origInst, builder->getIntValue(origInst->getDataType(), 0));
-        }
-
-        getSink()->diagnose(
-            origInst->sourceLoc,
-            Diagnostics::unimplemented,
-            "attempting to differentiate unhandled const type");
-
-        return InstPair(nullptr, nullptr);
-    }
-
-    InstPair transcribeSpecialize(IRBuilder*, IRSpecialize* origSpecialize)
-    {
-        // In general, we should not see any specialize insts at this stage.
-        // The exceptions are target intrinsics.
-        auto genericInnerVal = findInnerMostGenericReturnVal(as<IRGeneric>(origSpecialize->getBase()));
-        if (genericInnerVal->findDecoration<IRTargetIntrinsicDecoration>())
-        {
-            // Look for an IRForwardDerivativeDecoration on the specialize inst.
-            // (Normally, this would be on the inner IRFunc, but in this case only the JVP func
-            // can be specialized, so we put a decoration on the IRSpecialize)
-            //
-            if (auto jvpFuncDecoration = origSpecialize->findDecoration<IRForwardDerivativeDecoration>())
-            {
-                auto jvpFunc = jvpFuncDecoration->getForwardDerivativeFunc();
-
-                // Make sure this isn't itself a specialize .
-                SLANG_RELEASE_ASSERT(!as<IRSpecialize>(jvpFunc));
-
-                return InstPair(jvpFunc, jvpFunc);
-            }
-        }
-        else
-        {
-            getSink()->diagnose(origSpecialize->sourceLoc,
-                Diagnostics::unexpected,
-                "should not be attempting to differentiate anything specialized here.");
-        }
-
-        return InstPair(nullptr, nullptr);
-    }
-
-    InstPair transcibeLookupInterfaceMethod(IRBuilder* builder, IRLookupWitnessMethod* origLookup)
-    {
-        // This is slightly counter-intuitive, but we don't perform any differentiation
-        // logic here. We simple clone the original lookup which points to the original function,
-        // or the cloned version in case we're inside a generic scope.
-        // The differentiation logic is inserted later when this is used in an IRCall.
-        // This decision is mostly to maintain a uniform convention of ForwardDifferentiate(Lookup(Table))
-        // rather than have Lookup(ForwardDifferentiate(Table))
-        // 
-        auto diffLookup = cloneInst(&cloneEnv, builder, origLookup);
-        return InstPair(diffLookup, diffLookup);
-    }
-
-    // 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.
-    // The current implementation requires that types are defined linearly.
-    // 
-    IRInst* getDifferentialZeroOfType(IRBuilder* builder, IRType* primalType)
-    {
-        if (auto diffType = differentiateType(builder, primalType))
-        {
-            switch (diffType->getOp())
-            {
-            case kIROp_DifferentialPairType:
-                return builder->emitMakeDifferentialPair(
-                    diffType,
-                    getDifferentialZeroOfType(builder, as<IRDifferentialPairType>(diffType)->getValueType()),
-                    getDifferentialZeroOfType(builder, as<IRDifferentialPairType>(diffType)->getValueType()));
-            }
-            // Since primalType has a corresponding differential type, we can lookup the 
-            // definition for zero().
-            auto zeroMethod = differentiableTypeConformanceContext.getZeroMethodForType(builder, primalType);
-            SLANG_ASSERT(zeroMethod);
-
-            auto emptyArgList = List<IRInst*>();
-            return builder->emitCallInst((IRType*)diffType, zeroMethod, emptyArgList);
-        }
-        else
-        {
-            if (isScalarIntegerType(primalType))
-            {
-                return builder->getIntValue(primalType, 0);
-            }
-
-            getSink()->diagnose(primalType->sourceLoc,
-                Diagnostics::internalCompilerError,
-                "could not generate zero value for given type");
-            return nullptr;
-        }
-    }
-
-    InstPair transcribeBlock(IRBuilder* builder, IRBlock* origBlock)
-    {
-        IRBuilder subBuilder(builder->getSharedBuilder());
-        subBuilder.setInsertLoc(builder->getInsertLoc());
-        
-        IRInst* diffBlock = subBuilder.emitBlock();
-        
-        // Note: for blocks, we setup the mapping _before_
-        // processing the children since we could encounter
-        // a lookup while processing the children.
-        // 
-        mapPrimalInst(origBlock, diffBlock);
-        mapDifferentialInst(origBlock, diffBlock);
-
-        subBuilder.setInsertInto(diffBlock);
-
-        // First transcribe every parameter in the block.
-        for (auto param = origBlock->getFirstParam(); param; param = param->getNextParam())
-            this->transcribe(&subBuilder, param);
-
-        // Then, run through every instruction and use the transcriber to generate the appropriate
-        // derivative code.
-        //
-        for (auto child = origBlock->getFirstOrdinaryInst(); child; child = child->getNextInst())
-            this->transcribe(&subBuilder, child);
-
-        return InstPair(diffBlock, diffBlock);
-    }
-
-    InstPair transcribeFieldExtract(IRBuilder* builder, IRInst* originalInst)
-    {
-        SLANG_ASSERT(as<IRFieldExtract>(originalInst) || as<IRFieldAddress>(originalInst));
-
-        IRInst* origBase = originalInst->getOperand(0);
-        auto primalBase = findOrTranscribePrimalInst(builder, origBase);
-        auto field = originalInst->getOperand(1);
-        auto derivativeRefDecor = field->findDecoration<IRDerivativeMemberDecoration>();
-        auto primalType = (IRType*)lookupPrimalInst(originalInst->getDataType(), originalInst->getDataType());
-
-        IRInst* primalOperands[] = { primalBase, field };
-        IRInst* primalFieldExtract = builder->emitIntrinsicInst(
-            primalType,
-            originalInst->getOp(),
-            2,
-            primalOperands);
-
-        if (!derivativeRefDecor)
-        {
-            return InstPair(primalFieldExtract, nullptr);
-        }
-
-        IRInst* diffFieldExtract = nullptr;
-
-        if (auto diffType = differentiateType(builder, primalType))
-        {
-            if (auto diffBase = findOrTranscribeDiffInst(builder, origBase))
-            {
-                IRInst* diffOperands[] = { diffBase, derivativeRefDecor->getDerivativeMemberStructKey() };
-                diffFieldExtract = builder->emitIntrinsicInst(
-                    diffType,
-                    originalInst->getOp(),
-                    2,
-                    diffOperands);
-            }
-        }
-        return InstPair(primalFieldExtract, diffFieldExtract);
-    }
-
-    InstPair transcribeGetElement(IRBuilder* builder, IRInst* origGetElementPtr)
-    {
-        SLANG_ASSERT(as<IRGetElement>(origGetElementPtr) || as<IRGetElementPtr>(origGetElementPtr));
-
-        IRInst* origBase = origGetElementPtr->getOperand(0);
-        auto primalBase = findOrTranscribePrimalInst(builder, origBase);
-        auto primalIndex = findOrTranscribePrimalInst(builder, origGetElementPtr->getOperand(1));
-
-        auto primalType = (IRType*)lookupPrimalInst(origGetElementPtr->getDataType(), origGetElementPtr->getDataType());
-
-        IRInst* primalOperands[] = {primalBase, primalIndex};
-        IRInst* primalGetElementPtr = builder->emitIntrinsicInst(
-            primalType,
-            origGetElementPtr->getOp(),
-            2,
-            primalOperands);
-
-        IRInst* diffGetElementPtr = nullptr;
-
-        if (auto diffType = differentiateType(builder, primalType))
-        {
-            if (auto diffBase = findOrTranscribeDiffInst(builder, origBase))
-            {
-                IRInst* diffOperands[] = {diffBase, primalIndex};
-                diffGetElementPtr = builder->emitIntrinsicInst(
-                    diffType,
-                    origGetElementPtr->getOp(),
-                    2,
-                    diffOperands);
-            }
-        }
-
-        return InstPair(primalGetElementPtr, diffGetElementPtr);
-    }
-
-    InstPair transcribeLoop(IRBuilder* builder, IRLoop* origLoop)
-    {
-        // The loop comes with three blocks.. we just need to transcribe each one
-        // and assemble the new loop instruction.
-        
-        // Transcribe the target block (this is the 'condition' part of the loop, which
-        // will branch into the loop body)
-        auto diffTargetBlock = findOrTranscribeDiffInst(builder, origLoop->getTargetBlock());
-
-        // Transcribe the break block (this is the block after the exiting the loop)
-        auto diffBreakBlock = findOrTranscribeDiffInst(builder, origLoop->getBreakBlock());
-
-        // Transcribe the continue block (this is the 'update' part of the loop, which will
-        // branch into the condition block)
-        auto diffContinueBlock = findOrTranscribeDiffInst(builder, origLoop->getContinueBlock());
-
-        
-        List<IRInst*> diffLoopOperands;
-        diffLoopOperands.add(diffTargetBlock);
-        diffLoopOperands.add(diffBreakBlock);
-        diffLoopOperands.add(diffContinueBlock);
-
-        // If there are any other operands, use their primal versions.
-        for (UIndex ii = diffLoopOperands.getCount(); ii < origLoop->getOperandCount(); ii++)
-        {
-            auto primalOperand = findOrTranscribePrimalInst(builder, origLoop->getOperand(ii));
-            diffLoopOperands.add(primalOperand);
-        }
-
-        IRInst* diffLoop = builder->emitIntrinsicInst(
-            nullptr,
-            kIROp_loop,
-            diffLoopOperands.getCount(),
-            diffLoopOperands.getBuffer());
-
-        return InstPair(diffLoop, diffLoop);
-    }
-
-    InstPair transcribeIfElse(IRBuilder* builder, IRIfElse* origIfElse)
-    {
-        // IfElse Statements come with 4 blocks. We transcribe each block into it's
-        // linear form, and then wire them up in the same way as the original if-else
-        
-        // Transcribe condition block
-        auto primalConditionBlock = findOrTranscribePrimalInst(builder, origIfElse->getCondition());
-        SLANG_ASSERT(primalConditionBlock);
-
-        // Transcribe 'true' block (condition block branches into this if true)
-        auto diffTrueBlock = findOrTranscribeDiffInst(builder, origIfElse->getTrueBlock());
-        SLANG_ASSERT(diffTrueBlock);
-
-        // Transcribe 'false' block (condition block branches into this if true)
-        // TODO (sai): What happens if there's no false block?
-        auto diffFalseBlock = findOrTranscribeDiffInst(builder, origIfElse->getFalseBlock());
-        SLANG_ASSERT(diffFalseBlock);
-
-        // Transcribe 'after' block (true and false blocks branch into this)
-        auto diffAfterBlock = findOrTranscribeDiffInst(builder, origIfElse->getAfterBlock());
-        SLANG_ASSERT(diffAfterBlock);
-        
-        List<IRInst*> diffIfElseArgs;
-        diffIfElseArgs.add(primalConditionBlock);
-        diffIfElseArgs.add(diffTrueBlock);
-        diffIfElseArgs.add(diffFalseBlock);
-        diffIfElseArgs.add(diffAfterBlock);
-
-        // If there are any other operands, use their primal versions.
-        for (UIndex ii = diffIfElseArgs.getCount(); ii < origIfElse->getOperandCount(); ii++)
-        {
-            auto primalOperand = findOrTranscribePrimalInst(builder, origIfElse->getOperand(ii));
-            diffIfElseArgs.add(primalOperand);
-        }
-
-        IRInst* diffLoop = builder->emitIntrinsicInst(
-            nullptr,
-            kIROp_ifElse,
-            diffIfElseArgs.getCount(),
-            diffIfElseArgs.getBuffer());
-
-        return InstPair(diffLoop, diffLoop);
-    }
-
-    InstPair transcribeMakeDifferentialPair(IRBuilder* builder, IRMakeDifferentialPair* origInst)
-    {
-        auto primalVal = findOrTranscribePrimalInst(builder, origInst->getPrimalValue());
-        SLANG_ASSERT(primalVal);
-        auto diffPrimalVal = findOrTranscribePrimalInst(builder, origInst->getDifferentialValue());
-        SLANG_ASSERT(diffPrimalVal);
-        auto primalDiffVal = findOrTranscribeDiffInst(builder, origInst->getPrimalValue());
-        SLANG_ASSERT(primalDiffVal);
-        auto diffDiffVal = findOrTranscribeDiffInst(builder, origInst->getDifferentialValue());
-        SLANG_ASSERT(diffDiffVal);
-
-        auto primalPair = builder->emitMakeDifferentialPair(origInst->getDataType(), primalVal, diffPrimalVal);
-        auto diffPair = builder->emitMakeDifferentialPair(
-            differentiateType(builder, origInst->getDataType()),
-            primalDiffVal,
-            diffDiffVal);
-        return InstPair(primalPair, diffPair);
-    }
-
-    InstPair transcribeDifferentialPairGetElement(IRBuilder* builder, IRInst* origInst)
-    {
-        SLANG_ASSERT(
-            origInst->getOp() == kIROp_DifferentialPairGetDifferential ||
-            origInst->getOp() == kIROp_DifferentialPairGetPrimal);
-
-        auto primalVal = findOrTranscribePrimalInst(builder, origInst->getOperand(0));
-        SLANG_ASSERT(primalVal);
-
-        auto diffVal = findOrTranscribeDiffInst(builder, origInst->getOperand(0));
-        SLANG_ASSERT(diffVal);
-
-        auto primalResult = builder->emitIntrinsicInst(origInst->getFullType(), origInst->getOp(), 1, &primalVal);
-
-        auto diffValPairType = as<IRDifferentialPairType>(diffVal->getDataType());
-        IRInst* diffResultType = nullptr;
-        if (origInst->getOp() == kIROp_DifferentialPairGetDifferential)
-            diffResultType = pairBuilder->getDiffTypeFromPairType(builder, diffValPairType);
-        else
-            diffResultType = diffValPairType->getValueType();
-        auto diffResult = builder->emitIntrinsicInst((IRType*)diffResultType, origInst->getOp(), 1, &diffVal);
-        return InstPair(primalResult, diffResult);
-    }
-
-    // Create an empty func to represent the transcribed func of `origFunc`.
-    InstPair transcribeFuncHeader(IRBuilder* inBuilder, IRFunc* origFunc)
-    {
-        IRBuilder builder(inBuilder->getSharedBuilder());
-        builder.setInsertBefore(origFunc);
-
-        IRFunc* primalFunc = origFunc;
-
-        differentiableTypeConformanceContext.setFunc(origFunc);
-
-        primalFunc = origFunc;
-
-        auto diffFunc = builder.createFunc();
-
-        SLANG_ASSERT(as<IRFuncType>(origFunc->getFullType()));
-        IRType* diffFuncType = this->differentiateFunctionType(
-            &builder,
-            as<IRFuncType>(origFunc->getFullType()));
-        diffFunc->setFullType(diffFuncType);
-
-        if (auto nameHint = origFunc->findDecoration<IRNameHintDecoration>())
-        {
-            auto originalName = nameHint->getName();
-            StringBuilder newNameSb;
-            newNameSb << "s_fwd_" << originalName;
-            builder.addNameHintDecoration(diffFunc, newNameSb.getUnownedSlice());
-        }
-        builder.addForwardDerivativeDecoration(origFunc, diffFunc);
-
-        // Mark the generated derivative function itself as differentiable.
-        builder.addForwardDifferentiableDecoration(diffFunc);
-
-        // Find and clone `DifferentiableTypeDictionaryDecoration` to the new diffFunc.
-        if (auto dictDecor = origFunc->findDecoration<IRDifferentiableTypeDictionaryDecoration>())
-        {
-            cloneDecoration(dictDecor, diffFunc);
-        }
-
-        auto result = InstPair(primalFunc, diffFunc);
-        followUpFunctionsToTranscribe.add(result);
-        return result;
-    }
-
-    // Transcribe a function definition.
-    InstPair transcribeFunc(IRBuilder* inBuilder, IRFunc* primalFunc, IRFunc* diffFunc)
-    {
-        IRBuilder builder(inBuilder->getSharedBuilder());
-        builder.setInsertInto(diffFunc);
-
-        differentiableTypeConformanceContext.setFunc(primalFunc);
-        // Transcribe children from origFunc into diffFunc
-        for (auto block = primalFunc->getFirstBlock(); block; block = block->getNextBlock())
-            this->transcribe(&builder, block);
-
-        return InstPair(primalFunc, diffFunc);
-    }
-
-    // Transcribe a generic definition
-    InstPair transcribeGeneric(IRBuilder* inBuilder, IRGeneric* origGeneric)
-    {
-        auto innerVal = findInnerMostGenericReturnVal(origGeneric);
-        if (auto innerFunc = as<IRFunc>(innerVal))
-        {
-            differentiableTypeConformanceContext.setFunc(innerFunc);
-        }
-        else
-        {
-            return InstPair(origGeneric, nullptr);
-        }
-
-        // For now, we assume there's only one generic layer. So this inst must be top level
-        bool isTopLevel = (as<IRModuleInst>(origGeneric->getParent()) != nullptr);
-        SLANG_RELEASE_ASSERT(isTopLevel);
-
-        IRGeneric* primalGeneric = origGeneric;
-
-        IRBuilder builder(inBuilder->getSharedBuilder());
-        builder.setInsertBefore(origGeneric);
-
-        auto diffGeneric = builder.emitGeneric();
-
-        // Process type of generic. If the generic is a function, then it's type will also be a 
-        // generic and this logic will transcribe that generic first before continuing with the 
-        // function itself.
-        // 
-        auto primalType =  primalGeneric->getFullType();
-
-        IRType* diffType = nullptr;
-        if (primalType)
-        {
-            diffType = (IRType*) findOrTranscribeDiffInst(&builder, primalType);
-        }
-
-        diffGeneric->setFullType(diffType);
-
-        // TODO(sai): Replace naming scheme
-        // if (auto jvpName = this->getJVPFuncName(builder, primalFn))
-        //    builder->addNameHintDecoration(diffFunc, jvpName);
-        
-        // Transcribe children from origFunc into diffFunc.
-        builder.setInsertInto(diffGeneric);
-        for (auto block = origGeneric->getFirstBlock(); block; block = block->getNextBlock())
-            this->transcribe(&builder, block);
-
-        return InstPair(primalGeneric, diffGeneric);
-    }
-
-    IRInst* transcribe(IRBuilder* builder, IRInst* origInst)
-    {
-        // If a differential intstruction is already mapped for 
-        // this original inst, return that.
-        //
-        if (auto diffInst = lookupDiffInst(origInst, nullptr))
-        {
-            SLANG_ASSERT(lookupPrimalInst(origInst)); // Consistency check.
-            return diffInst;
-        }
-
-        // Otherwise, dispatch to the appropriate method 
-        // depending on the op-code.
-        // 
-        instsInProgress.Add(origInst);
-        InstPair pair = transcribeInst(builder, origInst);
-
-        if (auto primalInst = pair.primal)
-        {
-            mapPrimalInst(origInst, pair.primal);
-            mapDifferentialInst(origInst, pair.differential);
-            if (pair.differential)
-            {
-                switch (pair.differential->getOp())
-                {
-                case kIROp_Func:
-                case kIROp_Generic:
-                case kIROp_Block:
-                    // Don't generate again for these.
-                    // Functions already have their names generated in `transcribeFuncHeader`.
-                    break;
-                default:
-                    // Generate name hint for the inst.
-                    if (auto primalNameHint = primalInst->findDecoration<IRNameHintDecoration>())
-                    {
-                        StringBuilder sb;
-                        sb << "s_diff_" << primalNameHint->getName();
-                        builder->addNameHintDecoration(pair.differential, sb.getUnownedSlice());
-                    }
-                    break;
-                }
-            }
-            return pair.differential;
-        }
-        instsInProgress.Remove(origInst);
-
-        getSink()->diagnose(origInst->sourceLoc,
-                    Diagnostics::internalCompilerError,
-                    "failed to transcibe instruction");
-        return nullptr;
-    }
-
-    InstPair transcribeInst(IRBuilder* builder, IRInst* origInst)
-    {
-        // Handle common SSA-style operations
-        switch (origInst->getOp())
-        {
-        case kIROp_Param:
-            return transcribeParam(builder, as<IRParam>(origInst));
-
-        case kIROp_Var:
-            return transcribeVar(builder, as<IRVar>(origInst));
-
-        case kIROp_Load:
-            return transcribeLoad(builder, as<IRLoad>(origInst));
-
-        case kIROp_Store:
-            return transcribeStore(builder, as<IRStore>(origInst));
-
-        case kIROp_Return:
-            return transcribeReturn(builder, as<IRReturn>(origInst));
-
-        case kIROp_Add:
-        case kIROp_Mul:
-        case kIROp_Sub:
-        case kIROp_Div:
-            return transcribeBinaryArith(builder, origInst);
-        
-        case kIROp_Less:
-        case kIROp_Greater:
-        case kIROp_And:
-        case kIROp_Or:
-        case kIROp_Geq:
-        case kIROp_Leq:
-            return transcribeBinaryLogic(builder, origInst);
-
-        case kIROp_Construct:
-            return transcribeConstruct(builder, origInst);
-        
-        case kIROp_Call:
-            return transcribeCall(builder, as<IRCall>(origInst));
-        
-        case kIROp_swizzle:
-            return transcribeSwizzle(builder, as<IRSwizzle>(origInst));
-        
-        case kIROp_constructVectorFromScalar:
-        case kIROp_MakeTuple:
-            return transcribeByPassthrough(builder, origInst);
-
-        case kIROp_unconditionalBranch:
-            return transcribeControlFlow(builder, origInst);
-
-        case kIROp_FloatLit:
-        case kIROp_IntLit:
-        case kIROp_VoidLit:
-            return transcribeConst(builder, origInst);
-
-        case kIROp_Specialize:
-            return transcribeSpecialize(builder, as<IRSpecialize>(origInst));
-
-        case kIROp_lookup_interface_method:
-            return transcibeLookupInterfaceMethod(builder, as<IRLookupWitnessMethod>(origInst));
-
-        case kIROp_FieldExtract:
-        case kIROp_FieldAddress:
-            return transcribeFieldExtract(builder, origInst);
-        case kIROp_getElement:
-        case kIROp_getElementPtr:
-            return transcribeGetElement(builder, origInst);
-        
-        case kIROp_loop:
-            return transcribeLoop(builder, as<IRLoop>(origInst));
-
-        case kIROp_ifElse:
-            return transcribeIfElse(builder, as<IRIfElse>(origInst));
-
-        case kIROp_MakeDifferentialPair:
-            return transcribeMakeDifferentialPair(builder, as<IRMakeDifferentialPair>(origInst));
-        case kIROp_DifferentialPairGetPrimal:
-        case kIROp_DifferentialPairGetDifferential:
-            return transcribeDifferentialPairGetElement(builder, origInst);
-        }
-    
-        // If none of the cases have been hit, check if the instruction is a
-        // type. Only need to explicitly differentiate types if they appear inside a block.
-        // 
-        if (auto origType = as<IRType>(origInst))
-        {   
-            // If this is a generic type, transcibe the parent 
-            // generic and derive the type from the transcribed generic's
-            // return value.
-            // 
-            if (as<IRGeneric>(origType->getParent()->getParent()) && 
-                findInnerMostGenericReturnVal(as<IRGeneric>(origType->getParent()->getParent())) == origType && 
-                !instsInProgress.Contains(origType->getParent()->getParent()))
-            {
-                auto origGenericType = origType->getParent()->getParent();
-                auto diffGenericType = findOrTranscribeDiffInst(builder, origGenericType);
-                auto innerDiffGenericType = findInnerMostGenericReturnVal(as<IRGeneric>(diffGenericType));
-                return InstPair(
-                    origGenericType,
-                    innerDiffGenericType
-                );
-            }
-            else if (as<IRBlock>(origType->getParent()))
-                return InstPair(
-                    cloneInst(&cloneEnv, builder, origType),
-                    differentiateType(builder, origType));
-            else
-                return InstPair(
-                    cloneInst(&cloneEnv, builder, origType),
-                    nullptr);
-        }
-
-        // Handle instructions with children
-        switch (origInst->getOp())
-        {
-        case kIROp_Func:
-            return transcribeFuncHeader(builder, as<IRFunc>(origInst)); 
-
-        case kIROp_Block:
-            return transcribeBlock(builder, as<IRBlock>(origInst));
-        
-        case kIROp_Generic:
-            return transcribeGeneric(builder, as<IRGeneric>(origInst)); 
-        }
-
-        
-        // If we reach this statement, the instruction type is likely unhandled.
-        getSink()->diagnose(origInst->sourceLoc,
-                    Diagnostics::unimplemented,
-                    "this instruction cannot be differentiated");
-
-        return InstPair(nullptr, nullptr);
-    }
-};
-
-
-struct BackwardDiffTranscriber
-{
-
-    // Stores the mapping of arbitrary 'R-value' instructions to instructions that represent
-    // their differential values.
-    Dictionary<IRInst*, IRInst*> orginalToTranscribed;
-
-    // Set of insts currently being transcribed. Used to avoid infinite loops.
-    HashSet<IRInst*>                        instsInProgress;
-
-    // Cloning environment to hold mapping from old to new copies for the primal
-    // instructions.
-    IRCloneEnv                              cloneEnv;
-
-    // Diagnostic sink for error messages.
-    DiagnosticSink* sink;
-
-    // Type conformance information.
-    AutoDiffSharedContext* autoDiffSharedContext;
-
-    // Builder to help with creating and accessing the 'DifferentiablePair<T>' struct
-    DifferentialPairTypeBuilder* pairBuilder;
-
-    DifferentiableTypeConformanceContext    differentiableTypeConformanceContext;
-
-    List<InstPair>                          followUpFunctionsToTranscribe;
-
-    // Map that stores the upper gradient given an IRInst*
-    Dictionary<IRInst*, List<IRInst*>> upperGradients;
-    Dictionary<IRInst*, IRInst*> primalToDiffPair;
-
-    SharedIRBuilder* sharedBuilder;
-    // Witness table that `DifferentialBottom:IDifferential`.
-    IRWitnessTable* differentialBottomWitness = nullptr;
-    Dictionary<InstPair, IRInst*> differentialPairTypes;
-
-    BackwardDiffTranscriber(AutoDiffSharedContext* shared, SharedIRBuilder* inSharedBuilder, DiagnosticSink* inSink)
-        : autoDiffSharedContext(shared)
-        , sink(inSink)
-        , differentiableTypeConformanceContext(shared)
-        , sharedBuilder(inSharedBuilder)
-    {}
-
-    DiagnosticSink* getSink()
-    {
-        SLANG_ASSERT(sink);
-        return sink;
-    }
-
-    IRFuncType* differentiateFunctionType(IRBuilder* builder, IRFuncType* funcType)
-    {
-        List<IRType*> newParameterTypes;
-        IRType* diffReturnType;
-
-        for (UIndex i = 0; i < funcType->getParamCount(); i++)
-        {
-            auto origType = funcType->getParamType(i);
-            if (auto diffPairType = tryGetDiffPairType(builder, origType))
-            {
-                auto inoutDiffPairType = builder->getPtrType(kIROp_InOutType, diffPairType);
-                newParameterTypes.add(inoutDiffPairType);
-            }
-            else
-                newParameterTypes.add(origType);
-        }
-
-        newParameterTypes.add(funcType->getResultType());
-
-        diffReturnType = builder->getVoidType();
-
-        return builder->getFuncType(newParameterTypes, diffReturnType);
-    }
-
-    IRWitnessTable* getDifferentialBottomWitness()
-    {
-        IRBuilder builder(sharedBuilder);
-        builder.setInsertInto(sharedBuilder->getModule()->getModuleInst());
-        auto result =
-            as<IRWitnessTable>(differentiableTypeConformanceContext.lookUpConformanceForType(
-                builder.getDifferentialBottomType()));
-        SLANG_ASSERT(result);
-        return result;
-    }
-
-    // Get or construct `:IDifferentiable` conformance for a DifferentiablePair.
-    IRWitnessTable* getDifferentialPairWitness(IRInst* inDiffPairType)
-    {
-        IRBuilder builder(sharedBuilder);
-        builder.setInsertInto(inDiffPairType->parent);
-        auto diffPairType = as<IRDifferentialPairType>(inDiffPairType);
-        SLANG_ASSERT(diffPairType);
-        auto result =
-            as<IRWitnessTable>(differentiableTypeConformanceContext.lookUpConformanceForType(
-                builder.getDifferentialBottomType()));
-        if (result)
-            return result;
-
-        auto table = builder.createWitnessTable(autoDiffSharedContext->differentiableInterfaceType, diffPairType);
-        auto diffType = differentiateType(&builder, diffPairType->getValueType());
-        auto differentialType = builder.getDifferentialPairType(diffType, getDifferentialBottomWitness());
-        builder.createWitnessTableEntry(table, autoDiffSharedContext->differentialAssocTypeStructKey, differentialType);
-        // Omit the method synthesis here, since we can just intercept those directly at `getXXMethodForType`.
-
-        differentiableTypeConformanceContext.differentiableWitnessDictionary[diffPairType] = table;
-        return table;
-    }
-
-    IRType* getOrCreateDiffPairType(IRInst* primalType, IRInst* witness)
-    {
-        IRBuilder builder(sharedBuilder);
-        builder.setInsertInto(primalType->parent);
-        return builder.getDifferentialPairType(
-            (IRType*)primalType,
-            witness);
-    }
-
-    IRType* getOrCreateDiffPairType(IRInst* primalType)
-    {
-        IRBuilder builder(sharedBuilder);
-        builder.setInsertInto(primalType->parent);
-        auto witness = as<IRWitnessTable>(
-            differentiableTypeConformanceContext.lookUpConformanceForType((IRType*)primalType));
-        if (!witness)
-            witness = getDifferentialBottomWitness();
-        return builder.getDifferentialPairType(
-            (IRType*)primalType,
-            witness);
-    }
-
-    IRType* differentiateType(IRBuilder* builder, IRType* origType)
-    {
-        IRInst* diffType = nullptr;
-        if (!orginalToTranscribed.TryGetValue(origType, diffType))
-        {
-            diffType = _differentiateTypeImpl(builder, origType);
-            orginalToTranscribed[origType] = diffType;
-        }
-        return (IRType*)diffType;
-    }
-
-    IRType* _differentiateTypeImpl(IRBuilder* builder, IRType* origType)
-    {
-        if (auto ptrType = as<IRPtrTypeBase>(origType))
-            return builder->getPtrType(
-                origType->getOp(),
-                differentiateType(builder, ptrType->getValueType()));
-
-        // If there is an explicit primal version of this type in the local scope, load that
-        // otherwise use the original type. 
-        //
-        IRInst* primalType = origType;
-
-        // Special case certain compound types (PtrType, FuncType, etc..)
-        // otherwise try to lookup a differential definition for the given type.
-        // If one does not exist, then we assume it's not differentiable.
-        // 
-        switch (primalType->getOp())
-        {
-        case kIROp_Param:
-            if (as<IRTypeType>(primalType->getDataType()))
-                return (IRType*)(differentiableTypeConformanceContext.getDifferentialForType(
-                    builder,
-                    (IRType*)primalType));
-            else if (as<IRWitnessTableType>(primalType->getDataType()))
-                return (IRType*)primalType;
-
-        case kIROp_ArrayType:
-        {
-            auto primalArrayType = as<IRArrayType>(primalType);
-            if (auto diffElementType = differentiateType(builder, primalArrayType->getElementType()))
-                return builder->getArrayType(
-                    diffElementType,
-                    primalArrayType->getElementCount());
-            else
-                return nullptr;
-        }
-
-        case kIROp_DifferentialPairType:
-        {
-            auto primalPairType = as<IRDifferentialPairType>(primalType);
-            return getOrCreateDiffPairType(
-                pairBuilder->getDiffTypeFromPairType(builder, primalPairType),
-                pairBuilder->getDiffTypeWitnessFromPairType(builder, primalPairType));
-        }
-
-        case kIROp_FuncType:
-            return differentiateFunctionType(builder, as<IRFuncType>(primalType));
-
-        case kIROp_OutType:
-            if (auto diffValueType = differentiateType(builder, as<IROutType>(primalType)->getValueType()))
-                return builder->getOutType(diffValueType);
-            else
-                return nullptr;
-
-        case kIROp_InOutType:
-            if (auto diffValueType = differentiateType(builder, as<IRInOutType>(primalType)->getValueType()))
-                return builder->getInOutType(diffValueType);
-            else
-                return nullptr;
-
-        case kIROp_TupleType:
-        {
-            auto tupleType = as<IRTupleType>(primalType);
-            List<IRType*> diffTypeList;
-            // TODO: what if we have type parameters here?
-            for (UIndex ii = 0; ii < tupleType->getOperandCount(); ii++)
-                diffTypeList.add(
-                    differentiateType(builder, (IRType*)tupleType->getOperand(ii)));
-
-            return builder->getTupleType(diffTypeList);
-        }
-
-        default:
-            return (IRType*)(differentiableTypeConformanceContext.getDifferentialForType(builder, (IRType*)primalType));
-        }
-    }
-
-    IRType* tryGetDiffPairType(IRBuilder* builder, IRType* primalType)
-    {
-        // If this is a PtrType (out, inout, etc..), then create diff pair from
-        // value type and re-apply the appropropriate PtrType wrapper.
-        // 
-        if (auto origPtrType = as<IRPtrTypeBase>(primalType))
-        {
-            if (auto diffPairValueType = tryGetDiffPairType(builder, origPtrType->getValueType()))
-                return builder->getPtrType(primalType->getOp(), diffPairValueType);
-            else
-                return nullptr;
-        }
-        auto diffType = differentiateType(builder, primalType);
-        if (diffType)
-            return (IRType*)getOrCreateDiffPairType(primalType);
-        return nullptr;
-    }
-
-    InstPair transcribeParam(IRBuilder* builder, IRParam* origParam)
-    {
-        auto primalDataType = origParam->getDataType();
-        // Do not differentiate generic type (and witness table) parameters
-        if (as<IRTypeType>(primalDataType) || as<IRWitnessTableType>(primalDataType))
-        {
-            return InstPair(
-                cloneInst(&cloneEnv, builder, origParam),
-                nullptr);
-        }
-
-        if (auto diffPairType = tryGetDiffPairType(builder, (IRType*)primalDataType))
-        {
-            IRInst* diffPairParam = builder->emitParam(diffPairType);
-
-            auto diffPairVarName = makeDiffPairName(origParam);
-            if (diffPairVarName.getLength() > 0)
-                builder->addNameHintDecoration(diffPairParam, diffPairVarName.getUnownedSlice());
-
-            SLANG_ASSERT(diffPairParam);
-
-            if (auto pairType = as<IRDifferentialPairType>(diffPairParam->getDataType()))
-            {
-                return InstPair(
-                    builder->emitDifferentialPairGetPrimal(diffPairParam),
-                    builder->emitDifferentialPairGetDifferential(
-                        (IRType*)pairBuilder->getDiffTypeFromPairType(builder, pairType),
-                        diffPairParam));
-            }
-            // If this is an `in/inout DifferentialPair<>` parameter, we can't produce
-            // its primal and diff parts right now because they would represent a reference
-            // to a pair field, which doesn't make sense since pair types are considered mutable.
-            // We encode the result as if the param is non-differentiable, and handle it
-            // with special care at load/store.
-            return InstPair(diffPairParam, nullptr);
-        }
-
-
-        return InstPair(
-            cloneInst(&cloneEnv, builder, origParam),
-            nullptr);
-    }
-
-    // Returns "dp<var-name>" to use as a name hint for parameters.
-    // If no primal name is available, returns a blank string.
-    // 
-    String makeDiffPairName(IRInst* origVar)
-    {
-        if (auto namehintDecoration = origVar->findDecoration<IRNameHintDecoration>())
-        {
-            return ("dp" + String(namehintDecoration->getName()));
-        }
-
-        return String("");
-    }
-
-
-    // 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.
-    // The current implementation requires that types are defined linearly.
-    // 
-    IRInst* getDifferentialZeroOfType(IRBuilder* builder, IRType* primalType)
-    {
-        if (auto diffType = differentiateType(builder, primalType))
-        {
-            switch (diffType->getOp())
-            {
-            case kIROp_DifferentialPairType:
-                return builder->emitMakeDifferentialPair(
-                    diffType,
-                    getDifferentialZeroOfType(builder, as<IRDifferentialPairType>(diffType)->getValueType()),
-                    getDifferentialZeroOfType(builder, as<IRDifferentialPairType>(diffType)->getValueType()));
-            }
-            // Since primalType has a corresponding differential type, we can lookup the 
-            // definition for zero().
-            auto zeroMethod = differentiableTypeConformanceContext.getZeroMethodForType(builder, primalType);
-            SLANG_ASSERT(zeroMethod);
-
-            auto emptyArgList = List<IRInst*>();
-            return builder->emitCallInst((IRType*)diffType, zeroMethod, emptyArgList);
-        }
-        else
-        {
-            if (isScalarIntegerType(primalType))
-            {
-                return builder->getIntValue(primalType, 0);
-            }
-
-            getSink()->diagnose(primalType->sourceLoc,
-                Diagnostics::internalCompilerError,
-                "could not generate zero value for given type");
-            return nullptr;
-        }
-    }
-
-    InstPair transcribeBlock(IRBuilder* builder, IRBlock* origBlock)
-    {
-        IRBuilder subBuilder(builder->getSharedBuilder());
-        subBuilder.setInsertLoc(builder->getInsertLoc());
-
-        IRBlock* diffBlock = subBuilder.emitBlock();
-
-        subBuilder.setInsertInto(diffBlock);
-
-        // First transcribe every parameter in the block.
-        for (auto param = origBlock->getFirstParam(); param; param = param->getNextParam())
-            this->copyParam(&subBuilder, param);
-
-        // The extra param for input gradient
-        auto gradParam = subBuilder.emitParam(as<IRFuncType>(origBlock->getParent()->getFullType())->getResultType());
-
-        // Then, run through every instruction and use the transcriber to generate the appropriate
-        // derivative code.
-        //
-        for (auto child = origBlock->getFirstOrdinaryInst(); child; child = child->getNextInst())
-            this->copyInst(&subBuilder, child);
-
-        auto lastInst = diffBlock->getLastOrdinaryInst();
-        List<IRInst*> grads = { gradParam };
-        upperGradients.Add(lastInst, grads);
-        for (auto child = diffBlock->getLastOrdinaryInst(); child; child = child->getPrevInst())
-        {
-            auto upperGrads = upperGradients.TryGetValue(child);
-            if (!upperGrads)
-                continue;
-            if (upperGrads->getCount() > 1)
-            {
-                auto sumGrad = upperGrads->getFirst();
-                for (auto i = 1; i < upperGrads->getCount(); i++)
-                {
-                    sumGrad = subBuilder.emitAdd(sumGrad->getDataType(), sumGrad, (*upperGrads)[i]);
-                }
-                this->transcribeInstBackward(&subBuilder, child, sumGrad);
-            }
-            else
-                this->transcribeInstBackward(&subBuilder, child, upperGrads->getFirst());
-        }
-
-        subBuilder.emitReturn();
-
-        return InstPair(diffBlock, diffBlock);
-    }
-
-    // Create an empty func to represent the transcribed func of `origFunc`.
-    InstPair transcribeFuncHeader(IRBuilder* inBuilder, IRFunc* origFunc)
-    {
-        IRBuilder builder(inBuilder->getSharedBuilder());
-        builder.setInsertBefore(origFunc);
-
-        IRFunc* primalFunc = origFunc;
-
-        differentiableTypeConformanceContext.setFunc(origFunc);
-
-        primalFunc = origFunc;
-
-        auto diffFunc = builder.createFunc();
-
-        SLANG_ASSERT(as<IRFuncType>(origFunc->getFullType()));
-        IRType* diffFuncType = this->differentiateFunctionType(
-            &builder,
-            as<IRFuncType>(origFunc->getFullType()));
-        diffFunc->setFullType(diffFuncType);
-
-        if (auto nameHint = origFunc->findDecoration<IRNameHintDecoration>())
-        {
-            auto originalName = nameHint->getName();
-            StringBuilder newNameSb;
-            newNameSb << "s_bwd_" << originalName;
-            builder.addNameHintDecoration(diffFunc, newNameSb.getUnownedSlice());
-        }
-        builder.addBackwardDerivativeDecoration(origFunc, diffFunc);
-
-        // Mark the generated derivative function itself as differentiable.
-        builder.addBackwardDifferentiableDecoration(diffFunc);
-
-        // Find and clone `DifferentiableTypeDictionaryDecoration` to the new diffFunc.
-        if (auto dictDecor = origFunc->findDecoration<IRDifferentiableTypeDictionaryDecoration>())
-        {
-            cloneDecoration(dictDecor, diffFunc);
-        }
-
-        auto result = InstPair(primalFunc, diffFunc);
-        followUpFunctionsToTranscribe.add(result);
-        return result;
-    }
-
-    // Transcribe a function definition.
-    InstPair transcribeFunc(IRBuilder* inBuilder, IRFunc* primalFunc, IRFunc* diffFunc)
-    {
-        IRBuilder builder(inBuilder->getSharedBuilder());
-        builder.setInsertInto(diffFunc);
-
-        differentiableTypeConformanceContext.setFunc(primalFunc);
-        // Transcribe children from origFunc into diffFunc
-        for (auto block = primalFunc->getFirstBlock(); block; block = block->getNextBlock())
-            this->transcribeBlock(&builder, block);
-
-        return InstPair(primalFunc, diffFunc);
-    }
-
-    IRInst* copyParam(IRBuilder* builder, IRParam* origParam)
-    {
-        auto primalDataType = origParam->getDataType();
-
-        if (auto diffPairType = tryGetDiffPairType(builder, (IRType*)primalDataType))
-        {
-            auto inoutDiffPairType = builder->getPtrType(kIROp_InOutType, diffPairType);
-            IRInst* diffParam = builder->emitParam(inoutDiffPairType);
-
-            auto diffPairVarName = makeDiffPairName(origParam);
-            if (diffPairVarName.getLength() > 0)
-                builder->addNameHintDecoration(diffParam, diffPairVarName.getUnownedSlice());
-
-            SLANG_ASSERT(diffParam);
-            auto paramValue = builder->emitLoad(diffParam);
-            auto primal = builder->emitDifferentialPairGetPrimal(paramValue);
-            orginalToTranscribed.Add(origParam, primal);
-            primalToDiffPair.Add(primal, diffParam);
-
-            return diffParam;
-        }
-
-
-        return cloneInst(&cloneEnv, builder, origParam);
-    }
-
-    InstPair copyBinaryArith(IRBuilder* builder, IRInst* origArith)
-    {
-        SLANG_ASSERT(origArith->getOperandCount() == 2);
-
-        auto origLeft = origArith->getOperand(0);
-        auto origRight = origArith->getOperand(1);
-
-        IRInst* primalLeft;
-        if (!orginalToTranscribed.TryGetValue(origLeft, primalLeft))
-        {
-            primalLeft = origLeft;
-        }
-        IRInst* primalRight;
-        if (!orginalToTranscribed.TryGetValue(origRight, primalRight))
-        {
-            primalRight = origRight;
-        }
-
-        auto resultType = origArith->getDataType();
-        IRInst* newInst;
-        switch (origArith->getOp())
-        {
-        case kIROp_Add:
-            newInst = builder->emitAdd(resultType, primalLeft, primalRight);
-            break;
-        case kIROp_Mul:
-            newInst = builder->emitMul(resultType, primalLeft, primalRight);
-            break;
-        case kIROp_Sub:
-            newInst = builder->emitSub(resultType, primalLeft, primalRight);
-            break;
-        case kIROp_Div:
-            newInst = builder->emitDiv(resultType, primalLeft, primalRight);
-            break;
-        default:
-            newInst = nullptr;
-            getSink()->diagnose(origArith->sourceLoc,
-                Diagnostics::unimplemented,
-                "this arithmetic instruction cannot be differentiated");
-        }
-        orginalToTranscribed.Add(origArith, newInst);
-        return InstPair(newInst, nullptr);
-    }
-
-    IRInst* transcribeBinaryArithBackward(IRBuilder* builder, IRInst* origArith, IRInst* grad)
-    {
-        SLANG_ASSERT(origArith->getOperandCount() == 2);
-
-        auto lhs = origArith->getOperand(0);
-        auto rhs = origArith->getOperand(1);
-
-        if (as<IRInOutType>(lhs->getDataType()))
-        {
-            lhs = builder->emitLoad(lhs);
-            lhs = builder->emitDifferentialPairGetPrimal(lhs);
-        }
-        if (as<IRInOutType>(rhs->getDataType()))
-        {
-            rhs = builder->emitLoad(rhs);
-            rhs = builder->emitDifferentialPairGetPrimal(rhs);
-        }
-
-        IRInst* leftGrad;
-        IRInst* rightGrad;
-
-
-        switch (origArith->getOp())
-        {
-        case kIROp_Add:
-            leftGrad = grad;
-            rightGrad = grad;
-            break;
-        case kIROp_Mul:
-            leftGrad = builder->emitMul(grad->getDataType(), rhs, grad);
-            rightGrad = builder->emitMul(grad->getDataType(), lhs, grad);
-            break;
-        case kIROp_Sub:
-            leftGrad = grad;
-            rightGrad = builder->emitNeg(grad->getDataType(), grad);
-            break;
-        case kIROp_Div:
-            leftGrad = builder->emitMul(grad->getDataType(), rhs, grad);
-            rightGrad = builder->emitMul(grad->getDataType(), lhs, grad); // TODO 1.0 / Grad
-            break;
-        default:
-            getSink()->diagnose(origArith->sourceLoc,
-                Diagnostics::unimplemented,
-                "this arithmetic instruction cannot be differentiated");
-        }
-
-        lhs = origArith->getOperand(0);
-        rhs = origArith->getOperand(1);
-        if (auto leftGrads = upperGradients.TryGetValue(lhs))
-        {
-            leftGrads->add(leftGrad);
-        }
-        else
-        {
-            upperGradients.Add(lhs, leftGrad);
-        }
-        if (auto rightGrads = upperGradients.TryGetValue(rhs))
-        {
-            rightGrads->add(rightGrad);
-        }
-        else
-        {
-            upperGradients.Add(rhs, rightGrad);
-        }
-
-        return nullptr;
-    }
-
-    InstPair copyInst(IRBuilder* builder, IRInst* origInst)
-    {
-        // Handle common SSA-style operations
-        switch (origInst->getOp())
-        {
-        case kIROp_Param:
-            return transcribeParam(builder, as<IRParam>(origInst));
-
-        case kIROp_Return:
-            return InstPair(nullptr, nullptr);
-
-        case kIROp_Add:
-        case kIROp_Mul:
-        case kIROp_Sub:
-        case kIROp_Div:
-            return copyBinaryArith(builder, origInst);
-
-        default:
-            // Not yet implemented
-            SLANG_ASSERT(0);
-        }
-
-        return InstPair(nullptr, nullptr);
-    }
-
-    IRInst* transcribeParamBackward(IRBuilder* builder, IRInst* param, IRInst* grad)
-    {
-        IRInOutType* inoutParam = as<IRInOutType>(param->getDataType());
-        auto pairType = as<IRDifferentialPairType>(inoutParam->getValueType());
-        auto paramValue = builder->emitLoad(param);
-        auto primal = builder->emitDifferentialPairGetPrimal(paramValue);
-        auto diff = builder->emitDifferentialPairGetDifferential(
-            (IRType*)pairBuilder->getDiffTypeFromPairType(builder, pairType),
-            paramValue
-        );
-        auto newDiff = builder->emitAdd(grad->getDataType(), diff, grad);
-        auto updatedParam = builder->emitMakeDifferentialPair(pairType, primal, newDiff);
-        auto store = builder->emitStore(param, updatedParam);
-
-        return store;
-    }
-
-    IRInst* transcribeInstBackward(IRBuilder* builder, IRInst* origInst, IRInst* grad)
-    {
-        // Handle common SSA-style operations
-        switch (origInst->getOp())
-        {
-        case kIROp_Param:
-            return transcribeParamBackward(builder, as<IRParam>(origInst), grad);
-
-        case kIROp_Add:
-        case kIROp_Mul:
-        case kIROp_Sub:
-        case kIROp_Div:
-            return transcribeBinaryArithBackward(builder, origInst, grad);
-
-        case kIROp_DifferentialPairGetPrimal:
-        {
-            if (auto param = primalToDiffPair.TryGetValue(origInst))
-            {
-                if (auto leftGrads = upperGradients.TryGetValue(*param))
-                {
-                    leftGrads->add(grad);
-                }
-                else
-                {
-                    upperGradients.Add(*param, grad);
-                }
-            }
-            else
-                SLANG_ASSERT(0);
-            return nullptr;
-        }
-
-        default:
-            // Not yet implemented
-            SLANG_ASSERT(0);
-        }
-
-        return nullptr;
-    }
-};
-
-
-struct JVPDerivativeContext : public InstPassBase
-{
-
-    DiagnosticSink* getSink()
-    {
-        return sink;
-    }
-
-    bool processModule()
-    {
-        // We start by initializing our shared IR building state,
-        // since we will re-use that state for any code we
-        // generate along the way.
-        //
-        SharedIRBuilder* sharedBuilder = &sharedBuilderStorage;
-        sharedBuilder->init(module);
-        sharedBuilder->deduplicateAndRebuildGlobalNumberingMap();
-
-        // TODO(sai): Move this call.
-        transcriberStorage.differentiableTypeConformanceContext.buildGlobalWitnessDictionary();
-
-        IRBuilder builderStorage(sharedBuilderStorage);
-        IRBuilder* builder = &builderStorage;
-
-        // Process all ForwardDifferentiate instructions (kIROp_ForwardDifferentiate), by 
-        // generating derivative code for the referenced function.
-        //
-        bool modified = processReferencedFunctions(builder);
-
-        // Replaces IRDifferentialPairType with an auto-generated struct,
-        // IRDifferentialPairGetDifferential with 'differential' field access,
-        // IRDifferentialPairGetPrimal with 'primal' field access, and
-        // IRMakeDifferentialPair with an IRMakeStruct.
-        // 
-        modified |= simplifyDifferentialBottomType(builder);
-
-        // De-duplicate any remaining types.
-        sharedBuilder->deduplicateAndRebuildGlobalNumberingMap();
-
-        modified |= processPairTypes(builder, module->getModuleInst());
-
-        modified |= eliminateDifferentialBottomType(builder);
-
-        return modified;
-    }
-
-    IRInst* lookupJVPReference(IRInst* primalFunction)
-    {
-        if (auto jvpDefinition = primalFunction->findDecoration<IRForwardDerivativeDecoration>())
-            return jvpDefinition->getForwardDerivativeFunc();
-        return nullptr;
-    }
-
-    // Recursively process instructions looking for JVP calls (kIROp_ForwardDifferentiate),
-    // then check that the referenced function is marked correctly for differentiation.
-    //
-    bool processReferencedFunctions(IRBuilder* builder)
-    {
-        List<IRInst*> autoDiffWorkList;
-
-        for (;;)
-        {
-            // Collect all `ForwardDifferentiate` insts from the module.
-            autoDiffWorkList.clear();
-            processAllInsts([&](IRInst* inst)
-                {
-                    switch (inst->getOp())
-                    {
-                    case kIROp_ForwardDifferentiate:
-                    case kIROp_BackwardDifferentiate:
-                        // Only process now if the operand is a materialized function.
-                        switch (inst->getOperand(0)->getOp())
-                        {
-                        case kIROp_Func:
-                        case kIROp_Specialize:
-                            autoDiffWorkList.add(inst);
-                            break;
-                        default:
-                            break;
-                        }
-                        break;
-                    default:
-                        break;
-                    }
-                });
-
-            if (autoDiffWorkList.getCount() == 0)
-                break;
-
-            // Process collected `ForwardDifferentiate` insts and replace them with placeholders for
-            // differentiated functions.
-
-            transcriberStorage.followUpFunctionsToTranscribe.clear();
-            backwardTranscriberStorage.followUpFunctionsToTranscribe.clear();
-
-            for (auto differentiateInst : autoDiffWorkList)
-            {
-                IRInst* baseInst = differentiateInst->getOperand(0);
-                if (as<IRForwardDifferentiate>(differentiateInst))
-                {
-                    if (auto existingDiffFunc = lookupJVPReference(baseInst))
-                    {
-                        differentiateInst->replaceUsesWith(existingDiffFunc);
-                        differentiateInst->removeAndDeallocate();
-                    }
-                    else if (isMarkedForForwardDifferentiation(baseInst))
-                    {
-                        if (as<IRFunc>(baseInst) || as<IRGeneric>(baseInst))
-                        {
-                            IRInst* diffFunc = transcriberStorage.transcribe(builder, baseInst);
-                            SLANG_ASSERT(diffFunc);
-                            differentiateInst->replaceUsesWith(diffFunc);
-                            differentiateInst->removeAndDeallocate();
-                        }
-                        else
-                        {
-                            getSink()->diagnose(differentiateInst->sourceLoc,
-                                Diagnostics::internalCompilerError,
-                                "Unexpected instruction. Expected func or generic");
-                        }
-                    }
-                    else
-                    {
-                        getSink()->diagnose(differentiateInst->sourceLoc,
-                            Diagnostics::internalCompilerError,
-                            "Requested differentiation on a function that isn't marked as differentiable.");
-                    }
-
-                }
-                else if (as<IRBackwardDifferentiate>(differentiateInst))
-                {
-                    if (isMarkedForBackwardDifferentiation(baseInst))
-                    {
-                        if (as<IRFunc>(baseInst))
-                        {
-                            IRInst* diffFunc =
-                                backwardTranscriberStorage
-                                    .transcribeFuncHeader(builder, (IRFunc*)baseInst)
-                                    .differential;
-                            SLANG_ASSERT(diffFunc);
-                            differentiateInst->replaceUsesWith(diffFunc);
-                            differentiateInst->removeAndDeallocate();
-                        }
-                        else
-                        {
-                            getSink()->diagnose(differentiateInst->sourceLoc,
-                                Diagnostics::internalCompilerError,
-                                "Unexpected instruction. Expected func or generic");
-                        }
-                    }
-                }
-            }
-            // Actually synthesize the derivatives.
-            List<InstPair> followUpWorkList = _Move(transcriberStorage.followUpFunctionsToTranscribe);
-            for (auto task : followUpWorkList)
-            {
-                auto diffFunc = as<IRFunc>(task.differential);
-                SLANG_ASSERT(diffFunc);
-                auto primalFunc = as<IRFunc>(task.primal);
-                SLANG_ASSERT(primalFunc);
-
-                transcriberStorage.transcribeFunc(builder, primalFunc, diffFunc);
-            }
-            followUpWorkList = _Move(backwardTranscriberStorage.followUpFunctionsToTranscribe);
-            for (auto task : followUpWorkList)
-            {
-                auto diffFunc = as<IRFunc>(task.differential);
-                SLANG_ASSERT(diffFunc);
-                auto primalFunc = as<IRFunc>(task.primal);
-                SLANG_ASSERT(primalFunc);
-
-                backwardTranscriberStorage.transcribeFunc(builder, primalFunc, diffFunc);
-            }
-
-            // Transcribing the function body really shouldn't produce more follow up function body work.
-            // However it may produce new `ForwardDifferentiate` instructions, which we collect and process
-            // in the next iteration.
-            SLANG_RELEASE_ASSERT(transcriberStorage.followUpFunctionsToTranscribe.getCount() == 0);
-
-        }
-        return true;
-    }
-
-    IRInst* lowerPairType(IRBuilder* builder, IRType* pairType, bool* isTrivial = nullptr)
-    {
-        builder->setInsertBefore(pairType);
-        auto loweredPairTypeInfo = (&pairBuilderStorage)->lowerDiffPairType(
-            builder,
-            pairType);
-        if (isTrivial)
-            *isTrivial = loweredPairTypeInfo.isTrivial;
-        return loweredPairTypeInfo.loweredType;
-    }
-
-    IRInst* lowerMakePair(IRBuilder* builder, IRInst* inst)
-    {
-
-        if (auto makePairInst = as<IRMakeDifferentialPair>(inst))
-        {
-            bool isTrivial = false;
-            auto pairType = as<IRDifferentialPairType>(makePairInst->getDataType());
-            if (auto loweredPairType = lowerPairType(builder, pairType, &isTrivial))
-            {
-                builder->setInsertBefore(makePairInst);
-                IRInst* result = nullptr;
-                if (isTrivial)
-                {
-                    result = makePairInst->getPrimalValue();
-                }
-                else
-                {
-                    IRInst* operands[2] = { makePairInst->getPrimalValue(), makePairInst->getDifferentialValue() };
-                    result = builder->emitMakeStruct((IRType*)(loweredPairType), 2, operands);
-                }
-                makePairInst->replaceUsesWith(result);
-                makePairInst->removeAndDeallocate();
-                return result;
-            }
-        }
-
-        return nullptr;
-    }
-
-    IRInst* lowerPairAccess(IRBuilder* builder, IRInst* inst)
-    {
-        if (auto getDiffInst = as<IRDifferentialPairGetDifferential>(inst))
-        {
-            auto pairType = getDiffInst->getBase()->getDataType();
-            if (auto pairPtrType = as<IRPtrTypeBase>(pairType))
-            {
-                pairType = pairPtrType->getValueType();
-            }
-
-            if (lowerPairType(builder, pairType, nullptr))
-            {
-                builder->setInsertBefore(getDiffInst);
-                IRInst* diffFieldExtract = nullptr;
-                diffFieldExtract = (&pairBuilderStorage)->emitDiffFieldAccess(builder, getDiffInst->getBase());
-                getDiffInst->replaceUsesWith(diffFieldExtract);
-                getDiffInst->removeAndDeallocate();
-                return diffFieldExtract;
-            }
-        }
-        else if (auto getPrimalInst = as<IRDifferentialPairGetPrimal>(inst))
-        {
-            auto pairType = getPrimalInst->getBase()->getDataType();
-            if (auto pairPtrType = as<IRPtrTypeBase>(pairType))
-            {
-                pairType = pairPtrType->getValueType();
-            }
-
-            if (lowerPairType(builder, pairType, nullptr))
-            {
-                builder->setInsertBefore(getPrimalInst);
-
-                IRInst* primalFieldExtract = nullptr;
-                primalFieldExtract = (&pairBuilderStorage)->emitPrimalFieldAccess(builder, getPrimalInst->getBase());
-                getPrimalInst->replaceUsesWith(primalFieldExtract);
-                getPrimalInst->removeAndDeallocate();
-                return primalFieldExtract;
-            }
-        }
-
-        return nullptr;
-    }
-
-    bool processPairTypes(IRBuilder* builder, IRInst* instWithChildren)
-    {
-        bool modified = false;
-        // Hoist all pair types to global scope when possible.
-        auto moduleInst = module->getModuleInst();
-        processInstsOfType<IRDifferentialPairType>(kIROp_DifferentialPairType, [&](IRInst* originalPairType)
-            {
-                if (originalPairType->parent != moduleInst)
-                {
-                    originalPairType->removeFromParent();
-                    ShortList<IRInst*> operands;
-                    for (UInt i = 0; i < originalPairType->getOperandCount(); i++)
-                    {
-                        operands.add(originalPairType->getOperand(i));
-                    }
-                    auto newPairType = builder->findOrEmitHoistableInst(
-                        originalPairType->getFullType(),
-                        originalPairType->getOp(),
-                        originalPairType->getOperandCount(),
-                        operands.getArrayView().getBuffer());
-                    originalPairType->replaceUsesWith(newPairType);
-                    originalPairType->removeAndDeallocate();
-                }
-            });
-
-        sharedBuilderStorage.deduplicateAndRebuildGlobalNumberingMap();
-
-        processAllInsts([&](IRInst* inst)
-            {
-                // Make sure the builder is at the right level.
-                builder->setInsertInto(instWithChildren);
-
-                switch (inst->getOp())
-                {
-                case kIROp_DifferentialPairGetDifferential:
-                case kIROp_DifferentialPairGetPrimal:
-                    lowerPairAccess(builder, inst);
-                    modified = true;
-                    break;
-
-                case kIROp_MakeDifferentialPair:
-                    lowerMakePair(builder, inst);
-                    modified = true;
-                    break;
-
-                default:
-                    break;
-                }
-            });
-
-        processInstsOfType<IRDifferentialPairType>(kIROp_DifferentialPairType, [&](IRDifferentialPairType* inst)
-            {
-                if (auto loweredType = lowerPairType(builder, inst))
-                {
-                    inst->replaceUsesWith(loweredType);
-                    inst->removeAndDeallocate();
-                }
-            });
-        return modified;
-    }
-
-    bool simplifyDifferentialBottomType(IRBuilder* builder)
-    {
-        bool modified = false;
-        auto diffBottom = builder->getDifferentialBottom();
-
-        bool changed = true;
-        List<IRUse*> uses;
-        while (changed)
-        {
-            changed = false;
-            // Replace all insts whose type is `DifferentialBottomType` to `diffBottom`.
-            processAllInsts([&](IRInst* inst)
-                {
-                    if (inst->getDataType() && inst->getDataType()->getOp() == kIROp_DifferentialBottomType)
-                    {
-                        if (inst != diffBottom)
-                        {
-                            inst->replaceUsesWith(diffBottom);
-                            inst->removeAndDeallocate();
-                            modified = true;
-                        }
-                    }
-                });
-            // Go through all uses of diffBottom and run simplification.
-            processAllInsts([&](IRInst* inst)
-                {
-                    if (!inst->hasUses())
-                        return;
-
-                    builder->setInsertBefore(inst);
-                    IRInst* valueToReplace = nullptr;
-                    switch (inst->getOp())
-                    {
-                    case kIROp_Store:
-                        if (as<IRStore>(inst)->getVal() == diffBottom)
-                        {
-                            inst->removeAndDeallocate();
-                            changed = true;
-                        }
-                        return;
-                    case kIROp_MakeDifferentialPair:
-                        // Our simplification could lead to a situation where
-                        // bottom is used to make a pair that has a non-bottom differential type,
-                        // in this case we should use zero instead.
-                        if (inst->getOperand(1) == diffBottom)
-                        {
-                            // Only apply if we are the second operand.
-                            auto pairType = as<IRDifferentialPairType>(inst->getDataType());
-                            if (pairBuilderStorage.getDiffTypeFromPairType(builder, pairType)->getOp() != kIROp_DifferentialBottomType)
-                            {
-                                auto zero = transcriberStorage.getDifferentialZeroOfType(builder, pairType->getValueType());
-                                inst->setOperand(1, zero);
-                                changed = true;
-                            }
-                        }
-                        return;
-                    case kIROp_DifferentialPairGetDifferential:
-                        if (inst->getOperand(0)->getOp() == kIROp_MakeDifferentialPair)
-                        {
-                            valueToReplace = inst->getOperand(0)->getOperand(1);
-                        }
-                        break;
-                    case kIROp_DifferentialPairGetPrimal:
-                        if (inst->getOperand(0)->getOp() == kIROp_MakeDifferentialPair)
-                        {
-                            valueToReplace = inst->getOperand(0)->getOperand(0);
-                        }
-                        break;
-                    case kIROp_Add:
-                        if (inst->getOperand(0) == diffBottom)
-                        {
-                            valueToReplace = inst->getOperand(1);
-                        }
-                        else if (inst->getOperand(1) == diffBottom)
-                        {
-                            valueToReplace = inst->getOperand(0);
-                        }
-                        break;
-                    case kIROp_Sub:
-                        if (inst->getOperand(0) == diffBottom)
-                        {
-                            // If left is bottom, and right is not bottom, then we should return -right.
-                            // However we can't possibly run into that case since both side of - operator
-                            // must be at the same order of differentiation.
-                            valueToReplace = diffBottom;
-                        }
-                        else if (inst->getOperand(1) == diffBottom)
-                        {
-                            valueToReplace = inst->getOperand(0);
-                        }
-                        break;
-                    case kIROp_Mul:
-                    case kIROp_Div:
-                        if (inst->getOperand(0) == diffBottom)
-                        {
-                            valueToReplace = diffBottom;
-                        }
-                        else if (inst->getOperand(1) == diffBottom)
-                        {
-                            valueToReplace = diffBottom;
-                        }
-                        break;
-                    default:
-                        break;
-                    }
-                    if (valueToReplace)
-                    {
-                        inst->replaceUsesWith(valueToReplace);
-                        changed = true;
-                    }
-                });
-            modified |= changed;
-        }
-
-        return modified;
-    }
-
-    bool eliminateDifferentialBottomType(IRBuilder* builder)
-    {
-        simplifyDifferentialBottomType(builder);
-
-        bool modified = false;
-        auto diffBottom = builder->getDifferentialBottom();
-        auto diffBottomType = diffBottom->getDataType();
-        diffBottom->replaceUsesWith(builder->getVoidValue());
-        diffBottom->removeAndDeallocate();
-        diffBottomType->replaceUsesWith(builder->getVoidType());
-
-        return modified;
-    }
-
-    // Checks decorators to see if the function should
-    // be differentiated (kIROp_ForwardDifferentiableDecoration)
-    // 
-    bool isMarkedForForwardDifferentiation(IRInst* callable)
-    {
-        return callable->findDecoration<IRForwardDifferentiableDecoration>() != nullptr;
-    }
-
-    IRStringLit* getForwardDerivativeFuncName(IRInst* func)
-    {
-        IRBuilder builder(&sharedBuilderStorage);
-        builder.setInsertBefore(func);
-
-        IRStringLit* name = nullptr;
-        if (auto linkageDecoration = func->findDecoration<IRLinkageDecoration>())
-        {
-            name = builder.getStringValue((String(linkageDecoration->getMangledName()) + "_fwd_diff").getUnownedSlice());
-        }
-        else if (auto namehintDecoration = func->findDecoration<IRNameHintDecoration>())
-        {
-            name = builder.getStringValue((String(namehintDecoration->getName()) + "_fwd_diff").getUnownedSlice());
-        }
-
-        return name;
-    }
-
-    // Checks decorators to see if the function should
-    // be differentiated (kIROp_ForwardDifferentiableDecoration)
-    // 
-    bool isMarkedForBackwardDifferentiation(IRInst* callable)
-    {
-        return callable->findDecoration<IRBackwardDifferentiableDecoration>() != nullptr;
-    }
-
-    IRStringLit* getBackwardDerivativeFuncName(IRInst* func)
-    {
-        IRBuilder builder(&sharedBuilderStorage);
-        builder.setInsertBefore(func);
-
-        IRStringLit* name = nullptr;
-        if (auto linkageDecoration = func->findDecoration<IRLinkageDecoration>())
-        {
-            name = builder.getStringValue((String(linkageDecoration->getMangledName()) + "_bwd_diff").getUnownedSlice());
-        }
-        else if (auto namehintDecoration = func->findDecoration<IRNameHintDecoration>())
-        {
-            name = builder.getStringValue((String(namehintDecoration->getName()) + "_bwd_diff").getUnownedSlice());
-        }
-
-        return name;
-    }
-
-    JVPDerivativeContext(IRModule* module, DiagnosticSink* sink) :
-        InstPassBase(module),
-        sink(sink),
-        autoDiffSharedContextStorage(module->getModuleInst()),
-        transcriberStorage(&autoDiffSharedContextStorage, &sharedBuilderStorage),
-        backwardTranscriberStorage(&autoDiffSharedContextStorage, &sharedBuilderStorage, sink)
-    {
-        autoDiffSharedContextStorage.sharedBuilder = &sharedBuilderStorage;
-        pairBuilderStorage.sharedContext = &autoDiffSharedContextStorage;
-        transcriberStorage.sink = sink;
-        transcriberStorage.autoDiffSharedContext = &(autoDiffSharedContextStorage);
-        transcriberStorage.pairBuilder = &(pairBuilderStorage);
-        backwardTranscriberStorage.pairBuilder = &pairBuilderStorage;
-    }
-
-protected:
-    // A transcriber object that handles the main job of 
-    // processing instructions while maintaining state.
-    //
-    JVPTranscriber                  transcriberStorage;
-
-    BackwardDiffTranscriber         backwardTranscriberStorage;
-
-    // Diagnostic object from the compile request for
-    // error messages.
-    DiagnosticSink* sink;
-
-    // Context to find and manage the witness tables for types 
-    // implementing `IDifferentiable`
-    AutoDiffSharedContext           autoDiffSharedContextStorage;
-
-    // Builder for dealing with differential pair types.
-    DifferentialPairTypeBuilder     pairBuilderStorage;
-
-};
-
-// Set up context and call main process method.
-//
-bool processDifferentiableFuncs(
-    IRModule* module,
-    DiagnosticSink* sink,
-    IRJVPDerivativePassOptions const&)
-{
-    // Simplify module to remove dead code.
-    IRDeadCodeEliminationOptions options;
-    options.keepExportsAlive = true;
-    options.keepLayoutsAlive = true;
-    eliminateDeadCode(module, options);
-
-    JVPDerivativeContext context(module, sink);
-    bool changed = context.processModule();
-    return changed;
-}
-
-void stripAutoDiffDecorationsFromChildren(IRInst* parent)
-{
-    for (auto inst : parent->getChildren())
-    {
-        for (auto decor = inst->getFirstDecoration(); decor; )
-        {
-            auto next = decor->getNextDecoration();
-            switch (decor->getOp())
-            {
-            case kIROp_ForwardDerivativeDecoration:
-            case kIROp_DerivativeMemberDecoration:
-            case kIROp_DifferentiableTypeDictionaryDecoration:
-                decor->removeAndDeallocate();
-                break;
-            default:
-                break;
-            }
-            decor = next;
-        }
-
-        if (inst->getFirstChild() != nullptr)
-        {
-            stripAutoDiffDecorationsFromChildren(inst);
-        }
-    }
-}
-
-void stripAutoDiffDecorations(IRModule* module)
-{
-    stripAutoDiffDecorationsFromChildren(module->getModuleInst());
-}
-
-AutoDiffSharedContext::AutoDiffSharedContext(IRModuleInst* inModuleInst)
-    : moduleInst(inModuleInst)
-{
-    differentiableInterfaceType = as<IRInterfaceType>(findDifferentiableInterface());
-    if (differentiableInterfaceType)
-    {
-        differentialAssocTypeStructKey = findDifferentialTypeStructKey();
-        differentialAssocTypeWitnessStructKey = findDifferentialTypeWitnessStructKey();
-        zeroMethodStructKey = findZeroMethodStructKey();
-        addMethodStructKey = findAddMethodStructKey();
-        mulMethodStructKey = findMulMethodStructKey();
-
-        if (differentialAssocTypeStructKey)
-            isInterfaceAvailable = true;
-    }
-}
-
-IRInst* AutoDiffSharedContext::findDifferentiableInterface()
-{
-    if (auto module = as<IRModuleInst>(moduleInst))
-    {
-        for (auto globalInst : module->getGlobalInsts())
-        {
-            // TODO: This seems like a particularly dangerous way to look for an interface.
-            // See if we can lower IDifferentiable to a separate IR inst.
-            //
-            if (globalInst->getOp() == kIROp_InterfaceType &&
-                as<IRInterfaceType>(globalInst)->findDecoration<IRNameHintDecoration>()->getName() == "IDifferentiable")
-            {
-                return globalInst;
-            }
-        }
-    }
-    return nullptr;
-}
-
-IRStructKey* AutoDiffSharedContext::getIDifferentiableStructKeyAtIndex(UInt index)
-{
-    if (as<IRModuleInst>(moduleInst) && differentiableInterfaceType)
-    {
-        // Assume for now that IDifferentiable has exactly five fields.
-        SLANG_ASSERT(differentiableInterfaceType->getOperandCount() == 5);
-        if (auto entry = as<IRInterfaceRequirementEntry>(differentiableInterfaceType->getOperand(index)))
-            return as<IRStructKey>(entry->getRequirementKey());
-        else
-        {
-            SLANG_UNEXPECTED("IDifferentiable interface entry unexpected type");
-        }
-    }
-
-    return nullptr;
-}
-
-void DifferentiableTypeConformanceContext::setFunc(IRGlobalValueWithCode* func)
-{
-    parentFunc = func;
-
-    auto decor = func->findDecoration<IRDifferentiableTypeDictionaryDecoration>();
-    SLANG_RELEASE_ASSERT(decor);
-
-    // Build lookup dictionary for type witnesses.
-    for (auto child = decor->getFirstChild(); child; child = child->next)
-    {
-        if (auto item = as<IRDifferentiableTypeDictionaryItem>(child))
-        {
-            auto existingItem = differentiableWitnessDictionary.TryGetValue(item->getConcreteType());
-            if (existingItem)
-            {
-                if (auto witness = as<IRWitnessTable>(item->getWitness()))
-                {
-                    if (witness->getConcreteType()->getOp() == kIROp_DifferentialBottomType)
-                        continue;
-                }
-                *existingItem = item->getWitness();
-            }
-            else
-            {
-                differentiableWitnessDictionary.Add((IRType*)item->getConcreteType(), item->getWitness());
-            }
-        }
-    }
-}
-
-
-// Lookup a witness table for the concreteType. One should exist if concreteType
-// inherits (successfully) from IDifferentiable.
-// 
-
-IRInst* DifferentiableTypeConformanceContext::lookUpConformanceForType(IRInst* type)
-{
-    IRInst* foundResult = nullptr;
-    differentiableWitnessDictionary.TryGetValue(type, foundResult);
-    return foundResult;
-}
-
-IRInst* DifferentiableTypeConformanceContext::lookUpInterfaceMethod(IRBuilder* builder, IRType* origType, IRStructKey* key)
-{
-    if (auto conformance = lookUpConformanceForType(origType))
-    {
-        return _lookupWitness(builder, conformance, key);
-    }
-    return nullptr;
-}
-
-void DifferentiableTypeConformanceContext::buildGlobalWitnessDictionary()
-{
-    for (auto globalInst : sharedContext->moduleInst->getChildren())
-    {
-        if (auto pairType = as<IRDifferentialPairType>(globalInst))
-        {
-            differentiableWitnessDictionary.Add(pairType->getValueType(), pairType->getWitness());
-        }
-    }
-}
-}