[KernelGen][Nvidia] Add addr_ operator with Triton kernel

benchmark/test_addr_.py

@@ -0,0 +1,35 @@
+# Generated by KernelGen: https://github.com/flagos-ai/KernelGen
+from typing import Generator
+
+import pytest
+import torch
+
+from . import base, consts
+
+
+class AddrInplaceBenchmark(base.BlasBenchmark):
+    def set_more_shapes(self):
+        return []
+
+    def get_input_iter(self, dtype) -> Generator:
+        for shape in self.shapes:
+            m, n = shape[0], shape[1]
+            yield from self.input_fn(m, n, dtype, self.device)
+
+
+def _input_fn(m, n, cur_dtype, device):
+    inp1 = torch.randn([m, n], dtype=cur_dtype, device=device)
+    inp2 = torch.randn([m], dtype=cur_dtype, device=device)
+    inp3 = torch.randn([n], dtype=cur_dtype, device=device)
+    yield inp1, inp2, inp3, {"alpha": 0.5, "beta": 0.5}
+
+
+@pytest.mark.addr_
+def test_addr_():
+    bench = AddrInplaceBenchmark(
+        op_name="addr_",
+        input_fn=_input_fn,
+        torch_op=torch.Tensor.addr_,
+        dtypes=consts.FLOAT_DTYPES,
+    )
+    bench.run()

conf/operators.yaml

@@ -388,6 +388,19 @@ ops:
       - BLAS
     stages:
       - stable: '4.0'
+  - id: addr_
+    description: |
+      In-place version of addr. Performs the outer-product of vectors `vec1` and `vec2`
+      and adds it to the matrix `input` in-place.
+    for:
+      - addr_
+    labels:
+      - aten
+      - KernelGen
+    kind:
+      - BLAS
+    stages:
+      - stable: '4.0'
   - id: affine_grid_generator
     description: |
       Generates a 2D or 3D flow field (sampling grid), given a batch of affine matrices theta.

src/flag_gems/__init__.py

@@ -211,6 +211,7 @@ def torch_ge(v):
     ("addmv", addmv),
     ("addmv.out", addmv_out),
     ("addr", addr),
+    ("addr_", addr_),
     ("affine_grid_generator", affine_grid_generator),
     ("alias", alias),
     ("alias_copy", alias_copy),

src/flag_gems/ops/__init__.py

@@ -105,6 +105,7 @@
 from flag_gems.ops.addmm_ import addmm_
 from flag_gems.ops.addmv import addmv, addmv_out
 from flag_gems.ops.addr import addr
+from flag_gems.ops.addr_ import addr_
 from flag_gems.ops.affine_grid_generator import affine_grid_generator
 from flag_gems.ops.alias import alias
 from flag_gems.ops.alias_copy import alias_copy, alias_copy_out
@@ -835,6 +836,7 @@
     "addmv",
     "addmv_out",
     "addr",
+    "addr_",
     "affine_grid_generator",
     "alias",
     "alias_copy",

src/flag_gems/ops/addr_.py

