add addcdiv.out, addcmul.out, baddbmm.out, cat.out etc; refactor op name according to native_functions.yaml

Author: taozhiwei

src/flag_gems/__init__.py

@@ -42,10 +42,14 @@ def torch_ge(v):
     ("_grouped_mm", group_mm),
     ("_is_all_true", _is_all_true),
     ("_log_softmax", log_softmax),
+    ("_log_softmax.out", log_softmax_out),
     ("_log_softmax_backward_data", log_softmax_backward),
+    ("_log_softmax_backward_data.out", log_softmax_backward_out),
     ("_safe_softmax", _safe_softmax),
     ("_softmax", softmax),
+    ("_softmax.out", softmax_out),
     ("_softmax_backward_data", softmax_backward),
+    ("_softmax_backward_data.out", softmax_backward_out),
     (
         "_to_copy",
         to_copy,
@@ -64,11 +68,15 @@ def torch_ge(v):
     ("add.Tensor", add),
     ("add_.Tensor", add_),
     ("addcdiv", addcdiv),
+    ("addcdiv.out", addcdiv_out),
     ("addcmul", addcmul),
+    ("addcmul.out", addcmul_out),
     ("addmv", addmv),
     ("addmv.out", addmv_out),
     ("addmm", addmm),
     ("addmm.out", addmm_out),
+    ("addmm.dtype", addmm_dtype),
+    ("addmm.dtype_out", addmm_dtype_out),
     ("addr", addr),
     ("alias_copy", alias_copy),
     ("all", all),
@@ -117,6 +125,7 @@ def torch_ge(v):
     ("bmm", bmm),
     ("bmm.out", bmm_out),
     ("cat", cat),
+    ("cat.out", cat_out),
     ("celu", celu),
     ("celu_", celu_),
     ("ceil", ceil),
@@ -378,8 +387,8 @@ def torch_ge(v):
     ("rms_norm", rms_norm),
     ("roll", roll),
     ("round", round),
-    ("round.out", round_out),
     ("round_", round_),
+    ("round.out", round_out),
     ("rrelu_with_noise_backward", rrelu_with_noise_backward),
     ("rsqrt", rsqrt),
     ("rsqrt_", rsqrt_),
@@ -481,6 +490,16 @@ def torch_ge(v):
     func_name = fn.__name__ if hasattr(fn, "__name__") else str(fn)
     FULL_CONFIG_BY_FUNC.setdefault(func_name, []).append(_item)
 
+# Friendly names for only_enable(include=[...]) when the registered impl is *.out
+for _alias, _target in (
+    ("softmax", "softmax_out"),
+    ("softmax_backward", "softmax_backward_out"),
+    ("log_softmax", "log_softmax_out"),
+    ("log_softmax_backward", "log_softmax_backward_out"),
+):
+    if _target in FULL_CONFIG_BY_FUNC:
+        FULL_CONFIG_BY_FUNC.setdefault(_alias, []).extend(FULL_CONFIG_BY_FUNC[_target])
+
 
 def enable(
     lib=aten_lib,

src/flag_gems/ops/__init__.py

@@ -8,9 +8,9 @@
 from flag_gems.ops.absolute import absolute
 from flag_gems.ops.acos import acos
 from flag_gems.ops.add import add, add_
-from flag_gems.ops.addcdiv import addcdiv
-from flag_gems.ops.addcmul import addcmul
-from flag_gems.ops.addmm import addmm, addmm_out
+from flag_gems.ops.addcdiv import addcdiv, addcdiv_out
+from flag_gems.ops.addcmul import addcmul, addcmul_out
+from flag_gems.ops.addmm import addmm, addmm_dtype, addmm_dtype_out, addmm_out
 from flag_gems.ops.addmv import addmv, addmv_out
 from flag_gems.ops.addr import addr
 from flag_gems.ops.alias_copy import alias_copy, alias_copy_out
@@ -39,7 +39,7 @@
     scaled_dot_product_attention_forward,
 )
 from flag_gems.ops.avg_pool2d import avg_pool2d, avg_pool2d_backward
-from flag_gems.ops.baddbmm import baddbmm
+from flag_gems.ops.baddbmm import baddbmm, baddbmm_out
 from flag_gems.ops.batch_norm import batch_norm, batch_norm_backward
 from flag_gems.ops.bernoulli_ import bernoulli_
 from flag_gems.ops.bitwise_and import (
@@ -60,7 +60,7 @@
 )
 from flag_gems.ops.bitwise_right_shift import bitwise_right_shift
 from flag_gems.ops.bmm import bmm, bmm_out
-from flag_gems.ops.cat import cat
+from flag_gems.ops.cat import cat, cat_out
 from flag_gems.ops.ceil import ceil, ceil_, ceil_out
 from flag_gems.ops.celu import celu, celu_
 from flag_gems.ops.clamp import (
@@ -167,7 +167,12 @@
 from flag_gems.ops.log1p_ import log1p_
 from flag_gems.ops.log10 import log10, log10_, log10_out
 from flag_gems.ops.log_sigmoid import log_sigmoid
-from flag_gems.ops.log_softmax import log_softmax, log_softmax_backward
+from flag_gems.ops.log_softmax import (
+    log_softmax,
+    log_softmax_backward,
+    log_softmax_backward_out,
+    log_softmax_out,
+)
 from flag_gems.ops.logaddexp import logaddexp, logaddexp_out
 from flag_gems.ops.logical_and import logical_and, logical_and_
 from flag_gems.ops.logical_not import logical_not
@@ -274,7 +279,12 @@
 from flag_gems.ops.slice_backward import slice_backward
 from flag_gems.ops.slice_scatter import slice_scatter
 from flag_gems.ops.soft_margin_loss import soft_margin_loss, soft_margin_loss_out
-from flag_gems.ops.softmax import softmax, softmax_backward
+from flag_gems.ops.softmax import (
+    softmax,
+    softmax_backward,
+    softmax_backward_out,
+    softmax_out,
+)
 from flag_gems.ops.softplus import softplus
 from flag_gems.ops.softshrink import softshrink, softshrink_out
 from flag_gems.ops.sort import sort, sort_stable
@@ -344,8 +354,12 @@
     "add",
     "add_",
     "addcdiv",
+    "addcdiv_out",
     "addcmul",
+    "addcmul_out",
     "addmm",
+    "addmm_dtype",
+    "addmm_dtype_out",
     "addmm_out",
     "addmv",
     "addmv_out",
@@ -378,6 +392,7 @@
     "avg_pool2d",
     "avg_pool2d_backward",
     "baddbmm",
+    "baddbmm_out",
     "batch_norm",
     "batch_norm_backward",
     "bernoulli_",
@@ -398,6 +413,7 @@
     "bmm",
     "bmm_out",
     "cat",
+    "cat_out",
     "ceil",
     "ceil_",
     "ceil_out",
@@ -538,6 +554,8 @@
     "log_sigmoid",
     "log_softmax",
     "log_softmax_backward",
+    "log_softmax_backward_out",
+    "log_softmax_out",
     "log1p_",
     "logaddexp",
     "logaddexp_out",
@@ -676,6 +694,8 @@
     "soft_margin_loss_out",
     "softmax",
     "softmax_backward",
+    "softmax_backward_out",
+    "softmax_out",
     "softplus",
     "softshrink",
     "softshrink_out",

src/flag_gems/ops/addcdiv.py

@@ -16,12 +16,14 @@ def addcdiv_kernel(x, t1, t2, value):
     return x + value * (t1 / t2)
 
 
-def addcdiv(inp, tensor1, tensor2, value=1.0, out=None):
-    logger.debug("GEMS ADDCDIV FORWARD")
-
-    if out is None:
-        out = torch.empty_like(inp)
-
+def addcdiv_out(inp, tensor1, tensor2, *, value=1.0, out):
+    logger.debug("GEMS ADDCDIV_OUT")
     addcdiv_kernel(inp, tensor1, tensor2, value, out0=out)
-
     return out
+
+
+def addcdiv(inp, tensor1, tensor2, value=1.0):
+    """Functional entry; CUDA may dispatch here without hitting ``addcdiv.out``."""
+    logger.debug("GEMS ADDCDIV")
+    out = torch.empty_like(inp)
+    return addcdiv_kernel(inp, tensor1, tensor2, value, out0=out)

src/flag_gems/ops/addcmul.py

@@ -16,15 +16,21 @@ def addcmul_forward(x, t1, t2, value):
     return x + value * t1 * t2
 
 
-def addcmul(inp, tensor1, tensor2, *, value=1.0, out=None):
-    logger.debug("GEMS ADDCMUL FORWARD")
-    if out is not None:
-        broadcast_shape = torch.broadcast_shapes(
-            inp.shape, tensor1.shape, tensor2.shape
-        )
-        if list(out.shape) != list(broadcast_shape):
-            out.resize_(broadcast_shape)
-        addcmul_forward(inp, tensor1, tensor2, value, out0=out)
-        return out
-    else:
-        return addcmul_forward(inp, tensor1, tensor2, value)
+def addcmul_out(inp, tensor1, tensor2, *, value=1.0, out):
+    logger.debug("GEMS ADDCMUL_OUT")
+    broadcast_shape = torch.broadcast_shapes(inp.shape, tensor1.shape, tensor2.shape)
+    if list(out.shape) != list(broadcast_shape):
+        out.resize_(broadcast_shape)
+    addcmul_forward(inp, tensor1, tensor2, value, out0=out)
+    return out
+
+
+def addcmul(inp, tensor1, tensor2, *, value=1.0):
+    """Functional entry; keep alongside ``addcmul.out`` for dispatch coverage."""
+    logger.debug("GEMS ADDCMUL")
+    broadcast_shape = torch.broadcast_shapes(inp.shape, tensor1.shape, tensor2.shape)
+    dtype = torch.promote_types(
+        inp.dtype, torch.promote_types(tensor1.dtype, tensor2.dtype)
+    )
+    out = torch.empty(broadcast_shape, device=inp.device, dtype=dtype)
+    return addcmul_out(inp, tensor1, tensor2, value=value, out=out)

src/flag_gems/ops/addmm.py

@@ -184,3 +184,39 @@ def addmm_out(bias, mat1, mat2, *, beta=1, alpha=1, out=None):
             out.stride(1),
         )
     return out
+
+
+def addmm_dtype(bias, mat1, mat2, out_dtype, *, beta=1, alpha=1):
+    logger.debug("GEMS ADDMM_DTYPE")
+    out = torch.empty(
+        (mat1.shape[0], mat2.shape[1]),
+        device=mat1.device,
+        dtype=out_dtype,
+    )
+    return addmm_dtype_out(bias, mat1, mat2, out_dtype, beta=beta, alpha=alpha, out=out)
+
+
+def addmm_dtype_out(bias, mat1, mat2, out_dtype, *, beta=1, alpha=1, out):
+    logger.debug("GEMS ADDMM_DTYPE_OUT")
+    if mat1.dtype != mat2.dtype:
+        raise RuntimeError(
+            f"mat1 and mat2 must have the same dtype, but got {mat1.dtype} and {mat2.dtype}"
+        )
+    if out.dtype != out_dtype:
+        raise RuntimeError(
+            "out_dtype must be the same as the dtype of the provided out tensor"
+        )
+    if not (
+        out_dtype == mat1.dtype
+        or (
+            out_dtype == torch.float32 and mat1.dtype in (torch.float16, torch.bfloat16)
+        )
+    ):
+        raise RuntimeError(
+            "out_dtype must be the same as input dtype or fp32 for fp16/bf16 inputs"
+        )
+    if bias.dtype != out_dtype and bias.dtype != mat1.dtype:
+        raise RuntimeError("self dtype must match either out_dtype or mat1 dtype")
+
+    bias_c = bias.to(out_dtype)
+    return addmm_out(bias_c, mat1, mat2, beta=beta, alpha=alpha, out=out)

src/flag_gems/ops/baddbmm.py

@@ -134,6 +134,38 @@ def baddbmm_kernel(
     tl.store(o_ptrs, o, mask=mask_c)
 
 
+def _baddbmm_launch(bias, A, B, beta, alpha, out):
+    batch, M, K = A.shape
+    _, _, N = B.shape
+    A = A.contiguous()
+    B = B.contiguous()
+    bbias = torch.broadcast_to(bias, (batch, M, N)).contiguous()
+    bias_batch_stride = bbias.stride(0)
+    bias_M_stride = bbias.stride(1)
+    bias_N_stride = bbias.stride(-1)
+
+    grid = lambda meta: (
+        triton.cdiv(meta["M"], meta["TILE_M"]),
+        triton.cdiv(meta["N"], meta["TILE_N"]),
+        batch,
+    )
+    with torch_device_fn.device(A.device):
+        baddbmm_kernel[grid](
+            A,
+            B,
+            out,
+            bbias,
+            alpha,
+            beta,
+            M,
+            N,
+            K,
+            bias_batch_stride=bias_batch_stride,
+            bias_M_stride=bias_M_stride,
+            bias_N_stride=bias_N_stride,
+        )
+
+
 class BaddbmmFunction(torch.autograd.Function):
     @staticmethod
     def forward(ctx, bias, A, B, beta, alpha):
@@ -145,35 +177,8 @@ def forward(ctx, bias, A, B, beta, alpha):
 
         batch, M, K = A.shape
         _, _, N = B.shape
-        A = A.contiguous()
-        B = B.contiguous()
         out = torch.empty((batch, M, N), dtype=A.dtype, device=A.device)
-
-        bbias = torch.broadcast_to(bias, (batch, M, N)).contiguous()
-        bias_batch_stride = bbias.stride(0)
-        bias_M_stride = bbias.stride(1)
-        bias_N_stride = bbias.stride(-1)
-
-        grid = lambda meta: (
-            triton.cdiv(meta["M"], meta["TILE_M"]),
-            triton.cdiv(meta["N"], meta["TILE_N"]),
-            batch,
-        )
-        with torch_device_fn.device(A.device):
-            baddbmm_kernel[grid](
-                A,
-                B,
-                out,
-                bbias,
-                alpha,
-                beta,
-                M,
-                N,
-                K,
-                bias_batch_stride=bias_batch_stride,
-                bias_M_stride=bias_M_stride,
-                bias_N_stride=bias_N_stride,
-            )
+        _baddbmm_launch(bias, A, B, beta, alpha, out)
         return out
 
     @staticmethod
@@ -234,6 +239,24 @@ def compute_B_grad(A, d_output, alpha):
     return grad_B
 
 
+def baddbmm_out(bias, A, B, *, beta=1.0, alpha=1.0, out):
+    logger.debug("GEMS BADDBMM_OUT")
+    batch, M, K = A.shape
+    _, _, N = B.shape
+    assert (
+        out.shape == (batch, M, N) and out.dtype == A.dtype
+    ), "Incompatible output shape or dtype for baddbmm.out"
+    _baddbmm_launch(
+        bias.contiguous(),
+        A.contiguous(),
+        B.contiguous(),
+        beta,
+        alpha,
+        out,
+    )
+    return out
+
+
 def baddbmm(bias, A, B, beta=1.0, alpha=1.0):
     return BaddbmmFunction.apply(
         bias.contiguous(),