Add typed-channel PPR sampler

gigl-core/core/sampling/ppr_forward_push.h

@@ -25,6 +25,21 @@ struct SeedNodeTypeState {
     std::unordered_set<int32_t> queuedNodes;       // snapshot captured by drainQueue()
 };
 
+// Batched drain result for typed-PPR channels.
+struct DrainedTypedPPRQueues {
+    // Channels that drained queued nodes and need pushResiduals() this iteration.
+    std::vector<int32_t> activeChannelIndices;
+
+    // Channels that have non-empty uncached frontiers and remaining fetch budget.
+    std::vector<int32_t> fetchChannelIndices;
+
+    // Edge types requested by each fetch channel, aligned with fetchChannelIndices.
+    std::vector<std::vector<int32_t>> edgeTypeIdsByFetchChannel;
+
+    // Unioned node frontier for one shared distributed neighbor fetch.
+    std::unordered_map<int32_t, torch::Tensor> unionNodesByEdgeTypeId;
+};
+
 // C++ kernel for PPR Forward Push (Andersen et al., 2006).
 // Hot-loop state lives here; distributed neighbor fetches are driven from Python.
 //
@@ -75,6 +90,12 @@ class PPRForwardPush {
     std::unordered_map<int32_t, std::tuple<torch::Tensor, torch::Tensor, torch::Tensor>> extractTopKWithResidualTopUp(
         int32_t maxPPRNodes, bool enableResidualTopUp);
 
+    friend std::unordered_map<int32_t, std::tuple<torch::Tensor, torch::Tensor, torch::Tensor>>
+    extractTypedTopKWithResidualTopUp(const std::vector<PPRForwardPush*>& states,
+                                      const std::vector<int32_t>& channelQuotas,
+                                      int32_t maxPPRNodes,
+                                      bool enableResidualTopUp);
+
 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;
@@ -116,4 +137,18 @@ class PPRForwardPush {
     std::unordered_map<uint64_t, std::vector<int32_t>> _neighborCache;
 };
 
+// Drain several independent PPR states and union their fetch frontier by edge type.
+// Used by typed PPR to keep per-channel PPR state separate while issuing one
+// shared neighbor fetch for duplicate channel frontier requests.
+DrainedTypedPPRQueues drainTypedPPRChannelQueues(const std::vector<PPRForwardPush*>& states,
+                                                 const std::vector<int32_t>& fetchIterationCounts,
+                                                 int32_t maxFetchIterations);
+
+// Extract and merge typed-PPR channel states in one C++ step.
+std::unordered_map<int32_t, std::tuple<torch::Tensor, torch::Tensor, torch::Tensor>> extractTypedTopKWithResidualTopUp(
+    const std::vector<PPRForwardPush*>& states,
+    const std::vector<int32_t>& channelQuotas,
+    int32_t maxPPRNodes,
+    bool enableResidualTopUp);
+
 } // namespace gigl

gigl-core/core/sampling/python_ppr_forward_push.cpp

@@ -66,6 +66,56 @@ extractTopKWithResidualTopUpWrapper(PPRForwardPush& state, int32_t maxPPRNodes,
     return result;
 }
 
+static py::tuple drainTypedPPRChannelQueuesWrapper(const py::sequence& states,
+                                                   const std::vector<int32_t>& fetchIterationCounts,
+                                                   int32_t maxFetchIterations) {
+    std::vector<PPRForwardPush*> statePtrs;
+    statePtrs.reserve(py::len(states));
+    // Sequence iteration and casting touch Python objects, so keep the GIL
+    // while copying raw C++ state pointers out of the Python container.
+    for (py::handle stateObj : states) {
+        statePtrs.push_back(&stateObj.cast<PPRForwardPush&>());
+    }
+
+    // C++ typed drain only reads/mutates PPRForwardPush states and builds C++
+    // containers. Reacquire the GIL before constructing the Python tuple.
+    // REQUIREMENT: no other thread may read or mutate these channel states
+    // while the GIL is released. The typed sampler drains and pushes each
+    // channel in a single sequenced loop iteration.
+    DrainedTypedPPRQueues drained;
+    {
+        py::gil_scoped_release release;
+        drained = drainTypedPPRChannelQueues(statePtrs, fetchIterationCounts, maxFetchIterations);
+    }
+    return py::make_tuple(std::move(drained.activeChannelIndices),
+                          std::move(drained.fetchChannelIndices),
+                          std::move(drained.edgeTypeIdsByFetchChannel),
+                          std::move(drained.unionNodesByEdgeTypeId));
+}
+
+static std::unordered_map<int32_t, std::tuple<torch::Tensor, torch::Tensor, torch::Tensor>>
+extractTypedTopKWithResidualTopUpWrapper(const py::sequence& states,
+                                         const std::vector<int32_t>& channelQuotas,
+                                         int32_t maxPPRNodes,
+                                         bool enableResidualTopUp) {
+    std::vector<PPRForwardPush*> statePtrs;
+    statePtrs.reserve(py::len(states));
+    // Sequence iteration and casting touch Python objects, so keep the GIL
+    // while copying raw C++ state pointers out of the Python container.
+    for (py::handle stateObj : states) {
+        statePtrs.push_back(&stateObj.cast<PPRForwardPush&>());
+    }
+
+    // C++ extraction only reads the completed channel states and builds C++
+    // tensors/containers. Reacquire the GIL before pybind converts the return.
+    std::unordered_map<int32_t, std::tuple<torch::Tensor, torch::Tensor, torch::Tensor>> result;
+    {
+        py::gil_scoped_release release;
+        result = extractTypedTopKWithResidualTopUp(statePtrs, channelQuotas, maxPPRNodes, enableResidualTopUp);
+    }
+    return result;
+}
+
 } // namespace gigl
 
 // TORCH_EXTENSION_NAME is set by PyTorch's build system to match the Python
