def test_gpu(self): from torch.utils.cpp_extension import CUDA_HOME create_extension( name='gpulib', headers=[test_dir + '/ffi/src/cuda/cudalib.h'], sources=[ test_dir + '/ffi/src/cuda/cudalib.c', ], with_cuda=True, verbose=False, include_dirs=[os.path.join(CUDA_HOME, 'include')], ).build() import gpulib tensor = torch.ones(2, 2).float() gpulib.good_func(tensor, 2, 1.5) self.assertEqual(tensor, torch.ones(2, 2) * 2 + 1.5) ctensor = tensor.cuda().fill_(1) gpulib.cuda_func(ctensor, 2, 1.5) self.assertEqual(ctensor, torch.ones(2, 2) * 2 + 1.5) self.assertRaises(TypeError, lambda: gpulib.cuda_func(tensor, 2, 1.5)) self.assertRaises(TypeError, lambda: gpulib.cuda_func(ctensor.storage(), 2, 1.5))
def test_gpu(self): compile_extension( name='gpulib', header=test_dir + '/ffi/src/cuda/cudalib.h', sources=[ test_dir + '/ffi/src/cuda/cudalib.c', ], with_cuda=True, verbose=False, ) import gpulib tensor = torch.ones(2, 2).float() gpulib.good_func(tensor, 2, 1.5) self.assertEqual(tensor, torch.ones(2, 2) * 2 + 1.5) ctensor = tensor.cuda().fill_(1) gpulib.cuda_func(ctensor, 2, 1.5) self.assertEqual(ctensor, torch.ones(2, 2) * 2 + 1.5) self.assertRaises(TypeError, lambda: gpulib.cuda_func(tensor, 2, 1.5)) self.assertRaises(TypeError, lambda: gpulib.cuda_func(ctensor.storage(), 2, 1.5))
def test_gpu(self): create_extension( name='gpulib', headers=[test_dir + '/ffi/src/cuda/cudalib.h'], sources=[ test_dir + '/ffi/src/cuda/cudalib.c', ], with_cuda=True, verbose=False, ).build() import gpulib tensor = torch.ones(2, 2).float() gpulib.good_func(tensor, 2, 1.5) self.assertEqual(tensor, torch.ones(2, 2) * 2 + 1.5) ctensor = tensor.cuda().fill_(1) gpulib.cuda_func(ctensor, 2, 1.5) self.assertEqual(ctensor, torch.ones(2, 2) * 2 + 1.5) self.assertRaises(TypeError, lambda: gpulib.cuda_func(tensor, 2, 1.5)) self.assertRaises(TypeError, lambda: gpulib.cuda_func(ctensor.storage(), 2, 1.5))