Add resource tracking to NullQubit

.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-dev19
+pennylane=0.42.0-dev22
 
 # For a custom LQ/LK version, update the package version here and at
 # 'doc/requirements.txt'

doc/releases/changelog-dev.md

@@ -124,6 +124,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)
@@ -186,4 +189,5 @@ Tzung-Han Juang,
 Christina Lee,
 Erick Ochoa Lopez,
 Mehrdad Malekmohammadi,
-Paul Haochen Wang.
+Paul Haochen Wang,
+Jake Zaia.

doc/requirements.txt

@@ -32,4 +32,4 @@ lxml_html_clean
 --extra-index-url https://test.pypi.org/simple/
 pennylane-lightning-kokkos==0.42.0-dev11
 pennylane-lightning==0.42.0-dev11
-pennylane==0.42.0-dev19
+pennylane==0.42.0-dev22

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::cerr << "Error opening file: " << resources_fname.str();
+                // TODO: Should this throw?
+            }
+            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);
+}