[Roadmap] CPU Performance Optimization for PyG
See original GitHub issue🚀 The feature, motivation and pitch
The goal of this roadmap is to optimize CPU performance for PyG (including torch_scatter
, torch_sparse
).
For the first step, we will start with single node inference performance optimization on:
- Homogeneous Models: GCN or GAT, PNA, EdgeConv
- Heterogeneous Models: to_hetero, R-GCN, R-GAT
Next step will extend to optimization effort to (distributed) training.
Performance Profiling
CPU platform: Icelake Xeon
Generic benchmarking
- GCN + ogbn-products: (
torch_sparse::spmm_sum
96.04%) - GCN + reddits: layer=1, hidden=16 (
DataLoader
83.49%,aten::scatter_add_
8.47%) - GCN + reddits: layer=3, hidden=32 (
DataLoader
59.83%,aten::scatter_add_
24.76%) - SAGE + ogbn-products: (
aten::scatter_add_
27.61%,DataLoader
25.70%,aten::index
20.26%) - GAT + CiteSeer: (
aten::scatter_add_
30.91%,torch_scatter::scatter_max
24.54%,aten::mm
10.34%,aten::index_select
6.71%) most of models under pytorch_geometric/benchmark/citation have similar behavior from performance perspective. - to_hetero_mag: (
aten::addmm
21.69%,aten::scatter_add_
20.60%,aten::index_select
13.48%,DataLoader
12.31%) - PNA: (t
orch_scatter::scatter_max
39.34%,torch_scatter::scatter_min
39.25%); need follow up: need to get scatter_reduce tensor shape/stride (similar issue as aten::scatter_add_?) - dynamicEdgeConv: (
torch_scatter::scatter_max
66.91%,torch_cluster::knn
23.56%) sourcebenchmark/points/edge_cnn.py
- EdgeConv: ((
torch_scatter::scatter_max
torch_scatter::scatter_max 53.61%,aten::index_select
21.73%,DataLoader
16.11%) source from https://github.com/pyg-team/pytorch_geometric/pull/4915 - pytorch_geometric/benchmark/kernel
- pytorch_geometric/benchmark/points
Large dataset benchmarking
- GraphSAGE + mag240m profiling (table below)
- analysis of ratio of profiler recorded time against total runtime (evaluate other overhead such numpy calling if any)
- gather the input range for spmm_{sum|max|mean} for oneDNN RFC proposal (future plan)
------------------------------------------------------- ------------ ------------ ------------ ------------ ------------ ------------
Name Self CPU % Self CPU CPU total % CPU total CPU time avg # of Calls
------------------------------------------------------- ------------ ------------ ------------ ------------ ------------ ------------
enumerate(DataLoader)#_MultiProcessingDataLoaderIter... 71.27% 608.842s 71.39% 609.891s 70.223ms 8685
torch_sparse::spmm_mean 14.91% 127.390s 14.93% 127.522s 7.342ms 17370
aten::addmm 3.77% 32.166s 7.34% 62.727s 1.806ms 34740
aten::copy_ 3.60% 30.766s 3.60% 30.766s 161.007us 191082
aten::mm 2.29% 19.588s 2.30% 19.683s 1.133ms 17370
aten::native_batch_norm 0.94% 7.989s 1.01% 8.657s 332.256us 26055
DataLoader (with preprocess of input data) is the major bottleneck here, mostly from_numpy
(246s) and to
(169s) triggered by data type conversion, source from convert_batch.
Performance Hotspots
- DataLoader (mini-batch mode): mostly introduced by preprocessing by Samplers (e.g.
NeighborSampler
). - edge_index in CSR: spmm_sum or spmm_max from torch_sparse (memory format CSR).
- edge_index in COO: scatter_add, scatter_max(torch_scatter), index_select, index, etc.
Python level API upgrade in model scripts
The DataLoader is a major hotspot so the first step is to upgrade DataLoader from NeightborSampler
to NeighborLoader
which has native C++ impelemtation:
Native level kernel optimization
Phase One Optimizations
-
NeighborLoader
parallelization: the current impl is sequential (probably to avoid oversubscription with multiple workers on the data loader). Unlike GPU runs, asynchronously run data loading thread and computation thread does not always make sense. On some occasions, run data loading step and computation step sequentially while making each of the torch operator parallel on OpenMP (which is case of intra-parallelism) makes more sense. Hotspot ontorch_sparse::neighbor_sample
. -
aten::sort
: GCN + ogbn-products spent roughly 1/3 time on sort in the preprocessing step (which is not covered during profiler result for the model inference), introduced by indexing from sparse tensor at gnn.py#L123. Root cause is aten::sort(dim) could only be paralleled on dimensions != dim, and the grain size is not correctly set. Fixed by #74897. -
spmm_{max|mean|sum}
(torch_sparse). Add vectorization and prefetch (indirect memory access) and apply blocking on M and K (if necessary). -
scatte_add
andscatter_max
(torch_scatter). Optimizedscatter_add
(with extended index) with #82703. Still need more polishing work.
the current impl for scatter_add
will try to parallel on the inner dimension to avoid write conflict; while ideally we should try to parallel on the outer dimension and vectorize on the inner dimension, yet need to resolve the write conflict on the output tensor. Experiment different impls for the given input range.
-
index_select
, optimized via #76868. -
index
, directly optimizeindex
would be difficult, maybe we can change it to more performance ops likeindex_select
from NeighborLoader or customize its kernel from NeighborLoader.
Phase Two Optimizations
- enable ‘std’ reduce type in
scatter_reduce
- add MultiAggregation kernel
- optimize
segment_reduce
and align the reduce types between scatter_reduce, spmm_reduce, segment_reduce. - kernel fusion for GAS, #71300, maybe dispatch on
TensorTypeId
ofCPU
andSparseCPU
. -
scatter_add
: cache the sorted index. -
knn
(torch_cluster), need follow up shape info to determine proper method to parallel the kernel. Evaluate knn fromoneAPI
dal.
Design option for vectorization
To vectorize kernels from torch-sparse and torch-scatter, we have multiple options:
- vectorize inside torch-sparse/torch-scatter: the most simple way is to use
#pragma omp simd
and add a compiler flagmarch=skylake-avx512
but this won’t apply bfloat16 (bfloat16 is a overload of uint16 and won’t be vectorized properly by compiler) - vectorize inside torch-sparse/torch-scatter: use the wrapper of
at::vec::Vectorized<scalar_t>
, this will apply to bfloat16 but we need to customize the cmake scripts to make it compatible with PyTorch’s cpu build flags: _DEFAULT(scalar code), _AVX2 and _AVX512. - vectorize inside torch core: in this manner
at::vec::Vectorized<scalar_t>
will work without any change but need to move the operator from torch-sparse/torch-scatter to torch. Makes more sense for the fused kernel of GAS.
(current decision is to go with option 3 as much as we can)
Bfloat16 enabling in torch-sparse/torch-scatter
(highly related to the vectorization method choosn)
- Need more work to determine operator list for bfloat16 support.
Validation
- verify float32 accuracy
- verify bfloat16 accuracy
Issue Analytics
- State:
- Created a year ago
- Comments:16 (7 by maintainers)
Update on
spmm
optimizations, PR submitted at https://github.com/pytorch/pytorch/pull/83727.Port
spmm
reduction fromtorch-sparse
totorch
, the current PR is only for demonstrating performance gains, API definition needs more amendment.Now only
sum
is added, more will come in future (max, mean, min), the algorithm is pretty much the same.Select benchmark from
./ogb/examples/nodeproppred/products/gnn.py
, since originally this one spent majority of time ontorch_sparse::spmm_sum
. Thespmm
roughly got 5x speedup on my 20 core machine.To break down the optimization scheme a little bit:
The
balanced thread partition
is targeting at balancing the thread payload. Basically if we directly parallel on row direction, it will be (I collect number of edges for each thread):We can see that the first 2 threads have more payload than others, need to balance the thread payload here. Normally we can use dynamic scheduling for omp, but this won’t fit into pytorch’s
at::parallel_for
which is essentially a static scheduling, so I did manual partitioning here (the logic may be further refined, will do later).current benchmark profiling result uses the default setting. Some scripts, for example to_hetero_mag would explicitly set the
num_workers
, if not the pytorch default setting will be 4.DataLoader
time in the benchmark profile result actually comprises of two parts:The second part takes more time, so it is still possible to be improved with single worker + parallel openmp. If we use
num_workers>0
, need to make sure openmp in the worker have correct setting (omp_num_threads and core affinity binding) to avoid over-subscription.Actually the data loader optimization is a rather complexed issue, perhaps more complexed than optimizing the kernels 😦 since it is more likely a tuning job to achieve the most balanced situation between workload payload (memory footprint, computation complexity etc.) and hardware capacity (IO, memory bandwidth, ALU flops).
Usually we do not do data loading optimizations since the real case in deployment would probably be even more complexed (some venders have mechanisms like prefetching, batching to improve overall user experience and efficiency). But the thing is DGL has done some optimizations here so we need to at least something similar, otherwise out of box performance on PyG would look bad.
Anyway, we will make sure that openmp have correct settings either
num_workers=0
ornum_workers=N
, and also each of the sampler can be properly paralleled.num_workers=0
benefits more for the pre processing andnum_workers=N
benefits more for the IO. And let the users to decide which way to go (maybe we can give a BKM or some simple guideline).