@@ -85,7 +135,18 @@ PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) {
         .def("drain_queue", gigl::drainQueueWrapper)
         .def("push_residuals", gigl::pushResidualsWrapper)
         .def("extract_top_k_with_residual_top_up",
-             gigl::extractTopKWithResidualTopUpWrapper,
+             &gigl::extractTopKWithResidualTopUpWrapper,
              py::arg("max_ppr_nodes"),
              py::arg("enable_residual_topup"));
+    m.def("drain_typed_ppr_channel_queues",
+          &gigl::drainTypedPPRChannelQueuesWrapper,
+          py::arg("states"),
+          py::arg("fetch_iteration_counts"),
+          py::arg("max_fetch_iterations") = -1);
+    m.def("extract_typed_top_k_with_residual_top_up",
+          &gigl::extractTypedTopKWithResidualTopUpWrapper,
+          py::arg("states"),
+          py::arg("channel_quotas"),
+          py::arg("max_ppr_nodes"),
+          py::arg("enable_residual_topup"));
 }

gigl-core/src/gigl_core/__init__.py

@@ -1,3 +1,11 @@
-from gigl_core.ppr_forward_push import PPRForwardPush
+from gigl_core.ppr_forward_push import (
+    PPRForwardPush,
+    drain_typed_ppr_channel_queues,
+    extract_typed_top_k_with_residual_top_up,
+)
 
-__all__ = ["PPRForwardPush"]
+__all__ = [
+    "PPRForwardPush",
+    "drain_typed_ppr_channel_queues",
+    "extract_typed_top_k_with_residual_top_up",
+]

gigl-core/src/gigl_core/ppr_forward_push.pyi

@@ -1,3 +1,5 @@
+from typing import Sequence
+
 import torch
 
 class PPRForwardPush:
@@ -21,3 +23,15 @@ class PPRForwardPush:
         max_ppr_nodes: int,
         enable_residual_topup: bool,
     ) -> dict[int, tuple[torch.Tensor, torch.Tensor, torch.Tensor]]: ...
+
+def drain_typed_ppr_channel_queues(
+    states: Sequence[PPRForwardPush],
+    fetch_iteration_counts: Sequence[int],
+    max_fetch_iterations: int = -1,
+) -> tuple[list[int], list[int], list[list[int]], dict[int, torch.Tensor]]: ...
+def extract_typed_top_k_with_residual_top_up(
+    states: Sequence[PPRForwardPush],
+    channel_quotas: Sequence[int],
+    max_ppr_nodes: int,
+    enable_residual_topup: bool,
+) -> dict[int, tuple[torch.Tensor, torch.Tensor, torch.Tensor]]: ...

gigl-core/tests/ppr_forward_push_test.cpp

@@ -2,7 +2,10 @@
 #include "sampling/ppr_forward_push.h"
 
 #include <unordered_map>
+#include <unordered_set>
 
+using gigl::drainTypedPPRChannelQueues;
+using gigl::extractTypedTopKWithResidualTopUp;
 using gigl::PPRForwardPush;
 
 // Builds a single-edge-type, single-node-type PPRForwardPush.
@@ -32,6 +35,25 @@ static std::unordered_map<int32_t, std::tuple<torch::Tensor, torch::Tensor, torc
               torch::tensor(counts, torch::kLong)}}};
 }
 
