classdef (Sealed, ConstructOnLoad) AcceleratedFunction < handle & ...
        matlab.mixin.internal.indexing.Paren & ...
        matlab.mixin.internal.Scalar & ...
        matlab.mixin.internal.indexing.DotNotOverloaded
    %AcceleratedFunction  Accelerated deep learning function
    %
    %   To create an AcceleratedFunction, use the DLACCELERATE function.
    %
    %   To evaluate the accelerated function, evaluate the object directly, or,
    %   equivalently, use DLFEVAL or FEVAL. For example,
    %     accfun = dlaccelerate(@modelGradients);
    %     gradients = accfun(dlnet, dlX, T);
    %
    %   An AcceleratedFunction stores traces of the underlying function.
    %   A trace depends on the function inputs and outputs:
    %     - For dlarray inputs, the trace depends on the size, format, and
    %       underlying datatype of the dlarray.
    %     - For dlnetwork inputs, the trace depends on the size, format, and
    %       underlying datatype of the dlnetwork state and learnable parameters.
    %     - For other inputs, the trace depends on the values of the input.
    %     - The trace depends on the number of outputs.
    %   When a new trace is needed, the underlying function is evaluated and the
    %   resulting trace is cached in the AcceleratedFunction.
    %
    %   AcceleratedFunction properties:
    %      Function - Underlying traced function, specified as a function
    %                 handle. This property is read-only.
    %
    %      Enabled - Flag to enable tracing, specified as a logical. The
    %                default is true.
    %
    %      CacheSize  - Size of cache, specified as a positive integer.
    %                   This corresponds to the maximum number of input and
    %                   output combinations to cache. The default value is 50.
    %
    %      HitRate - Cache hit rate, specified as a positive
    %                scalar. This property corresponds to the percentage of
    %                reused evaluations. This property is read-only.
    %
    %      Occupancy - Cache occupancy, specified as a positive
    %                  scalar. This property corresponds to the percentage
    %                  of cache in use. This property is read-only.
    %
    %      CheckMode - Check mode, specified as 'none' or 'tolerance'.
    %                  Checking outputs will compare the accelerated
    %                  results with the results of evaluating the
    %                  underlying function. Warn if the results differ by
    %                  more than CheckTolerance. The default is
    %                  'none'.
    %
    %      CheckTolerance - Check tolerance specified as a scalar numeric
    %                       value. The tolerance is used for both absolute
    %                       and relative error. This value is only used
    %                       when CheckMode is 'tolerance'. The default is 1e-4.
    %
    %   AcceleratedFunction methods:
    %    clearCache - Clears the accelerated function trace cache.
    %
    %   See also: DLACCELERATE, DLFEVAL

    % Copyright 2020-2023 The MathWorks, Inc.

    properties (SetAccess = private)
        % FUNCTION  MATLAB function_handle referenced.
        Function function_handle;
    end

    properties
        % ENABLED  Logical value indicating state of the AcceleratedFunction. Use
        % this property to toggle using the cache.
        Enabled logical = true;
    end

    properties (Dependent)
        % CacheSize  Double value indicating maximum number of cache entries.
        CacheSize;
    end

    properties (Dependent, SetAccess = private)
        % HitRate  Percentage of evaluations reusing a cached trace.
        HitRate;

        % Occupancy  Percentage of available cache in use.
        Occupancy;
    end

    properties (Dependent)
        % CheckMode  Output checking mode. Values are
        % 'none' and 'tolerance'.
        CheckMode
    end

    properties
        % CHECKTOLERANCE  Scalar numeric value for absolute and
        % relative error tolerance while checking outputs. Only
        % used when CheckMode is 'tolerance'.
        CheckTolerance = 1e-4;
    end

    properties (Transient, Access = private)
        % Holds all cached data
        Cache;

        % Specify number of inputs to compute gradients with respect to. Used
        % to avoid backpropagating to state inputs in a network.
        MaxNumGradients = inf;
    end

    properties (Access = private)
        % Default cache size
        PrivateCacheSize double = 50;

        % Logical value controlling output checking.
        CheckOutputs logical = false;

        % Versioning to maintain backwards compatibility
        Version double = 1.0;

        % Logical value controlling whether to optimize the traced
        % function.
        EnableOptimization (1,1) logical = true;

        % Logical value controlling whether to trace the full
        % backward (if necessary) of the traced function.
        EnableFullBackward(1,1) logical = true;

        % Logical value controlling whether to use static analysis on the
        % supplied function.
        EnableStaticAnalysis (1,1) logical = false;

        % Force backward pass code generation in tracing mode.
        ForceBackwardPassCodegenInTracingMode (1,1) logical = false;

        % Logical value controlling whether to use C++ for MATLAB code
        % generation.
        EnableCPPCodeGeneration(1,1) logical = true;

        % Logical value controlling whether to use "no-grad" mode during
        % tracing of function. In no-grad mode, we trace the function but
        % don't save intermediates for a backward pass. Using no-grad mode
        % means that we will not run a backward pass on the trace. If this
        % AcceleratedFunction is called from a tracing context then this
        % mode has no effect.
        NoGradUnlessTracingMode(1,1) logical = true;

        % Integer specifying the maximum number of lines in the generated
        % function. If generating a single function would exceed the
        % maximum length, then AcceleratedFunction attempts to generate one
        % main function with calls to local functions of length smaller
        % than equal to MaxGeneratedFunctionLength.
        MaxGeneratedFunctionLength double {mustBeScalarOrEmpty, mustBeInteger, mustBePositive} = [];
    end

    methods
        function set.Enabled(obj, state)
            % Setter method for Enabled
            validateattributes(state, {'logical'}, {'scalar','nonnan'});
            obj.Enabled = state;
        end

        function cSize = get.CacheSize(obj)
            % Getter method for CacheSize
            cSize = obj.PrivateCacheSize;
        end
        function set.CacheSize(obj, val)
            % Setter method for CacheSize
            validateattributes(val, {'numeric'},{'scalar', 'positive', 'nonnan', 'integer'});
            updateCacheSize(obj, val);
        end

        function p = get.HitRate(obj)
            objCache = obj.Cache;
            p = objCache.TotalHits /(objCache.TotalHits + objCache.TotalMisses) * 100;
            if isnan(p)
                p = 0;
            end
        end

        function p = get.Occupancy(obj)
            objCache = obj.Cache;
            p = objCache.Load/obj.PrivateCacheSize*100;
        end

        function set.CheckMode(obj, state)
            % Setter method for CheckMode
            state = validatestring(state, {'none','tolerance'});
            obj.CheckOutputs = strcmp(state,'tolerance');
        end
        function p = get.CheckMode(obj)
            if obj.CheckOutputs
                p = 'tolerance';
            else
                p = 'none';
            end
        end

        function set.CheckTolerance(obj, val)
            % Setter method for CheckTolerance
            validateattributes(val, {'numeric'},{'scalar', 'positive', 'nonnan'});
            obj.CheckTolerance = val;
        end

        function clearCache(obj)
            %clearCache Clear accelerated function trace cache
            %   clearCache(accfun) clears the trace cache held by the specified AcceleratedFunction object.
            initializeCache(obj);
        end
    end

    methods (Access = {?deep.internal.AcceleratedFunctionManager})
        function obj = AcceleratedFunction(fun, maxCacheSize)
            % Constructor, only accessible from the manager
            if nargin > 0
                obj.Function = fun;
                if nargin == 2
                    obj.PrivateCacheSize = maxCacheSize;
                end
            end
            initializeCache(obj);
        end

        function delete(~)
            % Hiding deletion, as deleting one handle will invalidate all
            % other handles that exist for the same function
        end

    end

    methods (Hidden)

        
        function varargout = parenReference(obj,varargin)
            % Overloading parenReference to mimic function call syntax
            persistent tm;
            if isempty(tm)
                tm = deep.internal.recording.TapeManager;
            end
            if ~obj.Enabled
                [varargout{1:nargout}] = obj.execute(varargin);
                return;
            end

            % local variables for commonly used values.
            load = obj.Cache.Load;
            cachedSignatures = obj.Cache.Signatures;

            try
                [signature, parameters] = iEvaluationInputs(varargin);
            catch me
                throw(me);
            end
            isTracing = getTracingCount(tm) > 0;
            signature = [{isTracing} signature];
            retraceWithoutGradients = false;
            initialHitCount = 0;
            if load > 0
                for i = load : -1 : 1 % Iterate backwards
                    if isequal(signature, cachedSignatures{i}) ...
                            && nargout == obj.Cache.Nargout(1,i)
                        obj.Cache.TotalHits = obj.Cache.TotalHits + 1;
                        if obj.Cache.FlagForRetracing(i)
                            % Cache hit but we want to retrace
                            retraceWithoutGradients = true;
                            initialHitCount = 1;
                            obj.deleteCacheEntry(i);
                            break;
                        end
                        % Cache hit!
                        fun = obj.Cache.Functions{i};
                        nout = obj.Cache.Nargout(2,i);
                        if obj.CheckOutputs
                            rngState = rng;
                        end
                        if isempty(fun)
                            outputs = {};
                        else
                            [outputs{1:nout}] = fun(parameters{:});
                        end
                        cachedOutputs = obj.Cache.Outputs{i};
                        varargout = iEvaluationOutputs(cachedOutputs, outputs);
                        obj.Cache.HitCount(i) = obj.Cache.HitCount(i) + 1;
                        if i ~= load
                            obj.prioritize(i, load);
                        end
                        if obj.CheckOutputs
                            rng(rngState);
                            [underlyingOutputs{1:nargout}] = dlfeval(obj.Function, varargin{:});
                            if ~deep.internal.isEqual(varargout, underlyingOutputs, obj.CheckTolerance)
                                warning(message('deep:dlaccelerate:OutputsDiffer'));
                            end
                        end
                        return
                    end
                end
            end

            % Trace the function.
            if iIsThreadPoolWorker()
                error(message('deep:dlaccelerate:ThreadWorkerNotSupported'));
            end
            needsForward = isTracing || ~retraceWithoutGradients;
            [inputs, inputNodes, restore, priorTapeCount] = obj.setupTrace(varargin, tm, isTracing, needsForward); %#ok<ASGLU>
            restore{end+1} = setupNoGradModeIfNeeded(obj, isTracing, needsForward);
            restore{end+1} = setupToOptimizeTape(obj);
            restore = flip(restore); %#ok<NASGU> % Guarantees destruction order
            if obj.EnableStaticAnalysis
                [varargout{1:nargout}] = obj.transformAndExecuteWithFallback(inputs);
            else
                [varargout{1:nargout}] = obj.Function(inputs{:});
            end
            % If we did not encounter dlgradient, flag this for next time
            % so we can regenerate a more efficient function. If dlgradient
            % was encountered no-grad mode will have been disabled on the
            % tape.
            retraceWithoutGradients = ~isTracing && ...
                deep.internal.recording.tape.isNoGradMode() && ...
                iRetraceForInference() && ... % Is there any advantage in retracing?
                ~retraceWithoutGradients;
            [cacheData, varargout, illegalOutputs] = generateForward(obj, varargout, inputNodes, tm, priorTapeCount, isTracing);
            if illegalOutputs
                error(message('deep:dlaccelerate:ObjectOutput'));
            end
            obj.Cache.TotalMisses = obj.Cache.TotalMisses + 1 - initialHitCount;

            % Refresh load variable
            load = obj.Cache.Load;

            nargoutEntry = [nargout; cacheData.NumTraceOutputs];

            % Check for illegal recursion
            cachedSignatures = obj.Cache.Signatures;
            for i = 1:load
                if isequal(signature, cachedSignatures{i}) ...
                        && nargout == obj.Cache.Nargout(1,i)
                    error(message('deep:dlaccelerate:RecursiveDlarray'));
                end
            end
            % Cache result
            if load < obj.PrivateCacheSize
                obj.Cache.Signatures{end+1} = signature;
                obj.Cache.Nargout(:,end+1) = nargoutEntry;
                obj.Cache.Functions{end+1} = cacheData.Function;
                obj.Cache.Outputs{end+1} = cacheData.Output;
                obj.Cache.GeneratedCode{end+1} = cacheData.GeneratedCode;
                obj.Cache.TapeFunction{end+1} = cacheData.TapeFunction;
                obj.Cache.HitCount(end+1) = initialHitCount;
                obj.Cache.FlagForRetracing(end+1) = retraceWithoutGradients;
                obj.Cache.Load = load + 1;
            else
                obj.Cache.Signatures  = [obj.Cache.Signatures(2:end), {signature}];
                obj.Cache.Nargout = [obj.Cache.Nargout(:,2:end), nargoutEntry];
                obj.Cache.Functions = [obj.Cache.Functions(2:end), {cacheData.Function}];
                obj.Cache.Outputs = [obj.Cache.Outputs(2:end), {cacheData.Output}];
                obj.Cache.GeneratedCode = [obj.Cache.GeneratedCode(2:end), {cacheData.GeneratedCode}];
                obj.Cache.TapeFunction = [obj.Cache.TapeFunction(2:end), {cacheData.TapeFunction}];
                obj.Cache.HitCount = [obj.Cache.HitCount(2:end), initialHitCount];
                obj.Cache.FlagForRetracing = [obj.Cache.FlagForRetracing(2:end), retraceWithoutGradients];
            end
        end

        function varargout = feval(obj, varargin)
            [varargout{1:nargout}] = parenReference(obj,varargin{:});
        end

        function varargout = dlfeval(obj,varargin)
            % If we're already inside dlfeval, calling dlfeval is not
            % allowed. Instead of duplicating the error here, allow
            % dlfevalWithNestingCheck to handle it.
            tapeData = deep.internal.recording.tape.getTapeData();
            alreadyInsideDlfeval = ~isempty(tapeData) && tapeData.IsTracingForDlfeval;
            % If the first argument isn't an AcceleratedFunction then it is
            % a function input not the function to be evaluated.
            isFunctionToBeEvaluated = isa(obj,'deep.AcceleratedFunction');
            % Mark the tape as inside dlfeval
            tapeDataRestore = deep.internal.recording.TapeData.setTapeDataForDlfeval(); %#ok<NASGU>
            if ~isFunctionToBeEvaluated || alreadyInsideDlfeval
                [varargout{1:nargout}] = deep.internal.dlfevalWithNestingCheck(obj,varargin{:});
            else
                if ~obj.Enabled
                    [varargout{1:nargout}] = dlfeval(obj.Function,varargin{:});
                else
                    origMode = obj.NoGradUnlessTracingMode;
                    restoreFcn = @(o)setNoGradUnlessTracingMode(o, origMode);
                    restore = onCleanup(@()iProtectedMethodCall(restoreFcn, obj));
                    obj.NoGradUnlessTracingMode = false;
                    [varargout{1:nargout}] = parenReference(obj,varargin{:});
                end
            end
        end

        function varargout = repmat(obj, varargin) %#ok<STOUT>
            % REPMAT always errors
            try
                obj(2) = obj; %#ok Mimic array formation error message.
            catch err
                throwAsCaller(err)
            end
        end

    end

    % Internal APIs for debugging and controlling gradients.
    methods (Access = ?deep.internal.AcceleratedFunctionManager)
        function setMaxNumGradients(obj, n)
            obj.MaxNumGradients = n;
        end
        function cache = getCache(obj)
            cache = obj.Cache;
        end
        function setEnableOptimization(obj, tf)
            obj.EnableOptimization = tf;
        end
        function setEnableFullBackward(obj, tf)
            obj.EnableFullBackward = tf;
        end
        function setEnableStaticAnalysis(obj, tf)
            obj.EnableStaticAnalysis = tf;
        end
        function setForceBackwardPassCodegenInTracingMode(obj, tf)
            obj.ForceBackwardPassCodegenInTracingMode = tf;
        end
        function setEnableCPPCodeGeneration(obj, tf)
            obj.EnableCPPCodeGeneration = tf;
        end
        function setNoGradUnlessTracingMode(obj, tf)
            obj.NoGradUnlessTracingMode = tf;
        end
        function setMaxGeneratedFunctionLength(obj,value)
            obj.MaxGeneratedFunctionLength = value;
        end
    end

    methods (Static, Hidden)
        function obj = createStandaloneInstance(fun)
            obj = deep.AcceleratedFunction(fun);
        end

        function rval = loadobj(obj)
            % Ask the manager for a AcceleratedFunction with the same handle.
            manager = deep.internal.AcceleratedFunctionManager.getInstance();
            rval = manager.getAcceleratedFunction(obj.Function);

            % Set the state to the object being loaded.
            rval = copyState(obj, rval);

        end
    end

    methods (Access = private)

        function copy = copyState(obj, copy)
            % Copy state from one object to another.
            assert(isa(copy,'deep.AcceleratedFunction'));

            % To maintain backwards compatibility, each case in the switch
            % encodes conversion from the version number in the case
            % to the current version
            switch obj.Version
                case 1
                    % Conversion from Version 1 to Current Version
                    copy = copyStateFromVersion1(obj,copy);

                otherwise
                    % Add cases to handle conversions from older versions
            end
        end

        function copy = copyStateFromVersion1(obj, copy)
            % Helper to copy state information from objects of Version 1
            assert(obj.Version == 1);

            copy.Enabled = obj.Enabled;
            copy.CacheSize = obj.CacheSize;
            copy.CheckOutputs = obj.CheckOutputs;
            copy.CheckTolerance = obj.CheckTolerance;
        end

        function initializeCache(obj)
            % Initializes an empty cache.
            obj.Cache.Signatures = {};
            obj.Cache.Functions = {};
            obj.Cache.GeneratedCode = {};
            obj.Cache.Nargout = zeros(2,0);
            obj.Cache.Outputs = {};
            obj.Cache.TapeFunction = {};

            obj.Cache.HitCount = [];
            obj.Cache.FlagForRetracing = [];
            obj.Cache.Load = 0;
            obj.Cache.TotalHits = 0;
            obj.Cache.TotalMisses = 0;
        end

        function deleteCacheEntry(obj, idx)
            obj.Cache.Signatures(idx) = [];
            obj.Cache.Functions(idx) = [];
            obj.Cache.GeneratedCode(idx) = [];
            obj.Cache.Nargout(:,idx) = [];
            obj.Cache.Outputs(idx) = [];
            obj.Cache.TapeFunction(idx) = [];
            obj.Cache.HitCount(idx) = [];
            obj.Cache.FlagForRetracing(idx) = [];
            obj.Cache.Load = obj.Cache.Load - 1;
        end

        function varargout = execute(obj, inputs)
            % Execute without using cache
            [varargout{1:nargout}] = obj.Function(inputs{:});
        end

        function varargout = transformAndExecuteWithFallback(obj, inputs)
            % Execute without using cache along with an initial static
            % analysis. Any errors during static analysis will fallback to
            % the non-static function evaluation.
            try
                [varargout{1:nargout}] = deep.internal.recording.astvisitor.transformAndExecute(obj.Function,inputs{:});
            catch ex 
                [varargout{1:nargout}] = obj.Function(inputs{:});
            end
        end

        function [inputs, inputNodes, restore, priorTapeCount] = setupTrace(~, inputs, tm, isTracing, needsForward)
            traceCount = needsForward + 1;
            traceRestore = iStartTracing(tm, traceCount);
            [inputs, inputNodes, priorTapeCount] = iTracingInputs(inputs, isTracing, tm);
            if isempty(inputNodes)
                traceRestore = []; % Stop tracing if there is nothing to trace.
            end
            % Mark tape as accelerating
            tapeDataRestore = deep.internal.recording.TapeData.setTapeDataAndSetupCleanup( ...
                IsTracingForAcceleration=true);
            restore = {tapeDataRestore, traceRestore};
        end

        function [cacheData, outputs, illegalOutputs] = generateForward(obj, outputs, inputNodes, tm, priorTapeCount, wasTracing)
            [outputs, cacheOutput, outputNodes, illegalOutputs] = iTracingOutputs(outputs, wasTracing);
            if isempty(inputNodes) || isempty(outputNodes) || illegalOutputs
                fun = [];
                generatedCode = [];
                tapeFunction = [];
            else
                tapeStartForOptimization = priorTapeCount+1;
                if obj.EnableOptimization
                    iOptimizeTape(tm,inputNodes,outputNodes,tapeStartForOptimization);
                end
                inputIDs = iGetIDs(inputNodes);
                outputIDs = iGetIDs(outputNodes);
                % Get the tape after iTracingOutputs so that we capture any internal permutation
                tape = deep.internal.recording.tapeStruct(tm, priorTapeCount);
                allIndependentOutputs = deep.internal.recording.allIndependent(tape, outputIDs);
                traceBackward = wasTracing && (obj.ForceBackwardPassCodegenInTracingMode || allIndependentOutputs);
                if traceBackward
                    [intermediateIDs, backwardMask] = deep.internal.recording.getIntermediatesForBackward(tape, outputIDs, inputIDs);
                    outputIDs = [outputIDs intermediateIDs];
                    outputNodes = [outputNodes iGetNodes(tm,intermediateIDs)];
                    numIntermediates = length(intermediateIDs);
                end
                [fun, fileName] = iTapeToFunction(tape,inputIDs,outputIDs,[],obj.EnableCPPCodeGeneration,obj.EnableOptimization,obj.MaxGeneratedFunctionLength);
                generatedCode = deep.internal.GeneratedCode(fun, fileName);
                tapeFunction = iTapeFunction(tape,inputIDs,outputIDs);

                if traceBackward
                    % Generate backward functions and modify fun to record an AcceleratedOp.
                    args.InputIDs = inputIDs;
                    args.InputNodes = inputNodes;
                    args.OutputIDs = outputIDs;
                    args.OutputNodes = outputNodes;
                    args.TapeManager = tm;
                    args.ForwardTapeStart = priorTapeCount;
                    args.TapeStartForOptimization = tapeStartForOptimization;
                    args.BackwardMask = backwardMask;
                    args.AllIndependentOutputs = allIndependentOutputs;
                    args.IsNonDifferentiableOutput = iMarkOutputsThatDeriveOtherOutputs(tape,outputIDs,inputIDs,allIndependentOutputs);
                    [fun, backwardTapeFunctions] = augmentWithBackwardFunctions(obj, args, numIntermediates, generatedCode);
                    tapeFunction = [{tapeFunction} backwardTapeFunctions];
                end
            end
            cacheData.Function = fun;
            cacheData.GeneratedCode = generatedCode;
            cacheData.Output = cacheOutput;
            cacheData.TapeFunction = tapeFunction;
            cacheData.NumTraceOutputs = length(outputNodes);
        end

        function [fun, backwardTapeFunctions] = augmentWithBackwardFunctions(obj, args, numIntermediates, generatedCode)        
            restore = deep.internal.recording.setupNoGradMode(); %#ok<NASGU>
            numOutputs = length(args.OutputIDs);
            adjointSizes = cell(1,numOutputs);
            for i = 1:numOutputs
                adjointSizes{i} = size(args.OutputNodes{i});
            end
            numGradients = min(length(args.InputNodes), obj.MaxNumGradients);
            args.NumGradients = numGradients;
            args.AdjointSizes = adjointSizes;
            args.EnableOptimization = obj.EnableOptimization;
            args.Filename = [];
            args.EnableCPPCodeGeneration = obj.EnableCPPCodeGeneration;
            args.EnableOptimization = obj.EnableOptimization;
            args.MaxGeneratedFunctionLength = obj.MaxGeneratedFunctionLength;
            enableFullBackward = obj.EnableFullBackward;
            try
                [fastBackwardFun, fastBackwardFileName, fastBackwardTapeFunction] = iGenerateBackwardFunction(args, numIntermediates);
            catch err
                enableFullBackward = false;
                fastBackwardFun = @(varargin)iThrow(err);
                fastBackwardFileName = [];
                fastBackwardTapeFunction = [];
            end
            fullBackwardFileName = [];
            fullBackwardFun = [];
            fullBackwardTapeFunction = [];
            merr = [];
            
            if enableFullBackward
                try
                    [fullBackwardFun, fullBackwardFileName, fullBackwardTapeFunction] = iGenerateBackwardFunction(args, 0);
                catch merr
                    if ~iIsGPUOutOfMemory(merr)
                        rethrow(merr)
                    end
                end
            end
            args.OutputNodes = iMarkNonDifferentiableNodes(args.OutputNodes,args.IsNonDifferentiableOutput);
            appendBackwardFcns(generatedCode,fullBackwardFun,fastBackwardFun, ...
                               [string(fullBackwardFileName), string(fastBackwardFileName)]);
            args.GeneratedCode = generatedCode;
            backwardTapeFunctions = {fullBackwardTapeFunction, fastBackwardTapeFunction};
            args.Memory = args.BackwardMask;
            args.IntermediateAdjointInputRange = [numOutputs-numIntermediates+1  numOutputs];
            args.FullBackwardOutOfGPUMemoryError = merr;
            op = deep.internal.AcceleratedOp(args);
            % Assign field value to a variable and then pass it to the
            % anonymous function to avoid potential memory leak (g2842070)
            nonDiffArgs = args.IsNonDifferentiableOutput;
            fun = @(varargin)iApplyOp(op, nonDiffArgs, numOutputs, varargin{:});
        end

        function prioritize(obj, i, load)
            j = ceil((i + load)/2); % try moving half way to head
            hitcount = obj.Cache.HitCount;
            if hitcount(i) < hitcount(j)
                j = i + 1;  % fall back to moving by one
            end
            if hitcount(i) > hitcount(j)
                % swap i and j entries
                lhs = [i j];
                rhs = [j i];
                obj.Cache.Signatures(lhs) = obj.Cache.Signatures(rhs);
                obj.Cache.Functions(lhs) = obj.Cache.Functions(rhs);
                obj.Cache.Nargout(:,lhs) = obj.Cache.Nargout(:,rhs);
                obj.Cache.Outputs(lhs) = obj.Cache.Outputs(rhs);
                obj.Cache.GeneratedCode(lhs) = obj.Cache.GeneratedCode(rhs);
                obj.Cache.HitCount(lhs) = obj.Cache.HitCount(rhs);
                obj.Cache.FlagForRetracing(lhs) = obj.Cache.FlagForRetracing(rhs);
            end
        end

        function updateCacheSize(obj, cacheSize)
            % Updates cache size
            if cacheSize < obj.Cache.Load
                % If cacheSize is smaller, truncate cache by dropping
                % early entries
                N = length(obj.Cache.Signatures) - cacheSize + 1;
                obj.Cache.Signatures = obj.Cache.Signatures(N:end);
                obj.Cache.Functions = obj.Cache.Functions(N:end);
                obj.Cache.Nargout = obj.Cache.Nargout(:,N:end);
                obj.Cache.Outputs = obj.Cache.Outputs(N:end);
                obj.Cache.GeneratedCode = obj.Cache.GeneratedCode(N:end);
                obj.Cache.HitCount = obj.Cache.HitCount(N:end);
                obj.Cache.FlagForRetracing = obj.Cache.FlagForRetracing(N:end);
                obj.Cache.Load = cacheSize;
            end

            obj.PrivateCacheSize = cacheSize;
        end

        function optimizeTapeRestore = setupToOptimizeTape(obj)
            optimizeTapeRestore = [];
            if obj.EnableOptimization
                optimizeTapeRestore = iSetupToOptimizeTape();
            end
        end

        function nogradModeRestore = setupNoGradModeIfNeeded(obj, isTracing, needsForward)
            nogradModeRestore = [];
            % If this AcceleratedFunction was called from a tracing
            % context, we will not use the no-grad mode.
            if isTracing
                return;
            end
            if obj.NoGradUnlessTracingMode && ~iIsNoGradMode()
                % We want to trace in no-grad mode and are not currently
                % using no-grad mode. So turn no-grad mode on.
                nogradModeRestore = deep.internal.recording.setupNoGradMode(true);
                % In this situation since it is legitimate to call
                % dlgradient during a trace, turn on needs forward mode
                deep.internal.recording.tape.setTapeNeedsForward(needsForward);
            elseif ~obj.NoGradUnlessTracingMode && iIsNoGradMode()
                % We want to trace in normal mode and are currently using
                % no-grad mode. So turn no-grad mode off.
                nogradModeRestore = deep.internal.recording.setupNoGradMode(false);
            end
        end

    end