@@ -0,0 +1,99 @@
+# 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
+from flag_gems.utils import libentry
+
+logger = logging.getLogger(__name__)
+
+
+@libentry()
+@triton.jit(do_not_specialize=["beta", "alpha"])
+def addr_inplace_kernel(
+    input_ptr,
+    vec1_ptr,
+    vec2_ptr,
+    beta,
+    alpha,
+    M,
+    N,
+    stride_input_m,
+    stride_input_n,
+    stride_vec1,
+    stride_vec2,
+    BLOCK_SIZE_M: tl.constexpr,
+    BLOCK_SIZE_N: tl.constexpr,
+):
+    pid_m = tl.program_id(0)
+    pid_n = tl.program_id(1)
+
+    offs_m = pid_m * BLOCK_SIZE_M + tl.arange(0, BLOCK_SIZE_M)
+    offs_n = pid_n * BLOCK_SIZE_N + tl.arange(0, BLOCK_SIZE_N)
+
+    mask_m = offs_m < M
+    mask_n = offs_n < N
+
+    vec1_ptrs = vec1_ptr + offs_m * stride_vec1
+    vec2_ptrs = vec2_ptr + offs_n * stride_vec2
+
+    vec1 = tl.load(vec1_ptrs, mask=mask_m, other=0.0).to(tl.float32)
+    vec2 = tl.load(vec2_ptrs, mask=mask_n, other=0.0).to(tl.float32)
+
+    input_ptrs = (
+        input_ptr + offs_m[:, None] * stride_input_m + offs_n[None, :] * stride_input_n
+    )
+    mask_2d = mask_m[:, None] & mask_n[None, :]
+    input_val = tl.load(input_ptrs, mask=mask_2d, other=0.0).to(tl.float32)
+
+    # Fused computation: beta * input + alpha * outer(vec1, vec2)
+    result = beta * input_val + alpha * (vec1[:, None] * vec2[None, :])
+
+    tl.store(input_ptrs, result, mask=mask_2d)
+
+
+def addr_(input, vec1, vec2, *, beta=1, alpha=1):
+    logger.debug("GEMS ADDR_")
+    assert input.dtype in (
+        torch.float16,
+        torch.bfloat16,
+        torch.float32,
+        torch.float64,
+    ), f"addr_: unsupported dtype {input.dtype}, expected floating point"
+    if vec1.dim() != 1 or vec2.dim() != 1:
+        raise ValueError("addr_: expected 1-D vectors")
+
+    M, N = input.shape
+    if vec1.shape[0] != M or vec2.shape[0] != N:
+        raise ValueError(
+            f"addr_: vec1 size {vec1.shape[0]} must match input rows {M}, "
+            f"vec2 size {vec2.shape[0]} must match input cols {N}"
+        )
+
+    # Single fused kernel dispatch: reads input, computes, writes back in-place
+    BLOCK_SIZE_M = 32  # Tile size for rows
+    BLOCK_SIZE_N = 32  # Tile size for columns
+    grid = lambda META: (
+        triton.cdiv(M, BLOCK_SIZE_M),
+        triton.cdiv(N, BLOCK_SIZE_N),
+    )
+    with torch_device_fn.device(input.device):
+        addr_inplace_kernel[grid](
+            input,
+            vec1,
+            vec2,
+            beta,
+            alpha,
+            M,
+            N,
+            input.stride(0),
+            input.stride(1),
+            vec1.stride(0),
+            vec2.stride(0),
+            BLOCK_SIZE_M=BLOCK_SIZE_M,
+            BLOCK_SIZE_N=BLOCK_SIZE_N,
+        )
+    return input

tests/test_addr_.py

@@ -0,0 +1,45 @@
+# Generated by KernelGen: https://github.com/flagos-ai/KernelGen
+import pytest
+import torch
+
+import flag_gems
+
+from . import accuracy_utils as utils
+from . import conftest as cfg
+
+if cfg.QUICK_MODE:
+    # Quick mode: minimal shape for fast CI validation
+    MN_SHAPES = [
+        (1, 32),
+    ]
+    # Quick mode: only float32 for fast validation
+    FLOAT_DTYPES = [torch.float32]
+else:
+    # Cover edge case (1-row), medium, and non-aligned shapes
+    MN_SHAPES = [
+        (1, 32),
+        (160, 1024),
+        (5333, 497),
+    ]
+    FLOAT_DTYPES = utils.FLOAT_DTYPES
+
+
+@pytest.mark.addr_
+@pytest.mark.parametrize("M, N", MN_SHAPES)
+@pytest.mark.parametrize("dtype", FLOAT_DTYPES)
+def test_addr_(M, N, dtype):
+    input_tensor = torch.randn((M, N), dtype=dtype, device=flag_gems.device)
+    vec1 = torch.randn((M,), dtype=dtype, device=flag_gems.device)
+    vec2 = torch.randn((N,), dtype=dtype, device=flag_gems.device)
+    alpha = torch.randn((), dtype=dtype, device=flag_gems.device)
+    beta = torch.randn((), dtype=dtype, device=flag_gems.device)
+
+    ref_inp = utils.to_reference(input_tensor, True)
+    ref_vec1 = utils.to_reference(vec1, True)
+    ref_vec2 = utils.to_reference(vec2, True)
+
+    ref_out = torch.addr(ref_inp, ref_vec1, ref_vec2, alpha=alpha, beta=beta)
+    with flag_gems.use_gems():
+        res_out = input_tensor.addr_(vec1, vec2, alpha=alpha, beta=beta)
+
+    utils.gems_assert_close(res_out, ref_out, dtype, equal_nan=True)