Skip to content

Add residual top-up extraction to PPR forward push #699

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Closed
mkolodner-sc wants to merge 1 commit into from
+193 −106
Closed

Add residual top-up extraction to PPR forward push #699

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
122 changes: 71 additions & 51 deletions gigl-core/core/sampling/ppr_forward_push.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -146,11 +146,11 @@ std::optional<std::unordered_map<int32_t, torch::Tensor>> PPRForwardPush::drainQ
return result;
}

void PPRForwardPush::pushResiduals(
const std::unordered_map<int32_t, std::tuple<torch::Tensor, torch::Tensor, torch::Tensor>>& fetchedByEtypeId) {
// Step 1: Unpack the input map into a C++ map keyed by packKey(nodeId, edgeTypeId)
// for fast lookup during the residual-push loop below.
std::unordered_map<uint64_t, std::vector<int32_t>> fetched;
void PPRForwardPush::pushResiduals(const NeighborFetchMap& fetchedByEtypeId) {
// Step 1: Persist fetched neighbor lists in the per-state cache. drainQueue()
// consults this cache before requesting future lookups, so storing every
// fetched row here avoids re-fetching a (node, edge type) pair if it re-enters
// the frontier later in the same PPR channel.
for (const auto& [edgeTypeId, neighborTensors] : fetchedByEtypeId) {
const auto& nodeIdsTensor = std::get<0>(neighborTensors);
const auto& flatNeighborIdsTensor = std::get<1>(neighborTensors);
Expand All @@ -172,7 +172,10 @@ void PPRForwardPush::pushResiduals(
for (int64_t neighborIdx = 0; neighborIdx < count; ++neighborIdx) {
neighborIds[neighborIdx] = static_cast<int32_t>(flatNeighborIdsAccessor[offset + neighborIdx]);
}
fetched[packKey(nodeId, edgeTypeId)] = std::move(neighborIds);
uint64_t cacheKey = packKey(nodeId, edgeTypeId);
if (_neighborCache.find(cacheKey) == _neighborCache.end()) {
_neighborCache.emplace(cacheKey, std::move(neighborIds));
}
offset += count;
}
}
Expand All @@ -197,21 +200,16 @@ void PPRForwardPush::pushResiduals(
srcNodeTypeState.pprScores[sourceNodeId] += sourceResidual;
srcNodeTypeState.residuals[sourceNodeId] = 0.0;

// b. Count total fetched/cached neighbors across all edge types for
// b. Count total cached neighbors across all edge types for
// this source node. We normalise by the number of neighbors we
// actually retrieved, not the true degree, so residual is fully
// distributed among known neighbors rather than leaking to unfetched
// ones (which matters when num_neighbors_per_hop < true_degree).
int32_t totalFetched = 0;
int32_t totalCachedNeighbors = 0;
for (int32_t edgeTypeId : _nodeTypeToEdgeTypeIds[nodeTypeId]) {
auto fetchedEntry = fetched.find(packKey(sourceNodeId, edgeTypeId));
if (fetchedEntry != fetched.end()) {
totalFetched += static_cast<int32_t>(fetchedEntry->second.size());
} else {
auto cachedEntry = _neighborCache.find(packKey(sourceNodeId, edgeTypeId));
if (cachedEntry != _neighborCache.end()) {
totalFetched += static_cast<int32_t>(cachedEntry->second.size());
}
auto cachedEntry = _neighborCache.find(packKey(sourceNodeId, edgeTypeId));
if (cachedEntry != _neighborCache.end()) {
totalCachedNeighbors += static_cast<int32_t>(cachedEntry->second.size());
}
}
// Two cases reach here:
Expand All @@ -221,32 +219,23 @@ void PPRForwardPush::pushResiduals(
// This overstates src and understates its neighbors. This is expected
// behavior when max_fetch_iterations is set, which intentionally trades
// theoretical PPR correctness for better throughput.
if (totalFetched == 0) {
if (totalCachedNeighbors == 0) {
continue;
}

double residualPerNeighbor = (1.0 - _alpha) * sourceResidual / static_cast<double>(totalFetched);
double residualPerNeighbor =
(1.0 - _alpha) * sourceResidual / static_cast<double>(totalCachedNeighbors);

for (int32_t edgeTypeId : _nodeTypeToEdgeTypeIds[nodeTypeId]) {
// Invariant: fetched and _neighborCache are mutually exclusive for
// any given (node, etype) key within one iteration. drainQueue()
// only requests a fetch for nodes absent from _neighborCache, so a
// key is in at most one of the two.
//
// Neighbor list for this (src, edgeTypeId) pair, borrowed from whichever
// map holds it. reference_wrapper is used because std::optional cannot
// hold a reference directly, and we want to avoid copying the vector —
// the data already exists in fetched or _neighborCache and both outlive
// this loop body. Access via neighborList->get().
// the data already exists in _neighborCache and outlives this loop body.
// Access via neighborList->get().
std::optional<std::reference_wrapper<const std::vector<int32_t>>> neighborList;
auto fetchedEntry = fetched.find(packKey(sourceNodeId, edgeTypeId));
if (fetchedEntry != fetched.end()) {
neighborList = std::cref(fetchedEntry->second);
} else {
auto cachedEntry = _neighborCache.find(packKey(sourceNodeId, edgeTypeId));
if (cachedEntry != _neighborCache.end()) {
neighborList = std::cref(cachedEntry->second);
}
auto cachedEntry = _neighborCache.find(packKey(sourceNodeId, edgeTypeId));
if (cachedEntry != _neighborCache.end()) {
neighborList = std::cref(cachedEntry->second);
}
if (!neighborList || neighborList->get().empty()) {
continue;
Expand All @@ -267,18 +256,6 @@ void PPRForwardPush::pushResiduals(
dstNodeTypeState.residuals[neighborNodeId] >= threshold) {
dstNodeTypeState.queue.insert(neighborNodeId);
++_numNodesInQueue;

// Promote neighbor lists to the persistent cache: this node will
// be processed next iteration, so caching avoids a re-fetch.
for (int32_t neighborEdgeTypeId : _nodeTypeToEdgeTypeIds[dstNodeTypeId]) {
uint64_t packedKey = packKey(neighborNodeId, neighborEdgeTypeId);
if (_neighborCache.find(packedKey) == _neighborCache.end()) {
auto fetchedNeighborEntry = fetched.find(packedKey);
if (fetchedNeighborEntry != fetched.end()) {
_neighborCache[packedKey] = fetchedNeighborEntry->second;
}
}
}
}
}
}
Expand All @@ -287,9 +264,15 @@ void PPRForwardPush::pushResiduals(
}
}

std::unordered_map<int32_t, std::tuple<torch::Tensor, torch::Tensor, torch::Tensor>> PPRForwardPush::extractTopK(
int32_t maxPprNodes) {
std::unordered_map<int32_t, std::tuple<torch::Tensor, torch::Tensor, torch::Tensor>> result;
PPRExtractResult PPRForwardPush::extractTopK(int32_t maxPprNodes) {
return extractTopKWithResidualTopUp(maxPprNodes, /*maxResidualNodes=*/0);
}

PPRExtractResult PPRForwardPush::extractTopKWithResidualTopUp(int32_t maxPprNodes, int32_t maxResidualNodes) {
TORCH_CHECK(maxPprNodes >= 0, "maxPprNodes must be non-negative, got ", maxPprNodes, ".");
TORCH_CHECK(maxResidualNodes >= 0, "maxResidualNodes must be non-negative, got ", maxResidualNodes, ".");

PPRExtractResult result;
// Emit an entry for every node type, even if unreachable in this batch (empty tensors,
// all-zero valid_counts). This keeps the output shape consistent across batches so
// downstream model architectures see a fixed set of PPR edge types every iteration.
Expand All @@ -299,21 +282,58 @@ std::unordered_map<int32_t, std::tuple<torch::Tensor, torch::Tensor, torch::Tens
std::vector<int64_t> validCounts;

for (int32_t seedIdx = 0; seedIdx < _batchSize; ++seedIdx) {
const auto& scores = _state[seedIdx][nodeTypeId].pprScores;
const auto& nodeTypeState = _state[seedIdx][nodeTypeId];
const auto& scores = nodeTypeState.pprScores;
int32_t topK = std::min(maxPprNodes, static_cast<int32_t>(scores.size()));
int32_t residualTopK = 0;
std::unordered_set<int32_t> selectedNodeIds;
if (topK > 0) {
selectedNodeIds.reserve(static_cast<size_t>(topK));
std::vector<std::pair<int32_t, double>> scorePairs(scores.begin(), scores.end());
std::partial_sort(scorePairs.begin(),
scorePairs.begin() + topK,
scorePairs.end(),
[](const auto& a, const auto& b) { return a.second > b.second; });

for (int32_t rankIdx = 0; rankIdx < topK; ++rankIdx) {
flatIds.push_back(static_cast<int64_t>(scorePairs[rankIdx].first));
int32_t nodeId = scorePairs[rankIdx].first;
flatIds.push_back(static_cast<int64_t>(nodeId));
flatWeights.push_back(scorePairs[rankIdx].second);
selectedNodeIds.insert(nodeId);
}
}

if (maxResidualNodes > 0) {
std::vector<std::pair<int32_t, double>> residualPairs;
residualPairs.reserve(nodeTypeState.residuals.size());
for (const auto& [nodeId, residual] : nodeTypeState.residuals) {
if (residual <= 0.0 || selectedNodeIds.find(nodeId) != selectedNodeIds.end()) {
continue;
}

auto scoreIter = scores.find(nodeId);
double pprScore = (scoreIter != scores.end()) ? scoreIter->second : 0.0;
double calibratedScore = pprScore + residual;
if (calibratedScore <= 0.0) {
continue;
}
residualPairs.emplace_back(nodeId, calibratedScore);
}

residualTopK = std::min(maxResidualNodes, static_cast<int32_t>(residualPairs.size()));
if (residualTopK > 0) {
std::partial_sort(residualPairs.begin(),
residualPairs.begin() + residualTopK,
residualPairs.end(),
[](const auto& a, const auto& b) { return a.second > b.second; });

for (int32_t rankIdx = 0; rankIdx < residualTopK; ++rankIdx) {
flatIds.push_back(static_cast<int64_t>(residualPairs[rankIdx].first));
flatWeights.push_back(residualPairs[rankIdx].second);
}
}
}
validCounts.push_back(static_cast<int64_t>(topK));
validCounts.push_back(static_cast<int64_t>(topK + residualTopK));
}

result[nodeTypeId] = {torch::tensor(flatIds, torch::kLong),
Expand Down
55 changes: 32 additions & 23 deletions gigl-core/core/sampling/ppr_forward_push.h
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -20,12 +20,16 @@ namespace gigl {
// not co-located in memory — only the control-plane metadata (size, bucket pointer)
// lives inside the struct.
struct SeedNodeTypeState {
std::unordered_map<int32_t, double> pprScores; // absorbed PPR mass
std::unordered_map<int32_t, double> residuals; // unabsorbed mass waiting to push
std::unordered_set<int32_t> queue; // nodes queued for the next drain
std::unordered_set<int32_t> queuedNodes; // snapshot captured by drainQueue()
std::unordered_map<int32_t, double> pprScores; // absorbed PPR mass
std::unordered_map<int32_t, double> residuals; // unabsorbed mass waiting to push
std::unordered_set<int32_t> queue; // nodes queued for the next drain
std::unordered_set<int32_t> queuedNodes; // snapshot captured by drainQueue()
};

using TensorTriplet = std::tuple<torch::Tensor, torch::Tensor, torch::Tensor>;
using NeighborFetchMap = std::unordered_map<int32_t, TensorTriplet>;
using PPRExtractResult = std::unordered_map<int32_t, TensorTriplet>;

// C++ kernel for PPR Forward Push (Andersen et al., 2006).
// Hot-loop state lives here; distributed neighbor fetches are driven from Python.
//
Expand All @@ -37,14 +41,14 @@ struct SeedNodeTypeState {
// 4. pushResiduals(fetchedByEtypeId)
// 5. extractTopK(maxPprNodes)
class PPRForwardPush {
public:
public:
PPRForwardPush(const torch::Tensor& seedNodes,
int32_t seedNodeTypeId,
double alpha,
double requeueThresholdFactor,
std::vector<std::vector<int32_t>> nodeTypeToEdgeTypeIds,
std::vector<int32_t> edgeTypeToDstNtypeId,
std::vector<torch::Tensor> degreeTensors);
int32_t seedNodeTypeId,
double alpha,
double requeueThresholdFactor,
std::vector<std::vector<int32_t>> nodeTypeToEdgeTypeIds,
std::vector<int32_t> edgeTypeToDstNtypeId,
std::vector<torch::Tensor> degreeTensors);

// Drain queued nodes and return {etype_id: int64 node tensor} for neighbor lookup.
// Returns nullopt when the queue is empty (convergence). Empty map means all nodes
Expand All @@ -53,30 +57,36 @@ class PPRForwardPush {

// Push residuals given fetched neighbor data.
// fetchedByEtypeId: {etype_id: (node_ids[N], flat_nbrs[sum(counts)], counts[N])}
void pushResiduals(const std::unordered_map<
int32_t, std::tuple<torch::Tensor, torch::Tensor, torch::Tensor>>&
fetchedByEtypeId);
void pushResiduals(const NeighborFetchMap& fetchedByEtypeId);

// Return top-k PPR nodes per seed per node type.
// Result: {ntype_id: (flat_ids, flat_weights, valid_counts)} — one entry per node type,
// including types unreachable in this batch (empty tensors, all-zero valid_counts).
std::unordered_map<int32_t, std::tuple<torch::Tensor, torch::Tensor, torch::Tensor>>
extractTopK(int32_t maxPprNodes);
PPRExtractResult extractTopK(int32_t maxPprNodes);

private:
// Return top-k PPR nodes, then append residual-mass top-up nodes.
//
// Residual top-up does not issue new neighbor fetches. It only reads the
// residual table already built by Forward Push. Scores are emitted on the
// same mass scale as PPR scores: ppr_score(node) + residual(node), i.e. the
// score the node would have if the remaining residual at that node were
// absorbed locally.
PPRExtractResult extractTopKWithResidualTopUp(int32_t maxPprNodes, int32_t maxResidualNodes);

private:
// Total out-degree of a node across all edge types. Returns 0 for sink nodes.
[[nodiscard]] int32_t getTotalDegree(int32_t nodeId, int32_t nodeTypeId) const;

double _alpha;
double _requeueThresholdFactor; // alpha * eps; per-node requeue threshold = factor * degree
double _requeueThresholdFactor; // alpha * eps; per-node requeue threshold = factor * degree

// NOTE: int32_t is used for batch size, node IDs, and type IDs throughout this class.
// All of this code will break silently (overflow) if batch size or node IDs exceed ~2B
// (INT32_MAX = 2,147,483,647). This is not a realistic concern today, but if graph
// scale ever approaches that threshold, these should be widened to int64_t.
int32_t _batchSize; // number of seed nodes in the current batch
int32_t _numNodeTypes; // total distinct node types (1 for homogeneous graphs)
int32_t _numNodesInQueue{0}; // running count of queued nodes across all seeds and types
int32_t _batchSize; // number of seed nodes in the current batch
int32_t _numNodeTypes; // total distinct node types (1 for homogeneous graphs)
int32_t _numNodesInQueue{0}; // running count of queued nodes across all seeds and types

// Graph structure — set at construction, read-only during the algorithm.
// _nodeTypeToEdgeTypeIds[ntype_id] → list of edge type IDs that originate from that node type.
Expand All @@ -102,7 +112,6 @@ class PPRForwardPush {
// Hash map: nodeId is a sparse graph ID from a large graph, so a dense array is
// impractical (contrast with _state above). Populated incrementally; avoids re-fetching.
std::unordered_map<uint64_t, std::vector<int32_t>> _neighborCache;

};

} // namespace gigl
} // namespace gigl
5 changes: 3 additions & 2 deletions gigl-core/core/sampling/python_ppr_forward_push.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -21,7 +21,7 @@ namespace gigl {
// pybind11/stl.h handles all type conversions automatically; the other methods use
// direct member function pointers for the same reason.
static void pushResidualsWrapper(PPRForwardPush& state, const py::dict& fetchedByEtypeId) {
std::unordered_map<int32_t, std::tuple<torch::Tensor, torch::Tensor, torch::Tensor>> neighborTensorsByEtypeId;
NeighborFetchMap neighborTensorsByEtypeId;
// Dict iteration touches Python objects — GIL must be held here.
for (auto item : fetchedByEtypeId) {
auto edgeTypeId = item.first.cast<int32_t>();
Expand Down Expand Up @@ -57,5 +57,6 @@ PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) {
std::vector<torch::Tensor>>())
.def("drain_queue", &gigl::PPRForwardPush::drainQueue)
.def("push_residuals", gigl::pushResidualsWrapper)
.def("extract_top_k", &gigl::PPRForwardPush::extractTopK);
.def("extract_top_k", &gigl::PPRForwardPush::extractTopK)
.def("extract_top_k_with_residual_top_up", &gigl::PPRForwardPush::extractTopKWithResidualTopUp);
}
16 changes: 10 additions & 6 deletions gigl-core/src/gigl_core/ppr_forward_push.pyi
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -1,5 +1,9 @@
import torch

TensorTriplet = tuple[torch.Tensor, torch.Tensor, torch.Tensor]
ExtractResult = dict[int, TensorTriplet]


class PPRForwardPush:
def __init__(
self,
Expand All @@ -12,10 +16,10 @@ class PPRForwardPush:
degree_tensors: list[torch.Tensor],
) -> None: ...
def drain_queue(self) -> dict[int, torch.Tensor] | None: ...
def push_residuals(
def push_residuals(self, fetched_by_etype_id: dict[int, TensorTriplet]) -> None: ...
def extract_top_k(self, max_ppr_nodes: int) -> ExtractResult: ...
def extract_top_k_with_residual_top_up(
self,
fetched_by_etype_id: dict[int, tuple[torch.Tensor, torch.Tensor, torch.Tensor]],
) -> None: ...
def extract_top_k(
self, max_ppr_nodes: int
) -> dict[int, tuple[torch.Tensor, torch.Tensor, torch.Tensor]]: ...
max_ppr_nodes: int,
max_residual_nodes: int,
) -> ExtractResult: ...
Loading