end

function restore = iStartTracing(tm,order)
restore = cell(1,order);
for i = 1:order
    restore{i} = deep.internal.recording.TapeGuard;
    startTracing(tm);
end
end

function [signature, parameters] = iEvaluationInputs(x)
% Split the evaluation inputs X into a SIGNATURE and a cell array
% PARAMETERS of extracted dlarray data.
visitor = deep.internal.accelerate.ExtractSignatureAndParametersVisitor;
signature = deep.internal.container.accept(x, visitor);
parameters = visitor.Parameters;
end

function [x, nodes, priorTapeCount] = iTracingInputs(x, isTracing, tm)
% Enables tracing on dlarray values in x and returns RecordingNodes in
% the order expected by iEvaluationInputs.
visitor = deep.internal.accelerate.StartTracingAndExtractNodesVisitor;
x = deep.internal.container.accept(x, visitor);
if isTracing && ~iHaveUniqueIDs(visitor.Nodes)
    uniqueNodesVisitor = deep.internal.accelerate.EnsureUniqueNodesVisitor;
    x = deep.internal.container.accept(x, uniqueNodesVisitor);
    nodes = uniqueNodesVisitor.Nodes;
else
    nodes = visitor.Nodes;
end
priorTapeCount = getTapeSize(tm);

% In case custom layers have setters on their learnable properties, sync
% the value with what we are recording.
visitor = deep.internal.accelerate.SetNodesVisitor(nodes);
x = deep.internal.container.accept(x, visitor);
end

