GCNConv and parallel execution on a CPU

[l-hoang]

Hello.

I've been experimenting with PyTorch Geometric's GCNConv layer on CPUs, and it seems like
they do not leverage much parallelism when given multiple threads.

Here is the code I use:

class GCN(torch.nn.Module):
    def __init__(self, in_dim, out_dim):
        super().__init__()
        self.hidden_feat = 2000

        self.conv1 = GCNConv(in_dim, self.hidden_feat)
        self.conv2 = GCNConv(self.hidden_feat, self.hidden_feat)
        self.conv3 = GCNConv(self.hidden_feat, out_dim)

    def forward(self, data):
        x, edge_index = data.x, data.edge_index

        x = self.conv1(x, edge_index)
        x = self.conv2(x, edge_index)
        x = self.conv3(x, edge_index)

        return x.log_softmax(dim=1)

3 GCN layers back to back with a hidden feature size of 2000.

If I run this with 1 thread (confirmed with torch.get_num_threads()) the forward phases takes 68 seconds.
The data batch print shown is the number of rows and edges in the batch for reference.

Data(x=[96185, 128], edge_index=[2, 255628], y=[96185], seed_nodes=[10000])
model run done 68.32124781608582

If I run with 32 threads (the number of cores on the machine I am running on), the time improves, but
not by much (note the data batch is different since sampling occurs randomly, but it's in the same ballpark):

Data(x=[96091, 128], edge_index=[2, 254745], y=[96091], seed_nodes=[10000])
model run done 49.60090661048889

How much parallelism is expected from the GCNConv layer? It seems like going from 1 thread to 32 threads doesn't
improve the runtime much, and from my observations with top, most of the time seems to be spent
not using all the allocated threads (i.e. there doesn't seem to be much parallel execution going on).

Thank you.

[rusty1s]

As far as I know, torch.scatter_add does not have great multi-threading support on CPU, which probably explains your observed runtimes. That is, you only get minor benefits in efficiency on CPU through other ops (matmul, gather, ...), while most of the time is spent in single-threaded scatter ops. It might be a good idea to do a micro benchmark of torch.scatter_(..., reduce="sum") and torch_scatter.scatter and submit it to the PyTorch and torch-scatter repos.

[rusty1s]

Parallel scatter support should be available once rusty1s/pytorch_scatter#270 lands (thanks to @mananshah99).

The spmm_value_bw_cpu computes the gradients w.r.t. to the edge weights which is only used in case of attention. The gradients w.r.t. to node features use spmm under-the-hood and should make use of parallelization already.

[bowu]

Parallel scatter support should be available once rusty1s/pytorch_scatter#270 lands (thanks to @mananshah99).

The spmm_value_bw_cpu computes the gradients w.r.t. to the edge weights which is only used in case of attention. The gradients w.r.t. to node features use spmm under-the-hood and should make use of parallelization already.

After it lands, should we use COO edge tensor or SparseTensor for the best performance?

[rusty1s]

This is hard to tell and depends on the use-case. For dense graphs, it is definitely advisable to use SparseTensor.

[bowu]

I thought SparseTensor is optimized for sparse graphs. Could you let us know why it mainly works for dense graphs?

[rusty1s]

That's more due to reasons of memory consumption. If you have a dense graph, then the usage ofSparseTensor may be able to avoid materializing edge-level message representations (dependent on the GNN operator). See here for more details.