dmlc--dgl
Xin Yao
79b0a50afc
* Fix test_csrmm for tensor core * unset allow tf32 flag * update test unified tensor * skip fp16 for CPU
86 行
2.9 KiB
Python
86 行
2.9 KiB
Python
import dgl.multiprocessing as mp
|
|
import unittest, os
|
|
import pytest
|
|
|
|
import torch as th
|
|
import dgl
|
|
import backend as F
|
|
|
|
def start_unified_tensor_worker(dev_id, input, seq_idx, rand_idx, output_seq, output_rand):
|
|
device = th.device('cuda:'+str(dev_id))
|
|
th.cuda.set_device(device)
|
|
input_unified = dgl.contrib.UnifiedTensor(input, device=device)
|
|
output_seq.copy_(input_unified[seq_idx.to(device)].cpu())
|
|
output_rand.copy_(input_unified[rand_idx.to(device)].cpu())
|
|
|
|
@unittest.skipIf(os.name == 'nt', reason='Do not support windows yet')
|
|
@unittest.skipIf(F.ctx().type == 'cpu', reason='gpu only test')
|
|
def test_unified_tensor():
|
|
test_row_size = 65536
|
|
test_col_size = 128
|
|
rand_test_size = 8192
|
|
|
|
device = th.device('cuda:0')
|
|
|
|
input = th.rand((test_row_size, test_col_size))
|
|
input_unified = dgl.contrib.UnifiedTensor(input, device=device)
|
|
|
|
seq_idx = th.arange(0, test_row_size)
|
|
# CPU indexing
|
|
assert th.all(th.eq(input[seq_idx], input_unified[seq_idx]))
|
|
# GPU indexing
|
|
assert th.all(th.eq(input[seq_idx].to(device), input_unified[seq_idx.to(device)]))
|
|
|
|
rand_idx = th.randint(0, test_row_size, (rand_test_size,))
|
|
# CPU indexing
|
|
assert th.all(th.eq(input[rand_idx], input_unified[rand_idx]))
|
|
# GPU indexing
|
|
assert th.all(th.eq(input[rand_idx].to(device), input_unified[rand_idx.to(device)]))
|
|
|
|
@unittest.skipIf(os.name == 'nt', reason='Do not support windows yet')
|
|
@unittest.skipIf(F.ctx().type == 'cpu', reason='gpu only test')
|
|
@pytest.mark.parametrize("num_workers", [1, 2])
|
|
def test_multi_gpu_unified_tensor(num_workers):
|
|
if F.ctx().type == 'cuda' and th.cuda.device_count() < num_workers:
|
|
pytest.skip("Not enough number of GPUs to do this test, skip multi-gpu test.")
|
|
|
|
test_row_size = 65536
|
|
test_col_size = 128
|
|
|
|
rand_test_size = 8192
|
|
|
|
input = th.rand((test_row_size, test_col_size)).share_memory_()
|
|
seq_idx = th.arange(0, test_row_size).share_memory_()
|
|
rand_idx = th.randint(0, test_row_size, (rand_test_size,)).share_memory_()
|
|
|
|
output_seq = []
|
|
output_rand = []
|
|
|
|
output_seq_cpu = input[seq_idx]
|
|
output_rand_cpu = input[rand_idx]
|
|
|
|
worker_list = []
|
|
|
|
ctx = mp.get_context('spawn')
|
|
for i in range(num_workers):
|
|
output_seq.append(th.zeros((test_row_size, test_col_size)).share_memory_())
|
|
output_rand.append(th.zeros((rand_test_size, test_col_size)).share_memory_())
|
|
p = ctx.Process(target=start_unified_tensor_worker,
|
|
args=(i, input, seq_idx, rand_idx, output_seq[i], output_rand[i],))
|
|
p.start()
|
|
worker_list.append(p)
|
|
|
|
for p in worker_list:
|
|
p.join()
|
|
for p in worker_list:
|
|
assert p.exitcode == 0
|
|
for i in range(num_workers):
|
|
assert th.all(th.eq(output_seq_cpu, output_seq[i]))
|
|
assert th.all(th.eq(output_rand_cpu, output_rand[i]))
|
|
|
|
|
|
if __name__ == '__main__':
|
|
test_unified_tensor()
|
|
test_multi_gpu_unified_tensor(1)
|
|
test_multi_gpu_unified_tensor(2)
|