function [y, cacheOutput, nodes, illegalOutputs] = iTracingOutputs(y, wasTracing)
% Loop over arbitrary outputs y and extract information.
% Outputs:
%   y            y with recording off (if needed)
%   cacheOutput  y with recording dlarrays replaced with format information
%   nodes        flattened vector of RecordingArray nodes
%   illegalOutputs  true if y contained any unsupported containers
visitor = deep.internal.accelerate.ExtractOutputCacheRepresentationAndNodesVisitor;
cacheOutput = deep.internal.container.accept(y, visitor);
nodes = visitor.Nodes;
illegalOutputs = visitor.IllegalOutputs;
if ~wasTracing
    visitor = deep.internal.accelerate.StopTracingVisitor;
    y = deep.internal.container.accept(y, visitor);
end
end

function y = iEvaluationOutputs(y, newValues)
% Put back new values in place of the FormattedData objects.
visitor = deep.internal.accelerate.InsertComputedValuesVisitor(newValues);
y = deep.internal.container.accept(y, visitor);
end

function varargout = iApplyOp(op, isNonDifferentiableOutput, numOutputs, varargin)
[varargout{1:numOutputs}] = recordNary(varargin{:}, numOutputs, op);
if any(isNonDifferentiableOutput)
    varargout = iMarkNonDifferentiableNodes(varargout,isNonDifferentiableOutput);
