Merge branch 'fix_triton_extra_name_on_ascend_20260225' into ascend/support_bf16

Signed-off-by: ldwang <ftgreat@163.com>

Author: ftgreat

benchmark/test_binary_pointwise_perf.py

@@ -99,6 +99,97 @@ def test_general_binary_pointwise_perf(op_name, torch_op, dtypes):
     bench.run()
 
 
+class BinaryScalarPointwiseBenchmark(Benchmark):
+    """
+    Benchmark class for binary pointwise operations with a scalar operand.
+    """
+
+    DEFAULT_METRICS = DEFAULT_METRICS[:] + ["tflops"]
+
+    def set_more_shapes(self):
+        special_shapes_2d = [(1024, 2**i) for i in range(0, 20, 4)]
+        shapes_3d = [(64, 64, 2**i) for i in range(0, 20, 4)]
+        return special_shapes_2d + shapes_3d
+
+    def get_input_iter(self, cur_dtype) -> Generator:
+        for shape in self.shapes:
+            inp1 = generate_tensor_input(shape, cur_dtype, self.device)
+            inp2 = generate_tensor_input(shape, cur_dtype, self.device)
+            yield inp1, inp2
+
+    def get_tflops(self, op, *args, **kwargs):
+        shape1 = list(args[0].shape)
+        return torch.tensor(shape1).prod().item()
+
+
+@pytest.mark.parametrize(
+    "op_name, torch_op, dtypes",
+    [
+        pytest.param(
+            "ne_scalar",
+            torch.ne,
+            FLOAT_DTYPES,
+            marks=pytest.mark.ne,
+        ),
+        pytest.param(
+            "lt_scalar",
+            torch.lt,
+            FLOAT_DTYPES,
+            marks=pytest.mark.lt,
+        ),
+        pytest.param(
+            "ge_scalar",
+            torch.ge,
+            FLOAT_DTYPES,
+            marks=pytest.mark.ge,
+        ),
+    ],
+)
+def test_binary_scalar_pointwise_perf(op_name, torch_op, dtypes):
+    bench = BinaryScalarPointwiseBenchmark(
+        op_name=op_name, torch_op=torch_op, dtypes=dtypes
+    )
+    bench.run()
+
+
+class BinaryTensorScalarBenchmark(Benchmark):
+    """
+    Benchmark class for binary pointwise operations with tensor and scalar operands
+    (e.g., pow_tensor_scalar).
+    """
+
+    DEFAULT_METRICS = DEFAULT_METRICS[:] + ["tflops"]
+
+    def __init__(self, *args, scalar=2.0, **kwargs):
+        super().__init__(*args, **kwargs)
+        self.scalar = scalar
+
+    def set_more_shapes(self):
+        special_shapes_2d = [(1024, 2**i) for i in range(0, 20, 4)]
+        shapes_3d = [(64, 64, 2**i) for i in range(0, 20, 4)]
+        return special_shapes_2d + shapes_3d
+
+    def get_input_iter(self, cur_dtype) -> Generator:
+        for shape in self.shapes:
+            inp1 = generate_tensor_input(shape, cur_dtype, self.device)
+            yield inp1, self.scalar
+
+    def get_tflops(self, op, *args, **kwargs):
+        shape1 = list(args[0].shape)
+        return torch.tensor(shape1).prod().item()
+
+
+@pytest.mark.pow
+def test_pow_tensor_scalar_perf():
+    bench = BinaryTensorScalarBenchmark(
+        op_name="pow_tensor_scalar",
+        torch_op=torch.pow,
+        dtypes=FLOAT_DTYPES,
+        scalar=2.0,
+    )
+    bench.run()
+
+
 @pytest.mark.parametrize(
     "op_name, torch_op, dtypes",
     [

src/flag_gems/runtime/backend/_ascend/ops/__init__.py

@@ -1,3 +1,4 @@
+from .add import add, add_
 from .addmm import addmm
 from .all import all, all_dim, all_dims
 from .amax import amax
@@ -40,22 +41,27 @@
 from .full import full
 from .full_like import full_like
 from .gather import gather
+from .scatter_add_ import scatter_add_
 from .groupnorm import group_norm, group_norm_backward
 from .hstack import hstack
 from .index import index
 from .index_add import index_add
+from .index_put import index_put, index_put_
 from .index_select import index_select
 from .isin import isin
 from .linspace import linspace
 from .logical_and import logical_and, logical_and_
 from .log_softmax import log_softmax, log_softmax_backward
+from .lt import lt, lt_scalar
 from .masked_fill import masked_fill, masked_fill_
 from .masked_select import masked_select
 from .max import max, max_dim
 from .mean import mean, mean_dim
 from .min import min, min_dim
 from .mm import mm
 from .multinomial import multinomial
+from .ne import ne, ne_scalar
+from .nonzero import nonzero
 from .ones import ones
 from .ones_like import ones_like
 from .outer import outer
@@ -76,6 +82,7 @@
 from .softmax import softmax, softmax_backward
 from .sort import sort
 from .stack import stack
+from .topk import topk
 from .threshold import threshold, threshold_backward
 from .triu import triu
 from .unique import _unique2
@@ -88,6 +95,8 @@
 from .zeros_like import zeros_like
 
 __all__ = [
+    "add",
+    "add_",
     "addmm",
     "all",
     "all_dim",
@@ -177,8 +186,16 @@
     "bitwise_or_tensor",
     "bitwise_or_tensor_",
     "multinomial",
+    "lt",
+    "lt_scalar",
+    "ne",
+    "ne_scalar",
+    "nonzero",
     "index_add",
+    "index_put",
+    "index_put_",
     "_unique2",
+    "topk",
     "upsample_nearest2d",
     "randperm",
     "true_divide",
@@ -192,4 +209,5 @@
     "div_mode_",
     "logical_and",
     "logical_and_",
+    "scatter_add_",
 ]

src/flag_gems/runtime/backend/_ascend/ops/add.py

@@ -0,0 +1,116 @@
+import logging
+import math
+
+import torch
+import triton
+import triton.language as tl
+
+from flag_gems.runtime import torch_device_fn
+from flag_gems.utils import pointwise_dynamic
+
+logger = logging.getLogger(f'flag_gems.runtime._ascend.ops.{__name__.split(".")[-1]}')
+
+CORE_NUM = 40
+
+try:
+    import torch_npu  # noqa: F401
+    import triton.runtime.driver as driver
+
+    device = torch.npu.current_device()
+    props = driver.active.utils.get_device_properties(device)
+    CORE_NUM = props["num_vectorcore"]
+except Exception:
+    CORE_NUM = 40
+
+
+@triton.jit
+def add_kernel(
+    x_ptr,
+    y_ptr,
+    out_ptr,
+    alpha,
+    data_len,
+    BLOCK_SIZE: tl.constexpr,
+    TILE_SIZE: tl.constexpr,
+):
+    pid = tl.program_id(0)
+    iter_num = tl.cdiv(BLOCK_SIZE, TILE_SIZE)
+
+    for idx in tl.range(0, iter_num):
+        offsets = pid * BLOCK_SIZE + idx * TILE_SIZE + tl.arange(0, TILE_SIZE)
+        mask = offsets < data_len
+        x = tl.load(x_ptr + offsets, mask=mask, care_padding=False)
+        y = tl.load(y_ptr + offsets, mask=mask, care_padding=False)
+        out = x + y * alpha
+        tl.store(out_ptr + offsets, out, mask=mask)
+
+
+@pointwise_dynamic(
+    is_tensor=[True, False, False], promotion_methods=[(0, 1, "DEFAULT")]
+)
+@triton.jit
+def add_func_tensor_scalar(x, y, alpha):
+    return x + y * alpha
+
+
+@pointwise_dynamic(
+    is_tensor=[False, True, False], promotion_methods=[(0, 1, "DEFAULT")]
+)
+@triton.jit
+def add_func_scalar_tensor(x, y, alpha):
+    return x + y * alpha
+
+
+def _launch_add_kernel(x_flat, y_flat, out_flat, alpha, data_len, device):
+    TILE_SIZE = 8192
+    BLOCK_SIZE = math.ceil(data_len / CORE_NUM)
+    BLOCK_SIZE = max(BLOCK_SIZE, TILE_SIZE)
+    # Round up to multiple of TILE_SIZE for proper tiling
+    BLOCK_SIZE = triton.cdiv(BLOCK_SIZE, TILE_SIZE) * TILE_SIZE
+    grid = lambda meta: (triton.cdiv(data_len, meta["BLOCK_SIZE"]),)
+    with torch_device_fn.device(device):
+        add_kernel[grid](
+            x_flat, y_flat, out_flat, float(alpha), data_len, BLOCK_SIZE, TILE_SIZE
+        )
+
+
+def add(A, B, *, alpha=1):
+    logger.debug("GEMS_ASCEND ADD")
+    if isinstance(A, torch.Tensor) and isinstance(B, torch.Tensor):
+        if B.device != A.device:
+            B = B.to(A.device)
+        result_type = torch.result_type(A, B)
+        A_cont = A.contiguous()
+        B_cont = B.contiguous()
+        if A_cont.dtype != result_type:
+            A_cont = A_cont.to(result_type)
+        if B_cont.dtype != result_type:
+            B_cont = B_cont.to(result_type)
+        A_flat = A_cont.view(-1)
+        B_flat = B_cont.view(-1)
+        out = torch.empty_like(A_flat, dtype=result_type)
+        data_len = A_flat.numel()
+        _launch_add_kernel(A_flat, B_flat, out, alpha, data_len, A.device)
+        return out.view(A.shape)
+    elif isinstance(A, torch.Tensor):
+        return add_func_tensor_scalar(A, B, alpha)
+    elif isinstance(B, torch.Tensor):
+        return add_func_scalar_tensor(A, B, alpha)
+    else:
+        return torch.tensor(A + B * alpha)
+
+
+def add_(A, B, *, alpha=1):
+    logger.debug("GEMS_ASCEND ADD_")
+    if isinstance(A, torch.Tensor) and isinstance(B, torch.Tensor):
+        if B.device != A.device:
+            B = B.to(A.device)
+        A_flat = A.contiguous().view(-1)
+        B_flat = B.contiguous().view(-1)
+        data_len = A_flat.numel()
+        _launch_add_kernel(A_flat, B_flat, A_flat, alpha, data_len, A.device)
+        return A
+    elif isinstance(A, torch.Tensor):
+        return add_func_tensor_scalar(A, B, alpha, out0=A)
+    else:
+        raise ValueError("Unreachable.")

src/flag_gems/runtime/backend/_ascend/ops/cumsum.py

@@ -236,7 +236,7 @@ def cumsum_wrapper(inp, dim=1, dtype=None, out=None):
 
 
 def reduce_then_scan_row(x, out, M, N, compute_dtype):
-    if N <= 16384:  # persistent
+    if N <= 8192:  # persistent (capped to avoid UB overflow with int64 + multibuffer)
         TILE_SIZE = triton.next_power_of_2(N)
         reduce_then_scan_root_scan_kernel_row[(M, 1, 1)](
             x,
@@ -248,7 +248,12 @@ def reduce_then_scan_row(x, out, M, N, compute_dtype):
 
     TILE_SIZE = min(4096, triton.next_power_of_2(N))
     num_tiles = triton.cdiv(N, TILE_SIZE)
-    num_ctas = num_tiles
+    # Cap num_ctas so ROOT_SCAN_TILE_SIZE fits in UB (max 4096 for int64 + multibuffer)
+    MAX_ROOT_SCAN = 4096
+    num_ctas = min(num_tiles, MAX_ROOT_SCAN)
+    # Ensure total grid (M * num_ctas) doesn't exceed Ascend coreDim limit
+    max_ctas_for_grid = max(1, GRID_Y_LIMIT // M)
+    num_ctas = min(num_ctas, max_ctas_for_grid)
     ROOT_SCAN_TILE_SIZE = triton.next_power_of_2(num_ctas)
     tiles_per_cta = triton.cdiv(num_tiles, num_ctas)
     block_sums = torch.empty(