Fix: test/test_signal_pad (#432)

Author: joydddd

test/test_signal_wait.expected

@@ -72,29 +72,6 @@ def gmem_signal_cas_kernel(signal_pad: torch.Tensor, *, _launcher=_default_launc
     _launcher(_gmem_signal_cas_kernel_kernel, (n,), signal_pad, signal_pad.stride(0), num_warps=4, num_stages=3)
     return signal_pad
 
---- assertExpectedJournal(TestWait.test_signal_stack_signalpad)
-from __future__ import annotations
-
-import torch
-import helion
-import triton
-import triton.language as tl
-from helion.runtime import default_launcher as _default_launcher
-
-@triton.jit
-def _gmem_signal_pointers_kernel_kernel(signal_pad_ptrs, signal_pad_ptrs_size_0, example_stride_0, signal_pad_ptrs_stride_0, _RDIM_SIZE_1: tl.constexpr):
-    pid_0 = tl.program_id(0)
-    offset_0 = pid_0
-    indices_1 = tl.arange(0, _RDIM_SIZE_1).to(tl.int32)
-    mask_1 = indices_1 < signal_pad_ptrs_size_0
-    ptr_tile = tl.load(signal_pad_ptrs + indices_1 * signal_pad_ptrs_stride_0, mask_1, other=0)
-    helion.runtime.triton_send_signal(addr=ptr_tile.to(tl.pointer_type(tl.int32))[:] + (offset_0 * example_stride_0)[None], update=1, sem='release', scope='gpu', op='atomic_xchg', skip_sync=False)
-
-def gmem_signal_pointers_kernel(signal_pad_ptrs: torch.Tensor, example: torch.Tensor, *, _launcher=_default_launcher):
-    _RDIM_SIZE_1 = triton.next_power_of_2(signal_pad_ptrs.size(0))
-    _launcher(_gmem_signal_pointers_kernel_kernel, (example.size(0),), signal_pad_ptrs, signal_pad_ptrs.size(0), example.stride(0), signal_pad_ptrs.stride(0), _RDIM_SIZE_1, num_warps=4, num_stages=3)
-    return signal_pad_ptrs
-
 --- assertExpectedJournal(TestWait.test_signal_multiple)
 from __future__ import annotations
 
@@ -143,6 +120,31 @@ def gmem_signal_tensor_bar_kernel(signal_pad: torch.Tensor, *, _launcher=_defaul
     _launcher(_gmem_signal_tensor_bar_kernel_kernel, (triton.cdiv(n, _BLOCK_SIZE_0),), signal_pad, signal_pad.stride(0), _BLOCK_SIZE_0, num_warps=4, num_stages=3)
     return signal_pad
 
+--- assertExpectedJournal(TestWait.test_signal_stack_signalpad)
+from __future__ import annotations
+
+import torch
+import helion
+import triton
+import triton.language as tl
+from helion.runtime import default_launcher as _default_launcher
+
+helion.runtime.set_triton_allocator()
+
+@triton.jit
+def _gmem_signal_pointers_kernel_kernel(signal_pad_ptrs, signal_pad_ptrs_size_0, example_stride_0, signal_pad_ptrs_stride_0, _RDIM_SIZE_1: tl.constexpr):
+    pid_0 = tl.program_id(0)
+    offset_0 = pid_0
+    indices_1 = tl.arange(0, _RDIM_SIZE_1).to(tl.int32)
+    mask_1 = indices_1 < signal_pad_ptrs_size_0
+    ptr_tile = tl.load(signal_pad_ptrs + indices_1 * signal_pad_ptrs_stride_0, mask_1, other=0)
+    helion.runtime.triton_send_signal(addr=ptr_tile.to(tl.pointer_type(tl.int32))[:] + (offset_0 * example_stride_0)[None], update=1, sem='release', scope='gpu', op='atomic_xchg', skip_sync=False)
+
+def gmem_signal_pointers_kernel(signal_pad_ptrs: torch.Tensor, example: torch.Tensor, *, _launcher=_default_launcher):
+    _RDIM_SIZE_1 = triton.next_power_of_2(signal_pad_ptrs.size(0))
+    _launcher(_gmem_signal_pointers_kernel_kernel, (example.size(0),), signal_pad_ptrs, signal_pad_ptrs.size(0), example.stride(0), signal_pad_ptrs.stride(0), _RDIM_SIZE_1, num_warps=4, num_stages=3)
+    return signal_pad_ptrs
+
 --- assertExpectedJournal(TestWait.test_wait_2d_tile)
 from __future__ import annotations
 
@@ -256,31 +258,6 @@ def gmem_wait_multi_bar_kernel_cas(signal_pad: torch.Tensor, *, _launcher=_defau
     _launcher(_gmem_wait_multi_bar_kernel_cas_kernel, (triton.cdiv(N, _BLOCK_SIZE_0),), signal_pad, signal_pad.stride(0), _BLOCK_SIZE_0, num_warps=4, num_stages=3)
     return signal_pad
 
---- assertExpectedJournal(TestWait.test_wait_stack_signalpad)
-from __future__ import annotations
-
-import torch
-import helion
-import triton
-import triton.language as tl
-from helion.runtime import default_launcher as _default_launcher
-
-@triton.jit
-def _gmem_wait_pointers_kernel_kernel(signal_pad_ptrs, out, signal_pad_ptrs_size_0, example_stride_0, out_stride_0, signal_pad_ptrs_stride_0, _RDIM_SIZE_1: tl.constexpr):
-    pid_0 = tl.program_id(0)
-    offset_0 = pid_0
-    indices_1 = tl.arange(0, _RDIM_SIZE_1).to(tl.int32)
-    mask_1 = indices_1 < signal_pad_ptrs_size_0
-    dev_tile = tl.load(signal_pad_ptrs + indices_1 * signal_pad_ptrs_stride_0, mask_1, other=0)
-    helion.runtime.triton_wait_multiple_signal(addr=dev_tile.to(tl.pointer_type(tl.int32))[:] + (offset_0 * example_stride_0)[None], expect=1, update=0, sem='acquire', scope='gpu', op='ld', skip_sync=False)
-    tl.store(out + offset_0 * out_stride_0, offset_0, None)
-
-def gmem_wait_pointers_kernel(signal_pad_ptrs: torch.Tensor, example: torch.Tensor, *, _launcher=_default_launcher):
-    out = torch.empty_like(example)
-    _RDIM_SIZE_1 = triton.next_power_of_2(signal_pad_ptrs.size(0))
-    _launcher(_gmem_wait_pointers_kernel_kernel, (example.size(0),), signal_pad_ptrs, out, signal_pad_ptrs.size(0), example.stride(0), out.stride(0), signal_pad_ptrs.stride(0), _RDIM_SIZE_1, num_warps=4, num_stages=3)
-    return out
-
 --- assertExpectedJournal(TestWait.test_wait_pointers)
 from __future__ import annotations
 
@@ -311,3 +288,30 @@ def gmem_wait_pointers_kernel(signal_pad_ptrs: torch.Tensor, pad_shape: hl.const
     _BLOCK_SIZE_1 = N
     _launcher(_gmem_wait_pointers_kernel_kernel, (4,), signal_pad_ptrs, out, out.stride(0), signal_pad_ptrs.stride(0), N, _BLOCK_SIZE_1, num_warps=4, num_stages=3)
     return out
+
+--- assertExpectedJournal(TestWait.test_wait_stack_signalpad)
+from __future__ import annotations
+
+import torch
+import helion
+import triton
+import triton.language as tl
+from helion.runtime import default_launcher as _default_launcher
+
+helion.runtime.set_triton_allocator()
+
+@triton.jit
+def _gmem_wait_pointers_kernel_kernel(signal_pad_ptrs, out, signal_pad_ptrs_size_0, example_stride_0, out_stride_0, signal_pad_ptrs_stride_0, _RDIM_SIZE_1: tl.constexpr):
+    pid_0 = tl.program_id(0)
+    offset_0 = pid_0
+    indices_1 = tl.arange(0, _RDIM_SIZE_1).to(tl.int32)
+    mask_1 = indices_1 < signal_pad_ptrs_size_0
+    dev_tile = tl.load(signal_pad_ptrs + indices_1 * signal_pad_ptrs_stride_0, mask_1, other=0)
+    helion.runtime.triton_wait_multiple_signal(addr=dev_tile.to(tl.pointer_type(tl.int32))[:] + (offset_0 * example_stride_0)[None], expect=1, update=0, sem='acquire', scope='gpu', op='ld', skip_sync=False)
+    tl.store(out + offset_0 * out_stride_0, offset_0, None)
+
+def gmem_wait_pointers_kernel(signal_pad_ptrs: torch.Tensor, example: torch.Tensor, *, _launcher=_default_launcher):
+    out = torch.empty_like(example)
+    _RDIM_SIZE_1 = triton.next_power_of_2(signal_pad_ptrs.size(0))
+    _launcher(_gmem_wait_pointers_kernel_kernel, (example.size(0),), signal_pad_ptrs, out, signal_pad_ptrs.size(0), example.stride(0), out.stride(0), signal_pad_ptrs.stride(0), _RDIM_SIZE_1, num_warps=4, num_stages=3)
+    return out