xref: /aosp_15_r20/external/pytorch/c10/xpu/test/impl/XPUDeviceTest.cpp (revision da0073e96a02ea20f0ac840b70461e3646d07c45)
1 #include <gtest/gtest.h>
2 
3 #include <c10/xpu/XPUFunctions.h>
4 
has_xpu()5 bool has_xpu() {
6   return c10::xpu::device_count() > 0;
7 }
8 
TEST(XPUDeviceTest,DeviceBehavior)9 TEST(XPUDeviceTest, DeviceBehavior) {
10   if (!has_xpu()) {
11     return;
12   }
13 
14   c10::xpu::set_device(0);
15   EXPECT_EQ(c10::xpu::current_device(), 0);
16 
17   if (c10::xpu::device_count() <= 1) {
18     return;
19   }
20 
21   c10::xpu::set_device(1);
22   EXPECT_EQ(c10::xpu::current_device(), 1);
23   EXPECT_EQ(c10::xpu::exchange_device(0), 1);
24   EXPECT_EQ(c10::xpu::current_device(), 0);
25 }
26 
TEST(XPUDeviceTest,DeviceProperties)27 TEST(XPUDeviceTest, DeviceProperties) {
28   if (!has_xpu()) {
29     return;
30   }
31 
32   c10::xpu::DeviceProp device_prop{};
33   c10::xpu::get_device_properties(&device_prop, 0);
34 
35   EXPECT_TRUE(device_prop.max_compute_units > 0);
36   EXPECT_TRUE(device_prop.gpu_eu_count > 0);
37 }
38 
TEST(XPUDeviceTest,PointerGetDevice)39 TEST(XPUDeviceTest, PointerGetDevice) {
40   if (!has_xpu()) {
41     return;
42   }
43 
44   sycl::device& raw_device = c10::xpu::get_raw_device(0);
45   void* ptr =
46       sycl::malloc_device(8, raw_device, c10::xpu::get_device_context());
47 
48   EXPECT_EQ(c10::xpu::get_device_idx_from_pointer(ptr), 0);
49   sycl::free(ptr, c10::xpu::get_device_context());
50 
51   int dummy = 0;
52   ASSERT_THROW(c10::xpu::get_device_idx_from_pointer(&dummy), c10::Error);
53 }
54