end
end

function ids = iGetIDs(nodes)
ids = zeros(1, length(nodes), 'uint64');
for i = 1:length(nodes)
    ids(i) = getNodeID(nodes{i});
end
end

function nodes = iGetNodes(tm, ids)
nodes = cell(1, length(ids));
for i = 1:length(ids)
    nodes{i} = getRecordingArrayForID(tm, ids(i));
end
end

function tf = iIsGPUOutOfMemory(merr)
tf = "parallel:gpu:array:OOM" == merr.identifier;
end

function [backwardFun, backwardFileName, backwardTapeFunction] = iGenerateBackwardFunction(args,numIntermediateAdjointsToDrop)
tm = args.TapeManager;
rewindTapeLocation = getTapeSize(tm);
tmp = onCleanup(@()removeLastTapeEntries(tm, getTapeSize(tm) - rewindTapeLocation));
[backwardFun, backwardFileName, backwardTapeFunction] = iGenerateBackwardFunctionNoCleanup(args,numIntermediateAdjointsToDrop);
end

function [backwardFun, backwardFileName, backwardTapeFunction] = iGenerateBackwardFunctionNoCleanup(args,numIntermediateAdjointsToDrop)
tm = args.TapeManager;
numAdjoints = length(args.OutputIDs) - numIntermediateAdjointsToDrop;
adjointNodes = cell(1,numAdjoints);
adjointIDs = zeros(1,numAdjoints,'uint64');
for i = 1:numAdjoints
    value = gather(args.OutputNodes{i}.Value);
    if isnumeric(value)
        adjointValue = ones(args.AdjointSizes{i},'like',value);
    elseif islogical(value)
        adjointValue = ones(args.AdjointSizes{i},'single');
    else
        adjointValue = [];
    end
    adjointNodes{i} = deep.internal.recording.RecordingArray(adjointValue);
    adjointIDs(i) = getNodeID(adjointNodes{i});
