[AscendNPU-IR] Fix op document and tests.

docs/Tilelang.language/同步管道操作/T.sync_block_wait.md

@@ -35,21 +35,43 @@ T.sync_block_wait(id)
 以下代码示例了sync_block_wait同步指令的使用
 
 ```python
-def simple_sync_demo(A, Workspace, Output):
-    with T.Kernel(T.ceildiv(seq_len, block_m), is_npu=True) as (cid, _):
+def simple_sync(M, N, block_M, block_N, dtype="float16", inner_dtype="float32"):
+    m_num = M // block_M
+    n_num = N // block_N
+
+    @T.prim_func
+    def main(
+            A: T.Tensor((M, N), dtype),
+            B: T.Tensor((M, N), dtype),
+    ):
+        with T.Kernel(m_num * n_num, is_npu=True) as (cid, subid):
+
+            with T.Scope("Cube"):
+                bx = (cid // n_num) * block_M
+                by = (cid % n_num) * block_N
+
+                A_BUF = T.alloc_L1((block_M, block_N), dtype)
+                T.copy(A[bx, by], A_BUF)
+
+                C_BUF = T.alloc_L0C((block_M, block_N), inner_dtype)
 
-        with T.Scope("Cube"):
-            for i in T.serial(num_blocks):
                 with T.rs("PIPE_FIX"):
-                    T.npuir_store_fixpipe(l0_result, Workspace[cid * block_m, i * block_n])
-                    T.sync_block_set(i)
-
-        with T.Scope("Vector"):
-            for i in T.serial(num_blocks):
+                    T.sync_block_wait(1)  
+                    T.npuir_store_fixpipe(C_BUF, B[bx, by], [block_M, block_N], enable_nz2nd=True)
+                    T.sync_block_set(0)    
+
+            with T.Scope("Vector"):
+                bx = (cid // n_num) * block_M
+                by = (cid % n_num) * block_N
+
+                C_VEC = T.alloc_ub((block_M, block_N), dtype)
+
+                T.sync_block_set(1)
+
                 with T.rs("PIPE_MTE2"):
-                    T.sync_block_wait(i)
-                    T.copy(Workspace[cid * block_m, i * block_n], cross_kernel_ub)
-
+                    T.sync_block_wait(0)  
+                    T.copy(B[bx, by], C_VEC)
+                    T.sync_block_set(1)   
     return main
 ```