atomspace-gpu / opencog / persist / gpu / GpuStorageNode.cc

/*
 * opencog/persist/gpu/GpuStorageNode.cc
 *
 * BackingStore wrapper for GpuAtomTable. Bridges the AtomSpace
 * persistence API to the GPU atom table.
 *
 * Step 3 of Linas' plan: storeAtom/getAtom via BackingStore.h.
 *
 * Copyright (C) 2026 OpenCog Foundation
 * SPDX-License-Identifier: AGPL-3.0-or-later
 */

#include <algorithm>
#include <set>

#include <opencog/atoms/base/Link.h>
#include <opencog/atoms/base/Node.h>
#include <opencog/atoms/atom_types/atom_types.h>
#include <opencog/atoms/value/UnisetValue.h>
#include <opencog/util/exceptions.h>
#include <opencog/persist/gpu/GpuStorageNode.h>

using namespace opencog;

GpuStorageNode::GpuStorageNode(Type t, const std::string&& uri)
	: StorageNode(t, std::move(uri)), _is_open(false)
{
	memset(&_gpu_table, 0, sizeof(_gpu_table));
}

GpuStorageNode::GpuStorageNode(const std::string& uri)
	: StorageNode(GPU_STORAGE_NODE, uri), _is_open(false)
{
	memset(&_gpu_table, 0, sizeof(_gpu_table));
}

GpuStorageNode::~GpuStorageNode()
{
	if (_is_open) close();
}

// ================================================================
// Lifecycle

void GpuStorageNode::open(void)
{
	if (_is_open) return;

	int rc = gpu_table_alloc(&_gpu_table);
	if (rc != 0)
		throw IOException(TRACE_INFO, "GPU: failed to allocate atom table");

	_is_open = true;
}

void GpuStorageNode::close(void)
{
	if (!_is_open) return;

	gpu_table_barrier(&_gpu_table);
	gpu_table_free(&_gpu_table);

	_tlb.clear();
	_handle_to_slot.clear();
	_slot_is_node.clear();

	_is_open = false;
}

bool GpuStorageNode::connected(void)
{
	return _is_open;
}

void GpuStorageNode::destroy(void)
{
	if (!_is_open) return;
	gpu_table_clear(&_gpu_table);
	_gpu_table.atom_count = 0;
	_gpu_table.name_pool_used = 0;
	_gpu_table.out_pool_used = 0;
	_tlb.clear();
	_handle_to_slot.clear();
	_slot_is_node.clear();
}

void GpuStorageNode::erase(void)
{
	destroy();
}

// ================================================================
// Slot management
//
// Each atom gets a sequential GPU slot (0, 1, 2, ...).
// The TLB maps slot ↔ Handle (UUID = slot).
// _handle_to_slot provides fast forward lookup.

bool GpuStorageNode::has_slot(const Handle& h) const
{
	return _handle_to_slot.count(h) > 0;
}

uint32_t GpuStorageNode::assign_slot(const Handle& h)
{
	auto it = _handle_to_slot.find(h);
	if (it != _handle_to_slot.end())
		return it->second;

	uint32_t slot = _gpu_table.atom_count;
	_gpu_table.atom_count++;

	_handle_to_slot[h] = slot;
	_tlb.addAtom(h, (UUID)slot);
	_slot_is_node.push_back(h->is_node());

	return slot;
}

// ================================================================
// Store

void GpuStorageNode::do_store_atom(const Handle& h)
{
	// Already on GPU — nothing to update in Step 3 (no Values).
	if (has_slot(h)) return;

	// Links: recursively store outgoing atoms first.
	if (h->is_link())
	{
		for (const Handle& out : h->getOutgoingSet())
			do_store_atom(out);
	}

	uint32_t slot = assign_slot(h);
	int rc;

	if (h->is_node())
	{
		NodePtr np = NodeCast(h);
		const std::string& name = np->get_name();
		rc = gpu_store_node(&_gpu_table, slot,
			h->get_type(), name.c_str(), (uint16_t)name.size());
	}
	else
	{
		const HandleSeq& outgoing = h->getOutgoingSet();
		std::vector<uint32_t> out_slots;
		out_slots.reserve(outgoing.size());
		for (const Handle& out : outgoing)
			out_slots.push_back(_handle_to_slot.at(out));

		rc = gpu_store_link(&_gpu_table, slot,
			h->get_type(), out_slots.data(), (uint16_t)outgoing.size());
	}

	if (rc != 0)
		throw IOException(TRACE_INFO,
			"GPU: failed to store atom at slot %u", slot);

	// New atom invalidates the incoming-set index.
	_gpu_table.in_built = 0;
}

