issue/207: add sync in every test script

test/infiniop/attention.py

@@ -101,6 +101,7 @@ def test(
     v_stride=None,
     k_cache_stride=None,
     v_cache_stride=None,
+    sync=None
 ):
     print(
         f"Testing Attention on {torch_device} with n_q_head:{n_q_head} n_kv_head:{n_kv_head} seq_len:{seq_len} head_dim:{head_dim} pos:{pos} "
@@ -139,6 +140,9 @@ def test(
     v_tensor = to_tensor(v, lib)
     k_cache_tensor = to_tensor(k_cache, lib)
     v_cache_tensor = to_tensor(v_cache, lib)
+    
+    if sync is not None:
+        sync()
 
     descriptor = infiniopAttentionDescriptor_t()
     check_error(

test/infiniop/avg_pool.py

@@ -88,6 +88,7 @@ def test(
     padding,
     strides,
     tensor_dtype=torch.float16,
+    sync=None
 ):
     print(
         f"Testing AvgPool on {torch_device} with x_shape:{x_shape} kernel_shape:{k_shape} padding:{padding} strides:{strides} dtype:{tensor_dtype}"
@@ -109,6 +110,10 @@ def test(
 
     x_tensor = to_tensor(x, lib)
     y_tensor = to_tensor(y, lib)
+    
+    if sync is not None:
+        sync()
+
     descriptor = infiniopAvgPoolDescriptor_t()
 
     check_error(

test/infiniop/causal_softmax.py

@@ -87,6 +87,7 @@ def test(
     y_stride=None,
     inplace=Inplace.OUT_OF_PLACE,
     dtype=torch.float16,
+    sync=None
 ):
     print(
         f"Testing CausalSoftmax on {torch_device} with shape:{shape} x_stride:{x_stride} y_stride:{y_stride} dtype:{dtype} inplace:{inplace}"
@@ -107,6 +108,9 @@ def test(
         y = torch.zeros(shape, dtype=dtype).to(torch_device)
         y = rearrange_if_needed(y, y_stride)
         y_tensor = to_tensor(y, lib)
+        
+    if sync is not None:
+        sync()
 
     descriptor = infiniopCausalSoftmaxDescriptor_t()
     check_error(

test/infiniop/conv.py

@@ -95,6 +95,7 @@ def test(
     dilations,
     tensor_stride=None,
     tensor_dtype=torch.float16,
+    sync=None
 ):
     assert len(pads) == len(strides) == len(dilations)
     print(
@@ -118,8 +119,11 @@ def test(
     x_tensor = to_tensor(x, lib)
     w_tensor = to_tensor(w, lib)
     y_tensor = to_tensor(y, lib)
-    descriptor = infiniopConvDescriptor_t()
+    
+    if sync is not None:
+        sync()
 
+    descriptor = infiniopConvDescriptor_t()
     check_error(
         lib.infiniopCreateConvDescriptor(
             handle,

test/infiniop/expand.py

@@ -52,6 +52,7 @@ def test(
     y_stride=None,
     x_stride=None,
     tensor_dtype=torch.float16,
+    sync=None
 ):
     print(
         f"Testing Expand on {torch_device} with x_shape:{x_shape} y_shape:{y_shape} x_stride:{x_stride} y_stride:{y_stride} dtype:{tensor_dtype}"
@@ -76,8 +77,11 @@ def test(
 
     x_tensor = to_tensor(x, lib)
     y_tensor = to_tensor(y, lib)
-    descriptor = infiniopExpandDescriptor_t()
+    
+    if sync is not None:
+        sync()
 
+    descriptor = infiniopExpandDescriptor_t()
     check_error(
         lib.infiniopCreateExpandDescriptor(
             handle,

test/infiniop/global_avg_pool.py

@@ -51,6 +51,7 @@ def test(
     torch_device,
     x_shape,
     tensor_dtype=torch.float16,
+    sync=None
 ):
     print(
         f"Testing GlobalAvgPool on {torch_device} with input tensor_shape: {x_shape} dtype: {tensor_dtype}"
@@ -70,8 +71,11 @@ def test(
 
     x_tensor = to_tensor(x, lib)
     y_tensor = to_tensor(y, lib)
-    descriptor = infiniopGlobalAvgPoolDescriptor_t()
+    
+    if sync is not None:
+        sync()
 
+    descriptor = infiniopGlobalAvgPoolDescriptor_t()
     check_error(
         lib.infiniopCreateGlobalAvgPoolDescriptor(
             handle,

test/infiniop/max_pool.py

@@ -83,6 +83,7 @@ def test(
     padding,
     strides,
     tensor_dtype=torch.float16,
+    sync=None
 ):
     print(
         f"Testing MaxPool on {torch_device} with x_shape:{x_shape} kernel_shape:{k_shape} padding:{padding} strides:{strides} dtype:{tensor_dtype}"
@@ -104,8 +105,11 @@ def test(
 
     x_tensor = to_tensor(x, lib)
     y_tensor = to_tensor(y, lib)
-    descriptor = infiniopMaxPoolDescriptor_t()
+    
+    if sync is not None:
+        sync()
 
+    descriptor = infiniopMaxPoolDescriptor_t()
     check_error(
         lib.infiniopCreateMaxPoolDescriptor(
             handle,

test/infiniop/mlp.py

@@ -65,6 +65,7 @@ def test(
     y_stride=None,
     w12_stride=None,
     w3_stride=None,
+    sync=None
 ):
     print(
         f"Testing MLP on {torch_device} with num_tokens:{num_tokens} hidden_size:{hidden_size} intermediate_size:{intermediate_size}"
@@ -97,6 +98,10 @@ def test(
     x_tensor = to_tensor(x, lib)
     w12_tensor = to_tensor(w12, lib)
     w3_tensor = to_tensor(w3, lib)
+    
+    if sync is not None:
+        sync()
+
     descriptor = infiniopMLPDescriptor_t()
     check_error(
         lib.infiniopCreateMLPDescriptor(

test/infiniop/random_sample.py

@@ -103,6 +103,7 @@ def test(
     topk,
     temperature,
     dtype=torch.float16,
+    sync=None
 ):
     print(
         f"Testing RandomSample on {torch_device} with voc:{voc} random_val:{random_val} topp:{topp} topk:{topk} temperature:{temperature} dtype:{dtype}"
@@ -122,6 +123,9 @@ def test(
 
     indices_tensor.descriptor.contents.dt = InfiniDtype.U64  # treat int64 as uint64
 
+    if sync is not None:
+        sync()
+
     descriptor = infiniopRandomSampleDescriptor_t()
     check_error(
         lib.infiniopCreateRandomSampleDescriptor(

test/infiniop/rearrange.py

@@ -131,6 +131,7 @@ def test(
     x_stride,
     y_stride,
     dtype=torch.float16,
+    sync=None
 ):
     print(
         f"Testing Rerrange on {torch_device} with shape:{shape} x_stride:{x_stride} y_stride:{y_stride} dtype:{dtype}"
@@ -145,6 +146,9 @@ def test(
     ]
 
     x_tensor, y_tensor = [to_tensor(tensor, lib) for tensor in [x, y]]
+    
+    if sync is not None:
+        sync()
 
     descriptor = infiniopRearrangeDescriptor_t()
     check_error(

test/infiniop/relu.py

@@ -55,6 +55,7 @@ def test(
     tensor_shape,
     tensor_dtype=torch.float16,
     inplace=Inplace.OUT_OF_PLACE,
+    sync=None
 ):
     print(
         f"Testing Relu on {torch_device} with tensor_shape:{tensor_shape} dtype:{tensor_dtype} inplace: {inplace.name}"
@@ -78,8 +79,11 @@ def test(
 
     x_tensor = to_tensor(x, lib)
     y_tensor = to_tensor(y, lib) if inplace == Inplace.OUT_OF_PLACE else x_tensor
-    descriptor = infiniopReluDescriptor_t()
 
+    if sync is not None:
+        sync()    
+
+    descriptor = infiniopReluDescriptor_t()
     check_error(
         lib.infiniopCreateReluDescriptor(
             handle,

test/infiniop/rope.py

@@ -117,6 +117,7 @@ def test(
     y_strides=None,
     inplace=Inplace.OUT_OF_PLACE,
     dtype=torch.float32,
+    sync=None
 ):
     if inplace == Inplace.INPLACE_X:
         y_strides = x_strides
@@ -147,8 +148,8 @@ def test(
     else:
         y_tensor = to_tensor(y, lib)
 
-    if torch_device == "npu":
-        synchronize_device(torch_device)
+    if sync is not None:
+        sync()
 
     check_error(
         lib.infiniopCreateRoPEDescriptor(