Define device string as contant variable in python package

Author: EikanWang

intel_pytorch_extension_py/__init__.py

@@ -6,6 +6,8 @@
 from .ops import *
 import _torch_ipex as core
 
+DEVICE = 'dpcpp'
+
 def get_auto_optimization():
     return core.get_auto_dnnl()
 

setup.py

@@ -254,7 +254,7 @@ def make_relative_rpath(path):
       'intel_pytorch_extension',
       'intel_pytorch_extension.optim',
       'intel_pytorch_extension.ops'],
-    package_dir={'torch_ipex': 'torch_ipex_py', 'intel_pytorch_extension': 'intel_pytorch_extension_py'},
+    package_dir={'intel_pytorch_extension': 'intel_pytorch_extension_py'},
     zip_safe=False,
     ext_modules=[DPCPPExt('_torch_ipex')],
     cmdclass={

tests/cpu/common_device_type.py

@@ -17,23 +17,23 @@
 
 All contributions by Facebook:
 Copyright (c) 2016 Facebook Inc.
- 
+
 All contributions by Google:
 Copyright (c) 2015 Google Inc.
 All rights reserved.
- 
+
 All contributions by Yangqing Jia:
 Copyright (c) 2015 Yangqing Jia
 All rights reserved.
- 
+
 All contributions from Caffe:
 Copyright(c) 2013, 2014, 2015, the respective contributors
 All rights reserved.
- 
+
 All other contributions:
 Copyright(c) 2015, 2016 the respective contributors
 All rights reserved.
- 
+
 Caffe2 uses a copyright model similar to Caffe: each contributor holds
 copyright over their contributions to Caffe2. The project versioning records
 all such contribution and copyright details. If a contributor wants to further
@@ -49,7 +49,7 @@
 from functools import wraps
 import unittest
 import torch
-import _torch_ipex as ipex
+import intel_pytorch_extension as ipex
 import copy
 from common_utils import TestCase, TEST_WITH_ROCM, TEST_MKL, \
     skipCUDANonDefaultStreamIf
@@ -254,15 +254,15 @@ class CPUTestBase(DeviceTypeTestBase):
     device_type = 'cpu'
 
 class DPCPPTestBase(DeviceTypeTestBase):
-    device_type = 'dpcpp'
+    device_type = ipex.DEVICE
 
     @classmethod
     def get_primary_device(cls):
         return cls.primary_device
 
     @classmethod
     def get_all_devices(cls):
-        return ['dpcpp']
+        return [ipex.DEVICE]
 
     # Returns the dtypes the test has requested.
     # Prefers device-specific dtype specifications over generic ones.
@@ -276,7 +276,7 @@ def _get_dtypes(cls, test):
 
     @classmethod
     def setUpClass(cls):
-        cls.primary_device = 'dpcpp'
+        cls.primary_device = ipex.DEVICE
 
 class CUDATestBase(DeviceTypeTestBase):
     device_type = 'cuda'

tests/cpu/test_bf16_lazy_reorder.py

@@ -32,7 +32,7 @@
 def get_rand_seed():
     return int(time.time() * 1000000000)
 
-device = torch.device("dpcpp:0")
+device = ipex.DEVICE
 class TestConv(TestCase):
     def test_Conv2d_with_cpu(self):
         rand_seed = int(get_rand_seed())

tests/cpu/test_emb.py

@@ -9,14 +9,14 @@ class TestEMB(TestCase):
     def test_emb(self):
         #E = nn.EmbeddingBag(10, 5, mode="sum", sparse=True)
         cpu_emb = nn.EmbeddingBag(10, 3, mode='sum', sparse=True)
-        dpcpp_emb = copy.deepcopy(cpu_emb).to('dpcpp:0')
-        bf16_emb = copy.deepcopy(cpu_emb).to('dpcpp').bfloat16()
+        dpcpp_emb = copy.deepcopy(cpu_emb).to(ipex.DEVICE)
+        bf16_emb = copy.deepcopy(cpu_emb).to(ipex.DEVICE).bfloat16()
         # a batch of 2 samples of 4 indices each
         cpu_input = torch.LongTensor([1,2,4,5,4,3,2,9])
-        dpcpp_input = cpu_input.clone().detach().to('dpcpp:0')
+        dpcpp_input = cpu_input.clone().detach().to(ipex.DEVICE)
 
         cpu_offsets = torch.LongTensor([0,1,2,3,4,5,6,7])
-        dpcpp_offsets = cpu_offsets.clone().detach().to('dpcpp:0')
+        dpcpp_offsets = cpu_offsets.clone().detach().to(ipex.DEVICE)
 
         cpu_out = cpu_emb(cpu_input, cpu_offsets)
 

tests/cpu/test_interaction.py

@@ -16,21 +16,21 @@ def interact_features(x, ly):
     Z = torch.bmm(T, torch.transpose(T, 1, 2))
     _, ni, nj = Z.shape
     offset =  0
-    li = torch.tensor([i for i in range(ni) for j in range(i + offset)], device='dpcpp')
-    lj = torch.tensor([j for i in range(nj) for j in range(i + offset)], device='dpcpp')
+    li = torch.tensor([i for i in range(ni) for j in range(i + offset)], device=ipex.DEVICE)
+    lj = torch.tensor([j for i in range(nj) for j in range(i + offset)], device=ipex.DEVICE)
     Zflat = Z[:, li, lj]
     # concatenate dense features and interactions
     R = torch.cat([x] + [Zflat], dim=1)
     return R
 
 def run(dtype='float32'):
     print("##################### testing with %s"% str(dtype))
-    x1 = torch.randn([2048, 128], device='dpcpp').to(dtype).clone().detach().requires_grad_()
+    x1 = torch.randn([2048, 128], device=ipex.DEVICE).to(dtype).clone().detach().requires_grad_()
     x2 = x1.clone().detach().requires_grad_()
     ly1 = []
     ly2 = []
     for i in range(0, 26):
-        V = torch.randn([2048, 128], device='dpcpp').to(dtype).clone().detach().requires_grad_()
+        V = torch.randn([2048, 128], device=ipex.DEVICE).to(dtype).clone().detach().requires_grad_()
         ly1.append(V)
         ly2.append(V.clone().detach().requires_grad_())
 

tests/cpu/test_jit.py

@@ -59,7 +59,7 @@
 from torch.jit._recursive import wrap_cpp_module
 import copy
 
-import intel_pytorch_extension
+import intel_pytorch_extension as ipex
 from intel_pytorch_extension import core
 
 import torch.nn as nn
@@ -76,7 +76,7 @@
     IS_SANDCASTLE, load_tests, brute_pdist, brute_cdist, slowTest, \
     skipCUDANonDefaultStreamIf, skipCUDAMemoryLeakCheckIf
 
-device = 'dpcpp:0'
+device = ipex.DEVICE
 #device = 'cpu:0'
 SIZE = 100
 
@@ -210,8 +210,8 @@ def _test_output(self, model, x, kind=None):
         core.disable_jit()
         core.disable_mix_bf16_fp32()
 
-        model = model.to('dpcpp').eval()
-        x = x.to('dpcpp')
+        model = model.to(device).eval()
+        x = x.to(device)
         with torch.no_grad():
             result = model(x)
 
@@ -247,8 +247,8 @@ def _test_output_bf16(self, model, x, kind=None, prec=None):
         core.enable_jit()
         core.disable_mix_bf16_fp32()
 
-        model = model.to('dpcpp').eval()
-        x = x.to('dpcpp')
+        model = model.to(ipex.DEVICE).eval()
+        x = x.to(ipex.DEVICE)
         x2 = x.clone()
 
         fused_model = torch.jit.script(copy.deepcopy(model))