void GpuStorageNode::storeAtom(const Handle& h, bool synchronous)
{
	std::lock_guard<std::mutex> lk(_mtx);
	do_store_atom(h);
	if (synchronous)
		gpu_table_barrier(&_gpu_table);
}

// ================================================================
// Load
//
// Reconstruct atoms from GPU data. Two-pass: nodes first (no
// dependencies), then links (outgoing references resolved via
// the slot→Handle vector built during pass 1).

void GpuStorageNode::loadAtomSpace(AtomSpace* as)
{
	if (!_is_open) return;

	gpu_table_barrier(&_gpu_table);

	uint32_t n = _gpu_table.atom_count;
	std::vector<Handle> loaded(n);

	// Pass 1: nodes
	for (uint32_t slot = 0; slot < n; slot++)
	{
		if (!_slot_is_node[slot]) continue;

		uint16_t type;
		char name_buf[4096];
		uint16_t name_len = sizeof(name_buf);

		int rc = gpu_fetch_node(&_gpu_table, slot,
			&type, name_buf, &name_len);
		if (rc != 0) continue;

		std::string name(name_buf, name_len);
		loaded[slot] = as->add_node(type, std::move(name));
	}

	// Pass 2: links (outgoing slots are always < current slot
	// because do_store_atom stores outgoing first).
	for (uint32_t slot = 0; slot < n; slot++)
	{
		if (_slot_is_node[slot]) continue;

		uint16_t type;
		uint32_t out_buf[256];
		uint16_t arity = sizeof(out_buf) / sizeof(out_buf[0]);

		int rc = gpu_fetch_link(&_gpu_table, slot,
			&type, out_buf, &arity);
		if (rc != 0) continue;

		HandleSeq outgoing;
		outgoing.reserve(arity);
		for (uint16_t i = 0; i < arity; i++)
			outgoing.push_back(loaded[out_buf[i]]);

		loaded[slot] = as->add_link(type, std::move(outgoing));
	}
}

void GpuStorageNode::storeAtomSpace(const AtomSpace* as)
{
	HandleSeq all;
	as->get_handles_by_type(all, ATOM, true);

	std::lock_guard<std::mutex> lk(_mtx);
	for (const Handle& h : all)
		do_store_atom(h);
}

// ================================================================
// Fetch (Step 3: no Values stored yet)

void GpuStorageNode::getAtom(const Handle&)
{
	// Step 3: no Values on GPU. Nothing to load.
}

// Build the incoming-set CSR index on the GPU if it is stale.
void GpuStorageNode::ensure_incoming()
{
	if (!_gpu_table.in_built)
		gpu_table_build_incoming(&_gpu_table);
}

// Reconstruct one GPU link-slot as a Link in `as`. Returns the Handle
// (or undefined if any outgoing atom is unknown). type_filter == NOTYPE
// means "any type".
Handle GpuStorageNode::fetch_link_to(AtomSpace* as, uint32_t link_slot,
                                     Type type_filter)
{
	uint16_t ltype = 0;
	uint32_t out_buf[256];
	uint16_t arity = sizeof(out_buf) / sizeof(out_buf[0]);
	if (gpu_fetch_link(&_gpu_table, link_slot, &ltype, out_buf, &arity) != 0)
		return Handle::UNDEFINED;
	if (NOTYPE != type_filter && (Type)ltype != type_filter)
		return Handle::UNDEFINED;

	HandleSeq oset;
	oset.reserve(arity);
	for (uint16_t i = 0; i < arity; i++)
	{
		Handle oh = _tlb.getAtom((UUID)out_buf[i]);
		if (nullptr == oh) return Handle::UNDEFINED;
		oset.push_back(oh);
	}
	return as->add_link((Type)ltype, std::move(oset));
}

// Fetch the incoming set of `h` (optionally filtered by type) from the GPU
// and install the links into `as`.
void GpuStorageNode::fetch_incoming_impl(AtomSpace* as, const Handle& h,
                                         Type type_filter)
{
	std::lock_guard<std::mutex> lk(_mtx);
	if (!_is_open) return;

	auto it = _handle_to_slot.find(h);
	if (it == _handle_to_slot.end()) return;   // atom not on GPU

	ensure_incoming();

	uint32_t slot = it->second;
	uint32_t count = 0;
	gpu_fetch_incoming(&_gpu_table, slot, nullptr, 0, &count);
	if (0 == count) return;

	std::vector<uint32_t> links(count);
	gpu_fetch_incoming(&_gpu_table, slot, links.data(), count, &count);

	for (uint32_t ls : links)
		fetch_link_to(as, ls, type_filter);
}

