【KernelGen】Migrate replication_pad1d from experimental_ops to ops (flagos-ai#2005)

* Migrate replication_pad1d from experimental_ops to ops

Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>

* feat: add logger.debug for tracing

* fix: add missing blank lines in test_special_ops.py

* fix: complete truncated benchmark function (E999) and remove blank lines between decorators (E304)

---------

Co-authored-by: Claude Opus 4.6 <noreply@anthropic.com>

Author: factnn

benchmark/test_special_perf.py

@@ -1047,6 +1047,37 @@ def upsample_bicubic2d_input_fn(shape, dtype, device):
     bench.run()
 
 
+@pytest.mark.replication_pad1d
+def test_perf_replication_pad1d():
+    def replication_pad1d_input_fn(config, dtype, device):
+        shape, padding = config
+        x = torch.randn(shape, dtype=dtype, device=device)
+        yield x, list(padding)
+
+    class ReplicationPad1dBenchmark(Benchmark):
+        def set_shapes(self, shape_file_path=None):
+            self.shapes = [
+                ((2, 3, 7), (1, 2)),
+                ((4, 16, 64), (3, 1)),
+                ((8, 32, 256), (1, 2)),
+                ((32, 256), (3, 1)),
+            ]
+
+        def set_more_shapes(self):
+            return None
+
+        def get_input_iter(self, cur_dtype):
+            for config in self.shapes:
+                yield from replication_pad1d_input_fn(config, cur_dtype, self.device)
+
+    bench = ReplicationPad1dBenchmark(
+        op_name="replication_pad1d",
+        torch_op=torch.ops.aten.replication_pad1d,
+        dtypes=FLOAT_DTYPES,
+    )
+    bench.run()
+
+
 @pytest.mark.unfold
 def test_perf_unfold_backward():
     def unfold_backward_input_fn(config, dtype, device):

experimental_tests/performance/replication_pad1d_test.py

@@ -314,6 +314,8 @@ def torch_ge(v):
     ("repeat_interleave.self_int", repeat_interleave_self_int),
     ("repeat_interleave.self_Tensor", repeat_interleave_self_tensor),
     ("repeat_interleave.Tensor", repeat_interleave_tensor),
+    ("replication_pad1d", replication_pad1d),
+    ("replication_pad1d.out", replication_pad1d_out),
     ("replication_pad3d", replication_pad3d),
     ("resolve_conj", resolve_conj),
     ("resolve_neg", resolve_neg),

experimental_tests/unit/replication_pad1d_test.py

@@ -210,6 +210,7 @@
     repeat_interleave_self_tensor,
     repeat_interleave_tensor,
 )
+from flag_gems.ops.replication_pad1d import replication_pad1d, replication_pad1d_out
 from flag_gems.ops.replication_pad3d import replication_pad3d
 from flag_gems.ops.resolve_conj import resolve_conj
 from flag_gems.ops.resolve_neg import resolve_neg
@@ -519,6 +520,8 @@
     "repeat_interleave_self_int",
     "repeat_interleave_self_tensor",
     "repeat_interleave_tensor",
+    "replication_pad1d",
+    "replication_pad1d_out",
     "replication_pad3d",
     "resolve_conj",
     "resolve_neg",

src/flag_gems/__init__.py

@@ -1,7 +1,14 @@
+# Generated by KernelGen: https://github.com/flagos-ai/KernelGen
+import logging
+
 import torch
 import triton
 import triton.language as tl
 
+from flag_gems.runtime import torch_device_fn
+
+logger = logging.getLogger(__name__)
+
 
 @triton.jit
 def replication_pad1d_kernel(
@@ -46,9 +53,6 @@ def replication_pad1d_kernel(
 
 
 def _launch_replication_pad1d_kernel(input: torch.Tensor, padding, out: torch.Tensor):
-    if not input.is_cuda or not out.is_cuda:
-        raise RuntimeError("Triton kernels require CUDA tensors")
-
     if isinstance(padding, torch.Tensor):
         padding = tuple(padding.tolist())
     left, right = int(padding[0]), int(padding[1])
@@ -68,7 +72,6 @@ def _launch_replication_pad1d_kernel(input: torch.Tensor, padding, out: torch.Te
     else:
         C, W_in = input.shape
         B = 1
-        N = 1  # dummy
         in_s_c, in_s_w = input.stride()
         in_s_n = 0
         if out.dim() == 2:
@@ -89,26 +92,28 @@ def _launch_replication_pad1d_kernel(input: torch.Tensor, padding, out: torch.Te
         )
 
     grid = (triton.cdiv(W_out, 256), B * C)
-    replication_pad1d_kernel[grid](
-        input,
-        out,
-        B,
-        C,
-        W_in,
-        W_out,
-        left,
-        in_s_n if dim == 3 else in_s_n,
-        in_s_c,
-        in_s_w,
-        out_s_n if (dim == 3 or out.dim() == 3) else 0,
-        out_s_c,
-        out_s_w,
-        BLOCK_SIZE=256,
-    )
+    with torch_device_fn.device(input.device):
+        replication_pad1d_kernel[grid](
+            input,
+            out,
+            B,
+            C,
+            W_in,
+            W_out,
+            left,
+            in_s_n if dim == 3 else in_s_n,
+            in_s_c,
+            in_s_w,
+            out_s_n if (dim == 3 or out.dim() == 3) else 0,
+            out_s_c,
+            out_s_w,
+            BLOCK_SIZE=256,
+        )
     return out
 
 
 def replication_pad1d(input: torch.Tensor, padding):
+    logger.debug("GEMS REPLICATION_PAD1D")
     if isinstance(padding, torch.Tensor):
         padding = tuple(padding.tolist())
     left, right = int(padding[0]), int(padding[1])
@@ -134,6 +139,7 @@ def replication_pad1d(input: torch.Tensor, padding):
 
 
 def replication_pad1d_out(input: torch.Tensor, padding, out: torch.Tensor):
+    logger.debug("GEMS REPLICATION_PAD1D_OUT")
     if isinstance(padding, torch.Tensor):
         padding = tuple(padding.tolist())
     left, right = int(padding[0]), int(padding[1])