Add resource tracking to NullQubit (#1619)

**Context:**
For benchmarking, we need a method of precise resource tracking. This
method has to be compatible with code run directly from PennyLane and
should be able to track precise resources after various transformation
passes have already occurred.

See also: PennyLaneAI/pennylane-benchmarks#76
and PennyLaneAI/pennylane#7226

**Description of the Change:**
Adds 2 new class attributes to `NullQubit`, as well as an optional
argument to its constructor. When this option is toggled on, `NullQubit`
will track precise gate counts while "executing" a circuit and save them
to a file when all qubits are released.

**Benefits:**
It is now possible to track precise gate counts (with or without gate
decomposition or other transformation passes) for measuring circuit
complexity.

**Possible Drawbacks:**
This implementation currently write the resources JSON to a file on disk
as opposed to neatly returning an object back through to the Python
caller in PennyLane. In the future, a more robust solution may be
desired.

**Related GitHub Issues:**
[sc-91366]

---------

Co-authored-by: Ali Asadi <[email protected]>
Co-authored-by: David Ittah <[email protected]>

Author: 3 people

.dep-versions

@@ -8,7 +8,7 @@ enzyme=v0.0.149
 
 # For a custom PL version, update the package version here and at
 # 'doc/requirements.txt
-pennylane=0.42.0-dev27
+pennylane=0.42.0-dev29
 
 # For a custom LQ/LK version, update the package version here and at
 # 'doc/requirements.txt'

.gitignore

@@ -9,6 +9,7 @@ pennylane_catalyst.egg-info
 # Testing files and directories
 .lit
 .pytest_cache
+__pennylane_resources_data_*.json
 
 # Packaging directories
 dist

doc/releases/changelog-dev.md

@@ -154,6 +154,9 @@
 
 <h3>Internal changes ⚙️</h3>
 
+* `null.qubit` can now support an optional `track_resources` argument which allows it to record which gates are executed.
+  [(#1619)](https://github.com/PennyLaneAI/catalyst/pull/1619)
+
 * Add an xDSL MLIR plugin to denote whether we will be using xDSL to execute some passes.
   This changelog entry may be moved to new features once all branches are merged together.
   [(#1707)](https://github.com/PennyLaneAI/catalyst/pull/1707)
@@ -239,4 +242,5 @@ Mehrdad Malekmohammadi,
 Anton Naim Ibrahim,
 Erick Ochoa Lopez,
 Ritu Thombre,
-Paul Haochen Wang.
+Paul Haochen Wang,
+Jake Zaia.

doc/requirements.txt

@@ -33,4 +33,4 @@ lxml_html_clean
 --extra-index-url https://test.pypi.org/simple/
 pennylane-lightning-kokkos==0.42.0-dev16
 pennylane-lightning==0.42.0-dev16
-pennylane==0.42.0-dev27
+pennylane==0.42.0-dev29

frontend/test/pytest/test_device_api.py

@@ -137,5 +137,16 @@ def circuit():
     assert circuit.mlir
 
 
+def test_track_resources():
+    """Test that resource tracking settings get passed to the device."""
+    dev = NullQubit(wires=2)
+    assert "track_resources" in QJITDevice.extract_backend_info(dev, dev.capabilities).kwargs
+    assert QJITDevice.extract_backend_info(dev, dev.capabilities).kwargs["track_resources"] is False
+
+    dev = NullQubit(wires=2, track_resources=True)
+    assert "track_resources" in QJITDevice.extract_backend_info(dev, dev.capabilities).kwargs
+    assert QJITDevice.extract_backend_info(dev, dev.capabilities).kwargs["track_resources"] is True
+
+
 if __name__ == "__main__":
     pytest.main(["-x", __file__])

runtime/lib/backend/common/Utils.hpp

@@ -126,6 +126,25 @@ static inline auto parse_kwargs(std::string kwargs) -> std::unordered_map<std::s
     return map;
 }
 
+template <class K, class V>
+void pretty_print_dict(const std::unordered_map<K, V> &map, size_t leadingSpaces = 0,
+                       std::ostream &out = std::cout)
+{
+    const std::string indent(leadingSpaces, ' ');
+    const std::string innerIndent = indent + "  ";
+
+    out << indent << "{\n";
+    auto it = map.begin();
+    while (it != map.end()) {
+        out << innerIndent << "\"" << it->first << "\": " << it->second;
+        if (++it != map.end()) {
+            out << ",";
+        }
+        out << "\n";
+    }
+    out << indent << "}";
+}
+
 enum class MeasurementsT : uint8_t {
     None, // = 0
     Expval,

runtime/lib/backend/null_qubit/NullQubit.hpp

@@ -15,16 +15,21 @@
 #pragma once
 
 #include <algorithm> // generate_n
+#include <chrono>
 #include <complex>
+#include <cstdio>
+#include <fstream>
 #include <memory>
 #include <optional>
 #include <random>
+#include <unordered_map>
 #include <vector>
 
 #include "DataView.hpp"
 #include "QuantumDevice.hpp"
 #include "QubitManager.hpp"
 #include "Types.h"
+#include "Utils.hpp"
 
 namespace Catalyst::Runtime::Devices {
 
@@ -40,14 +45,48 @@ namespace Catalyst::Runtime::Devices {
  *   of the device; these are used to implement Quantum Instruction Set (QIS) instructions.
  */
 struct NullQubit final : public Catalyst::Runtime::QuantumDevice {
-    NullQubit(const std::string &kwargs = "{}") {}
-    ~NullQubit() = default; // LCOV_EXCL_LINE
+    NullQubit(const std::string &kwargs = "{}")
+    {
+        auto device_kwargs = Catalyst::Runtime::parse_kwargs(kwargs);
+        if (device_kwargs.find("track_resources") != device_kwargs.end()) {
+            track_resources_ = device_kwargs["track_resources"] == "True";
+        }
+    }
+    ~NullQubit() {} // LCOV_EXCL_LINE
 
     NullQubit &operator=(const NullQubit &) = delete;
     NullQubit(const NullQubit &) = delete;
     NullQubit(NullQubit &&) = delete;
     NullQubit &operator=(NullQubit &&) = delete;
 
+    /**
+     * @brief Prints resources that would be used to execute this circuit as a JSON
+     */
+    void PrintResourceUsage(FILE *resources_file)
+    {
+        // Store the 2 special variables and clear them from the map to make
+        // pretty-printing easier
+        const size_t num_qubits = resource_data_["num_qubits"];
+        const size_t num_gates = resource_data_["num_gates"];
+        resource_data_.erase("num_gates");
+        resource_data_.erase("num_qubits");
+
+        std::stringstream resources;
+
+        resources << "{\n";
+        resources << "  \"num_qubits\": " << num_qubits << ",\n";
+        resources << "  \"num_gates\": " << num_gates << ",\n";
+        resources << "  \"gate_types\": ";
+        pretty_print_dict(resource_data_, 2, resources);
+        resources << "\n}" << std::endl;
+
+        fwrite(resources.str().c_str(), 1, resources.str().size(), resources_file);
+
+        // Restore 2 special variables
+        resource_data_["num_qubits"] = num_qubits;
+        resource_data_["num_gates"] = num_gates;
+    }
+
     /**
      * @brief Allocate a "null" qubit.
      *
@@ -56,6 +95,10 @@ struct NullQubit final : public Catalyst::Runtime::QuantumDevice {
     auto AllocateQubit() -> QubitIdType
     {
         num_qubits_++; // next_id
+        if (this->track_resources_) {
+            // Store the highest number of qubits allocated at any time since device creation
+            resource_data_["num_qubits"] = std::max(num_qubits_, resource_data_["num_qubits"]);
+        }
         return this->qubit_manager.Allocate(num_qubits_);
     }
 
@@ -94,6 +137,27 @@ struct NullQubit final : public Catalyst::Runtime::QuantumDevice {
     {
         num_qubits_ = 0;
         this->qubit_manager.ReleaseAll();
+        if (this->track_resources_) {
+            auto time = std::chrono::high_resolution_clock::now();
+            auto timestamp =
+                std::chrono::duration_cast<std::chrono::nanoseconds>(time.time_since_epoch())
+                    .count();
+            std::stringstream resources_fname;
+            resources_fname << "__pennylane_resources_data_" << timestamp << ".json";
+
+            // Need to use FILE* instead of ofstream since ofstream has no way to atomically open a
+            // file only if it does not already exist
+            FILE *resources_file = fopen(resources_fname.str().c_str(), "wx");
+            if (resources_file == nullptr) {
+                std::string err_msg = "Error opening file '" + resources_fname.str() + "'.";
+                RT_FAIL(err_msg.c_str());
+            }
+            else {
+                PrintResourceUsage(resources_file);
+                fclose(resources_file);
+            }
+            this->resource_data_.clear();
+        }
     }
 
     /**
@@ -170,17 +234,47 @@ struct NullQubit final : public Catalyst::Runtime::QuantumDevice {
                         const std::vector<QubitIdType> &controlled_wires = {},
                         const std::vector<bool> &controlled_values = {})
     {
+        if (this->track_resources_) {
+            std::string prefix = "";
+            std::string suffix = "";
+            if (!controlled_wires.empty()) {
+                if (controlled_wires.size() > 1) {
+                    prefix += std::to_string(controlled_wires.size());
+                }
+                prefix += "C(";
+                suffix += ")";
+            }
+            if (inverse) {
+                prefix += "Adj(";
+                suffix += ")";
+            }
+            resource_data_["num_gates"]++;
+            resource_data_[prefix + name + suffix]++;
+        }
     }
 
     /**
      * @brief Doesn't apply a given matrix directly to the state vector of a device.
      *
      */
     void MatrixOperation(const std::vector<std::complex<double>> &,
-                         const std::vector<QubitIdType> &, bool,
+                         const std::vector<QubitIdType> &, bool inverse,
                          const std::vector<QubitIdType> &controlled_wires = {},
                          const std::vector<bool> &controlled_values = {})
     {
+        if (this->track_resources_) {
+            resource_data_["num_gates"]++;
+
+            std::string op_name = "QubitUnitary";
+
+            if (!controlled_wires.empty()) {
+                op_name = "Controlled" + op_name;
+            }
+            if (inverse) {
+                op_name = "Adj(" + op_name + ")";
+            }
+            resource_data_[op_name]++;
+        }
     }
 
     /**
@@ -360,9 +454,28 @@ struct NullQubit final : public Catalyst::Runtime::QuantumDevice {
         return {0, 0, 0, {}, {}};
     }
 
+    /**
+     * @brief Returns the number of gates used since the last time all qubits were released. Only
+     * works if resource tracking is enabled
+     */
+    auto ResourcesGetNumGates() -> std::size_t { return resource_data_["num_gates"]; }
+
+    /**
+     * @brief Returns the maximum number of qubits used since the last time all qubits were
+     * released. Only works if resource tracking is enabled
+     */
+    auto ResourcesGetNumQubits() -> std::size_t { return resource_data_["num_qubits"]; }
+
+    /**
+     * @brief Returns whether the device is tracking resources or not.
+     */
+    auto IsTrackingResources() const -> bool { return track_resources_; }
+
   private:
+    bool track_resources_{false};
     std::size_t num_qubits_{0};
     std::size_t device_shots_{0};
+    std::unordered_map<std::string, std::size_t> resource_data_;
     Catalyst::Runtime::QubitManager<QubitIdType, size_t> qubit_manager{};
 
     // static constants for RESULT values

runtime/tests/Test_NullQubit.cpp

@@ -16,6 +16,7 @@
 #include <catch2/catch_approx.hpp>
 #include <catch2/catch_test_macros.hpp>
 #include <catch2/matchers/catch_matchers_string.hpp>
+#include <cstdio>
 
 #include "ExecutionContext.hpp"
 #include "QuantumDevice.hpp"
@@ -598,3 +599,104 @@ TEST_CASE("Test NullQubit device shots methods", "[NullQubit]")
         CHECK(sim->GetDeviceShots() == i);
     }
 }
+
+TEST_CASE("Test NullQubit device resource tracking", "[NullQubit]")
+{
+    // The name of the file where the resource usage data is stored
+    constexpr char RESOURCES_FNAME[] = "__pennylane_resources_data.json";
+
+    // Open a file for writing the resources JSON
+    FILE *resource_file_w = fopen(RESOURCES_FNAME, "wx");
+    if (resource_file_w == nullptr) {                            // LCOV_EXCL_LINE
+        FAIL("Failed to open resource usage file for writing."); // LCOV_EXCL_LINE
+    }
+
+    std::unique_ptr<NullQubit> dummy = std::make_unique<NullQubit>();
+    CHECK(dummy->IsTrackingResources() == false);
+
+    std::unique_ptr<NullQubit> sim = std::make_unique<NullQubit>("{'track_resources':True}");
+    CHECK(sim->IsTrackingResources() == true);
+    CHECK(sim->ResourcesGetNumGates() == 0);
+    CHECK(sim->ResourcesGetNumQubits() == 0);
+
+    std::vector<QubitIdType> Qs = sim->AllocateQubits(4);
+
+    CHECK(sim->ResourcesGetNumGates() == 0);
+    CHECK(sim->ResourcesGetNumQubits() == 4);
+
+    // Apply named gates to test all possible name modifiers
+    sim->NamedOperation("PauliX", {}, {Qs[0]}, false);
+    sim->NamedOperation("T", {}, {Qs[0]}, true);
+    sim->NamedOperation("S", {}, {Qs[0]}, false, {Qs[2]});
+    sim->NamedOperation("S", {}, {Qs[0]}, false, {Qs[1], Qs[2]});
+    sim->NamedOperation("T", {}, {Qs[0]}, true, {Qs[2]});
+    sim->NamedOperation("CNOT", {}, {Qs[0], Qs[1]}, false);
+
+    CHECK(sim->ResourcesGetNumGates() == 6);
+    CHECK(sim->ResourcesGetNumQubits() == 4);
+
+    // Applying an empty matrix is fine for NullQubit
+    sim->MatrixOperation({}, {Qs[0]}, false);
+    sim->MatrixOperation({}, {Qs[0]}, false, {Qs[1]});
+    sim->MatrixOperation({}, {Qs[0]}, true);
+    sim->MatrixOperation({}, {Qs[0]}, true, {Qs[1], Qs[2]});
+
+    CHECK(sim->ResourcesGetNumGates() == 10);
+    CHECK(sim->ResourcesGetNumQubits() == 4);
+
+    // Capture resources usage
+    sim->PrintResourceUsage(resource_file_w);
+    fclose(resource_file_w);
+
+    // Open the file of resource data
+    std::ifstream resource_file_r(RESOURCES_FNAME);
+    CHECK(resource_file_r.is_open()); // fail-fast if file failed to create
+
+    std::vector<std::string> resource_names = {"PauliX",
+                                               "C(Adj(T))",
+                                               "Adj(T)",
+                                               "C(S)",
+                                               "2C(S)",
+                                               "S",
+                                               "CNOT",
+                                               "Adj(ControlledQubitUnitary)",
+                                               "ControlledQubitUnitary",
+                                               "Adj(QubitUnitary)",
+                                               "QubitUnitary"};
+
+    // Check all fields have the correct value
+    std::string full_json;
+    while (resource_file_r) {
+        std::string line;
+        std::getline(resource_file_r, line);
+        if (line.find("num_qubits") != std::string::npos) {
+            CHECK(line.find("4") != std::string::npos);
+        }
+        if (line.find("num_gates") != std::string::npos) {
+            CHECK(line.find("10") != std::string::npos);
+        }
+        // If one of the resource names is in the line, check that there is precisely 1
+        for (const auto &name : resource_names) {
+            if (line.find(name) != std::string::npos) {
+                CHECK(line.find("1") != std::string::npos);
+                break;
+            }
+        }
+        full_json += line + "\n";
+    }
+    resource_file_r.close();
+    std::remove(RESOURCES_FNAME);
+
+    // Ensure all expected fields are present
+    CHECK(full_json.find("num_qubits") != std::string::npos);
+    CHECK(full_json.find("num_gates") != std::string::npos);
+    for (const auto &name : resource_names) {
+        CHECK(full_json.find(name) != std::string::npos);
+    }
+
+    // Check that releasing resets
+    sim->ReleaseAllQubits();
+
+    CHECK(sim->ResourcesGetNumGates() == 0);
+    CHECK(sim->ResourcesGetNumQubits() == 0);
+}