mirror of
https://github.com/pytorch/pytorch.git
synced 2025-10-21 05:34:18 +08:00
8888565597
Summary: Pull Request resolved: https://github.com/pytorch/pytorch/pull/58178 https://www.internalfb.com/T90561249: change the test to enforce Test Plan: buck test //caffe2/test:kernel_launch_checks before fixing LinearAlgebra.cu and file close: https://www.internalfb.com/intern/testinfra/testconsole/testrun/1970324893386017/ after: https://www.internalfb.com/intern/testinfra/testconsole/testrun/2814749824394650/ Reviewed By: r-barnes Differential Revision: D28390166 fbshipit-source-id: 8a217ce8c0b204922005c731aa38bc03f70fabcc
48 lines
1.9 KiB
Python
48 lines
1.9 KiB
Python
from torch.testing._internal.common_utils import TestCase, run_tests
|
|
from torch.testing import check_cuda_kernel_launches, check_code_for_cuda_kernel_launches
|
|
|
|
|
|
class AlwaysCheckCudaLaunchTest(TestCase):
|
|
def test_check_code(self):
|
|
"""Verifies that the regex works for a few different situations"""
|
|
|
|
# Try some different spacings
|
|
self.assertEqual(2, check_code_for_cuda_kernel_launches("""
|
|
some_function_call<TemplateArg><<<1,2,0,stream>>>(arg1,arg2,arg3);
|
|
C10_CUDA_KERNEL_LAUNCH_CHECK();
|
|
some_function_call<TemplateArg><<<1,2,0,stream>>>(arg1,arg2,arg3);
|
|
|
|
some_function_call<TemplateArg><<<1,2,0,stream>>>(arg1,arg2,arg3);
|
|
C10_CUDA_KERNEL_LAUNCH_CHECK();
|
|
some_function_call<TemplateArg><<<1,2,0,stream>>>(arg1,arg2,arg3);
|
|
some_other_stuff;
|
|
some_function_call<TemplateArg><<<1,2,0,stream>>>(arg1,arg2,arg3);
|
|
C10_CUDA_KERNEL_LAUNCH_CHECK();
|
|
some_function_call<TemplateArg><<<1,2,0,stream>>> (arg1,arg2,arg3);
|
|
C10_CUDA_KERNEL_LAUNCH_CHECK();
|
|
some_function_call<TemplateArg><<<1,2,0,stream>>> ( arg1 , arg2 , arg3 ) ;
|
|
|
|
C10_CUDA_KERNEL_LAUNCH_CHECK();
|
|
"""))
|
|
|
|
# Does it work for macros?
|
|
self.assertEqual(0, check_code_for_cuda_kernel_launches(r"""
|
|
#define SOME_MACRO(x) some_function_call<<<1,2>>> ( x ) ; \
|
|
C10_CUDA_KERNEL_LAUNCH_CHECK();
|
|
|
|
#define SMALL_INDEX(TENSOR_TYPE, INDICES_TYPE, TYPE, SELF_DIM, SOURCE_DIM, IDX_DIM) \
|
|
indexAddSmallIndex<TENSOR_TYPE, INDICES_TYPE, TYPE, SELF_DIM, SOURCE_DIM, IDX_DIM> \
|
|
<<<smallIndexGrid, smallIndexBlock, 0, stream>>>( \
|
|
selfInfo, sourceInfo, indexInfo, \
|
|
selfAddDim, sourceAddDim, sliceSize, selfAddDimSize); \
|
|
C10_CUDA_KERNEL_LAUNCH_CHECK();
|
|
"""))
|
|
|
|
def test_check_cuda_launches(self):
|
|
unsafeLaunchesCount = check_cuda_kernel_launches()
|
|
self.assertTrue(unsafeLaunchesCount == 0)
|
|
|
|
|
|
if __name__ == '__main__':
|
|
run_tests()
|