Added a tracker util so that we can track the allocations of QIR type objs (ebeee843) · Commits · ORNL Quantum Computing Institute / qcor

mlir/qir_qrt/qir-types.cpp

0 → 100644

+15 −0

Original line number	Diff line number	Diff line
		#include "qir-types.hpp"

		namespace qcor {
		namespace internal {

		AllocationTracker *AllocationTracker::m_globalTracker = nullptr;

		AllocationTracker &AllocationTracker::get() {
		if (!m_globalTracker) {
		m_globalTracker = new AllocationTracker();
		}
		return *m_globalTracker;
		}
		} // namespace internal
		} // namespace qcor
		No newline at end of file

mlir/qir_qrt/qir-types.hpp

+63 −13

Original line number	Diff line number	Diff line
		#pragma once

		#include <atomic>
		#include <cassert>
		#include <memory>
		#include <stdexcept>
		#include <atomic>
		#include <unordered_map>
		#include <vector>

		// Defines implementations of QIR Opaque types

		namespace qcor {
		namespace internal {
		// Internal tracker to make sure we're doing proper
		// ref. counting, e.g. when generating QIR LLVM IR.
		class AllocationTracker {
		public:
		static AllocationTracker &get();

		void onAllocate(void *objPtr) {
		assert(m_refCountMap.find(objPtr) == m_refCountMap.end());
		m_refCountMap[objPtr] = 1;
		}

		void updateCount(void *objPtr, int newCount) {
		assert(m_refCountMap.find(objPtr) != m_refCountMap.end());
		assert(newCount >= 0);
		m_refCountMap[objPtr] = newCount;
		}

		// Check if we have any leakage.
		// Returns false if no leak, true otherwise.
		// Can be use at shut-down (Finalize) to detect leakage.
		bool checkLeak() const {
		for (const auto &[ptr, count] : m_refCountMap) {
		if (count > 0) {
		return true;
		}
		}
		// No leak, all objects have been released.
		return false;
		}

		private:
		AllocationTracker(){};
		static AllocationTracker *m_globalTracker;
		std::unordered_map<void *, int> m_refCountMap;
		};
		} // namespace internal
		} // namespace qcor

		// FIXME - Qubit should be a struct that keeps track of idx
		// qreg name, array it comes from, and associated accelerator buffer
		// Make this a struct now so that we can upgrade the API later
		@@ -42,15 +86,17 @@ struct Array {
		Array(int64_t nbItems, int itemSizeInBytes = sizeof(int8_t *))
		: m_itemSizeInBytes(itemSizeInBytes),
		// Initialized to zero
		m_storage(nbItems * itemSizeInBytes, 0),
		m_refCount(1) {
		m_storage(nbItems * itemSizeInBytes, 0), m_refCount(1) {
		assert(m_itemSizeInBytes > 0);
		qcor::internal::AllocationTracker::get().onAllocate(this);
		};
		// Copy ctor:
		// note: we copy the Storage vector, hence set ref count to 1
		Array(const Array &other)
		: m_itemSizeInBytes(other.m_itemSizeInBytes), m_storage(other.m_storage),
		m_refCount(1) {}
		m_refCount(1) {
		qcor::internal::AllocationTracker::get().onAllocate(this);
		}

		void append(const Array &other) {
		if (other.m_itemSizeInBytes != m_itemSizeInBytes) {
		@@ -66,12 +112,16 @@ struct Array {
		int64_t element_size() const { return m_itemSizeInBytes; }

		// Ref. counting:
		void add_ref() { m_refCount += 1; }
		void add_ref() {
		m_refCount += 1;
		qcor::internal::AllocationTracker::get().updateCount(this, m_refCount);
		}
		// Release a single ref.
		// Returns true if this Array should be deleted
		// if heap allocated (via new)
		bool release_ref() {
		m_refCount -= 1;
		qcor::internal::AllocationTracker::get().updateCount(this, m_refCount);
		return (m_refCount == 0);
		}

mlir/qir_qrt/quantum_rt_array.cpp

+2 −2

Original line number	Diff line number	Diff line
		@@ -55,7 +55,7 @@ void __quantum__rt__array_update_alias_count(Array *array, int64_t increment) {
		// Looks like alias count has no functional significance, hence ignored.
		}

		void __quantum__rt__array_update_reference_count(Array *aux, int64_t count) {
		void __quantum__rt__array_update_reference_count(Array *array, int64_t increment) {
		// Spec:
		// Deallocates the array if the reference count becomes 0.
		// The behavior is undefined if the reference count becomes negative.
		@@ -63,7 +63,7 @@ void __quantum__rt__array_update_reference_count(Array *aux, int64_t count) {
		if (verbose)
		std::cout << "CALL: " << __PRETTY_FUNCTION__ << "\n";

		if (!aux) {
		if (!array) {
		// The call should be ignored if the given %Array* is a null pointer.
		return;
		}