Caching in QDK's resource estimator

CachedEstimator.cs

namespace CodeCachedEstimates;

using System;
using System.Collections.Generic;
using System.Linq;
using Microsoft.Quantum.Simulation.Core;
using Microsoft.Quantum.Simulation.Simulators;
using Microsoft.Quantum.Simulation.Simulators.QCTraceSimulators;
using Microsoft.Quantum.Simulation.QCTraceSimulatorRuntime;

public class CachedEstimator : ResourcesEstimator {
    // Keeps track of operation call stack
    private Stack<string> operationStack = new();
    // Tracks cached CNOT counts for cached operations
    private Dictionary<string, double> cachedResources = new();
    // Keeps track of operations on the call stack that needs caching
    private Stack<string> needsCaching = new();
    // Reference to gate count results
    private Microsoft.Quantum.Simulation.QCTraceSimulatorRuntime.StatisticsCollector<Microsoft.Quantum.Simulation.QCTraceSimulatorRuntime.CallGraphEdge> operationCounts = null;

    // We overload the recommended config to have access to the call stack. This
    // number might need to be increased depending on the depth of operation
    // calls.
    public static new QCTraceSimulatorConfiguration RecommendedConfig() {
        var config = ResourcesEstimator.RecommendedConfig();
        config.CallStackDepthLimit = 32;
        return config;
    }

    public CachedEstimator() : base(RecommendedConfig()) {
        // Access primitive operations counter from resource estimator
        foreach (var listener in tCoreConfig.Listeners) {
            if (listener is PrimitiveOperationsCounter) {
                operationCounts = (listener as PrimitiveOperationsCounter).Results as Microsoft.Quantum.Simulation.QCTraceSimulatorRuntime.StatisticsCollector<Microsoft.Quantum.Simulation.QCTraceSimulatorRuntime.CallGraphEdge>;
                break;
            }
        }

        // When an operation start, add most recent operation name to operation stack
        OnOperationStart += (callable, _) => {
            operationStack.Push(callable.FullName);
        };

        // When an operation ends, update the call stack and check whether the
        // operation needs to be cached
        OnOperationEnd += (callable, _) => {
            operationStack.Pop();

            if (needsCaching.Count > 0 && needsCaching.Peek() == callable.FullName) {
                needsCaching.Pop();
                // If operation needs caching, obtain CNOT count from top of the call stack (i.e., last row in the operation counts results table)
                var operationCNOTCount = operationCounts.GetStatistic(operationCounts.Data.Last().Key, PrimitiveOperationsGroupsNames.CNOT, StatisticsNames.Sum);
                cachedResources.Add(callable.FullName, operationCNOTCount);
            }
        };
    }

    // Updates resource counts if operation is cached with cached resource
    // counts, otherwise adds operation to stack for caching
    public bool IsCached() {
        var top = operationStack.Peek();

        double cachedCNOTCount;
        if (cachedResources.TryGetValue(top, out cachedCNOTCount)) {
            var listener = this.tCoreConfig.Listeners[0] as Microsoft.Quantum.Simulation.QCTraceSimulatorRuntime.PrimitiveOperationsCounter;
            for (var i = 0; i < cachedCNOTCount; ++i) {
                listener.OnPrimitiveOperation((int)PrimitiveOperationsGroups.CNOT, null, 0);
            }
            return true;
        } else {
            needsCaching.Push(top);
            return false;
        }
    }

    // Update implementation for `IsCached` for `CachedEstimator` to
    // `CachedEstimator.IsCached` function
    public class IsCachedImpl : CodeCachedEstimates.IsCached {
        private CachedEstimator estimator = null;

        public IsCachedImpl(CachedEstimator m) : base(m) {
            estimator = m;
        }

        public override Func<QVoid, bool> __Body__ => _ => {
            return estimator.IsCached();
        };
    }
}

CachedEstimator.csproj

<Project Sdk="Microsoft.Quantum.Sdk/0.24.210930">

  <PropertyGroup>
    <OutputType>Exe</OutputType>
    <TargetFramework>net6.0</TargetFramework>
  </PropertyGroup>

</Project>

Program.cs

using CodeCachedEstimates;

// you can pick any of the two estimators to see differences in runtime
var simulator = new CachedEstimator();
// var simulator = new Microsoft.Quantum.Simulation.Simulators.ResourcesEstimator();

RunProgram.Run(simulator).Wait();

System.Console.WriteLine(simulator.ToTSV());

Program.qs

namespace CodeCachedEstimates {

    open Microsoft.Quantum.Canon;
    open Microsoft.Quantum.Intrinsic;
    open Microsoft.Quantum.Random;

    // This operation returns false by default, but we can modify the
    // functionality for the CachedEstimator
    function IsCached() : Bool {
        return false;
    }

    // This operation should be cached
    operation Circuit(a : Qubit, b : Qubit) : Unit {
        // We add this if-statement to the operation
        if not IsCached() {
            // this should simulate some heavy computation
            for _ in 1..100000 {
                let _ = DrawRandomDouble(-1.0, 1.0);
            }

            // apply gates
            CNOT(a, b);
            CNOT(b, a);
            CNOT(a, b);
        }
    }

    operation RunProgram() : Unit {
        use a = Qubit();
        use b = Qubit();

        // Call the operation many times, only the first one should be executed,
        // all other calls use cached values
        for _ in 1..1000 {
            Circuit(a, b);
        }
    }
}