Resource tracking can now handle statepreps (#2230)

**Context:**
The `quantum.set_state` and `quantum.set_basis_state` instructions are
not currently tracked via runtime resource tracking.

**Description of the Change:**
* Adds 2 new functions, `SetState` and `SetBasisState` to the
ResourceTracker class, allowing these instructions to be tracked.
* Reports uses of these functions as `StatePrep` and `BasisState`
operators respectively

**Benefits:**
Device-level `qml.specs` for qjit'd circuits now reports the number of
initial `StatePrep` operators.

**Possible Drawbacks:**

**Related GitHub Issues:**
PennyLaneAI/pennylane#8650
[sc-104048]

---------

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

Author: jzaia18

doc/releases/changelog-dev.md

@@ -69,6 +69,10 @@
 
 <h3>Improvements 🛠</h3>
 
+* Resource tracking now tracks calls to `SetState` and `SetBasisState`, and can report results
+  that include `qml.StatePrep` operations.
+  [(#2230)](https://github.com/PennyLaneAI/catalyst/pull/2230)
+  
 * Remove the hardcoded list of runtime operations in the frontend.
   This will allow arbitrary PL gates to be represented without hyperparameters in MLIR.
   For gates that do not have a QIR representation, a runtime error will be raised at execution.

runtime/lib/backend/null_qubit/NullQubit.hpp

@@ -224,7 +224,12 @@ struct NullQubit final : public Catalyst::Runtime::QuantumDevice {
      * @param state The state vector data (ignored)
      * @param wires The qubits to prepare (ignored)
      */
-    void SetState(DataView<std::complex<double>, 1> &, std::vector<QubitIdType> &) {}
+    void SetState(DataView<std::complex<double>, 1> &, std::vector<QubitIdType> &wires)
+    {
+        if (this->track_resources_) {
+            this->resource_tracker_.SetState(wires);
+        }
+    }
 
     /**
      * @brief No-op implementation for computational basis state preparation
@@ -235,7 +240,12 @@ struct NullQubit final : public Catalyst::Runtime::QuantumDevice {
      * @param basis_state The computational basis state (ignored)
      * @param wires The qubits to prepare (ignored)
      */
-    void SetBasisState(DataView<int8_t, 1> &, std::vector<QubitIdType> &) {}
+    void SetBasisState(DataView<int8_t, 1> &, std::vector<QubitIdType> &wires)
+    {
+        if (this->track_resources_) {
+            this->resource_tracker_.SetBasisState(wires);
+        }
+    }
 
     /**
      * @brief No-op implementation for a named quantum operation

runtime/lib/backend/null_qubit/ResourceTracker.hpp

@@ -318,6 +318,26 @@ struct ResourceTracker final {
         RecordOperation(op_name, wires, controlled_wires);
     }
 
+    /**
+     * @brief Records a state preparation operation for resource tracking
+     *
+     * @param wires The target wires the operation acts upon
+     */
+    void SetState(const std::vector<QubitIdType> &wires)
+    {
+        RecordOperation("StatePrep", wires, {});
+    }
+
+    /**
+     * @brief Records a basis state preparation operation for resource tracking
+     *
+     * @param wires The target wires the operation acts upon
+     */
+    void SetBasisState(const std::vector<QubitIdType> &wires)
+    {
+        RecordOperation("BasisState", wires, {});
+    }
+
     /**
      * @brief Prints resource usage statistics in JSON format to the specified file
      *

runtime/tests/Test_NullQubit.cpp

@@ -924,6 +924,20 @@ TEST_CASE("Test NullQubit device resource tracking integration", "[NullQubit]")
 
     std::vector<QubitIdType> Qs = sim->AllocateQubits(4);
 
+    // Apply set state operations
+    {
+        std::vector<std::complex<double>> data = {{0.5, 0.5}, {0.0, 0.0}};
+        DataView<std::complex<double>, 1> data_view(data);
+        std::vector<QubitIdType> wires = {1};
+        sim->SetState(data_view, wires);
+    }
+    {
+        std::vector<int8_t> data = {0};
+        DataView<int8_t, 1> data_view(data);
+        std::vector<QubitIdType> wires = {0};
+        sim->SetBasisState(data_view, wires);
+    }
+
     // Apply named gates to test all possible name modifiers
     sim->NamedOperation("PauliX", {}, {Qs[0]}, false);
     sim->NamedOperation("T", {}, {Qs[0]}, true);
@@ -953,16 +967,20 @@ TEST_CASE("Test NullQubit device resource tracking integration", "[NullQubit]")
     resource_file_r.open(RESOURCES_FILENAME);
     CHECK(resource_file_r.is_open()); // fail-fast if file failed to create
 
-    std::vector<std::string> resource_names = {"PauliX",
-                                               "C(Adjoint(T))",
-                                               "Adjoint(T)",
-                                               "C(S)",
-                                               "2C(S)",
-                                               "CNOT",
-                                               "Adjoint(ControlledQubitUnitary)",
-                                               "ControlledQubitUnitary",
-                                               "Adjoint(QubitUnitary)",
-                                               "QubitUnitary"};
+    std::vector<std::string> resource_names = {
+        "PauliX",
+        "C(Adjoint(T))",
+        "Adjoint(T)",
+        "C(S)",
+        "2C(S)",
+        "CNOT",
+        "Adjoint(ControlledQubitUnitary)",
+        "ControlledQubitUnitary",
+        "Adjoint(QubitUnitary)",
+        "QubitUnitary",
+        "StatePrep",
+        "BasisState",
+    };
 
     // Read full Json, check if num_wires and num_gates are correct
     std::string full_json;
@@ -973,10 +991,10 @@ TEST_CASE("Test NullQubit device resource tracking integration", "[NullQubit]")
             CHECK(line.find("4") != std::string::npos);
         }
         if (line.find("num_gates") != std::string::npos) {
-            CHECK(line.find("10") != std::string::npos);
+            CHECK(line.find("12") != std::string::npos);
         }
         if (line.find("depth") != std::string::npos) {
-            CHECK(line.find("10") != std::string::npos);
+            CHECK(line.find("11") != std::string::npos);
         }
         full_json += line + "\n";
     }

runtime/tests/Test_ResourceTracker.cpp

@@ -427,3 +427,25 @@ TEST_CASE("Test Resource Tracker WriteOut", "[resourcetracking]")
         CHECK(full_json.find(name) != std::string::npos);
     }
 }
+
+TEST_CASE("Test Resource Tracker SetState Operations", "[resourcetracking]")
+{
+    ResourceTracker tracker;
+    tracker.SetComputeDepth(true);
+    for (size_t i = 0; i < 5; i++) {
+        tracker.AllocateQubit(i);
+    }
+    CHECK(tracker.GetNumGates() == 0);
+    CHECK(tracker.GetNumGates("StatePrep") == 0);
+    CHECK(tracker.GetNumGates("BasisState") == 0);
+
+    tracker.SetState({0});
+    tracker.SetBasisState({0, 1, 2});
+    CHECK(tracker.GetNumGates() == 2);
+    CHECK(tracker.GetNumGates("StatePrep") == 1);
+    CHECK(tracker.GetNumGates("BasisState") == 1);
+
+    CHECK(tracker.GetNumGatesBySize(1) == 1);
+    CHECK(tracker.GetNumGatesBySize(3) == 1);
+    CHECK(tracker.GetDepth() == 2);
+}