+static std::unordered_set<int64_t> tensorValues(const torch::Tensor& values) {
+    std::unordered_set<int64_t> result;
+    auto accessor = values.accessor<int64_t, 1>();
+    for (int64_t index = 0; index < values.size(0); ++index) {
+        result.insert(accessor[index]);
+    }
+    return result;
+}
+
+static std::vector<int64_t> tensorToInt64Vector(const torch::Tensor& values) {
+    std::vector<int64_t> result;
+    auto accessor = values.accessor<int64_t, 1>();
+    result.reserve(static_cast<size_t>(values.size(0)));
+    for (int64_t index = 0; index < values.size(0); ++index) {
+        result.push_back(accessor[index]);
+    }
+    return result;
+}
+
 // After construction, drainQueue() returns the seed node under etype 0.
 TEST(PPRForwardPush, DrainQueueReturnsSeedNodeInitially) {
     auto state = makeState(/*seeds=*/{0}, /*alpha=*/0.15, /*requeueThresholdFactor=*/1e-6, /*degrees=*/{1});
@@ -110,6 +132,106 @@ TEST(PPRForwardPush, NeighborCacheAvoidsRefetchingPreviouslyFetchedNode) {
     EXPECT_TRUE(iter3->empty());
 }
 
+TEST(PPRForwardPush, DrainTypedPPRChannelQueuesUnionsChannelFrontiers) {
+    auto channel0 = makeState(/*seeds=*/{0}, /*alpha=*/0.5, /*requeueThresholdFactor=*/1e-9, /*degrees=*/{1, 1});
+    auto channel1 = makeState(/*seeds=*/{1}, /*alpha=*/0.5, /*requeueThresholdFactor=*/1e-9, /*degrees=*/{1, 1});
+
+    auto drained = drainTypedPPRChannelQueues(std::vector<PPRForwardPush*>{&channel0, &channel1},
+                                              /*fetchIterationCounts=*/{0, 0},
+                                              /*maxFetchIterations=*/-1);
+
+    EXPECT_EQ(drained.activeChannelIndices, std::vector<int32_t>({0, 1}));
+    EXPECT_EQ(drained.fetchChannelIndices, std::vector<int32_t>({0, 1}));
+    ASSERT_EQ(drained.edgeTypeIdsByFetchChannel.size(), 2);
+    EXPECT_EQ(drained.edgeTypeIdsByFetchChannel[0], std::vector<int32_t>({0}));
+    EXPECT_EQ(drained.edgeTypeIdsByFetchChannel[1], std::vector<int32_t>({0}));
+
+    ASSERT_NE(drained.unionNodesByEdgeTypeId.find(0), drained.unionNodesByEdgeTypeId.end());
+    EXPECT_EQ(tensorValues(drained.unionNodesByEdgeTypeId.at(0)), std::unordered_set<int64_t>({0, 1}));
+}
+
+TEST(PPRForwardPush, DrainTypedPPRChannelQueuesHonorsPerChannelFetchBudget) {
+    auto channel0 = makeState(/*seeds=*/{0}, /*alpha=*/0.5, /*requeueThresholdFactor=*/1e-9, /*degrees=*/{1, 1});
+    auto channel1 = makeState(/*seeds=*/{1}, /*alpha=*/0.5, /*requeueThresholdFactor=*/1e-9, /*degrees=*/{1, 1});
+
+    auto drained = drainTypedPPRChannelQueues(std::vector<PPRForwardPush*>{&channel0, &channel1},
+                                              /*fetchIterationCounts=*/{1, 0},
+                                              /*maxFetchIterations=*/1);
+
+    EXPECT_EQ(drained.activeChannelIndices, std::vector<int32_t>({0, 1}));
+    EXPECT_EQ(drained.fetchChannelIndices, std::vector<int32_t>({1}));
+    ASSERT_EQ(drained.edgeTypeIdsByFetchChannel.size(), 1);
+    EXPECT_EQ(drained.edgeTypeIdsByFetchChannel[0], std::vector<int32_t>({0}));
+
+    ASSERT_NE(drained.unionNodesByEdgeTypeId.find(0), drained.unionNodesByEdgeTypeId.end());
+    EXPECT_EQ(tensorValues(drained.unionNodesByEdgeTypeId.at(0)), std::unordered_set<int64_t>({1}));
+}
+
+TEST(PPRForwardPush, ExtractTypedTopKWithResidualTopUpMergesChannelsInCpp) {
+    std::vector<int32_t> degrees(21, 1);
+    auto channel0 = makeState(/*seeds=*/{0}, /*alpha=*/0.5, /*requeueThresholdFactor=*/1.0, degrees);
+    auto channel1 = makeState(/*seeds=*/{0}, /*alpha=*/0.5, /*requeueThresholdFactor=*/1.0, degrees);
+
+    channel0.drainQueue();
+    channel0.pushResiduals(makeFetched(/*edgeTypeId=*/0, /*nodeIds=*/{0}, /*flatNeighborIds=*/{10}, /*counts=*/{1}));
+    channel1.drainQueue();
+    channel1.pushResiduals(makeFetched(/*edgeTypeId=*/0, /*nodeIds=*/{0}, /*flatNeighborIds=*/{20}, /*counts=*/{1}));
+
+    auto result = extractTypedTopKWithResidualTopUp(std::vector<PPRForwardPush*>{&channel0, &channel1},
+                                                    /*channelQuotas=*/{2, 1},
+                                                    /*maxPPRNodes=*/3,
+                                                    /*enableResidualTopUp=*/true);
+
+    ASSERT_NE(result.find(0), result.end());
+    const auto& [ids, features, counts] = result.at(0);
+    EXPECT_EQ(tensorToInt64Vector(ids), std::vector<int64_t>({0, 10, 20}));
+    EXPECT_EQ(tensorToInt64Vector(counts), std::vector<int64_t>({3}));
+    ASSERT_EQ(features.size(0), 3);
+    ASSERT_EQ(features.size(1), 5);
+
+    auto featureAccessor = features.accessor<double, 2>();
+    EXPECT_NEAR(featureAccessor[0][0], 1.0, 1e-9);
+    EXPECT_NEAR(featureAccessor[0][1], 1.0, 1e-9);
+    EXPECT_NEAR(featureAccessor[0][2], 1.0, 1e-9);
+    EXPECT_NEAR(featureAccessor[0][3], 1.0, 1e-9);
+    EXPECT_NEAR(featureAccessor[0][4], 1.0, 1e-9);
+
+    EXPECT_NEAR(featureAccessor[1][0], 0.5, 1e-9);
+    EXPECT_NEAR(featureAccessor[1][1], 0.5, 1e-9);
+    EXPECT_NEAR(featureAccessor[1][2], 0.0, 1e-9);
+    EXPECT_NEAR(featureAccessor[1][3], 1.0, 1e-9);
+    EXPECT_NEAR(featureAccessor[1][4], 0.0, 1e-9);
+
+    EXPECT_NEAR(featureAccessor[2][0], 0.5, 1e-9);
+    EXPECT_NEAR(featureAccessor[2][1], 0.0, 1e-9);
+    EXPECT_NEAR(featureAccessor[2][2], 0.5, 1e-9);
+    EXPECT_NEAR(featureAccessor[2][3], 0.0, 1e-9);
+    EXPECT_NEAR(featureAccessor[2][4], 1.0, 1e-9);
+}
+
+TEST(PPRForwardPush, ExtractTypedTopKWithResidualTopUpCanDisableTopUp) {
+    std::vector<int32_t> degrees(21, 1);
+    auto channel0 = makeState(/*seeds=*/{0}, /*alpha=*/0.5, /*requeueThresholdFactor=*/1.0, degrees);
+    auto channel1 = makeState(/*seeds=*/{0}, /*alpha=*/0.5, /*requeueThresholdFactor=*/1.0, degrees);
+
+    channel0.drainQueue();
+    channel0.pushResiduals(makeFetched(/*edgeTypeId=*/0, /*nodeIds=*/{0}, /*flatNeighborIds=*/{10}, /*counts=*/{1}));
+    channel1.drainQueue();
+    channel1.pushResiduals(makeFetched(/*edgeTypeId=*/0, /*nodeIds=*/{0}, /*flatNeighborIds=*/{20}, /*counts=*/{1}));
+
+    auto result = extractTypedTopKWithResidualTopUp(std::vector<PPRForwardPush*>{&channel0, &channel1},
+                                                    /*channelQuotas=*/{2, 1},
+                                                    /*maxPPRNodes=*/3,
+                                                    /*enableResidualTopUp=*/false);
+
+    ASSERT_NE(result.find(0), result.end());
+    const auto& [ids, features, counts] = result.at(0);
+    EXPECT_EQ(tensorToInt64Vector(ids), std::vector<int64_t>({0}));
+    EXPECT_EQ(tensorToInt64Vector(counts), std::vector<int64_t>({1}));
+    ASSERT_EQ(features.size(0), 1);
+    ASSERT_EQ(features.size(1), 5);
+}
+
 // Two seeds (0 and 1) both push residual to sink node 2.  The neighbor-lookup
 // request must deduplicate to one entry for node 2, yet both seeds must still
 // accumulate a PPR score for it.