void GpuStorageNode::fetchIncomingSet(AtomSpace* as, const Handle& h)
{
	fetch_incoming_impl(as, h, NOTYPE);
}

void GpuStorageNode::fetchIncomingByType(AtomSpace* as, const Handle& h, Type t)
{
	fetch_incoming_impl(as, h, t);
}

void GpuStorageNode::removeAtom(AtomSpace*, const Handle&, bool) {}
void GpuStorageNode::loadType(AtomSpace*, Type) {}

// Values are not persisted on the GPU yet (Step 4). No-ops keep the default
// BackingStore::runQuery (which caches results via storeValue) from throwing.
void GpuStorageNode::storeValue(const Handle&, const Handle&) {}
void GpuStorageNode::loadValue(const Handle&, const Handle&) {}

// ================================================================
// GPU-resident query (the runQuery override) — GENERAL multi-clause.
//
// Handles  (Meet <vardecl> body)  where there is exactly ONE variable and
// body is a single clause or an AndLink of clauses, each clause one of:
//   (a) a direct binary link   (T $X c) / (T c $X)        [e.g. InheritanceLink]
//   (b) an Eval relation  (Evaluation (Predicate P) (List $X c)) / (List c $X)
// Each clause's candidate var-bindings are computed via the GPU incoming
// index; the answer is their intersection. Anything else (multi-variable
// joins, nesting, etc.) returns false → default BackingStore::runQuery.

// Single query variable, or UNDEFINED if not exactly one.
static Handle gpu_single_var(const Handle& vardecl)
{
	Type t = vardecl->get_type();
	if (VARIABLE_NODE == t) return vardecl;
	if (TYPED_VARIABLE_LINK == t) return vardecl->getOutgoingAtom(0);
	if (VARIABLE_LIST == t or VARIABLE_SET == t)
	{
		if (1 != vardecl->get_arity()) return Handle::UNDEFINED;
		return gpu_single_var(vardecl->getOutgoingAtom(0));
	}
	return Handle::UNDEFINED;
}

bool GpuStorageNode::gpu_clause_candidates(const Handle& clause, const Handle& var,
                                           std::set<uint32_t>& out)
{
	Type ct = clause->get_type();

	// --- shape (b): Eval relation (Predicate, List) ---
	if (EVALUATION_LINK == ct and 2 == clause->get_arity() and
	    LIST_LINK == clause->getOutgoingAtom(1)->get_type())
	{
		const Handle& pred = clause->getOutgoingAtom(0);
		const Handle& lst  = clause->getOutgoingAtom(1);
		if (2 != lst->get_arity()) return false;
		int varpos = (lst->getOutgoingAtom(0) == var) ? 0
		           : (lst->getOutgoingAtom(1) == var) ? 1 : -1;
		if (varpos < 0) return false;
		const Handle& cst = lst->getOutgoingAtom(1 - varpos);
		if (not cst->is_node() or not pred->is_node()) return false;
		auto pit = _handle_to_slot.find(pred); if (pit == _handle_to_slot.end()) return false;
		auto cit = _handle_to_slot.find(cst);  if (cit == _handle_to_slot.end()) return false;
		uint32_t pslot = pit->second, cslot = cit->second, cstpos = 1 - varpos;

		uint32_t n = 0;
		gpu_fetch_incoming(&_gpu_table, cslot, nullptr, 0, &n);
		if (0 == n) return true;
		std::vector<uint32_t> cin(n);
		gpu_fetch_incoming(&_gpu_table, cslot, cin.data(), n, &n);
		for (uint32_t lslot : cin)            // candidate ListLink slots
		{
			uint16_t lt; uint32_t lob[8]; uint16_t lar = 8;
			if (gpu_fetch_link(&_gpu_table, lslot, &lt, lob, &lar) != 0) continue;
			if ((Type)lt != LIST_LINK or lar != 2) continue;
			if (lob[cstpos] != cslot) continue;       // const in the right list position
			uint32_t binding = lob[varpos];
			// require an EvaluationLink(pred, thisList)
			uint32_t m = 0;
			gpu_fetch_incoming(&_gpu_table, lslot, nullptr, 0, &m);
			if (0 == m) continue;
			std::vector<uint32_t> lin(m);
			gpu_fetch_incoming(&_gpu_table, lslot, lin.data(), m, &m);
			for (uint32_t eslot : lin)
			{
				uint16_t et; uint32_t eob[8]; uint16_t ear = 8;
				if (gpu_fetch_link(&_gpu_table, eslot, &et, eob, &ear) != 0) continue;
				if ((Type)et == EVALUATION_LINK and ear == 2 and
				    eob[0] == pslot and eob[1] == lslot)
				{ out.insert(binding); break; }
			}
		}
		return true;
	}

	// --- shape (a): direct binary link (T $X c) / (T c $X) ---
	if (2 == clause->get_arity())
	{
		const Handle& o0 = clause->getOutgoingAtom(0);
		const Handle& o1 = clause->getOutgoingAtom(1);
		int varpos = (o0 == var) ? 0 : (o1 == var) ? 1 : -1;
		if (varpos < 0) return false;
		const Handle& cst = (0 == varpos) ? o1 : o0;
		if (not cst->is_node()) return false;
		auto it = _handle_to_slot.find(cst); if (it == _handle_to_slot.end()) return false;
		uint32_t cslot = it->second, cstpos = 1 - varpos;

		uint32_t n = 0;
		gpu_fetch_incoming(&_gpu_table, cslot, nullptr, 0, &n);
		if (0 == n) return true;
		std::vector<uint32_t> cin(n);
		gpu_fetch_incoming(&_gpu_table, cslot, cin.data(), n, &n);
		for (uint32_t ls : cin)
		{
			uint16_t lt; uint32_t ob[8]; uint16_t ar = 8;
			if (gpu_fetch_link(&_gpu_table, ls, &lt, ob, &ar) != 0) continue;
			if ((Type)lt != ct or ar != 2) continue;
			if (ob[cstpos] != cslot) continue;
			out.insert(ob[varpos]);
		}
		return true;
	}

	return false;
}

