Skip to content

GitLab

Commit 653428bd authored by Lingfan Yu's avatar Lingfan Yu Committed by Minjie Wang
Browse files

[Feature][Kernel] DGL kernel support (#596) 

* [Kernel] Minigun integration and fused kernel support (#519)

* kernel interface

* add minigun

* Add cuda build

* functors

* working on binary elewise

* binary reduce

* change kernel interface

* WIP

* wip

* fix minigun

* compile

* binary reduce kernels

* compile

* simple test passed

* more reducers

* fix thrust problem

* fix cmake

* fix cmake; add proper guard for atomic

* WIP: bcast

* WIP

* bcast kernels

* update to new minigun pass-by-value practice

* broadcasting dim

* add copy src and copy edge

* fix linking

* fix none array problem

* fix copy edge

* add device_type and device_id to backend operator

* cache csr adj, remove cache for adjmat and incmat

* custom ops in backend and pytorch impl

* change dgl-mg kernel python interface

* add id_mapping var

* clean up plus v2e spmv schedule

* spmv schedule & clean up fall back

* symbolic message and reduce func, remove bundle func

* new executors

* new backend interface for dgl kernels and pytorch impl

* minor fix

* fix

* fix docstring, comments, func names

* nodeflow

* fix message id mapping and bugs...

* pytorch test case & fix

* backward binary reduce

* fix bug

* WIP: cusparse

* change to int32 csr for cusparse workaround

* disable cusparse

* change back to int64

* broadcasting backward

* cusparse; WIP: add rev_csr

* unit test for kernels

* pytorch backward with dgl kernel

* edge softmax

* fix backward

* improve softmax

* cache edge on device

* cache mappings on device

* fix partial forward code

* cusparse done

* copy_src_sum with cusparse

* rm id getter

* reduce grad for broadcast

* copy edge reduce backward

* kernel unit test for broadcasting

* full kernel unit test

* add cpu kernels

* edge softmax unit test

* missing ref

* fix compile and small bugs

* fix bug in bcast

* Add backward both

* fix torch utests

* expose infershape

* create out tensor in python

* fix c++ lint

* [Kernel] Add GPU utest and kernel utest (#524)

* fix gpu utest

* cuda utest runnable

* temp disable test nodeflow; unified test for kernel

* cuda test kernel done

* [Kernel] Update kernel branch (#550)

* [Model] add multiprocessing training with sampling. (#484)

* reorganize sampling code.

* add multi-process training.

* speed up gcn_cv

* fix graphsage_cv.

* add new API in graph store.

* update barrier impl.

* support both local and distributed training.

* fix multiprocess train.

* fix.

* fix barrier.

* add script for loading data.

* multiprocessing sampling.

* accel training.

* replace pull with spmv for speedup.

* nodeflow copy from parent with context.

* enable GPU.

* fix a bug in graph store.

* enable multi-GPU training.

* fix lint.

* add comments.

* rename to run_store_server.py

* fix gcn_cv.

* fix a minor bug in sampler.

* handle error better in graph store.

* improve graphsage_cv for distributed mode.

* update README.

* fix.

* update.

* [Tutorial] add sampling tutorial. (#522)

* add sampling tutorial.

* add readme

* update author list.

* fix indent in the code.

* rename the file.

* update tutorial.

* fix the last API.

* update image.

* [BUGFIX] fix the problems in the sampling tutorial. (#523)

* add index.

* update.

* update tutorial.

* fix gpu utest

* cuda utest runnable

* temp disable test nodeflow; unified test for kernel

* cuda test kernel done

* Fixing typo in JTNN after interface change (#536)

* [BugFix] Fix getting src and dst id of ALL edges in NodeFlow.apply_block (#515)

* [Bug Fix] Fix inplace op at backend (#546)

* Fix inplace operation

* fix line seprator

* [Feature] Add batch and unbatch for immutable graph (#539)

* Add batch and unbatch for immutable graph

* fix line seprator

* fix lintr

* remove unnecessary include

* fix code review

* [BUGFix] Improve multi-processing training (#526)

* fix.

* add comment.

* remove.

* temp fix.

* initialize for shared memory.

* fix graphsage.

* fix gcn.

* add more unit tests.

* add more tests.

* avoid creating shared-memory exclusively.

* redefine remote initializer.

* improve initializer.

* fix unit test.

* fix lint.

* fix lint.

* initialize data in the graph store server properly.

* fix test.

* fix test.

* fix test.

* small fix.

* add comments.

* cleanup server.

* test graph store with a random port.

* print.

* print to stderr.

* test1

* test2

* remove comment.

* adjust the initializer signature.

* [API] update graph store API. (#549)

* add init_ndata and init_edata in DGLGraph.

* adjust SharedMemoryGraph API.

* print warning.

* fix comment.

* update example

* fix.

* fix examples.

* add unit tests.

* add comments.

* [Refactor] Immutable graph index (#543)

* WIP

* header

* WIP .cc

* WIP

* transpose

* wip

* immutable graph .h and .cc

* WIP: nodeflow.cc

* compile

* remove all tmp dl managed ctx; they caused refcount issue

* one simple test

* WIP: testing

* test_graph

* fix graph index

* fix bug in sampler; pass pytorch utest

* WIP on mxnet

* fix lint

* fix mxnet unittest w/ unfortunate workaround

* fix msvc

* fix lint

* SliceRows and test_nodeflow

* resolve reviews

* resolve reviews

* try fix win ci

* try fix win ci

* poke win ci again

* poke

* lazy multigraph flag; stackoverflow error

* revert node subgraph test

* lazy object

* try fix win build

* try fix win build

* poke ci

* fix build script

* fix compile

* add a todo

* fix reviews

* fix compile

* [Kernel] Update kernel branch (#576)

* [Model] add multiprocessing training with sampling. (#484)

* reorganize sampling code.

* add multi-process training.

* speed up gcn_cv

* fix graphsage_cv.

* add new API in graph store.

* update barrier impl.

* support both local and distributed training.

* fix multiprocess train.

* fix.

* fix barrier.

* add script for loading data.

* multiprocessing sampling.

* accel training.

* replace pull with spmv for speedup.

* nodeflow copy from parent with context.

* enable GPU.

* fix a bug in graph store.

* enable multi-GPU training.

* fix lint.

* add comments.

* rename to run_store_server.py

* fix gcn_cv.

* fix a minor bug in sampler.

* handle error better in graph store.

* improve graphsage_cv for distributed mode.

* update README.

* fix.

* update.

* [Tutorial] add sampling tutorial. (#522)

* add sampling tutorial.

* add readme

* update author list.

* fix indent in the code.

* rename the file.

* update tutorial.

* fix the last API.

* update image.

* [BUGFIX] fix the problems in the sampling tutorial. (#523)

* add index.

* update.

* update tutorial.

* fix gpu utest

* cuda utest runnable

* temp disable test nodeflow; unified test for kernel

* cuda test kernel done

* Fixing typo in JTNN after interface change (#536)

* [BugFix] Fix getting src and dst id of ALL edges in NodeFlow.apply_block (#515)

* [Bug Fix] Fix inplace op at backend (#546)

* Fix inplace operation

* fix line seprator

* [Feature] Add batch and unbatch for immutable graph (#539)

* Add batch and unbatch for immutable graph

* fix line seprator

* fix lintr

* remove unnecessary include

* fix code review

* [BUGFix] Improve multi-processing training (#526)

* fix.

* add comment.

* remove.

* temp fix.

* initialize for shared memory.

* fix graphsage.

* fix gcn.

* add more unit tests.

* add more tests.

* avoid creating shared-memory exclusively.

* redefine remote initializer.

* improve initializer.

* fix unit test.

* fix lint.

* fix lint.

* initialize data in the graph store server properly.

* fix test.

* fix test.

* fix test.

* small fix.

* add comments.

* cleanup server.

* test graph store with a random port.

* print.

* print to stderr.

* test1

* test2

* remove comment.

* adjust the initializer signature.

* [API] update graph store API. (#549)

* add init_ndata and init_edata in DGLGraph.

* adjust SharedMemoryGraph API.

* print warning.

* fix comment.

* update example

* fix.

* fix examples.

* add unit tests.

* add comments.

* [Refactor] Immutable graph index (#543)

* WIP

* header

* WIP .cc

* WIP

* transpose

* wip

* immutable graph .h and .cc

* WIP: nodeflow.cc

* compile

* remove all tmp dl managed ctx; they caused refcount issue

* one simple test

* WIP: testing

* test_graph

* fix graph index

* fix bug in sampler; pass pytorch utest

* WIP on mxnet

* fix lint

* fix mxnet unittest w/ unfortunate workaround

* fix msvc

* fix lint

* SliceRows and test_nodeflow

* resolve reviews

* resolve reviews

* try fix win ci

* try fix win ci

* poke win ci again

* poke

* lazy multigraph flag; stackoverflow error

* revert node subgraph test

* lazy object

* try fix win build

* try fix win build

* poke ci

* fix build script

* fix compile

* add a todo

* fix reviews

* fix compile

* all demo use python-3 (#555)

* [DEMO] Reproduce numbers of distributed training in AMLC giant graph paper (#556)

* update

* update

* update

* update num_hops

* fix bug

* update

* report numbers of distributed training in AMLC giant graph paper

* [DEMO] Remove duplicate code for sampling (#557)

* update

* update

* re-use single-machine code

* update

* use relative path

* update

* update

* update

* add __init__.py

* add __init__.py

* import sys, os

* fix typo

* update

* [Perf] Improve performance of graph store. (#554)

* fix.

* use inplace.

* move to shared memory graph store.

* fix.

* add more unit tests.

* fix.

* fix test.

* fix test.

* disable test.

* fix.

* [BUGIFX] fix a bug in edge_ids (#560)

* add test.

* fix compute.

* fix test.

* turn on test.

* fix a bug.

* add test.

* fix.

* disable test.

* [DEMO] Add Pytorch demo for distributed sampler (#562)

* update

* update

* update

* add sender

* update

* remove duplicate cpde

* [Test] Add gtest to project (#547)

* add gtest module

* add gtest

* fix

* Update CMakeLists.txt

* Update README.md

* [Perf] lazily create msg_index. (#563)

* lazily create msg_index.

* update test.

* [BUGFIX] fix bugs for running GCN on giant graphs. (#561)

* load mxnet csr.

* enable load large csr.

* fix

* fix.

* fix int overflow.

* fix test.

* [BugFix] Fix error when bfs_level = 0 in Entity Classification with RGCN (#559)

* [DEMO] Update demo of distributed sampler (#564)

* update

* update

* update demo

* add network cpp test (#565)

* Add unittest for C++ RPC (#566)

* [CI] Fix CI for cpp test (#570)

* fix CI for cpp test

* update port number

* [Docker] update docker image (#575)

* update docker image

* specify lint version

* rm torch import from unified tests

* [Kernel][Scheduler][MXNet] Scheduler for DGL kernels and MXNet backend support (#541)

* [Model] add multiprocessing training with sampling. (#484)

* reorganize sampling code.

* add multi-process training.

* speed up gcn_cv

* fix graphsage_cv.

* add new API in graph store.

* update barrier impl.

* support both local and distributed training.

* fix multiprocess train.

* fix.

* fix barrier.

* add script for loading data.

* multiprocessing sampling.

* accel training.

* replace pull with spmv for speedup.

* nodeflow copy from parent with context.

* enable GPU.

* fix a bug in graph store.

* enable multi-GPU training.

* fix lint.

* add comments.

* rename to run_store_server.py

* fix gcn_cv.

* fix a minor bug in sampler.

* handle error better in graph store.

* improve graphsage_cv for distributed mode.

* update README.

* fix.

* update.

* [Tutorial] add sampling tutorial. (#522)

* add sampling tutorial.

* add readme

* update author list.

* fix indent in the code.

* rename the file.

* update tutorial.

* fix the last API.

* update image.

* [BUGFIX] fix the problems in the sampling tutorial. (#523)

* add index.

* update.

* update tutorial.

* fix gpu utest

* cuda utest runnable

* temp disable test nodeflow; unified test for kernel

* cuda test kernel done

* edge softmax module

* WIP

* Fixing typo in JTNN after interface change (#536)

* mxnet backend support

* improve reduce grad

* add max to unittest backend

* fix kernel unittest

* [BugFix] Fix getting src and dst id of ALL edges in NodeFlow.apply_block (#515)

* lint

* lint

* win build

* [Bug Fix] Fix inplace op at backend (#546)

* Fix inplace operation

* fix line seprator

* [Feature] Add batch and unbatch for immutable graph (#539)

* Add batch and unbatch for immutable graph

* fix line seprator

* fix lintr

* remove unnecessary include

* fix code review

* [BUGFix] Improve multi-processing training (#526)

* fix.

* add comment.

* remove.

* temp fix.

* initialize for shared memory.

* fix graphsage.

* fix gcn.

* add more unit tests.

* add more tests.

* avoid creating shared-memory exclusively.

* redefine remote initializer.

* improve initializer.

* fix unit test.

* fix lint.

* fix lint.

* initialize data in the graph store server properly.

* fix test.

* fix test.

* fix test.

* small fix.

* add comments.

* cleanup server.

* test graph store with a random port.

* print.

* print to stderr.

* test1

* test2

* remove comment.

* adjust the initializer signature.

* try

* fix

* fix

* fix

* fix

* fix

* try

* test

* test

* test

* try

* try

* try

* test

* fix

* try gen_target

* fix gen_target

* fix msvc var_args expand issue

* fix

* [API] update graph store API. (#549)

* add init_ndata and init_edata in DGLGraph.

* adjust SharedMemoryGraph API.

* print warning.

* fix comment.

* update example

* fix.

* fix examples.

* add unit tests.

* add comments.

* [Refactor] Immutable graph index (#543)

* WIP

* header

* WIP .cc

* WIP

* transpose

* wip

* immutable graph .h and .cc

* WIP: nodeflow.cc

* compile

* remove all tmp dl managed ctx; they caused refcount issue

* one simple test

* WIP: testing

* test_graph

* fix graph index

* fix bug in sampler; pass pytorch utest

* WIP on mxnet

* fix lint

* fix mxnet unittest w/ unfortunate workaround

* fix msvc

* fix lint

* SliceRows and test_nodeflow

* resolve reviews

* resolve reviews

* try fix win ci

* try fix win ci

* poke win ci again

* poke

* lazy multigraph flag; stackoverflow error

* revert node subgraph test

* lazy object

* try fix win build

* try fix win build

* poke ci

* fix build script

* fix compile

* add a todo

* fix reviews

* fix compile

* WIP

* WIP

* all demo use python-3 (#555)

* ToImmutable and CopyTo

* [DEMO] Reproduce numbers of distributed training in AMLC giant graph paper (#556)

* update

* update

* update

* update num_hops

* fix bug

* update

* report numbers of distributed training in AMLC giant graph paper

* [DEMO] Remove duplicate code for sampling (#557)

* update

* update

* re-use single-machine code

* update

* use relative path

* update

* update

* update

* add __init__.py

* add __init__.py

* import sys, os

* fix typo

* update

* [Perf] Improve performance of graph store. (#554)

* fix.

* use inplace.

* move to shared memory graph store.

* fix.

* add more unit tests.

* fix.

* fix test.

* fix test.

* disable test.

* fix.

* [BUGIFX] fix a bug in edge_ids (#560)

* add test.

* fix compute.

* fix test.

* turn on test.

* fix a bug.

* add test.

* fix.

* disable test.

* DGLRetValue DGLContext conversion

* [DEMO] Add Pytorch demo for distributed sampler (#562)

* update

* update

* update

* add sender

* update

* remove duplicate cpde

* [Test] Add gtest to project (#547)

* add gtest module

* add gtest

* fix

* Update CMakeLists.txt

* Update README.md

* Add support to convert immutable graph to 32 bits

* [Perf] lazily create msg_index. (#563)

* lazily create msg_index.

* update test.

* fix binary reduce following new minigun template

* enable both int64 and int32 kernels

* [BUGFIX] fix bugs for running GCN on giant graphs. (#561)

* load mxnet csr.

* enable load large csr.

* fix

* fix.

* fix int overflow.

* fix test.

* new kernel interface done for CPU

* docstring

* rename & docstring

* copy reduce and backward

* [BugFix] Fix error when bfs_level = 0 in Entity Classification with RGCN (#559)

* [DEMO] Update demo of distributed sampler (#564)

* update

* update

* update demo

* adapt cuda kernels to the new interface

* add network cpp test (#565)

* fix bug

* Add unittest for C++ RPC (#566)

* [CI] Fix CI for cpp test (#570)

* fix CI for cpp test

* update port number

* [Docker] update docker image (#575)

* update docker image

* specify lint version

* rm torch import from unified tests

* remove pytorch-specific test_function

* fix unittest

* fix

* fix unittest backend bug in converting tensor to numpy array

* fix

* mxnet version

* [BUGFIX] fix for MXNet 1.5. (#552)

* remove clone.

* turn on numpy compatible.

* Revert "remove clone."

This reverts commit 17bbf76ed72ff178df6b3f35addc428048672457.

* revert format changes

* fix mxnet api name

* revert mistakes in previous revert

* roll back CI to 20190523 build

* fix unittest

* disable test_shared_mem_store.py for now

* remove mxnet/test_specialization.py

* sync win64 test script

* fix lowercase

* missing backend in gpu unit test

* transpose to get forward graph

* pass update all

* add sanity check

* passing test_specialization.py

* fix and pass test_function

* fix check

* fix pytorch softmax

* mxnet kernels

* c++ lint

* pylint

* try

* win build

* fix

* win

* ci enable gpu build

* init submodule recursively

* backend docstring

* try

* test win dev

* doc string

* disable pytorch test_nn

* try to fix windows issue

* bug fixed, revert changes

* [Test] fix CI. (#586)

* disable unit test in mxnet tutorial.

* retry socket connection.

* roll back to set_np_compat

* try to fix multi-processing test hangs when it fails.

* fix test.

* fix.

* doc string

* doc string and clean up

* missing field in ctypes

* fix node flow schedule and unit test

* rename

* pylint

* copy from parent default context

* fix unit test script

* fix

* demo bug in nodeflow gpu test

* [Kernel][Bugfix] fix nodeflow bug (#604)

* fix nodeflow bug

* remove debug code

* add build gtest option

* fix cmake; fix graph index bug in spmv.py

* remove clone

* fix div rhs grad bug

* [Kernel] Support full builtin method, edge softmax and unit tests (#605)

* add full builtin support

* unit test

* unit test backend

* edge softmax

* apply edge with builtin

* fix kernel unit test

* disable mxnet test_shared_mem_store

* gen builtin reduce

* enable mxnet gpu unittest

* revert some changes

* docstring

* add note for the hack

* [Kernel][Unittest][CI] Fix MXNet GPU CI (#607)

* update docker image for MXNet GPU CI

* force all dgl graph input and output on CPU

* fix gpu unittest

* speedup compilation

* add some comments

* lint

* add more comments

* fix as requested

* add some comments

* comment

* lint

* lint

* update pylint

* fix as requested

* lint

* lint

* lint

* docstrings of python DGL kernel entries

* disable lint warnings on arguments in kernel.py

* fix docstring in scheduler

* fix some bug in unittest; try again

* Revert "Merge branch 'kernel' of github.com:zzhang-cn/dgl into kernel"

This reverts commit 1d2299e68b004182ea6130b088de1f1122b18a49, reversing
changes made to ddc97fbf1bec2b7815c0da7c74f7ecb2f428889b.

* Revert "fix some bug in unittest; try again"

This reverts commit ddc97fbf1bec2b7815c0da7c74f7ecb2f428889b.

* more comprehensive kernel test

* remove shape check in test_specialization
parent da0c92a2
Hide whitespace changes
Inline Side-by-side
Showing with 663 additions and 537 deletions
src/runtime/cuda/cuda_common.h 0 → 100644
View file @ 653428bd
/*!
* Copyright (c) 2017 by Contributors
* \file cuda_common.h
* \brief Common utilities for CUDA
*/
#ifndef DGL_RUNTIME_CUDA_CUDA_COMMON_H_
#define DGL_RUNTIME_CUDA_CUDA_COMMON_H_
#include <cublas_v2.h>
#include <cusparse.h>
#include <cuda_runtime.h>
#include <dgl/runtime/packed_func.h>
#include <string>
#include "../workspace_pool.h"
namespace dgl {
namespace runtime {
#define CUDA_DRIVER_CALL(x) \
{ \
CUresult result = x; \
if (result != CUDA_SUCCESS && result != CUDA_ERROR_DEINITIALIZED) { \
const char *msg; \
cuGetErrorName(result, &msg); \
LOG(FATAL) \
<< "CUDAError: " #x " failed with error: " << msg; \
} \
}
#define CUDA_CALL(func) \
{ \
cudaError_t e = (func); \
CHECK(e == cudaSuccess || e == cudaErrorCudartUnloading) \
<< "CUDA: " << cudaGetErrorString(e); \
}
#define CUSPARSE_CALL(func) \
{ \
cusparseStatus_t e = (func); \
CHECK(e == CUSPARSE_STATUS_SUCCESS) \
<< "CUSPARSE ERROR: " << e; \
}
#define CUBLAS_CALL(func) \
{ \
cublasStatus_t e = (func); \
CHECK(e == CUBLAS_STATUS_SUCCESS) << "CUBLAS ERROR: " << e; \
}
/*! \brief Thread local workspace */
class CUDAThreadEntry {
public:
/*! \brief The cuda stream */
cudaStream_t stream{nullptr};
/*! \brief The cusparse handler */
cusparseHandle_t cusparse_handle{nullptr};
/*! \brief The cublas handler */
cublasHandle_t cublas_handle{nullptr};
/*! \brief thread local pool*/
WorkspacePool pool;
/*! \brief constructor */
CUDAThreadEntry();
// get the threadlocal workspace
static CUDAThreadEntry* ThreadLocal();
};
} // namespace runtime
} // namespace dgl
#endif // DGL_RUNTIME_CUDA_CUDA_COMMON_H_
src/runtime/cuda/cuda_device_api.cc 0 → 100644
View file @ 653428bd
/*!
* Copyright (c) 2017 by Contributors
* \file cuda_device_api.cc
* \brief GPU specific API
*/
#include <dgl/runtime/device_api.h>
#include <dmlc/thread_local.h>
#include <dgl/runtime/registry.h>
#include <cuda_runtime.h>
#include "cuda_common.h"
namespace dgl {
namespace runtime {
class CUDADeviceAPI final : public DeviceAPI {
public:
void SetDevice(DGLContext ctx) final {
CUDA_CALL(cudaSetDevice(ctx.device_id));
}
void GetAttr(DGLContext ctx, DeviceAttrKind kind, DGLRetValue* rv) final {
int value = 0;
switch (kind) {
case kExist:
value = (
cudaDeviceGetAttribute(
&value, cudaDevAttrMaxThreadsPerBlock, ctx.device_id)
== cudaSuccess);
break;
case kMaxThreadsPerBlock: {
CUDA_CALL(cudaDeviceGetAttribute(
&value, cudaDevAttrMaxThreadsPerBlock, ctx.device_id));
break;
}
case kWarpSize: {
CUDA_CALL(cudaDeviceGetAttribute(
&value, cudaDevAttrWarpSize, ctx.device_id));
break;
}
case kMaxSharedMemoryPerBlock: {
CUDA_CALL(cudaDeviceGetAttribute(
&value, cudaDevAttrMaxSharedMemoryPerBlock, ctx.device_id));
break;
}
case kComputeVersion: {
std::ostringstream os;
CUDA_CALL(cudaDeviceGetAttribute(
&value, cudaDevAttrComputeCapabilityMajor, ctx.device_id));
os << value << ".";
CUDA_CALL(cudaDeviceGetAttribute(
&value, cudaDevAttrComputeCapabilityMinor, ctx.device_id));
os << value;
*rv = os.str();
return;
}
case kDeviceName: {
cudaDeviceProp props;
CUDA_CALL(cudaGetDeviceProperties(&props, ctx.device_id));
*rv = std::string(props.name);
return;
}
case kMaxClockRate: {
CUDA_CALL(cudaDeviceGetAttribute(
&value, cudaDevAttrClockRate, ctx.device_id));
break;
}
case kMultiProcessorCount: {
CUDA_CALL(cudaDeviceGetAttribute(
&value, cudaDevAttrMultiProcessorCount, ctx.device_id));
break;
}
case kMaxThreadDimensions: {
int dims[3];
CUDA_CALL(cudaDeviceGetAttribute(
&dims[0], cudaDevAttrMaxBlockDimX, ctx.device_id));
CUDA_CALL(cudaDeviceGetAttribute(
&dims[1], cudaDevAttrMaxBlockDimY, ctx.device_id));
CUDA_CALL(cudaDeviceGetAttribute(
&dims[2], cudaDevAttrMaxBlockDimZ, ctx.device_id));
std::stringstream ss; // use json string to return multiple int values;
ss << "[" << dims[0] <<", " << dims[1] << ", " << dims[2] << "]";
*rv = ss.str();
return;
}
}
*rv = value;
}
void* AllocDataSpace(DGLContext ctx,
size_t nbytes,
size_t alignment,
DGLType type_hint) final {
CUDA_CALL(cudaSetDevice(ctx.device_id));
CHECK_EQ(256 % alignment, 0U)
<< "CUDA space is aligned at 256 bytes";
void *ret;
CUDA_CALL(cudaMalloc(&ret, nbytes));
return ret;
}
void FreeDataSpace(DGLContext ctx, void* ptr) final {
CUDA_CALL(cudaSetDevice(ctx.device_id));
CUDA_CALL(cudaFree(ptr));
}
void CopyDataFromTo(const void* from,
size_t from_offset,
void* to,
size_t to_offset,
size_t size,
DGLContext ctx_from,
DGLContext ctx_to,
DGLType type_hint,
DGLStreamHandle stream) final {
cudaStream_t cu_stream = static_cast<cudaStream_t>(stream);
from = static_cast<const char*>(from) + from_offset;
to = static_cast<char*>(to) + to_offset;
if (ctx_from.device_type == kDLGPU && ctx_to.device_type == kDLGPU) {
CUDA_CALL(cudaSetDevice(ctx_from.device_id));
if (ctx_from.device_id == ctx_to.device_id) {
GPUCopy(from, to, size, cudaMemcpyDeviceToDevice, cu_stream);
} else {
cudaMemcpyPeerAsync(to, ctx_to.device_id,
from, ctx_from.device_id,
size, cu_stream);
}
} else if (ctx_from.device_type == kDLGPU && ctx_to.device_type == kDLCPU) {
CUDA_CALL(cudaSetDevice(ctx_from.device_id));
GPUCopy(from, to, size, cudaMemcpyDeviceToHost, cu_stream);
} else if (ctx_from.device_type == kDLCPU && ctx_to.device_type == kDLGPU) {
CUDA_CALL(cudaSetDevice(ctx_to.device_id));
GPUCopy(from, to, size, cudaMemcpyHostToDevice, cu_stream);
} else {
LOG(FATAL) << "expect copy from/to GPU or between GPU";
}
}
DGLStreamHandle CreateStream(DGLContext ctx) {
CUDA_CALL(cudaSetDevice(ctx.device_id));
cudaStream_t retval;
CUDA_CALL(cudaStreamCreate(&retval));
return static_cast<DGLStreamHandle>(retval);
}
void FreeStream(DGLContext ctx, DGLStreamHandle stream) {
CUDA_CALL(cudaSetDevice(ctx.device_id));
cudaStream_t cu_stream = static_cast<cudaStream_t>(stream);
CUDA_CALL(cudaStreamDestroy(cu_stream));
}
void SyncStreamFromTo(DGLContext ctx, DGLStreamHandle event_src, DGLStreamHandle event_dst) {
CUDA_CALL(cudaSetDevice(ctx.device_id));
cudaStream_t src_stream = static_cast<cudaStream_t>(event_src);
cudaStream_t dst_stream = static_cast<cudaStream_t>(event_dst);
cudaEvent_t evt;
CUDA_CALL(cudaEventCreate(&evt));
CUDA_CALL(cudaEventRecord(evt, src_stream));
CUDA_CALL(cudaStreamWaitEvent(dst_stream, evt, 0));
CUDA_CALL(cudaEventDestroy(evt));
}
void StreamSync(DGLContext ctx, DGLStreamHandle stream) final {
CUDA_CALL(cudaSetDevice(ctx.device_id));
CUDA_CALL(cudaStreamSynchronize(static_cast<cudaStream_t>(stream)));
}
void SetStream(DGLContext ctx, DGLStreamHandle stream) final {
CUDAThreadEntry::ThreadLocal()
->stream = static_cast<cudaStream_t>(stream);
}
void* AllocWorkspace(DGLContext ctx, size_t size, DGLType type_hint) final {
return CUDAThreadEntry::ThreadLocal()->pool.AllocWorkspace(ctx, size);
}
void FreeWorkspace(DGLContext ctx, void* data) final {
CUDAThreadEntry::ThreadLocal()->pool.FreeWorkspace(ctx, data);
}
static const std::shared_ptr<CUDADeviceAPI>& Global() {
static std::shared_ptr<CUDADeviceAPI> inst =
std::make_shared<CUDADeviceAPI>();
return inst;
}
private:
static void GPUCopy(const void* from,
void* to,
size_t size,
cudaMemcpyKind kind,
cudaStream_t stream) {
if (stream != 0) {
CUDA_CALL(cudaMemcpyAsync(to, from, size, kind, stream));
} else {
CUDA_CALL(cudaMemcpy(to, from, size, kind));
}
}
};
typedef dmlc::ThreadLocalStore<CUDAThreadEntry> CUDAThreadStore;
CUDAThreadEntry::CUDAThreadEntry()
: pool(kDLGPU, CUDADeviceAPI::Global()) {
}
CUDAThreadEntry* CUDAThreadEntry::ThreadLocal() {
return CUDAThreadStore::Get();
}
DGL_REGISTER_GLOBAL("device_api.gpu")
.set_body([](DGLArgs args, DGLRetValue* rv) {
DeviceAPI* ptr = CUDADeviceAPI::Global().get();
*rv = static_cast<void*>(ptr);
});
} // namespace runtime
} // namespace dgl
tests/backend/__init__.py
View file @ 653428bd
......@@ -30,7 +30,7 @@ _softmax = softmax
_default_context_str = os.getenv('DGLTESTDEV', 'cpu')
_context_dict = {
'cpu': cpu(),
'cuda': cuda(),
'gpu': cuda(),
}
_default_context = _context_dict[_default_context_str]
......@@ -45,7 +45,10 @@ def randn(shape):
def tensor(data, dtype=None):
if dtype is None:
data = np.array(data)
if is_tensor(data):
data = zerocopy_to_numpy(data)
else:
data = np.array(data)
dtype = int64 if np.issubdtype(data.dtype, np.integer) else float32
return copy_to(_tensor(data, dtype), _default_context)
......@@ -59,4 +62,4 @@ def full_1d(length, fill_value, dtype, ctx=_default_context):
return _full_1d(length, fill_value, dtype, ctx)
def softmax(x, dim):
return _softmax(x, dim)
\ No newline at end of file
return _softmax(x, dim)
tests/backend/backend_unittest.py
View file @ 653428bd
......@@ -37,7 +37,7 @@ def attach_grad(x):
def backward(x, head_gradient=None):
"""Invoke backward computation with an optional head gradient.
Returns nothing."""
pass
......@@ -71,6 +71,41 @@ def softmax(x, dim):
"""Softmax Operation on Tensors"""
pass
def spmm(x, y):
"""Sparse dense matrix multiply"""
pass
def add(a, b):
"""Compute a + b"""
pass
def sub(a, b):
"""Compute a - b"""
pass
def mul(a, b):
"""Compute a * b"""
pass
def div(a, b):
"""Compute a / b"""
pass
def sum(x, dim):
"""Computes the sum of array elements over given axes"""
pass
def max(x, dim):
"""Computes the max of array elements over given axes"""
pass
def min(x, dim):
"""Computes the min of array elements over given axes"""
pass
def prod(x, dim):
"""Computes the prod of array elements over given axes"""
pass
###############################################################################
# Tensor functions used *only* on index tensor
# ----------------
......
tests/backend/mxnet/__init__.py
View file @ 653428bd
......@@ -48,6 +48,33 @@ def reduce_sum(x):
def softmax(x, dim):
return nd.softmax(x, dim)
def spmm(x, y):
return nd.dot(x, y)
def add(a, b):
return a + b
def sub(a, b):
return a - b
def mul(a, b):
return a * b
def div(a, b):
return a / b
def sum(x, dim):
return x.sum(dim)
def max(x, dim):
return x.max(dim)
def min(x, dim):
return x.min(dim)
def prod(x, dim):
return x.prod(dim)
record_grad = autograd.record
......
tests/backend/pytorch/__init__.py
View file @ 653428bd
......@@ -3,13 +3,14 @@ from __future__ import absolute_import
import torch as th
def cuda():
return th.device('cuda')
return th.device('cuda:0')
def array_equal(a, b):
return th.equal(a, b)
return th.equal(a.cpu(), b.cpu())
def allclose(a, b):
return th.allclose(a.float(), b.float(), rtol=1e-4, atol=1e-4)
return th.allclose(a.float().cpu(),
b.float().cpu(), rtol=1e-4, atol=1e-4)
def randn(shape):
return th.randn(*shape)
......@@ -48,6 +49,33 @@ def reduce_sum(x):
def softmax(x, dim):
return th.softmax(x, dim)
def spmm(x, y):
return th.spmm(x, y)
def add(a, b):
return a + b
def sub(a, b):
return a - b
def mul(a, b):
return a * b
def div(a, b):
return a / b
def sum(x, dim):
return x.sum(dim)
def max(x, dim):
return x.max(dim)[0]
def min(x, dim):
return x.min(dim)[0]
def prod(x, dim):
return x.prod(dim)
class record_grad(object):
def __init__(self):
pass
......
tests/compute/test_basics.py
View file @ 653428bd
......@@ -200,7 +200,7 @@ def test_nx_conversion():
assert F.allclose(g.ndata['n1'], n1)
# with id in nx edge feature, e1 should follow original order
assert F.allclose(g.edata['e1'], e1)
assert F.array_equal(g.get_e_repr()['id'], F.arange(0, 4))
assert F.array_equal(g.get_e_repr()['id'], F.copy_to(F.arange(0, 4), F.cpu()))
# test conversion after modifying DGLGraph
g.pop_e_repr('id') # pop id so we don't need to provide id when adding edges
......@@ -314,7 +314,7 @@ def test_apply_edges():
u = F.tensor([0, 0, 0, 4, 5, 6])
v = F.tensor([1, 2, 3, 9, 9, 9])
g.apply_edges(lambda edges : {'w' : edges.data['w'] * 0.}, (u, v))
eid = g.edge_ids(u, v)
eid = F.tensor(g.edge_ids(u, v))
assert F.allclose(F.gather_row(g.edata['w'], eid), F.zeros((6, D)))
def test_update_routines():
......@@ -643,8 +643,8 @@ def test_group_apply_edges():
u, v, eid = g.out_edges(1, form='all')
else:
u, v, eid = g.in_edges(5, form='all')
out_feat = g.edata['norm_feat'][eid]
result = (g.ndata['h'][u] + g.ndata['h'][v]) * g.edata['feat'][eid]
out_feat = g.edges[eid].data['norm_feat']
result = (g.nodes[u].data['h'] + g.nodes[v].data['h']) * g.edges[eid].data['feat']
result = F.softmax(F.sum(result, dim=1), dim=0)
assert F.allclose(out_feat, result)
......
tests/compute/test_dis_sampler.py
View file @ 653428bd
from dgl import backend as F
import backend as F
import numpy as np
import scipy as sp
import dgl
......
tests/compute/test_frame.py
View file @ 653428bd
......@@ -113,7 +113,7 @@ def test_append2():
assert not f.is_span_whole_column()
assert f.num_rows == 3 * N
new_idx = list(range(N)) + list(range(2*N, 4*N))
assert F.array_equal(f._index.tousertensor(), F.tensor(new_idx, dtype=F.int64))
assert F.array_equal(f._index.tousertensor(), F.copy_to(F.tensor(new_idx, dtype=F.int64), F.cpu()))
assert data.num_rows == 4 * N
def test_append3():
......@@ -144,13 +144,13 @@ def test_row1():
rows = f[rowid]
for k, v in rows.items():
assert tuple(F.shape(v)) == (len(rowid), D)
assert F.allclose(v, F.gather_row(data[k], rowid.tousertensor()))
assert F.allclose(v, F.gather_row(data[k], F.tensor(rowid.tousertensor())))
# test duplicate keys
rowid = Index(F.tensor([8, 2, 2, 1]))
rows = f[rowid]
for k, v in rows.items():
assert tuple(F.shape(v)) == (len(rowid), D)
assert F.allclose(v, F.gather_row(data[k], rowid.tousertensor()))
assert F.allclose(v, F.gather_row(data[k], F.tensor(rowid.tousertensor())))
# setter
rowid = Index(F.tensor([0, 2, 4]))
......@@ -282,7 +282,7 @@ def test_slicing():
'a3': F.zeros([2, D]),
}
assert F.allclose(f2['a1'], f2_a1)
f1[Index(F.tensor([2, 3]))] = {
'a1': F.ones([2, D]),
'a2': F.ones([2, D]),
......
tests/compute/test_graph.py
View file @ 653428bd
import time
import math
import numpy as np
import scipy.sparse as sp
......@@ -308,7 +307,7 @@ def test_readonly():
assert g.number_of_edges() == 4
g.readonly()
assert g._graph.is_readonly() == True
assert g._graph.is_readonly() == True
assert g.number_of_nodes() == 5
assert g.number_of_edges() == 4
......@@ -321,7 +320,7 @@ def test_readonly():
assert fail
g.readonly()
assert g._graph.is_readonly() == True
assert g._graph.is_readonly() == True
assert g.number_of_nodes() == 5
assert g.number_of_edges() == 4
......
tests/compute/test_kernel.py 0 → 100644
View file @ 653428bd
import dgl
import dgl.function as fn
import networkx as nx
import backend as F
from itertools import product
def udf_copy_src(edges):
return {'m': edges.src['u']}
def udf_copy_edge(edges):
return {'m': edges.data['e']}
def udf_sum(nodes):
return {'r2': nodes.mailbox['m'].sum(1)}
def udf_max(nodes):
return {'r2': F.max(nodes.mailbox['m'], 1)}
D1 = 5
D2 = 3
D3 = 4
builtin = {'sum': fn.sum, 'max': fn.max}
udf_reduce = {'sum': udf_sum, 'max': udf_max}
fill_value = {'sum': 0, 'max': float("-inf")}
def generate_feature(g, broadcast='none'):
"""Create graph with src, edge, dst feature. broadcast can be 'u',
'e', 'v', 'none'
"""
nv = g.number_of_nodes()
ne = g.number_of_edges()
if broadcast == 'e':
u = F.randn((nv, D1, D2, D3))
e = F.randn((ne, D2, 1))
v = F.randn((nv, D1, D2, D3))
elif broadcast == 'u':
u = F.randn((nv, D2, 1))
e = F.randn((ne, D1, D2, D3))
v = F.randn((nv, D1, D2, D3))
elif broadcast == 'v':
u = F.randn((nv, D1, D2, D3))
e = F.randn((ne, D1, D2, D3))
v = F.randn((nv, D2, 1))
else:
u = F.randn((nv, D1, D2, D3))
e = F.randn((ne, D1, D2, D3))
v = F.randn((nv, D1, D2, D3))
return u, v, e
def test_copy_src_reduce():
def _test(red):
g = dgl.DGLGraph(nx.erdos_renyi_graph(100, 0.1))
hu, hv, he = generate_feature(g, 'none')
g.ndata['u'] = F.attach_grad(F.clone(hu))
g.ndata['v'] = F.attach_grad(F.clone(hv))
g.edata['e'] = F.attach_grad(F.clone(he))
with F.record_grad():
g.update_all(fn.copy_src(src='u', out='m'),
builtin[red](msg='m', out='r1'))
r1 = g.ndata['r1']
F.backward(r1.sum())
n_grad1 = F.grad(g.ndata['u'])
# reset grad
g.ndata['u'] = F.attach_grad(F.clone(hu))
g.ndata['v'] = F.attach_grad(F.clone(hv))
g.edata['e'] = F.attach_grad(F.clone(he))
with F.record_grad():
g.update_all(udf_copy_src, udf_reduce[red])
r2 = g.ndata['r2']
F.backward(r2.sum())
n_grad2 = F.grad(g.ndata['u'])
assert F.allclose(r1, r2)
assert(F.allclose(n_grad1, n_grad2))
_test('sum')
_test('max')
def test_copy_edge_reduce():
def _test(red):
g = dgl.DGLGraph(nx.erdos_renyi_graph(100, 0.1))
hu, hv, he = generate_feature(g, 'none')
g.ndata['u'] = F.attach_grad(F.clone(hu))
g.ndata['v'] = F.attach_grad(F.clone(hv))
g.edata['e'] = F.attach_grad(F.clone(he))
with F.record_grad():
g.update_all(fn.copy_edge(edge='e', out='m'),
builtin[red](msg='m', out='r1'))
r1 = g.ndata['r1']
F.backward(r1.sum())
e_grad1 = F.grad(g.edata['e'])
# reset grad
g.ndata['u'] = F.attach_grad(F.clone(hu))
g.ndata['v'] = F.attach_grad(F.clone(hv))
g.edata['e'] = F.attach_grad(F.clone(he))
with F.record_grad():
g.update_all(udf_copy_edge, udf_reduce[red])
r2 = g.ndata['r2']
F.backward(r2.sum())
e_grad2 = F.grad(g.edata['e'])
assert F.allclose(r1, r2)
assert(F.allclose(e_grad1, e_grad2))
_test('sum')
_test('max')
def test_all_binary_builtins():
def _test(g, lhs, rhs, binary_op, reducer, broadcast='none'):
hu, hv, he = generate_feature(g, broadcast)
g.ndata['u'] = F.attach_grad(F.clone(hu))
g.ndata['v'] = F.attach_grad(F.clone(hv))
g.edata['e'] = F.attach_grad(F.clone(he))
builtin_msg_name = "{}_{}_{}".format(lhs, binary_op, rhs)
builtin_msg = getattr(fn, builtin_msg_name)
builtin_red = getattr(fn, reducer)
def target_feature_switch(g, target):
if target == "u":
return g.ndata["u"]
elif target == "v":
return g.ndata["v"]
else:
return g.edata["e"]
with F.record_grad():
g.update_all(builtin_msg(lhs, rhs, 'm'), builtin_red('m', 'r1'))
r1 = g.ndata['r1']
F.backward(r1.sum())
lhs_grad_1 = F.grad(target_feature_switch(g, lhs))
rhs_grad_1 = F.grad(target_feature_switch(g, rhs))
# reset grad
g.ndata['u'] = F.attach_grad(F.clone(hu))
g.ndata['v'] = F.attach_grad(F.clone(hv))
g.edata['e'] = F.attach_grad(F.clone(he))
def target_switch(edges, target):
if target == "u":
return edges.src
elif target == "v":
return edges.dst
elif target == "e":
return edges.data
else:
assert(0), "Unknown target {}".format(target)
def mfunc(edges):
op = getattr(F, binary_op)
lhs_data = target_switch(edges, lhs)
rhs_data = target_switch(edges, rhs)
return {"m": op(lhs_data[lhs], rhs_data[rhs])}
def rfunc(nodes):
op = getattr(F, reducer)
return {"r2": op(nodes.mailbox['m'], 1)}
with F.record_grad():
g.update_all(mfunc, rfunc)
r2 = g.ndata['r2']
F.backward(r2.sum())
lhs_grad_2 = F.grad(target_feature_switch(g, lhs))
rhs_grad_2 = F.grad(target_feature_switch(g, rhs))
def _print_error(a, b):
print("Test {}_{}_{}_{} {}".
format(lhs, binary_op, rhs, reducer, broadcast))
print(a)
print(b)
if not F.allclose(r1, r2):
_print_error(r1, r2)
assert F.allclose(r1, r2)
if not F.allclose(rhs_grad_1, rhs_grad_2):
print("left grad")
_print_error(lhs_grad_1, lhs_grad_2)
assert(F.allclose(lhs_grad_1, lhs_grad_2))
if not F.allclose(rhs_grad_1, rhs_grad_2):
print("right grad")
_print_error(rhs_grad_1, rhs_grad_2)
assert(F.allclose(rhs_grad_1, rhs_grad_2))
g = dgl.DGLGraph()
g.add_nodes(20)
for i in range(2, 18):
g.add_edge(0, i)
g.add_edge(1, i)
g.add_edge(i, 18)
g.add_edge(i, 19)
g.add_edge(18, 0)
g.add_edge(18, 1)
g.add_edge(19, 0)
g.add_edge(19, 1)
target = ["u", "v", "e"]
for lhs, rhs in product(target, target):
if lhs == rhs:
continue
for binary_op in ["add", "sub", "mul", "div"]:
for reducer in ["sum", "max", "min", "prod"]:
for broadcast in ["none", lhs, rhs]:
_test(g, lhs, rhs, binary_op, reducer)
if __name__ == '__main__':
test_copy_src_reduce()
test_copy_edge_reduce()
test_all_binary_builtins()
tests/compute/test_multi_send_recv.py
View file @ 653428bd
......@@ -60,7 +60,7 @@ def test_multi_send():
g.send((u, v))
# check if message indicator is as expected
expected = F.zeros((g.number_of_edges(),), dtype=F.int64)
expected = F.copy_to(F.zeros((g.number_of_edges(),), dtype=F.int64), F.cpu())
eid = g.edge_ids([0, 0, 0, 0, 0, 1, 2, 3, 4, 5],
[1, 2, 3, 4, 5, 9, 9, 9, 9, 9])
expected[eid] = 1
......@@ -73,7 +73,7 @@ def test_multi_recv():
g.register_message_func(message_func)
g.register_reduce_func(reduce_func)
g.register_apply_node_func(apply_node_func)
expected = F.zeros((g.number_of_edges(),), dtype=F.int64)
expected = F.copy_to(F.zeros((g.number_of_edges(),), dtype=F.int64), F.cpu())
# two separate round of send and recv
u = [4, 5, 6]
v = [9]
......@@ -249,7 +249,7 @@ def test_dynamic_addition():
g.edata.update({'h1': F.randn((2, D)),
'h2': F.randn((2, D))})
g.send()
expected = F.ones((g.number_of_edges(),), dtype=F.int64)
expected = F.copy_to(F.ones((g.number_of_edges(),), dtype=F.int64), F.cpu())
assert F.array_equal(g._get_msg_index().tousertensor(), expected)
# add more edges
......@@ -279,7 +279,7 @@ def test_recv_no_send():
g.set_n_initializer(dgl.init.zero_initializer)
g.ndata['h'] = F.randn((3, D))
g.send((1, 2), message_func)
expected = F.zeros((2,), dtype=F.int64)
expected = F.copy_to(F.zeros(2, dtype=F.int64), F.cpu())
expected[1] = 1
assert F.array_equal(g._get_msg_index().tousertensor(), expected)
g.recv(2, reduce_func)
......
tests/compute/test_nodeflow.py
View file @ 653428bd
......@@ -35,7 +35,7 @@ def test_self_loop():
nf = create_mini_batch(g, num_hops, add_self_loop=True)
for i in range(1, nf.num_layers):
in_deg = nf.layer_in_degree(i)
deg = F.ones(in_deg.shape, dtype=F.int64) * n
deg = F.copy_to(F.ones(in_deg.shape, dtype=F.int64), F.cpu()) * n
assert F.array_equal(in_deg, deg)
def create_mini_batch(g, num_hops, add_self_loop=False):
......@@ -57,9 +57,9 @@ def check_basic(g, nf):
assert nf.number_of_edges() == num_edges
deg = nf.layer_in_degree(0)
assert F.array_equal(deg, F.zeros((nf.layer_size(0)), F.int64))
assert F.array_equal(deg, F.copy_to(F.zeros((nf.layer_size(0)), F.int64), F.cpu()))
deg = nf.layer_out_degree(-1)
assert F.array_equal(deg, F.zeros((nf.layer_size(-1)), F.int64))
assert F.array_equal(deg, F.copy_to(F.zeros((nf.layer_size(-1)), F.int64), F.cpu()))
for i in range(1, nf.num_layers):
in_deg = nf.layer_in_degree(i)
out_deg = nf.layer_out_degree(i - 1)
......@@ -77,7 +77,7 @@ def test_basic():
assert nf.layer_size(1) == g.number_of_nodes()
check_basic(g, nf)
parent_nids = F.arange(0, g.number_of_nodes())
parent_nids = F.copy_to(F.arange(0, g.number_of_nodes()), F.cpu())
nids = nf.map_from_parent_nid(0, parent_nids)
assert F.array_equal(nids, parent_nids)
......@@ -138,7 +138,7 @@ def check_apply_edges(create_node_flow):
eids = nf.block_parent_eid(i)
srcs, dsts = g.find_edges(eids)
expected_f_sum = g.ndata["f"][srcs] + g.ndata["f"][dsts]
expected_f_sum = g.nodes[srcs].data["f"] + g.nodes[dsts].data["f"]
assert F.array_equal(nf.blocks[i].data['f2'], expected_f_sum)
......@@ -161,7 +161,7 @@ def check_flow_compute(create_node_flow, use_negative_block_id=False):
lambda nodes: {'h' : nodes.data['t'] + 1})
g.update_all(fn.copy_src(src='h', out='m'), fn.sum(msg='m', out='t'),
lambda nodes: {'h' : nodes.data['t'] + 1})
assert F.array_equal(nf.layers[i + 1].data['h'], g.ndata['h'][nf.layer_parent_nid(i + 1)])
assert F.allclose(nf.layers[i + 1].data['h'], g.nodes[nf.layer_parent_nid(i + 1)].data['h'])
# Test the computation when only a few nodes are active in a layer.
g.ndata['h'] = g.ndata['h1']
......@@ -173,8 +173,8 @@ def check_flow_compute(create_node_flow, use_negative_block_id=False):
g.update_all(fn.copy_src(src='h', out='m'), fn.sum(msg='m', out='t'),
lambda nodes: {'h' : nodes.data['t'] + 1})
data1 = nf.layers[i + 1].data['h'][0:4]
data2 = g.ndata['h'][nf.map_to_parent_nid(vs)]
assert F.array_equal(data1, data2)
data2 = g.nodes[nf.map_to_parent_nid(vs)].data['h']
assert F.allclose(data1, data2)
def test_flow_compute():
......@@ -198,7 +198,7 @@ def check_prop_flows(create_node_flow):
# Test the computation on all layers.
nf2.prop_flow(fn.copy_src(src='h', out='m'), fn.sum(msg='m', out='t'),
lambda nodes: {'h' : nodes.data['t'] + 1})
assert F.array_equal(nf2.layers[-1].data['h'], g.ndata['h'][nf2.layer_parent_nid(-1)])
assert F.allclose(nf2.layers[-1].data['h'], g.nodes[nf2.layer_parent_nid(-1)].data['h'])
def test_prop_flows():
......@@ -216,12 +216,12 @@ def test_copy():
assert len(g.ndata.keys()) == len(nf.layers[i].data.keys())
for key in g.ndata.keys():
assert key in nf.layers[i].data.keys()
assert F.array_equal(nf.layers[i].data[key], g.ndata[key][nf.layer_parent_nid(i)])
assert F.array_equal(nf.layers[i].data[key], g.nodes[nf.layer_parent_nid(i)].data[key])
for i in range(nf.num_blocks):
assert len(g.edata.keys()) == len(nf.blocks[i].data.keys())
for key in g.edata.keys():
assert key in nf.blocks[i].data.keys()
assert F.array_equal(nf.blocks[i].data[key], g.edata[key][nf.block_parent_eid(i)])
assert F.array_equal(nf.blocks[i].data[key], g.edges[nf.block_parent_eid(i)].data[key])
nf = create_mini_batch(g, num_layers)
node_embed_names = [['h'], ['h1'], ['h']]
......@@ -231,12 +231,12 @@ def test_copy():
assert len(node_embed_names[i]) == len(nf.layers[i].data.keys())
for key in node_embed_names[i]:
assert key in nf.layers[i].data.keys()
assert F.array_equal(nf.layers[i].data[key], g.ndata[key][nf.layer_parent_nid(i)])
assert F.array_equal(nf.layers[i].data[key], g.nodes[nf.layer_parent_nid(i)].data[key])
for i in range(nf.num_blocks):
assert len(edge_embed_names[i]) == len(nf.blocks[i].data.keys())
for key in edge_embed_names[i]:
assert key in nf.blocks[i].data.keys()
assert F.array_equal(nf.blocks[i].data[key], g.edata[key][nf.block_parent_eid(i)])
assert F.array_equal(nf.blocks[i].data[key], g.edges[nf.block_parent_eid(i)].data[key])
nf = create_mini_batch(g, num_layers)
g.ndata['h0'] = F.clone(g.ndata['h'])
......@@ -247,12 +247,12 @@ def test_copy():
lambda nodes: {'h%d' % (i+1) : nodes.data['t'] + 1})
g.update_all(fn.copy_src(src='h', out='m'), fn.sum(msg='m', out='t'),
lambda nodes: {'h' : nodes.data['t'] + 1})
assert F.array_equal(nf.layers[i + 1].data['h%d' % (i+1)],
g.ndata['h'][nf.layer_parent_nid(i + 1)])
assert F.allclose(nf.layers[i + 1].data['h%d' % (i+1)],
g.nodes[nf.layer_parent_nid(i + 1)].data['h'])
nf.copy_to_parent(node_embed_names=[['h0'], ['h1'], ['h2']])
for i in range(num_layers + 1):
assert F.array_equal(nf.layers[i].data['h%d' % i],
g.ndata['h%d' % i][nf.layer_parent_nid(i)])
g.nodes[nf.layer_parent_nid(i)].data['h%d' % i])
nf = create_mini_batch(g, num_layers)
g.ndata['h0'] = F.clone(g.ndata['h'])
......@@ -278,10 +278,10 @@ def test_block_edges():
nf = create_mini_batch(g, num_layers)
assert nf.num_layers == num_layers + 1
for i in range(nf.num_blocks):
src, dst, eid = nf.block_edges(i)
src, dst, eid = nf.block_edges(i, remap=True)
# should also work for negative block ids
src_by_neg, dst_by_neg, eid_by_neg = nf.block_edges(-nf.num_blocks + i)
src_by_neg, dst_by_neg, eid_by_neg = nf.block_edges(-nf.num_blocks + i, remap=True)
assert F.array_equal(src, src_by_neg)
assert F.array_equal(dst, dst_by_neg)
assert F.array_equal(eid, eid_by_neg)
......@@ -300,7 +300,7 @@ def test_block_adj_matrix():
nf = create_mini_batch(g, num_layers)
assert nf.num_layers == num_layers + 1
for i in range(nf.num_blocks):
u, v, _ = nf.block_edges(i)
u, v, _ = nf.block_edges(i, remap=True)
adj, _ = nf.block_adjacency_matrix(i, F.cpu())
adj = F.sparse_to_numpy(adj)
......@@ -337,7 +337,7 @@ def test_block_incidence_matrix():
adj_by_neg = F.sparse_to_numpy(adj_by_neg)
adjs_by_neg.append(adj_by_neg)
u, v, e = nf.block_edges(i)
u, v, e = nf.block_edges(i, remap=True)
u = utils.toindex(u)
v = utils.toindex(v)
e = utils.toindex(e)
......@@ -367,4 +367,4 @@ if __name__ == '__main__':
test_prop_flows()
test_self_loop()
test_block_edges()
test_block_incidence_matrix()
\ No newline at end of file
test_block_incidence_matrix()
tests/compute/test_pickle.py
View file @ 653428bd
......@@ -148,7 +148,6 @@ def test_pickling_graph():
assert new_g._message_func == _global_message_func
assert isinstance(new_g._reduce_func, type(reduce_func))
assert new_g._reduce_func._name == 'sum'
assert new_g._reduce_func.reduce_op == F.sum
assert new_g._reduce_func.msg_field == 'x'
assert new_g._reduce_func.out_field == 'x'
......
tests/compute/test_sampler.py
View file @ 653428bd
......@@ -106,7 +106,7 @@ def test_10neighbor_sampler():
check_10neighbor_sampler(g, seeds=np.unique(np.random.randint(0, g.number_of_nodes(),
size=int(g.number_of_nodes() / 10))))
def test_layer_sampler(prefetch=False):
def _test_layer_sampler(prefetch=False):
g = generate_rand_graph(100)
nid = g.nodes()
src, dst, eid = g.all_edges(form='all', order='eid')
......@@ -157,5 +157,5 @@ if __name__ == '__main__':
test_10neighbor_sampler_all()
test_1neighbor_sampler()
test_10neighbor_sampler()
test_layer_sampler()
test_layer_sampler(prefetch=True)
#test_layer_sampler()
#test_layer_sampler(prefetch=True)
tests/compute/test_specialization.py
View file @ 653428bd
......@@ -51,14 +51,9 @@ def test_v2v_update_all():
fn.sum(msg='m', out=fld), apply_func)
v2 = g.ndata[fld]
g.set_n_repr({fld : v1})
g.update_all(fn.src_mul_edge(src=fld, edge='e2', out='m'),
fn.sum(msg='m', out=fld), apply_func)
v3 = g.ndata[fld]
g.set_n_repr({fld : v1})
g.update_all(message_func_edge, reduce_func, apply_func)
v4 = g.ndata[fld]
assert F.allclose(v2, v3)
assert F.allclose(v3, v4)
assert F.allclose(v2, v4)
# test 1d node features
_test('f1')
# test 2d node features
......@@ -98,14 +93,9 @@ def test_v2v_snr():
fn.sum(msg='m', out=fld), apply_func)
v2 = g.ndata[fld]
g.set_n_repr({fld : v1})
g.send_and_recv((u, v), fn.src_mul_edge(src=fld, edge='e2', out='m'),
fn.sum(msg='m', out=fld), apply_func)
v3 = g.ndata[fld]
g.set_n_repr({fld : v1})
g.send_and_recv((u, v), message_func_edge, reduce_func, apply_func)
v4 = g.ndata[fld]
assert F.allclose(v2, v3)
assert F.allclose(v3, v4)
assert F.allclose(v2, v4)
# test 1d node features
_test('f1')
# test 2d node features
......@@ -141,17 +131,12 @@ def test_v2v_pull():
# send and recv with edge weights
v1 = g.ndata[fld]
g.pull(nodes, fn.src_mul_edge(src=fld, edge='e1', out='m'),
fn.sum(msg='m', out=fld), apply_func)
fn.sum(msg='m', out=fld), apply_func)
v2 = g.ndata[fld]
g.ndata[fld] = v1
g.pull(nodes, fn.src_mul_edge(src=fld, edge='e2', out='m'),
fn.sum(msg='m', out=fld), apply_func)
v3 = g.ndata[fld]
g.ndata[fld] = v1
g.pull(nodes, message_func_edge, reduce_func, apply_func)
v4 = g.ndata[fld]
assert F.allclose(v2, v3)
assert F.allclose(v3, v4)
assert F.allclose(v2, v4)
# test 1d node features
_test('f1')
# test 2d node features
......@@ -401,11 +386,6 @@ def test_update_all_multi_fallback():
fn.sum(msg='m2', out='o2'),
_afunc)
assert F.allclose(o2, g.ndata.pop('o2'))
# v2v fallback to degree bucketing
g.update_all(fn.src_mul_edge(src='h', edge='w1', out='m1'),
fn.max(msg='m1', out='o3'),
_afunc)
assert F.allclose(o3, g.ndata.pop('o3'))
# multi builtins, both v2v spmv
g.update_all([fn.src_mul_edge(src='h', edge='w1', out='m1'), fn.src_mul_edge(src='h', edge='w1', out='m2')],
[fn.sum(msg='m1', out='o1'), fn.sum(msg='m2', out='o2')],
......@@ -418,18 +398,6 @@ def test_update_all_multi_fallback():
_afunc)
assert F.allclose(o1, g.ndata.pop('o1'))
assert F.allclose(o2, g.ndata.pop('o2'))
# multi builtins, one v2v spmv, one fallback to e2v, one fallback to degree-bucketing
g.update_all([fn.src_mul_edge(src='h', edge='w1', out='m1'),
fn.src_mul_edge(src='h', edge='w2', out='m2'),
fn.src_mul_edge(src='h', edge='w1', out='m3')],
[fn.sum(msg='m1', out='o1'),
fn.sum(msg='m2', out='o2'),
fn.max(msg='m3', out='o3')],
_afunc)
assert F.allclose(o1, g.ndata.pop('o1'))
assert F.allclose(o2, g.ndata.pop('o2'))
assert F.allclose(o3, g.ndata.pop('o3'))
def test_pull_multi_fallback():
# create a graph with zero in degree nodes
......@@ -476,11 +444,6 @@ def test_pull_multi_fallback():
fn.sum(msg='m2', out='o2'),
_afunc)
assert F.allclose(o2, g.ndata.pop('o2'))
# v2v fallback to degree bucketing
g.pull(nodes, fn.src_mul_edge(src='h', edge='w1', out='m1'),
fn.max(msg='m1', out='o3'),
_afunc)
assert F.allclose(o3, g.ndata.pop('o3'))
# multi builtins, both v2v spmv
g.pull(nodes,
[fn.src_mul_edge(src='h', edge='w1', out='m1'), fn.src_mul_edge(src='h', edge='w1', out='m2')],
......@@ -495,18 +458,6 @@ def test_pull_multi_fallback():
_afunc)
assert F.allclose(o1, g.ndata.pop('o1'))
assert F.allclose(o2, g.ndata.pop('o2'))
# multi builtins, one v2v spmv, one fallback to e2v, one fallback to degree-bucketing
g.pull(nodes,
[fn.src_mul_edge(src='h', edge='w1', out='m1'),
fn.src_mul_edge(src='h', edge='w2', out='m2'),
fn.src_mul_edge(src='h', edge='w1', out='m3')],
[fn.sum(msg='m1', out='o1'),
fn.sum(msg='m2', out='o2'),
fn.max(msg='m3', out='o3')],
_afunc)
assert F.allclose(o1, g.ndata.pop('o1'))
assert F.allclose(o2, g.ndata.pop('o2'))
assert F.allclose(o3, g.ndata.pop('o3'))
# test#1: non-0deg nodes
nodes = [1, 2, 9]
_pull_nodes(nodes)
......
tests/compute/test_subgraph.py
View file @ 653428bd
......@@ -30,7 +30,7 @@ def test_basics():
sg = g.subgraph(nid)
eid = {2, 3, 4, 5, 10, 11, 12, 13, 16}
assert set(F.zerocopy_to_numpy(sg.parent_eid)) == eid
eid = sg.parent_eid
eid = F.tensor(sg.parent_eid)
# the subgraph is empty initially
assert len(sg.ndata) == 0
assert len(sg.edata) == 0
......
tests/compute/test_traversal.py
View file @ 653428bd
......@@ -15,7 +15,7 @@ np.random.seed(42)
def toset(x):
return set(F.zerocopy_to_numpy(x).tolist())
def test_bfs(n=1000):
def test_bfs(n=100):
def _bfs_nx(g_nx, src):
edges = nx.bfs_edges(g_nx, src)
layers_nx = [set([src])]
......@@ -31,14 +31,12 @@ def test_bfs(n=1000):
edges_nx.append(edge_frontier)
frontier = set([v])
edge_frontier = set([g.edge_id(u, v)])
# avoids case of no successors
if len(frontier) > 0 and len(edge_frontier) > 0:
layers_nx.append(frontier)
edges_nx.append(edge_frontier)
layers_nx.append(frontier)
edges_nx.append(edge_frontier)
return layers_nx, edges_nx
g = dgl.DGLGraph()
a = sp.random(n, n, 10 / n, data_rvs=lambda n: np.ones(n))
a = sp.random(n, n, 3 / n, data_rvs=lambda n: np.ones(n))
g.from_scipy_sparse_matrix(a)
g_nx = g.to_networkx()
src = random.choice(range(n))
......@@ -56,9 +54,9 @@ def test_bfs(n=1000):
assert len(edges_dgl) == len(edges_nx)
assert all(toset(x) == y for x, y in zip(edges_dgl, edges_nx))
def test_topological_nodes(n=1000):
def test_topological_nodes(n=100):
g = dgl.DGLGraph()
a = sp.random(n, n, 10 / n, data_rvs=lambda n: np.ones(n))
a = sp.random(n, n, 3 / n, data_rvs=lambda n: np.ones(n))
b = sp.tril(a, -1).tocoo()
g.from_scipy_sparse_matrix(b)
......@@ -67,13 +65,13 @@ def test_topological_nodes(n=1000):
adjmat = g.adjacency_matrix()
def tensor_topo_traverse():
n = g.number_of_nodes()
mask = F.ones((n, 1))
mask = F.copy_to(F.ones((n, 1)), F.cpu())
degree = F.spmm(adjmat, mask)
while F.reduce_sum(mask) != 0.:
v = F.astype((degree == 0.), F.float32)
v = v * mask
mask = mask - v
frontier = F.nonzero_1d(F.squeeze(v, 1))
frontier = F.copy_to(F.nonzero_1d(F.squeeze(v, 1)), F.cpu())
yield frontier
degree -= F.spmm(adjmat, v)
......@@ -83,7 +81,7 @@ def test_topological_nodes(n=1000):
assert all(toset(x) == toset(y) for x, y in zip(layers_dgl, layers_spmv))
DFS_LABEL_NAMES = ['forward', 'reverse', 'nontree']
def test_dfs_labeled_edges(n=1000, example=False):
def test_dfs_labeled_edges(example=False):
dgl_g = dgl.DGLGraph()
dgl_g.add_nodes(6)
dgl_g.add_edges([0, 1, 0, 3, 3], [1, 2, 2, 4, 5])
......
tests/mxnet/test_specialization.py deleted 100644 → 0
View file @ da0c92a2
import os
os.environ['DGLBACKEND'] = 'mxnet'
import mxnet as mx
from mxnet import autograd
import scipy as sp
import numpy as np
import dgl
import dgl.function as fn
D = 5
mx.random.seed(1)
np.random.seed(1)
def generate_graph(n):
arr = (sp.sparse.random(n, n, density=0.1, format='coo') != 0).astype(np.int64)
g = dgl.DGLGraph(arr, readonly=True)
num_nodes = g.number_of_nodes()
g.set_n_repr({'f1' : mx.nd.random.normal(shape=(num_nodes,)),
'f2' : mx.nd.random.normal(shape=(num_nodes, D))})
weights = mx.nd.random.normal(shape=(g.number_of_edges(),))
g.set_e_repr({'e1': weights, 'e2': mx.nd.expand_dims(weights, axis=1)})
return g
def generate_graph2(n):
arr = (sp.sparse.random(n, n, density=0.1, format='coo') != 0).astype(np.int64)
g1 = dgl.DGLGraph(arr, readonly=True)
g2 = dgl.DGLGraph(arr, readonly=True)
num_nodes = g1.number_of_nodes()
g1.set_n_repr({'f1' : mx.nd.random.normal(shape=(num_nodes,)),
'f2' : mx.nd.random.normal(shape=(num_nodes, D))})
weights = mx.nd.random.normal(shape=(g1.number_of_edges(),))
g1.set_e_repr({'e1': weights, 'e2': mx.nd.expand_dims(weights, axis=1)})
g2.set_n_repr({'f1' : g1.ndata['f1'].copy(), 'f2' : g1.ndata['f2'].copy()})
g2.set_e_repr({'e1': g1.edata['e1'].copy(), 'e2': g1.edata['e2'].copy()})
return g1, g2
def test_update_all():
def _test(fld):
def message_func(edges):
return {'m' : edges.src[fld]}
def message_func_edge(edges):
if len(edges.src[fld].shape) == 1:
return {'m' : edges.src[fld] * edges.data['e1']}
else:
return {'m' : edges.src[fld] * edges.data['e2']}
def reduce_func(nodes):
return {fld : mx.nd.sum(nodes.mailbox['m'], axis=1)}
def apply_func(nodes):
return {fld : 2 * nodes.data[fld]}
g1, g2 = generate_graph2(100)
# update all
g1_data = g1.ndata[fld]
g2_data = g2.ndata[fld]
g1_data.attach_grad()
g2_data.attach_grad()
with mx.autograd.record():
g1.update_all(fn.copy_src(src=fld, out='m'), fn.sum(msg='m', out=fld), apply_func)
g2.update_all(message_func, reduce_func, apply_func)
g1_res = g1.ndata[fld]
g2_res = g2.ndata[fld]
assert np.allclose(g1_res.asnumpy(), g2_res.asnumpy(), rtol=1e-05, atol=1e-05)
g1_res.backward()
g2_res.backward()
assert np.allclose(g1_data.grad.asnumpy(), g2_data.grad.asnumpy(), rtol=1e-05, atol=1e-05)
# update all with edge weights
g1_data = g1.ndata[fld]
g1.update_all(fn.src_mul_edge(src=fld, edge='e1', out='m'),
fn.sum(msg='m', out=fld), apply_func)
v2 = g1.ndata[fld]
g1.set_n_repr({fld : g1_data})
g1.update_all(fn.src_mul_edge(src=fld, edge='e2', out='m'),
fn.sum(msg='m', out=fld), apply_func)
v3 = g1.ndata[fld]
assert np.allclose(v2.asnumpy(), v3.asnumpy(), rtol=1e-05, atol=1e-05)
g1.set_n_repr({fld : g1_data})
g2_data = g2.ndata[fld]
g1_data.attach_grad()
g2_data.attach_grad()
with mx.autograd.record():
g1.update_all(fn.src_mul_edge(src=fld, edge='e2', out='m'),
fn.sum(msg='m', out=fld), apply_func)
g2.update_all(message_func_edge, reduce_func, apply_func)
g1_res = g1.ndata[fld]
g2_res = g2.ndata[fld]
assert np.allclose(g1_res.asnumpy(), g2_res.asnumpy(), rtol=1e-05, atol=1e-05)
g1_res.backward()
g2_res.backward()
assert np.allclose(g1_data.grad.asnumpy(), g2_data.grad.asnumpy(), rtol=1e-05, atol=1e-05)
# test 1d node features
_test('f1')
# test 2d node features
_test('f2')
def test_pull():
def _test(fld):
def message_func(edges):
return {'m' : edges.src[fld]}
def message_func_edge(edges):
if len(edges.src[fld].shape) == 1:
return {'m' : edges.src[fld] * edges.data['e1']}
else:
return {'m' : edges.src[fld] * edges.data['e2']}
def reduce_func(nodes):
return {fld : mx.nd.sum(nodes.mailbox['m'], axis=1)}
def apply_func(nodes):
return {fld : 2 * nodes.data[fld]}
g1, g2 = generate_graph2(100)
num_nodes = g1.number_of_nodes()
u = np.unique(np.random.randint(0, num_nodes, size=(int(num_nodes/10))))
# pull in DGL
g1_data = g1.ndata[fld]
g2_data = g2.ndata[fld]
if len(g1_data.shape) == 1:
g1_data = mx.nd.expand_dims(g1_data, axis=1)
g1.ndata[fld] = g1_data
if len(g2_data.shape) == 1:
g2_data = mx.nd.expand_dims(g2_data, axis=1)
g2.ndata[fld] = g2_data
g1_data.attach_grad()
g2_data.attach_grad()
with mx.autograd.record():
g1.pull(u, fn.copy_src(src=fld, out='m'), fn.sum(msg='m', out=fld), apply_func)
spm = mx.nd.take(g2.adjacency_matrix(), mx.nd.array(u, dtype=np.int64))
g2_res = mx.nd.dot(spm, g2_data) * 2
g1_res = g1.ndata[fld][u]
assert np.allclose(g1_res.asnumpy(), g2_res.asnumpy(), rtol=1e-05, atol=1e-05)
g1_res.backward()
g2_res.backward()
assert np.allclose(g1_data.grad.asnumpy(), g2_data.grad.asnumpy(), rtol=1e-05, atol=1e-05)
# test 1d node features
_test('f1')
# test 2d node features
_test('f2')
def test_send_and_recv():
def _test(fld):
def message_func(edges):
return {'m' : edges.src[fld]}
def message_func_edge(edges):
if len(edges.src[fld].shape) == 1:
return {'m' : edges.src[fld] * edges.data['e1']}
else:
return {'m' : edges.src[fld] * edges.data['e2']}
def reduce_func(nodes):
return {fld : mx.nd.sum(nodes.mailbox['m'], axis=1)}
def apply_func(nodes):
return {fld : 2 * nodes.data[fld]}
g1, g2 = generate_graph2(100)
u, v = g1.all_edges()
idxs = np.unique(np.random.randint(0, len(u), size=(int(len(u)/10))))
u = u[idxs]
v = v[idxs]
# send and recv
g1_data = g1.ndata[fld]
g2_data = g2.ndata[fld]
g1_data.attach_grad()
g2_data.attach_grad()
with mx.autograd.record():
g1.send_and_recv((u, v), fn.copy_src(src=fld, out='m'),
fn.sum(msg='m', out=fld), apply_func)
g2.send_and_recv((u, v), message_func, reduce_func, apply_func)
g1_res = g1.ndata[fld]
g2_res = g2.ndata[fld]
assert np.allclose(g1_res.asnumpy(), g2_res.asnumpy(), rtol=1e-05, atol=1e-05)
g1_res.backward()
g2_res.backward()
assert np.allclose(g1_data.grad.asnumpy(), g2_data.grad.asnumpy(), rtol=1e-05, atol=1e-05)
# send and recv with edge weights
g1_data = g1.ndata[fld]
g1.send_and_recv((u, v), fn.src_mul_edge(src=fld, edge='e1', out='m'),
fn.sum(msg='m', out=fld), apply_func)
v2 = g1.ndata[fld]
g1.set_n_repr({fld : g1_data})
g1.send_and_recv((u, v), fn.src_mul_edge(src=fld, edge='e2', out='m'),
fn.sum(msg='m', out=fld), apply_func)
v3 = g1.ndata[fld]
assert np.allclose(v2.asnumpy(), v3.asnumpy(), rtol=1e-05, atol=1e-05)
g1.set_n_repr({fld : g1_data})
g2_data = g2.ndata[fld]
g1_data.attach_grad()
g2_data.attach_grad()
with mx.autograd.record():
g1.send_and_recv((u, v), fn.src_mul_edge(src=fld, edge='e2', out='m'),
fn.sum(msg='m', out=fld), apply_func)
g2.send_and_recv((u, v), message_func_edge, reduce_func, apply_func)
g1_res = g1.ndata[fld]
g2_res = g2.ndata[fld]
assert np.allclose(g1_res.asnumpy(), g2_res.asnumpy(), rtol=1e-05, atol=1e-05)
g1_res.backward()
g2_res.backward()
assert np.allclose(g1_data.grad.asnumpy(), g1_data.grad.asnumpy(), rtol=1e-05, atol=1e-05)
# test 1d node features
# TODO for some reason, this test doesn't pass in MXNet.
# somehow, it fails in backward.
#_test('f1')
# test 2d node features
_test('f2')
def test_update_all_multi_fn():
def message_func(edges):
return {'m2': edges.src['f2']}
def message_func_edge(edges):
return {'m2': edges.src['f2'] * edges.data['e2']}
def reduce_func(nodes):
return {'v2': mx.nd.sum(nodes.mailbox['m2'], axis=1)}
g = generate_graph(100)
g.set_n_repr({'v1' : mx.nd.zeros(shape=(g.number_of_nodes(),)),
'v2' : mx.nd.zeros(shape=(g.number_of_nodes(),))})
fld = 'f2'
# run builtin with single message and reduce
g.update_all(fn.copy_src(src=fld, out='m'), fn.sum(msg='m', out='v1'), None)
v1 = g.ndata['v1']
# 1 message, 2 reduces
g.update_all(fn.copy_src(src=fld, out='m'), [fn.sum(msg='m', out='v2'), fn.sum(msg='m', out='v3')], None)
v2 = g.ndata['v2']
v3 = g.ndata['v3']
assert np.allclose(v1.asnumpy(), v2.asnumpy(), rtol=1e-05, atol=1e-05)
assert np.allclose(v1.asnumpy(), v3.asnumpy(), rtol=1e-05, atol=1e-05)
# update all with edge weights, 2 message, 3 reduces
g.update_all([fn.src_mul_edge(src=fld, edge='e1', out='m1'), fn.src_mul_edge(src=fld, edge='e2', out='m2')],
[fn.sum(msg='m1', out='v1'), fn.sum(msg='m2', out='v2'), fn.sum(msg='m1', out='v3')],
None)
v1 = g.ndata['v1']
v2 = g.ndata['v2']
v3 = g.ndata['v3']
assert np.allclose(v1.asnumpy(), v2.asnumpy(), rtol=1e-05, atol=1e-05)
assert np.allclose(v1.asnumpy(), v3.asnumpy(), rtol=1e-05, atol=1e-05)
# run UDF with single message and reduce
g.update_all(message_func_edge, reduce_func, None)
v2 = g.ndata['v2']
assert np.allclose(v1.asnumpy(), v2.asnumpy(), rtol=1e-05, atol=1e-05)
def test_send_and_recv_multi_fn():
u = mx.nd.array([0, 0, 0, 3, 4, 9], dtype=np.int64)
v = mx.nd.array([1, 2, 3, 9, 9, 0], dtype=np.int64)
def message_func(edges):
return {'m2': edges.src['f2']}
def message_func_edge(edges):
return {'m2': edges.src['f2'] * edges.data['e2']}
def reduce_func(nodes):
return {'v2' : mx.nd.sum(nodes.mailbox['m2'], axis=1)}
g = generate_graph(100)
g.set_n_repr({'v1' : mx.nd.zeros(shape=(g.number_of_nodes(), D)),
'v2' : mx.nd.zeros(shape=(g.number_of_nodes(), D)),
'v3' : mx.nd.zeros(shape=(g.number_of_nodes(), D))})
fld = 'f2'
# run builtin with single message and reduce
g.send_and_recv((u, v), fn.copy_src(src=fld, out='m'), fn.sum(msg='m', out='v1'),
None)
v1 = g.ndata['v1']
# 1 message, 2 reduces
g.send_and_recv((u, v),
fn.copy_src(src=fld, out='m'),
[fn.sum(msg='m', out='v2'), fn.sum(msg='m', out='v3')],
None)
v2 = g.ndata['v2']
v3 = g.ndata['v3']
assert np.allclose(v1.asnumpy(), v2.asnumpy(), rtol=1e-05, atol=1e-05)
assert np.allclose(v1.asnumpy(), v3.asnumpy(), rtol=1e-05, atol=1e-05)
# send and recv with edge weights, 2 message, 3 reduces
g.send_and_recv((u, v),
[fn.src_mul_edge(src=fld, edge='e1', out='m1'), fn.src_mul_edge(src=fld, edge='e2', out='m2')],
[fn.sum(msg='m1', out='v1'), fn.sum(msg='m2', out='v2'), fn.sum(msg='m1', out='v3')],
None)
v1 = g.ndata['v1']
v2 = g.ndata['v2']
v3 = g.ndata['v3']
assert np.allclose(v1.asnumpy(), v2.asnumpy(), rtol=1e-05, atol=1e-05)
assert np.allclose(v1.asnumpy(), v3.asnumpy(), rtol=1e-05, atol=1e-05)
# run UDF with single message and reduce
g.send_and_recv((u, v), message_func_edge,
reduce_func, None)
v2 = g.ndata['v2']
assert np.allclose(v1.asnumpy(), v2.asnumpy(), rtol=1e-05, atol=1e-05)
############################ Copy from torch
D = 5
def simple_graph():
g = dgl.DGLGraph()
g.add_nodes(10)
# create a graph where 0 is the source and 9 is the sink
for i in range(1, 9):
g.add_edge(0, i)
g.add_edge(i, 9)
# add a back flow from 9 to 0
g.add_edge(9, 0)
g.set_n_repr({'f1' : mx.nd.random.normal(shape=(10,)), 'f2' : mx.nd.random.normal(shape=(10, D))})
weights = mx.nd.random.normal(shape=(17,))
g.set_e_repr({'e1': weights, 'e2': mx.nd.expand_dims(weights, 1)})
return g
def test_v2v_update_all_sum():
def _test(fld):
def message_func(edges):
return {'m' : edges.src[fld]}
def message_func_edge(edges):
if len(edges.src[fld].shape) == 1:
return {'m' : edges.src[fld] * edges.data['e1']}
else:
return {'m' : edges.src[fld] * edges.data['e2']}
def reduce_func(nodes):
return {fld : mx.nd.sum(nodes.mailbox['m'], axis=1)}
def apply_func(nodes):
return {fld : 2 * nodes.data[fld]}
g = simple_graph()
# update all
v1 = g.ndata[fld]
g.update_all(fn.copy_src(src=fld, out='m'), fn.sum(msg='m', out=fld), apply_func)
v2 = g.ndata[fld]
g.set_n_repr({fld : v1})
g.update_all(message_func, reduce_func, apply_func)
v3 = g.ndata[fld]
assert np.allclose(v2.asnumpy(), v3.asnumpy(), rtol=1e-05, atol=1e-05)
# update all with edge weights
v1 = g.ndata[fld]
g.update_all(fn.src_mul_edge(src=fld, edge='e1', out='m'),
fn.sum(msg='m', out=fld), apply_func)
v2 = g.ndata[fld]
g.set_n_repr({fld : v1})
g.update_all(fn.src_mul_edge(src=fld, edge='e2', out='m'),
fn.sum(msg='m', out=fld), apply_func)
v3 = g.ndata[fld].squeeze()
g.set_n_repr({fld : v1})
g.update_all(message_func_edge, reduce_func, apply_func)
v4 = g.ndata[fld]
assert np.allclose(v2.asnumpy(), v3.asnumpy(), rtol=1e-05, atol=1e-05)
assert np.allclose(v3.asnumpy(), v4.asnumpy(), rtol=1e-05, atol=1e-05)
# test 1d node features
_test('f1')
# test 2d node features
_test('f2')
def test_v2v_update_all_max():
def _test(fld):
def message_func(edges):
return {'m' : edges.src[fld]}
def message_func_edge(edges):
if len(edges.src[fld].shape) == 1:
return {'m' : edges.src[fld] * edges.data['e1']}
else:
return {'m' : edges.src[fld] * edges.data['e2']}
def reduce_func(nodes):
return {fld : mx.nd.max(nodes.mailbox['m'], axis=1)}
def apply_func(nodes):
return {fld : 2 * nodes.data[fld]}
g = simple_graph()
# update all
v1 = g.ndata[fld]
g.update_all(fn.copy_src(src=fld, out='m'), fn.max(msg='m', out=fld), apply_func)
v2 = g.ndata[fld]
g.set_n_repr({fld : v1})
g.update_all(message_func, reduce_func, apply_func)
v3 = g.ndata[fld]
assert np.allclose(v2.asnumpy(), v3.asnumpy(), rtol=1e-05, atol=1e-05)
# update all with edge weights
v1 = g.ndata[fld]
g.update_all(fn.src_mul_edge(src=fld, edge='e1', out='m'),
fn.max(msg='m', out=fld), apply_func)
v2 = g.ndata[fld]
g.set_n_repr({fld : v1})
g.update_all(fn.src_mul_edge(src=fld, edge='e2', out='m'),
fn.max(msg='m', out=fld), apply_func)
v3 = g.ndata[fld].squeeze()
g.set_n_repr({fld : v1})
g.update_all(message_func_edge, reduce_func, apply_func)
v4 = g.ndata[fld]
assert np.allclose(v2.asnumpy(), v3.asnumpy(), rtol=1e-05, atol=1e-05)
assert np.allclose(v3.asnumpy(), v4.asnumpy(), rtol=1e-05, atol=1e-05)
# test 1d node features
_test('f1')
# test 2d node features
_test('f2')
############################ Copy from torch
if __name__ == '__main__':
test_update_all()
test_pull()
test_send_and_recv()
test_update_all_multi_fn()
test_send_and_recv_multi_fn()
test_v2v_update_all_sum()
test_v2v_update_all_max()
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