Hetero subgraph with dispatching

CHANGELOG.md

@@ -5,6 +5,7 @@ The format is based on [Keep a Changelog](http://keepachangelog.com/en/1.0.0/).
 
 ## [Unreleased]
 ### Added
+- Added hetero subgraph kernel ([#43](https://github.com/pyg-team/pyg-lib/pull/43)
 - Added download script for benchmark data ([#44](https://github.com/pyg-team/pyg-lib/pull/44)
 - Added `biased sampling` utils ([#38](https://github.com/pyg-team/pyg-lib/pull/38))
 - Added `CHANGELOG.md` ([#39](https://github.com/pyg-team/pyg-lib/pull/39))

pyg_lib/csrc/sampler/subgraph.cpp

@@ -1,8 +1,11 @@
 #include "subgraph.h"
+#include <pyg_lib/csrc/utils/hetero_dispatch.h>
 
 #include <ATen/core/dispatch/Dispatcher.h>
 #include <torch/library.h>
 
+#include <functional>
+
 namespace pyg {
 namespace sampler {
 
@@ -11,7 +14,7 @@ std::tuple<at::Tensor, at::Tensor, c10::optional<at::Tensor>> subgraph(
     const at::Tensor& col,
     const at::Tensor& nodes,
     const bool return_edge_id) {
-  at::TensorArg rowptr_t{rowptr, "rowtpr", 1};
+  at::TensorArg rowptr_t{rowptr, "rowptr", 1};
   at::TensorArg col_t{col, "col", 1};
   at::TensorArg nodes_t{nodes, "nodes", 1};
 
@@ -25,10 +28,42 @@ std::tuple<at::Tensor, at::Tensor, c10::optional<at::Tensor>> subgraph(
   return op.call(rowptr, col, nodes, return_edge_id);
 }
 
+c10::Dict<utils::edge_t,
+          std::tuple<at::Tensor, at::Tensor, c10::optional<at::Tensor>>>
+hetero_subgraph(const utils::edge_tensor_dict_t& rowptr,
+                const utils::edge_tensor_dict_t& col,
+                const utils::node_tensor_dict_t& nodes,
+                const c10::Dict<utils::edge_t, bool>& return_edge_id) {
+  // Define the homogeneous implementation as a std function to pass the type
+  // check
+  std::function<std::tuple<at::Tensor, at::Tensor, c10::optional<at::Tensor>>(
+      const at::Tensor&, const at::Tensor&, const at::Tensor&, bool)>
+      func = subgraph;
+
+  // Construct an operator
+  utils::HeteroDispatchOp<decltype(func)> op(rowptr, col, func);
+
+  // Construct dispatchable arguments
+  // TODO: We filter source node by assuming hetero graph is a dict of homo
+  // graph here; both source and destination nodes should be considered when
+  // filtering a bipartite graph
+  utils::HeteroDispatchArg<utils::node_tensor_dict_t, at::Tensor,
+                           utils::NodeSrcMode>
+      nodes_arg(nodes);
+  utils::HeteroDispatchArg<c10::Dict<utils::edge_t, bool>, bool,
+                           utils::EdgeMode>
+      edge_id_arg(return_edge_id);
+  return op(nodes_arg, edge_id_arg);
+}
+
 TORCH_LIBRARY_FRAGMENT(pyg, m) {
   m.def(TORCH_SELECTIVE_SCHEMA(
       "pyg::subgraph(Tensor rowptr, Tensor col, Tensor "
       "nodes, bool return_edge_id) -> (Tensor, Tensor, Tensor?)"));
+  m.def(TORCH_SELECTIVE_SCHEMA(
+      "pyg::hetero_subgraph(Dict(str, Tensor) rowptr, Dict(str, "
+      "Tensor) col, Dict(str, Tensor) nodes, Dict(str, bool) "
+      "return_edge_id) -> Dict(str, (Tensor, Tensor, Tensor?))"));
 }
 
 }  // namespace sampler

pyg_lib/csrc/sampler/subgraph.h

@@ -2,6 +2,7 @@
 
 #include <ATen/ATen.h>
 #include "pyg_lib/csrc/macros.h"
+#include "pyg_lib/csrc/utils/types.h"
 
 namespace pyg {
 namespace sampler {
@@ -15,5 +16,14 @@ PYG_API std::tuple<at::Tensor, at::Tensor, c10::optional<at::Tensor>> subgraph(
     const at::Tensor& nodes,
     const bool return_edge_id = true);
 
+// A heterogeneous version of the above function.
+// Returns a dict from each relation type to its result
+PYG_API c10::Dict<utils::edge_t,
+                  std::tuple<at::Tensor, at::Tensor, c10::optional<at::Tensor>>>
+hetero_subgraph(const utils::edge_tensor_dict_t& rowptr,
+                const utils::edge_tensor_dict_t& col,
+                const utils::node_tensor_dict_t& nodes,
+                const c10::Dict<utils::edge_t, bool>& return_edge_id);
+
 }  // namespace sampler
 }  // namespace pyg

pyg_lib/csrc/utils/hetero_dispatch.h

@@ -0,0 +1,191 @@
+#pragma once
+
+#include "types.h"
+
+#include <type_traits>
+
+namespace pyg {
+
+namespace utils {
+
+// Base class for easier type check
+struct HeteroDispatchMode {};
+
+// List hetero dispatch mode as different types to avoid non-type template
+// specialization.
+struct SkipMode : public HeteroDispatchMode {};
+
+struct NodeSrcMode : public HeteroDispatchMode {};
+
+struct NodeDstMode : public HeteroDispatchMode {};
+
+struct EdgeMode : public HeteroDispatchMode {};
+
+// Check if the argument is a c10::dict so that is could be filtered by an edge
+// type.
+template <typename... T>
+struct is_c10_dict : std::false_type {};
+
+template <typename T, typename N>
+struct is_c10_dict<c10::Dict<T, N>> : std::true_type {};
+
+// TODO: Should specialize as if-constexpr when in C++17
+template <typename T, typename V, typename MODE>
+class HeteroDispatchArg {};
+
+// In SkipMode we do not filter this arg
+template <typename T, typename V>
+class HeteroDispatchArg<T, V, SkipMode> {
+ public:
+  HeteroDispatchArg(const T& val) : val_(val) {}
+
+  // If we pass the filter, we will obtain the value of the argument.
+  template <typename K>
+  V value_by_edge(const K& key) {
+    return val_;
+  }
+
+  bool filter_by_edge(const edge_t& edge) { return true; }
+
+ private:
+  T val_;
+};
+
+// In NodeSrcMode we check if source node is in the dict
+template <typename T, typename V>
+class HeteroDispatchArg<T, V, NodeSrcMode> {
+ public:
+  HeteroDispatchArg(const T& val) : val_(val) {
+    static_assert(is_c10_dict<T>::value, "Should be a c10::dict");
+  }
+
+  // Dict value lookup
+  template <typename K>
+  V value_by_edge(const K& key) {
+    return val_.at(get_src(key));
+  }
+
+  // Dict if key exists
+  bool filter_by_edge(const edge_t& edge) {
+    return val_.contains(get_src(edge));
+  }
+
+ private:
+  T val_;
+};
+
+// In NodeDstMode we check if destination node is in the dict
+template <typename T, typename V>
+class HeteroDispatchArg<T, V, NodeDstMode> {
+ public:
+  HeteroDispatchArg(const T& val) : val_(val) {
+    static_assert(is_c10_dict<T>::value, "Should be a c10::dict");
+  }
+
+  template <typename K>
+  V value_by_edge(const K& key) {
+    return val_.at(get_dst(key));
+  }
+
+  bool filter_by_edge(const edge_t& edge) {
+    return val_.contains(get_dst(edge));
+  }
+
+ private:
+  T val_;
+};
+
+// In EdgeMode we check if edge is in the dict
+template <typename T, typename V>
+class HeteroDispatchArg<T, V, EdgeMode> {
+ public:
+  HeteroDispatchArg(const T& val) : val_(val) {
+    static_assert(is_c10_dict<T>::value, "Should be a c10::dict");
+  }
+
+  template <typename K>
+  V value_by_edge(const K& key) {
+    return val_.at(key);
+  }
+
+  bool filter_by_edge(const edge_t& edge) { return val_.contains(edge); }
+
+ private:
+  T val_;
+};
+
+// The following will help static type checks:
+template <typename... T>
+struct is_hetero_arg : std::false_type {};
+
+// Just check inheritance, a workaround without introducing concepts
+template <typename T, typename V, typename Mode>
+struct is_hetero_arg<HeteroDispatchArg<T, V, Mode>> : std::true_type {
+  static_assert(std::is_base_of<HeteroDispatchMode, Mode>::value,
+                "Must pass a mode for dispatching");
+};
+
+// Specialize
+template <typename... Args>
+bool filter_args_by_edge(const edge_t& edge, Args&&... args) {}
+
+// Stop condition of argument filtering
+template <>
+bool filter_args_by_edge(const edge_t& edge) {
+  return true;
+}
+
+// We filter each argument individually by the given edge using a variadic
+// template
+template <typename T, typename... Args>
+bool filter_args_by_edge(const edge_t& edge, T&& t, Args&&... args) {
+  static_assert(
+      is_hetero_arg<std::remove_const_t<std::remove_reference_t<T>>>::value,
+      "args should be HeteroDispatchArg");
+  return t.filter_by_edge(edge) && filter_args_by_edge(edge, args...);
+}
+
+// Check if a callable is wrapped by std::function
+template <typename... T>
+struct is_std_function : std::false_type {};
+
+template <typename T, typename... Args>
+struct is_std_function<std::function<T(Args...)>> : std::true_type {};
+
+template <typename T>
+class HeteroDispatchOp {
+ public:
+  using result_type = typename T::result_type;
+  HeteroDispatchOp(const edge_tensor_dict_t& rowptr,
+                   const edge_tensor_dict_t& col,
+                   T op)
+      : rowptr_(rowptr), col_(col), op_(op) {
+    // Check early
+    static_assert(is_std_function<T>::value, "Must pass a function");
+  }
+
+  template <typename... Args>
+  c10::Dict<edge_t, result_type> operator()(Args&&... args) {
+    c10::Dict<edge_t, result_type> dict;
+    for (const auto& kv : rowptr_) {
+      auto edge = kv.key();
+      auto rowptr = kv.value();
+      auto col = col_.at(edge);
+      bool pass = filter_args_by_edge(edge, args...);
+      if (pass) {
+        result_type res = op_(rowptr, col, args.value_by_edge(edge)...);
+        dict.insert(edge, res);
+      }
+    }
+    return dict;
+  }
+
+ private:
+  edge_tensor_dict_t rowptr_;
+  edge_tensor_dict_t col_;
+  T op_;
+};
+
+}  // namespace utils
+
+}  // namespace pyg

pyg_lib/csrc/utils/types.h

@@ -0,0 +1,34 @@
+#pragma once
+
+#include <string>
+
+#include <ATen/ATen.h>
+
+namespace pyg {
+namespace utils {
+
+const std::string SPLIT_TOKEN = "__";
+
+using edge_t = std::string;
+using node_t = std::string;
+using rel_t = std::string;
+
+using edge_tensor_dict_t = c10::Dict<edge_t, at::Tensor>;
+using node_tensor_dict_t = c10::Dict<node_t, at::Tensor>;
+
+node_t get_src(const edge_t& e) {
+  return e.substr(0, e.find_first_of(SPLIT_TOKEN));
+}
+
+rel_t get_rel(const edge_t& e) {
+  auto beg = e.find_first_of(SPLIT_TOKEN) + SPLIT_TOKEN.size();
+  return e.substr(beg,
+                  e.find_last_of(SPLIT_TOKEN) - SPLIT_TOKEN.size() + 1 - beg);
+}
+
+node_t get_dst(const edge_t& e) {
+  return e.substr(e.find_last_of(SPLIT_TOKEN) + 1);
+}
+}  // namespace utils
+
+}  // namespace pyg

test/csrc/sampler/test_subgraph.cpp

@@ -20,3 +20,59 @@ TEST(SubgraphTest, BasicAssertions) {
   auto expected_edge_id = at::tensor({3, 4, 5, 6, 7, 8}, options);
   EXPECT_TRUE(at::equal(std::get<2>(out).value(), expected_edge_id));
 }
+
+TEST(HeteroSubgraphPassFilterTest, BasicAssertions) {
+  auto options = at::TensorOptions().dtype(at::kLong);
+
+  auto nodes = at::arange(1, 5, options);
+  auto graph = cycle_graph(/*num_nodes=*/6, options);
+
+  pyg::utils::node_t node_name = "node";
+  pyg::utils::edge_t edge_name = "node__to__node";
+
+  pyg::utils::edge_tensor_dict_t rowptr_dict;
+  rowptr_dict.insert(edge_name, std::get<0>(graph));
+  pyg::utils::edge_tensor_dict_t col_dict;
+  col_dict.insert(edge_name, std::get<1>(graph));
+  pyg::utils::edge_tensor_dict_t nodes_dict;
+  nodes_dict.insert(node_name, nodes);
+  c10::Dict<pyg::utils::edge_t, bool> edge_id_dict;
+  edge_id_dict.insert(edge_name, true);
+
+  auto res = pyg::sampler::hetero_subgraph(rowptr_dict, col_dict, nodes_dict,
+                                           edge_id_dict);
+
+  EXPECT_EQ(res.size(), 1);
+  auto out = res.at(edge_name);
+
+  auto expected_rowptr = at::tensor({0, 1, 3, 5, 6}, options);
+  EXPECT_TRUE(at::equal(std::get<0>(out), expected_rowptr));
+  auto expected_col = at::tensor({1, 0, 2, 1, 3, 2}, options);
+  EXPECT_TRUE(at::equal(std::get<1>(out), expected_col));
+  auto expected_edge_id = at::tensor({3, 4, 5, 6, 7, 8}, options);
+  EXPECT_TRUE(at::equal(std::get<2>(out).value(), expected_edge_id));
+}
+
+TEST(HeteroSubgraphFailFilterTest, BasicAssertions) {
+  auto options = at::TensorOptions().dtype(at::kLong);
+
+  auto nodes = at::arange(1, 5, options);
+  auto graph = cycle_graph(/*num_nodes=*/6, options);
+
+  pyg::utils::node_t node_name = "node";
+  pyg::utils::edge_t edge_name = "node123__to456__node321";
+
+  pyg::utils::edge_tensor_dict_t rowptr_dict;
+  rowptr_dict.insert(edge_name, std::get<0>(graph));
+  pyg::utils::edge_tensor_dict_t col_dict;
+  col_dict.insert(edge_name, std::get<1>(graph));
+  pyg::utils::edge_tensor_dict_t nodes_dict;
+  nodes_dict.insert(node_name, nodes);
+  c10::Dict<pyg::utils::edge_t, bool> edge_id_dict;
+  edge_id_dict.insert(edge_name, true);
+
+  auto res = pyg::sampler::hetero_subgraph(rowptr_dict, col_dict, nodes_dict,
+                                           edge_id_dict);
+
+  EXPECT_EQ(res.size(), 0);
+}

test/csrc/utils/test_utils.cpp

@@ -0,0 +1,15 @@
+#include <gtest/gtest.h>
+
+#include "pyg_lib/csrc/sampler/subgraph.h"
+
+TEST(UtilsTypeTest, BasicAssertions) {
+  pyg::utils::edge_t edge = "node1__to__node2";
+
+  auto src = pyg::utils::get_src(edge);
+  auto dst = pyg::utils::get_dst(edge);
+  auto rel = pyg::utils::get_rel(edge);
+
+  EXPECT_EQ(src, std::string("node1"));
+  EXPECT_EQ(dst, std::string("node2"));
+  EXPECT_EQ(rel, std::string("to"));
+}