end
backwardTapeStart = getTapeSize(tm);
grads = backwardTape(tm, args.OutputNodes(1:numAdjoints), adjointNodes, args.InputNodes(1:args.NumGradients), false, false, args.ForwardTapeStart);
gradIDs = zeros(1,length(grads),'uint64');
for i = 1:length(grads)
    if isempty(grads{i})
        grads{i} = zeros(size(args.InputNodes{i}), 'like', adjointNodes{1});
    end
    gradIDs(i) = getNodeID(grads{i});
end
if args.EnableOptimization
    iOptimizeTape(tm,{},{},args.TapeStartForOptimization);
end
tape = deep.internal.recording.tapeStruct(tm, backwardTapeStart);
inputsAndOutputForBackward = [args.InputIDs args.OutputIDs];
inputsAndOutputForBackward = inputsAndOutputForBackward(args.BackwardMask);
backwardInputIDs = [adjointIDs inputsAndOutputForBackward];
[backwardFun, backwardFileName] = iTapeToFunction(tape, backwardInputIDs, gradIDs, args.Filename, args.EnableCPPCodeGeneration, args.EnableOptimization, args.MaxGeneratedFunctionLength);
backwardTapeFunction = iTapeFunction(tape, backwardInputIDs, gradIDs);
end

function restore = iSetupToOptimizeTape()
restore = deep.internal.recording.TapeData.setTapeDataAndSetupCleanup(OptimizeForwardWithAcceleratedFunction=true);
end

