[PinSAGE samper] Adjust the APIs for PinSAGESamper

[lixiaobai09]

Description

Compact the "to_simple", "select_topk" and "gather_row" to one API "randomwalk_topk" in class RandomWalkNeighborSampler. The new API can speed up to about 2X than old three APIs.

Checklist

Please feel free to remove inapplicable items for your PR.

• The PR title starts with [$CATEGORY] (such as [NN], [Model], [Doc], [Feature]])
• Changes are complete (i.e. I finished coding on this PR)
• All changes have test coverage
• Code is well-documented
• To the best of my knowledge, examples are either not affected by this change,
  or have been fixed to be compatible with this change
• Related issue is referred in this PR
• If the PR is for a new model/paper, I've updated the example index here.

Changes

[dgl-bot]

To trigger regression tests:

• @dgl-bot run [instance-type] [which tests] [compare-with-branch];
  For example: @dgl-bot run g4dn.4xlarge all dmlc/master or @dgl-bot run c5.9xlarge kernel,api dmlc/master

[lixiaobai09]

A simple performance test script is as follow:

import dgl
import torch
import scipy
import time

if __name__ == '__main__':
    g = scipy.sparse.random(6000, 8000, 0.3)
    G = dgl.heterograph({
            ('A', 'AB', 'B'): g.nonzero(),
            ('B', 'BA', 'A'): g.T.nonzero()
        })
    print('edge size AB and BA: {}, {}'.format(G.num_edges('AB'), G.num_edges('BA')))
    sampler = dgl.sampling.PinSAGESampler(G, 'A', 'B', 3, 0.5, 200, 10)
    seeds = torch.LongTensor([i for i in range(1000)])
    repeat = 10
    print('graph edges: {}'.format(G.num_edges()))
    t1 = time.time()
    for _ in range(repeat):
        frontier = sampler(seeds)
        print('frontier edges: {}'.format(frontier.num_edges()))
    time_cost = (time.time() - t1)
    # frontier.all_edges(form='uv')
    print('time cost: {:.5f}'.format(time_cost / repeat))

[BarclayII]

From what I understand, this is fusing to_simple, select_topk, and counting the number of occurrences together. The docstring did not show that. Could you rewrite the docstring? Also, I don't think randomwalk_topk is a suitable name since this function is not related to random walk.

[lixiaobai09]

Hello, thank you for your reply. I have modified the docstring according to your requirement. I modified these codes to optimize the random walk algorithm at first, so I just give this API with the name "randomwalk_topk". I tried to think of another name but did not get a suitable name, because of my poor English. Could you give me a suggestion about this name?

[BarclayII]

I think this function is only used for PinSAGE. I'll change it to an internal function instead.