graclus_cpu.cpp 2.17 KB
Newer Older
rusty1s's avatar
rusty1s committed
1
2
3
4
#include "graclus_cpu.h"

#include "utils.h"

rusty1s's avatar
rusty1s committed
5
6
7
torch::Tensor graclus_cpu(torch::Tensor rowptr, torch::Tensor col,
                          torch::optional<torch::Tensor> optional_weight) {
  CHECK_CPU(rowptr);
rusty1s's avatar
rusty1s committed
8
  CHECK_CPU(col);
rusty1s's avatar
rusty1s committed
9
  CHECK_INPUT(rowptr.dim() == 1 && col.dim() == 1);
rusty1s's avatar
rusty1s committed
10
11
  if (optional_weight.has_value()) {
    CHECK_CPU(optional_weight.value());
rusty1s's avatar
rusty1s committed
12
    CHECK_INPUT(optional_weight.value().dim() == 1);
rusty1s's avatar
rusty1s committed
13
14
15
    CHECK_INPUT(optional_weight.value().numel() == col.numel());
  }

rusty1s's avatar
rusty1s committed
16
17
18
  int64_t num_nodes = rowptr.numel() - 1;
  auto out = torch::full(num_nodes, -1, rowptr.options());
  auto node_perm = torch::randperm(num_nodes, rowptr.options());
rusty1s's avatar
rusty1s committed
19
20
21

  auto rowptr_data = rowptr.data_ptr<int64_t>();
  auto col_data = col.data_ptr<int64_t>();
rusty1s's avatar
rusty1s committed
22
  auto node_perm_data = node_perm.data_ptr<int64_t>();
rusty1s's avatar
rusty1s committed
23
24
25
  auto out_data = out.data_ptr<int64_t>();

  if (!optional_weight.has_value()) {
rusty1s's avatar
rusty1s committed
26
27
    for (int64_t n = 0; n < num_nodes; n++) {
      auto u = node_perm_data[n];
rusty1s's avatar
rusty1s committed
28
29
30
31
32
33

      if (out_data[u] >= 0)
        continue;

      out_data[u] = u;

rusty1s's avatar
rusty1s committed
34
35
36
      int64_t row_start = rowptr_data[u], row_end = rowptr_data[u + 1];

      for (auto e = 0; e < row_end - row_start; e++) {
rusty1s's avatar
rusty1s committed
37
        auto v = col_data[row_start + e];
rusty1s's avatar
rusty1s committed
38
39
40
41
42
43
44
45
46
47
48

        if (out_data[v] >= 0)
          continue;

        out_data[u] = std::min(u, v);
        out_data[v] = std::min(u, v);
        break;
      }
    }
  } else {
    auto weight = optional_weight.value();
Matthias Fey's avatar
Matthias Fey committed
49
    auto scalar_type = weight.scalar_type();
50
    AT_DISPATCH_ALL_TYPES_AND2(at::ScalarType::Half, at::ScalarType::BFloat16, scalar_type, "graclus_cpu", [&] {
rusty1s's avatar
rusty1s committed
51
52
      auto weight_data = weight.data_ptr<scalar_t>();

rusty1s's avatar
rusty1s committed
53
54
      for (auto n = 0; n < num_nodes; n++) {
        auto u = node_perm_data[n];
rusty1s's avatar
rusty1s committed
55
56
57
58
59
60
61

        if (out_data[u] >= 0)
          continue;

        auto v_max = u;
        scalar_t w_max = (scalar_t)0.;

rusty1s's avatar
rusty1s committed
62
63
        for (auto e = rowptr_data[u]; e < rowptr_data[u + 1]; e++) {
          auto v = col_data[e];
rusty1s's avatar
rusty1s committed
64
65
66
67

          if (out_data[v] >= 0)
            continue;

rusty1s's avatar
rusty1s committed
68
          if (weight_data[e] >= w_max) {
rusty1s's avatar
rusty1s committed
69
            v_max = v;
rusty1s's avatar
rusty1s committed
70
            w_max = weight_data[e];
rusty1s's avatar
rusty1s committed
71
72
73
74
75
76
77
78
79
80
81
          }
        }

        out_data[u] = std::min(u, v_max);
        out_data[v_max] = std::min(u, v_max);
      }
    });
  }

  return out;
}