You need to sign in or sign up before continuing.
Unverified Commit c94d2a5f authored by Andrzej Kotłowski's avatar Andrzej Kotłowski Committed by GitHub
Browse files

[Performance] Do not fuse neighbor sampler for 1 thread (#6421)

parent ad4df9c5
......@@ -3,6 +3,7 @@ from .. import backend as F
from ..base import EID, NID
from ..heterograph import DGLGraph
from ..transforms import to_block
from ..utils import get_num_threads
from .base import BlockSampler
......@@ -150,8 +151,9 @@ class NeighborSampler(BlockSampler):
def sample_blocks(self, g, seed_nodes, exclude_eids=None):
output_nodes = seed_nodes
blocks = []
if self.fused:
# sample_neighbors_fused function requires multithreading to be more efficient
# than sample_neighbors
if self.fused and get_num_threads() > 1:
cpu = F.device_type(g.device) == "cpu"
if isinstance(seed_nodes, dict):
for ntype in list(seed_nodes.keys()):
......
Markdown is supported
0% or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment