dmlc--dgl
Minjie Wang
0227ddfb66
* WIP: TypedLinear and new RelGraphConv * wip * further simplify RGCN * a bunch of tweak for performance; add basic cpu support * update on segmm * wip: segment.cu * new backward kernel works * fix a bunch of bugs in kernel; leave idx_a for future * add nn test for typed_linear * rgcn nn test * bugfix in corner case; update RGCN README * doc * fix cpp lint * fix lint * fix ut * wip: hgtconv; presorted flag for rgcn * hgt code and ut; WIP: some fix on reorder graph * better typed linear init * fix ut * fix lint; add docstring
43 行
1.3 KiB
Python
43 行
1.3 KiB
Python
"""dgl gather_mm operator module."""
|
|
from .. import backend as F
|
|
|
|
__all__ = ['gather_mm']
|
|
|
|
def gather_mm(a, b, *, idx_b):
|
|
r"""Gather data according to the given indices and perform matrix multiplication.
|
|
|
|
Let the result tensor be ``c``, the operator conducts the following computation:
|
|
|
|
c[i] = a[i] @ b[idx_b[i]]
|
|
, where len(c) == len(idx_b)
|
|
|
|
|
|
Parameters
|
|
----------
|
|
a : Tensor
|
|
A 2-D tensor of shape ``(N, D1)``
|
|
b : Tensor
|
|
A 3-D tensor of shape ``(R, D1, D2)``
|
|
idx_b : Tensor, optional
|
|
An 1-D integer tensor of shape ``(N,)``.
|
|
|
|
Returns
|
|
-------
|
|
Tensor
|
|
The output dense matrix of shape ``(N, D2)``
|
|
"""
|
|
N, D1 = F.shape(a)
|
|
R, _, D2 = F.shape(b)
|
|
if N > 1000000 or D1 > 8 or D2 > 8:
|
|
# Use segment_mm for large workload
|
|
import torch
|
|
sorted_idx_b, perm = torch.sort(idx_b)
|
|
_, rev_perm = torch.sort(perm)
|
|
sorted_a = torch.index_select(a, 0, perm)
|
|
pos_l = torch.searchsorted(sorted_idx_b, torch.arange(R, device=a.device))
|
|
pos_r = torch.cat([pos_l[1:], torch.tensor([len(idx_b)], device=a.device)])
|
|
seglen = (pos_r - pos_l).cpu() # XXX(minjie): cause device synchronize
|
|
return torch.index_select(F.segment_mm(sorted_a, b, seglen), 0, rev_perm)
|
|
else:
|
|
return F.gather_mm(a, b, None, idx_b)
|