[Device] Use torch.device(npu)

YWHyuk · YWHyuk · commit d444ea7392c7 · 2026-01-22T14:04:12.000Z
diff --git a/experiments/gemm.py b/experiments/gemm.py
@@ -48,7 +48,4 @@ def custom_matmul(a, b):
     if 'pytorchsim_functional_mode' in os.environ:
         del os.environ['pytorchsim_functional_mode']
 
-    from Scheduler.scheduler import PyTorchSimRunner
-    module = PyTorchSimRunner.setup_device()
-    device = module.custom_device()
     run_matmul(size[0], size[1], size[2], config)
diff --git a/scripts/ILS_experiment/test_matmul.py b/scripts/ILS_experiment/test_matmul.py
@@ -60,7 +60,5 @@ def custom_matmul(bias, a, b):
     args = parser.parse_args()
     shape = tuple(map(int, args.shape.strip('()').split(',')))
 
-    from Scheduler.scheduler import PyTorchSimRunner
-    module = PyTorchSimRunner.setup_device()
-    device = module.custom_device()
-    test_matmul(device, *shape)
+    device = torch.device("npu:0")
+    test_matmul(device, *shape)
diff --git a/scripts/chiplet_prep.py b/scripts/chiplet_prep.py
@@ -64,9 +64,7 @@ def modify_file(dump_path, name, address_numa_stride=None, subgraph_map=None):
     import sys
     sys.path.append(os.environ.get('TORCHSIM_DIR', default='/workspace/PyTorchSim'))
 
-    from Scheduler.scheduler import PyTorchSimRunner
-    module = PyTorchSimRunner.setup_device()
-    device = module.custom_device()
+    device = torch.device("npu:0")
     parser = argparse.ArgumentParser(description='Process folder argument.')
     parser.add_argument('size', type=int, help='Folder value', default=256)
     args = parser.parse_args()