bool GpuStorageNode::gpu_run_meet(const Handle& query, const Handle& key)
{
	if (query->get_type() != MEET_LINK || query->get_arity() != 2) return false;
	Handle var = gpu_single_var(query->getOutgoingAtom(0));
	if (nullptr == var || var->get_type() != VARIABLE_NODE) return false;  // single-var only

	const Handle& body = query->getOutgoingAtom(1);
	HandleSeq clauses;
	if (AND_LINK == body->get_type()) clauses = body->getOutgoingSet();
	else clauses.push_back(body);
	if (clauses.empty()) return false;

	AtomSpace* as = query->getAtomSpace();
	if (nullptr == as) return false;

	std::lock_guard<std::mutex> lk(_mtx);
	ensure_incoming();

	std::set<uint32_t> result;
	bool first = true;
	for (const Handle& c : clauses)
	{
		std::set<uint32_t> cand;
		if (not gpu_clause_candidates(c, var, cand)) return false;  // unsupported → fallback
		if (first) { result = std::move(cand); first = false; }
		else
		{
			std::set<uint32_t> inter;
			std::set_intersection(result.begin(), result.end(),
			                      cand.begin(), cand.end(),
			                      std::inserter(inter, inter.begin()));
			result = std::move(inter);
		}
		if (result.empty()) break;
	}

	UnisetValuePtr svp(createUnisetValue());
	for (uint32_t s : result)
	{
		Handle h = _tlb.getAtom((UUID)s);
		if (h) svp->add(as->add_atom(h));
	}
	svp->close();
	query->setValue(key, svp);
	return true;
}

void GpuStorageNode::runQuery(const Handle& query, const Handle& key,
                             const Handle& meta, bool fresh)
{
	if (gpu_run_meet(query, key)) return;             // GPU-resident fast-path
	BackingStore::runQuery(query, key, meta, fresh);  // default: CPU crawl, GPU-served incoming
}

// ================================================================
// Synchronization

void GpuStorageNode::barrier(AtomSpace*)
{
	if (_is_open)
		gpu_table_barrier(&_gpu_table);
}

// ================================================================
// Diagnostics

std::string GpuStorageNode::monitor(void)
{
	std::string s;
	s += "GPU Storage Monitor\n";
	s += "Connected: ";
	s += (_is_open ? "yes" : "no");
	s += "\n";
	if (_is_open)
	{
		s += "Backend: ";
		switch (_gpu_table.backend) {
			case GPU_BACKEND_CUDA:   s += "CUDA\n"; break;
			case GPU_BACKEND_OPENCL: s += "OpenCL\n"; break;
			default:                 s += "None\n"; break;
		}
		s += "Atoms: " + std::to_string(_gpu_table.atom_count) + "\n";
		s += "Name pool: " + std::to_string(_gpu_table.name_pool_used) + " bytes\n";
		s += "Out pool: " + std::to_string(_gpu_table.out_pool_used) + " slots\n";
	}
	return s;
}

// ================================================================
// Factory

Handle GpuStorageNode::factory(const Handle& base)
{
	Handle h(createGpuStorageNode(base->get_name()));
	return h;
}