function tf = iRetraceForInference()
data = deep.internal.recording.tape.getTapeData();
tf = ~isempty(data) && data.RetraceForInference;
end

function iOptimizeTape(tm,x,y,tapeStart)
passes = deep.internal.recording.OptimizationPass.passesForAcceleratedFunction;
deep.internal.recording.optimizeTape(tm,x,y,passes,tapeStart);
end

function nodes = iMarkNonDifferentiableNodes(nodes,isNonDifferentiable)
numNodes = numel(nodes);
for i = 1:numNodes
    if isNonDifferentiable(i)
        setDifferentiable(nodes{i},false);
    end
end
end

function tf = iMarkOutputsThatDeriveOtherOutputs(tape,outputIDs,inputIDs,allIndependentOutputs)
% Returns a logical vector of the same length as numel(outputIDs). Element
% i of this vector is true if outputIDs(i) is not a member of inputIDs and
% outputIDs(i) derives another element of outputIDs. Such outputIDs(i) are
% non-differentiable as outputs from the AcceleratedOp.
numOutputs = numel(outputIDs);
tf = false(numOutputs,1);
if ~allIndependentOutputs
    for i = 1:numOutputs
        tf(i) = ~ismember(outputIDs(i),inputIDs) && ...
            deep.internal.recording.anyOutputDerivedFromInput(tape,outputIDs,outputIDs(i));
    end