diff --git a/tests/Diffusion/test_diffusion.py b/tests/Diffusion/test_diffusion.py
@@ -637,9 +637,7 @@ def test_timesteps(
     args = parser.parse_args()
 
     sys.path.append(os.environ.get("TORCHSIM_DIR", "/workspace/PyTorchSim"))
-    from Scheduler.scheduler import PyTorchSimRunner
-    module = PyTorchSimRunner.setup_device()
-    device = module.custom_device()
+    device = torch.device("npu:0")
 
     #test_upsample2d(device)
     #test_groupnorm(device)
diff --git a/tests/Fusion/test_addmm_residual.py b/tests/Fusion/test_addmm_residual.py
@@ -43,9 +43,7 @@ def addmm_residual(a, b, c, d):
     import sys
     sys.path.append(os.environ.get('TORCHSIM_DIR', default='/workspace/PyTorchSim'))
 
-    from Scheduler.scheduler import PyTorchSimRunner
-    module = PyTorchSimRunner.setup_device()
-    device = module.custom_device()
+    device = torch.device("npu:0")
     test_addmm_residual(device, 32, 32, 32)
     test_addmm_residual(device, 128, 128, 128)
     test_addmm_residual(device, 512, 512, 512)
diff --git a/tests/Fusion/test_attention_fusion.py b/tests/Fusion/test_attention_fusion.py
@@ -75,9 +75,7 @@ def test_MHA(device, num_heads=12, embed_dim=768, input_seq=512):
     import sys
     sys.path.append(os.environ.get('TORCHSIM_DIR', default='/workspace/PyTorchSim'))
 
-    from Scheduler.scheduler import PyTorchSimRunner
-    module = PyTorchSimRunner.setup_device()
-    device = module.custom_device()
+    device = torch.device("npu:0")
     test_MHA(device)
     # test_Attention(device, head=16, seq=512, d_k=64)
     # test_MHA(device, num_heads=12, embed_dim=768)
diff --git a/tests/Fusion/test_bmm_reduction.py b/tests/Fusion/test_bmm_reduction.py
@@ -42,9 +42,7 @@ def bmm(a, b):
     import sys
     sys.path.append(os.environ.get('TORCHSIM_DIR', default='/workspace/PyTorchSim'))
 
-    from Scheduler.scheduler import PyTorchSimRunner
-    module = PyTorchSimRunner.setup_device()
-    device = module.custom_device()
+    device = torch.device("npu:0")
     #test_bmm_reduce(device)
     test_bmm_reduce(device, 12, 512)
     test_bmm_reduce(device, 4, 256)
diff --git a/tests/Fusion/test_conv_fusion.py b/tests/Fusion/test_conv_fusion.py
@@ -101,9 +101,7 @@ def custom_conv_bn_relu(a, b, bias, c, d, e, f):
     import sys
     sys.path.append(os.environ.get('TORCHSIM_DIR', default='/workspace/PyTorchSim'))
 
-    from Scheduler.scheduler import PyTorchSimRunner
-    module = PyTorchSimRunner.setup_device()
-    device = module.custom_device()
+    device = torch.device("npu:0")
 
     # Vanila test
     test_conv_residual(device, batch_size=3, in_channels=64, out_channels=64, input_size=28, kernel_size=3, stride=1, padding=1)
diff --git a/tests/Fusion/test_matmul_activation.py b/tests/Fusion/test_matmul_activation.py
@@ -73,9 +73,7 @@ def test_matmul_activation(device, batch_size=16, input_size=32, output_size=8,
     import sys
     sys.path.append(os.environ.get('TORCHSIM_DIR', default='/workspace/PyTorchSim'))
 
-    from Scheduler.scheduler import PyTorchSimRunner
-    module = PyTorchSimRunner.setup_device()
-    device = module.custom_device()
+    device = torch.device("npu:0")
     test_matmul_activation(device)
     test_matmul_activation(device, batch_size=32, input_size=32, output_size=32, activation_fn="sigmoid")
     test_matmul_activation(device, batch_size=42, input_size=42, output_size=42, activation_fn="sigmoid")
diff --git a/tests/Fusion/test_matmul_reduction.py b/tests/Fusion/test_matmul_reduction.py
@@ -89,9 +89,7 @@ def matmul_fused(a, b, c, d):
     import sys
     sys.path.append(os.environ.get('TORCHSIM_DIR', default='/workspace/PyTorchSim'))
 
-    from Scheduler.scheduler import PyTorchSimRunner
-    module = PyTorchSimRunner.setup_device()
-    device = module.custom_device()
+    device = torch.device("npu:0")
     test_matmul_reduce(device, 3072, 512, 768)
     test_matmul_var_mean(device)
     test_matmul_add_var_mean(device)
diff --git a/tests/Fusion/test_matmul_scalar.py b/tests/Fusion/test_matmul_scalar.py
@@ -39,7 +39,5 @@ def matmul_fused(a, b, c):
     import sys
     sys.path.append(os.environ.get('TORCHSIM_DIR', default='/workspace/PyTorchSim'))
 
-    from Scheduler.scheduler import PyTorchSimRunner
-    module = PyTorchSimRunner.setup_device()
-    device = module.custom_device()
+    device = torch.device("npu:0")
     test_matmul_scalar(device)
diff --git a/tests/Fusion/test_prologue_fusion.py b/tests/Fusion/test_prologue_fusion.py
@@ -88,9 +88,7 @@ def bmm(a, b, c, d):
     import sys
     sys.path.append(os.environ.get('TORCHSIM_DIR', default='/workspace/PyTorchSim'))
 
-    from Scheduler.scheduler import PyTorchSimRunner
-    module = PyTorchSimRunner.setup_device()
-    device = module.custom_device()
+    device = torch.device("npu:0")
     test_elem_broadcast_fusion(device)
     test_elem_fusion(device)
     test_elem_bmm_input_fusion(device, batch_size=4, m=512, n=512, k=64)
diff --git a/tests/Fusion/test_transformer_fusion.py b/tests/Fusion/test_transformer_fusion.py
@@ -203,9 +203,7 @@ def test_EncoderBlock_validation(head=12, embed_dim=768, input_seq=512):
     import sys
     sys.path.append(os.environ.get('TORCHSIM_DIR', default='/workspace/PyTorchSim'))
 
-    from Scheduler.scheduler import PyTorchSimRunner
-    module = PyTorchSimRunner.setup_device()
-    device = module.custom_device()
+    device = torch.device("npu:0")
     #test_MHA(device)
     test_EncoderBlock(device)
     # test_EncoderBlock_validation()
diff --git a/tests/Llama/test_llama.py b/tests/Llama/test_llama.py
@@ -369,9 +369,7 @@ def run_llama_model_test(
     args = parser.parse_args()
 
     sys.path.append(os.environ.get("PYTORCHSIM_ROOT_PATH", "/workspace/PyTorchSim"))
-    from Scheduler.scheduler import PyTorchSimRunner
-    module = PyTorchSimRunner.setup_device()
-    device = module.custom_device()
+    device = torch.device("npu:0")
     #test_triu(device, size=(32, 128), diagonal=1)
     torch.compiler.is_compiling = lambda: True # FIXME. How to fix this?
     #run_rmsnorm_test(device)
diff --git a/tests/Mixtral_8x7B/test_attention.py b/tests/Mixtral_8x7B/test_attention.py
@@ -163,9 +163,7 @@ def test_rmsnorm(device, seq=32):
     import sys
     sys.path.append(os.environ.get('TORCHSIM_DIR', default='/workspace/PyTorchSim'))
 
-    from Scheduler.scheduler import PyTorchSimRunner
-    module = PyTorchSimRunner.setup_device()
-    device = module.custom_device()
+    device = torch.device("npu:0")
     #test_rmsnorm(device, seq=1)
     #test_concat(device, size1=(1, 8, 64, 64), size2=(1,8,1,64), dim=2)
     test_decode(device, 32, 3)
diff --git a/tests/test_activation.py b/tests/test_activation.py
@@ -89,9 +89,7 @@ def test_SwiGLU(device, size=(128, 128)):
     args = parser.parse_args()
     shape = tuple(map(int, args.shape.strip('()').split(',')))
 
-    from Scheduler.scheduler import PyTorchSimRunner
-    module = PyTorchSimRunner.setup_device()
-    device = module.custom_device()
+    device = torch.device("npu:0")
     test_ReLU(device, (47, 10))
     test_ReLU(device, (128, 128))
     test_ReLU(device, (4071, 429))
diff --git a/tests/test_add.py b/tests/test_add.py
@@ -58,9 +58,7 @@ def vectoradd(a, b):
     args = parser.parse_args()
     shape = tuple(map(int, args.shape.strip('()').split(',')))
 
-    from Scheduler.scheduler import PyTorchSimRunner
-    module = PyTorchSimRunner.setup_device()
-    device = module.custom_device()
+    device = torch.device("npu:0")
     test_vectoradd(device, (1, 1))
     test_vectoradd(device, (47, 10))
     test_vectoradd(device, (128, 128))
diff --git a/tests/test_batchnorm.py b/tests/test_batchnorm.py
@@ -37,9 +37,7 @@ def test_BatchNorm(device, size=(1, 16, 64, 64)):
     import sys
     sys.path.append(os.environ.get('TORCHSIM_DIR', default='/workspace/PyTorchSim'))
 
-    from Scheduler.scheduler import PyTorchSimRunner
-    module = PyTorchSimRunner.setup_device()
-    device = module.custom_device()
+    device = torch.device("npu:0")
     test_BatchNorm(device)
     test_BatchNorm(device, size=(1,64, 32, 32))
     test_BatchNorm(device, size=(1, 8, 4, 4))
diff --git a/tests/test_bmm.py b/tests/test_bmm.py
@@ -46,9 +46,7 @@ def bmm(a, b, bias):
     import sys
     sys.path.append(os.environ.get('TORCHSIM_DIR', default='/workspace/PyTorchSim'))
 
-    from Scheduler.scheduler import PyTorchSimRunner
-    module = PyTorchSimRunner.setup_device()
-    device = module.custom_device()
+    device = torch.device("npu:0")
     test_BMM(device)
     test_BMM(device, 2, 256, 128, 256)
     test_BMM(device, 2, 128, 256, 256)
diff --git a/tests/test_cnn.py b/tests/test_cnn.py
@@ -53,7 +53,5 @@ def test_CNN(device):
     import sys
     sys.path.append(os.environ.get('TORCHSIM_DIR', default='/workspace/PyTorchSim'))
 
-    from Scheduler.scheduler import PyTorchSimRunner
-    module = PyTorchSimRunner.setup_device()
-    device = module.custom_device()
+    device = torch.device("npu:0")
     test_CNN(device)
diff --git a/tests/test_conv2d.py b/tests/test_conv2d.py
@@ -40,9 +40,7 @@ def custom_conv2d(a, b, bias):
     import sys
     sys.path.append(os.environ.get('TORCHSIM_DIR', default='/workspace/PyTorchSim'))
 
-    from Scheduler.scheduler import PyTorchSimRunner
-    module = PyTorchSimRunner.setup_device()
-    device = module.custom_device()
+    device = torch.device("npu:0")
     torch._dynamo.config.cache_size_limit = 64
     with torch.no_grad():
         test_conv2d(device, batch_size=8, in_channels=3, out_channels=32, input_size=32, kernel_size=1, stride=1, padding=0)
diff --git a/tests/test_exponent.py b/tests/test_exponent.py
@@ -31,7 +31,5 @@ def exponent(a):
     import sys
     sys.path.append(os.environ.get('TORCHSIM_DIR', default='/workspace/PyTorchSim'))
 
-    from Scheduler.scheduler import PyTorchSimRunner
-    module = PyTorchSimRunner.setup_device()
-    device = module.custom_device()
+    device = torch.device("npu:0")
     test_exponent(device, size=(32, 32))
diff --git a/tests/test_gqa.py b/tests/test_gqa.py
@@ -301,9 +301,7 @@ def scaled_dot_product_attention(query, key, value, attn_mask=None, dropout_p=0.
     args = parser.parse_args()
 
     sys.path.append(os.environ.get('TORCHSIM_DIR', default='/workspace/PyTorchSim'))
-    from Scheduler.scheduler import PyTorchSimRunner
-    module = PyTorchSimRunner.setup_device()
-    device = module.custom_device()
+    device = torch.device("npu:0")
     
     test_repeat_interleave_compilation(
         device=device,
diff --git a/tests/test_indirect_access.py b/tests/test_indirect_access.py
@@ -83,9 +83,7 @@ def vectoradd(a, idx, b):
     import sys
     sys.path.append(os.environ.get('TORCHSIM_DIR', default='/workspace/PyTorchSim'))
 
-    from Scheduler.scheduler import PyTorchSimRunner
-    module = PyTorchSimRunner.setup_device()
-    device = module.custom_device()
+    device = torch.device("npu:0")
     test_scatter_full(device)
     test_scatter_full(device, size=(2048, 2048))
     test_scatter_add(device)
diff --git a/tests/test_layernorm.py b/tests/test_layernorm.py
@@ -41,9 +41,7 @@ def test_LayerNorm(device, size=(64, 64)):
     args = parser.parse_args()
     shape = tuple(map(int, args.shape.strip('()').split(',')))
 
-    from Scheduler.scheduler import PyTorchSimRunner
-    module = PyTorchSimRunner.setup_device()
-    device = module.custom_device()
+    device = torch.device("npu:0")
     with torch.no_grad():
         #test_LayerNorm(device)
         test_LayerNorm(device, shape)
diff --git a/tests/test_matmul.py b/tests/test_matmul.py
@@ -94,9 +94,7 @@ def custom_linear(a, b, bias):
     import sys
     sys.path.append(os.environ.get('TORCHSIM_DIR', default='/workspace/PyTorchSim'))
 
-    from Scheduler.scheduler import PyTorchSimRunner
-    module = PyTorchSimRunner.setup_device()
-    device = module.custom_device()
+    device = torch.device("npu:0")
     test_matmul(device, 32, 32, 32)
     test_matmul(device, 128, 128, 128)
     test_matmul(device, 256, 256, 256)
diff --git a/tests/test_mlp.py b/tests/test_mlp.py
@@ -109,9 +109,7 @@ def test_optimizer(device):
     import sys
     sys.path.append(os.environ.get('TORCHSIM_DIR', default='/workspace/PyTorchSim'))
 
-    from Scheduler.scheduler import PyTorchSimRunner
-    module = PyTorchSimRunner.setup_device()
-    device = module.custom_device()
+    device = torch.device("npu:0")
     test_mlp(device)
     test_mlp_inf(device, batch_size=1, input_size=256, hidden_size=512, output_size=256)
     test_mlp_inf(device, batch_size=8, input_size=256, hidden_size=512, output_size=256)
diff --git a/tests/test_pool.py b/tests/test_pool.py
@@ -47,9 +47,7 @@ def avgpool(a):
     import sys
     sys.path.append(os.environ.get('TORCHSIM_DIR', default='/workspace/PyTorchSim'))
 
-    from Scheduler.scheduler import PyTorchSimRunner
-    module = PyTorchSimRunner.setup_device()
-    device = module.custom_device()
+    device = torch.device("npu:0")
     #test_maxpool(device, b=1, c=8, h=16, w=16)
     #test_maxpool(device, b=1, c=8, h=112, w=112)
     test_avgpool(device, b=1, c=512, h=7, w=7)
diff --git a/tests/test_reduce.py b/tests/test_reduce.py
@@ -47,9 +47,7 @@ def reduce_sum(a, dim, keepdim):
     args = parser.parse_args()
     shape = tuple(map(int, args.shape.strip('()').split(',')))
 
-    from Scheduler.scheduler import PyTorchSimRunner
-    module = PyTorchSimRunner.setup_device()
-    device = module.custom_device()
+    device = torch.device("npu:0")
     test_reduce_sum(device, (29, 47), 1, keepdim=True)
     test_reduce_sum(device, (17, 68), 0, keepdim=True)
     test_reduce_sum(device, (327, 447), 1, keepdim=True)
diff --git a/tests/test_resnet.py b/tests/test_resnet.py
@@ -49,7 +49,5 @@ def test_resnet(device, batch=1, model_type='resnet18'):
     args = args.parse_args()
     sys.path.append(os.environ.get('TORCHSIM_DIR', default='/workspace/PyTorchSim'))
 
-    from Scheduler.scheduler import PyTorchSimRunner
-    module = PyTorchSimRunner.setup_device()
-    device = module.custom_device()
+    device = torch.device("npu:0")
     test_resnet(device, model_type=args.model_type)
diff --git a/tests/test_single_perceptron.py b/tests/test_single_perceptron.py
@@ -82,7 +82,5 @@ def weight_update(a, b, lr):
     import sys
     sys.path.append(os.environ.get('TORCHSIM_DIR', default='/workspace/PyTorchSim'))
 
-    from Scheduler.scheduler import PyTorchSimRunner
-    module = PyTorchSimRunner.setup_device()
-    device = module.custom_device()
+    device = torch.device("npu:0")
     test_single_perceptron(device)
diff --git a/tests/test_softmax.py b/tests/test_softmax.py
@@ -67,9 +67,7 @@ def forward(self, x):
     args = parser.parse_args()
     shape = tuple(map(int, args.shape.strip('()').split(',')))
 
-    from Scheduler.scheduler import PyTorchSimRunner
-    module = PyTorchSimRunner.setup_device()
-    device = module.custom_device()
+    device = torch.device("npu:0")
     test_softmax(device, size=(64, 128))
     test_softmax(device, size=(64, 128), dim=0)
     test_softmax(device, size=(256, 128))
diff --git a/tests/test_sparsity.py b/tests/test_sparsity.py
@@ -96,9 +96,7 @@ def test_mlp_inf(device, batch_size=64, input_size=64, hidden_size=32, output_si
     )
     args = parser.parse_args()
 
-    from Scheduler.scheduler import PyTorchSimRunner
-    module = PyTorchSimRunner.setup_device()
-    device = module.custom_device()
+    device = torch.device("npu:0")
 
     #test_dec_inf(device, sparsity=args.sparsity, block=args.block)
     test_mlp_inf(device, batch_size=32, input_size=784, hidden_size=512, output_size=256, sparsity=args.sparsity, block=args.block)
diff --git a/tests/test_stonne.py b/tests/test_stonne.py
@@ -54,7 +54,5 @@ def test_sparse_mm(device, input_size=128, hidden_size=128, output_size=128, spa
     args = parser.parse_args()
     sys.path.append(os.environ.get('TORCHSIM_DIR', default='/root/workspace/PyTorchSim'))
  
-    from Scheduler.scheduler import PyTorchSimRunner
-    module = PyTorchSimRunner.setup_device()
-    device = module.custom_device()
+    device = torch.device("npu:0")
     test_sparse_mm(device, args.sz, args.sz, args.sz, args.sparsity)
diff --git a/tests/test_transcendental.py b/tests/test_transcendental.py
@@ -73,9 +73,7 @@ def cos(a):
     args = parser.parse_args()
     shape = tuple(map(int, args.shape.strip('()').split(',')))
 
-    from Scheduler.scheduler import PyTorchSimRunner
-    module = PyTorchSimRunner.setup_device()
-    device = module.custom_device()
+    device = torch.device("npu:0")
     test_tanh(device)
     test_exp(device)
     test_erf(device)
diff --git a/tests/test_transformer.py b/tests/test_transformer.py
@@ -119,9 +119,7 @@ def test_MHA(device, num_heads=12, embed_dim=768, input_seq=512):
     import sys
     sys.path.append(os.environ.get('TORCHSIM_DIR', default='/workspace/PyTorchSim'))
 
-    from Scheduler.scheduler import PyTorchSimRunner
-    module = PyTorchSimRunner.setup_device()
-    device = module.custom_device()
+    device = torch.device("npu:0")
     test_EncoderBlock(device)
     # test_Attention(device, head=16, seq=512, d_k=64)
     # test_MHA(device, num_heads=12, embed_dim=768)
diff --git a/tests/test_transpose2D.py b/tests/test_transpose2D.py
@@ -46,9 +46,7 @@ def transpose(a, b):
     import sys
     sys.path.append(os.environ.get('TORCHSIM_DIR', default='/workspace/PyTorchSim'))
 
-    from Scheduler.scheduler import PyTorchSimRunner
-    module = PyTorchSimRunner.setup_device()
-    device = module.custom_device()
+    device = torch.device("npu:0")
     test_Transpose2D(device, [64, 156])
     test_Transpose2D_2(device, [16, 64])
     test_Transpose2D(device, [640, 256])
diff --git a/tests/test_transpose3D.py b/tests/test_transpose3D.py
diff --git a/tests/test_vectorops.py b/tests/test_vectorops.py
diff --git a/tests/test_view3D_2D.py b/tests/test_view3D_2D.py
diff --git a/tests/test_vit.py b/tests/test_vit.py
