cleanup add

Signed-off-by: Mayank Mishra <[email protected]>
mayank31398 · Dec 19, 2024 · 1c31acd · 1c31acd
1 parent 13da3e4
commit 1c31acd
Show file tree

Hide file tree

Showing 4 changed files with 30 additions and 32 deletions.
diff --git a/cute_kernels/kernels/add/add_scalar/cuda_implementation/__init__.py b/cute_kernels/kernels/add/add_scalar/cuda_implementation/__init__.py
@@ -1 +1,15 @@
-from .kernels_forward import add_scalar_forward_cuda
+import torch
+
+from .....constants import LIBRARY_NAME
+from .....kernel_registry import KernelRegistry
+from .....utils import cute_op
+
+
+_KERNEL_NAME = "add_scalar_forward_cuda"
+
+
+@cute_op(f"{LIBRARY_NAME}::{_KERNEL_NAME}", mutates_args={"output"})
+def add_scalar_forward_cuda(
+    x: torch.Tensor, y: float, output: torch.Tensor, vector_instruction_width: int, BLOCK_SIZE: int
+) -> None:
+    KernelRegistry.get_kernel(_KERNEL_NAME)(x, y, output, vector_instruction_width, BLOCK_SIZE)
diff --git a/cute_kernels/kernels/add/add_scalar/cuda_implementation/kernels_forward.py b/cute_kernels/kernels/add/add_scalar/cuda_implementation/kernels_forward.py
diff --git a/cute_kernels/kernels/add/add_tensor/cuda_implementation/__init__.py b/cute_kernels/kernels/add/add_tensor/cuda_implementation/__init__.py
@@ -1 +1,15 @@
-from .kernels_forward import add_tensor_forward_cuda
+import torch
+
+from .....constants import LIBRARY_NAME
+from .....kernel_registry import KernelRegistry
+from .....utils import cute_op
+
+
+_KERNEL_NAME = "add_tensor_forward_cuda"
+
+
+@cute_op(f"{LIBRARY_NAME}::{_KERNEL_NAME}", mutates_args={"output"})
+def add_tensor_forward_cuda(
+    x: torch.Tensor, y: torch.Tensor, output: torch.Tensor, vector_instruction_width: int, BLOCK_SIZE: int
+) -> None:
+    KernelRegistry.get_kernel(_KERNEL_NAME)(x, y, output, vector_instruction_width, BLOCK_SIZE)
diff --git a/cute_kernels/kernels/add/add_tensor/cuda_implementation/kernels_forward.py b/cute_kernels/kernels/add/add_tensor/cuda_implementation/kernels_forward.py