end
end

function [fun, fileName] = iTapeToFunction(tape,inputIDs,outputIDs,fileName,useCPPTapeToFunction,doOptimize,maxGeneratedFunctionLength)
[fun, fileName] = deep.internal.recording.convert.tapeToFunctionWithFallback(tape,inputIDs,outputIDs,fileName,useCPPTapeToFunction,doOptimize,maxGeneratedFunctionLength);
end

function tapeFunction = iTapeFunction(tape,inputIDs,outputIDs)
tapeFunction = deep.internal.recording.TapeFunction(tape,inputIDs,outputIDs);
tapeFunction = iRemoveInputValuesFromTapeFunction(tapeFunction);
end

function tapeFunction = iRemoveInputValuesFromTapeFunction(tapeFunction)
inputSet = dictionary(tapeFunction.InputIDs,0);
entries = tapeFunction.Tape.Data;
tapeFunction.Tape.Data = [];
for i = 1:numel(entries)
    isTapeFunctionInput = isKey(inputSet,entries(i).Inputs);
    entries(i).InputValues(isTapeFunctionInput) = {[]};
end
tapeFunction.Tape.Data = entries;
end

function tf = iIsThreadPoolWorker()
tf = parallel.internal.pool.isPoolThreadWorker;
end

function tf = iIsNoGradMode()
tf = deep.internal.recording.tape.isNoGradMode();
end

function tf = iHaveUniqueIDs(nodes)
numNodes = numel(nodes);
IDs = zeros(1,numNodes,'uint64');
for i = 1:numNodes
    IDs(i) = nodes{i}.NodeID;
end
uniqueIDs = unique(IDs);
tf = numel(uniqueIDs) == numNodes;
end

function varargout = iThrow(err)
throw(err);
end

function iProtectedMethodCall(fcn, obj)
% Call method only if handle object is valid
if isvalid(obj)
    fcn(obj);
end
end