xref: /aosp_15_r20/external/pytorch/test/test_itt.py (revision da0073e96a02ea20f0ac840b70461e3646d07c45)
1*da0073e9SAndroid Build Coastguard Worker# Owner(s): ["module: intel"]
2*da0073e9SAndroid Build Coastguard Worker
3*da0073e9SAndroid Build Coastguard Workerimport torch
4*da0073e9SAndroid Build Coastguard Workerimport unittest
5*da0073e9SAndroid Build Coastguard Workerfrom torch.testing._internal.common_utils import TestCase, run_tests, load_tests
6*da0073e9SAndroid Build Coastguard Worker
7*da0073e9SAndroid Build Coastguard Worker# load_tests from common_utils is used to automatically filter tests for
8*da0073e9SAndroid Build Coastguard Worker# sharding on sandcastle. This line silences flake warnings
9*da0073e9SAndroid Build Coastguard Workerload_tests = load_tests
10*da0073e9SAndroid Build Coastguard Worker
11*da0073e9SAndroid Build Coastguard Worker@unittest.skipIf(not torch.profiler.itt.is_available(), "ITT is required")
12*da0073e9SAndroid Build Coastguard Workerclass TestItt(TestCase):
13*da0073e9SAndroid Build Coastguard Worker    def test_itt(self):
14*da0073e9SAndroid Build Coastguard Worker        # Just making sure we can see the symbols
15*da0073e9SAndroid Build Coastguard Worker        torch.profiler.itt.range_push("foo")
16*da0073e9SAndroid Build Coastguard Worker        torch.profiler.itt.mark("bar")
17*da0073e9SAndroid Build Coastguard Worker        torch.profiler.itt.range_pop()
18*da0073e9SAndroid Build Coastguard Worker
19*da0073e9SAndroid Build Coastguard Workerif __name__ == '__main__':
20*da0073e9SAndroid Build Coastguard Worker    run_tests()
21