mirror of
https://github.com/zebrajr/pytorch.git
synced 2025-12-07 12:21:27 +01:00
1e05e5e0ae
Summary: Changes in PR https://github.com/pytorch/pytorch/issues/39759 broke HIP caffe2. hipify for caffe2 renames CUDA to HIP; torch does not. If caffe2 calls into torch, it needs to use CUDA-named functions. CC ezyang xw285cornell sunway513 houseroad dzhulgakov Pull Request resolved: https://github.com/pytorch/pytorch/pull/39801 Differential Revision: D21982493 Pulled By: xw285cornell fbshipit-source-id: 8e88e0fb80c71f0342e23ef0214a42d5542bdc70
32 lines
1.2 KiB
Python
32 lines
1.2 KiB
Python
from __future__ import absolute_import
|
|
from __future__ import division
|
|
from __future__ import print_function
|
|
from __future__ import unicode_literals
|
|
|
|
import unittest
|
|
|
|
import torch
|
|
from caffe2.python import core, workspace
|
|
|
|
# This is a standalone test that doesn't use test_util as we're testing
|
|
# initialization and thus we should be the ones calling GlobalInit
|
|
@unittest.skipIf(not workspace.has_cuda_support,
|
|
"THC pool testing is obscure and doesn't work on HIP yet")
|
|
class TestGPUInit(unittest.TestCase):
|
|
def testTHCAllocator(self):
|
|
cuda_or_hip = 'hip' if workspace.has_hip_support else 'cuda'
|
|
flag = '--caffe2_{}_memory_pool=thc'.format(cuda_or_hip)
|
|
core.GlobalInit(['caffe2', flag])
|
|
# just run one operator
|
|
# it's importantant to not call anything here from Torch API
|
|
# even torch.cuda.memory_allocated would initialize CUDA context
|
|
workspace.RunOperatorOnce(core.CreateOperator(
|
|
'ConstantFill', [], ["x"], shape=[5, 5], value=1.0,
|
|
device_option=core.DeviceOption(workspace.GpuDeviceType)
|
|
))
|
|
# make sure we actually used THC allocator
|
|
self.assertGreater(torch.cuda.memory_allocated(), 0)
|
|
|
|
if __name__ == '__main__':
|
|
unittest.main()
|