Merge branch 'master' into fix/tensordescriptor

Signed-off-by: wuhulalala <114969601+wuhulalala@users.noreply.github.com>

Author: wuhulalala

benchmark/conftest.py

@@ -137,14 +137,19 @@ def pytest_addoption(parser):
         help="Specify the shape file name for benchmarks. If not specified, a default shape list will be used.",
     )
 
-    parser.addoption(
-        "--record",
-        action="store",
-        default="none",
-        required=False,
-        choices=["none", "log"],
-        help="Benchmark info recorded in log files or not",
-    )
+    try:
+        parser.addoption(
+            "--record",
+            action="store",
+            default="none",
+            required=False,
+            choices=["none", "log"],
+            help="Benchmark info recorded in log files or not",
+        )
+    except ValueError:
+        # Mixed test+benchmark pytest runs may already register --record in
+        # tests/conftest.py. Reuse the existing option in that case.
+        pass
 
     parser.addoption(
         "--parallel",
@@ -158,11 +163,16 @@ def pytest_addoption(parser):
         ),
     )
 
-    parser.addoption(
-        "--collect-marks",
-        action="store_true",
-        help="Collect the tests with marker information without executing them",
-    )
+    try:
+        parser.addoption(
+            "--collect-marks",
+            action="store_true",
+            help="Collect the tests with marker information without executing them",
+        )
+    except ValueError:
+        # Mixed test+benchmark pytest runs may already register this option in
+        # tests/conftest.py. Reuse the existing option in that case.
+        pass
 
 
 def pytest_configure(config):

benchmark/core_shapes.yaml

@@ -111,6 +111,19 @@ MvAndOuterBenchmark:
     - [8192, 8192]
     - [10240, 10240]  #from perf
 
+mm_self_transpose:
+  shapes:
+    - [257, 96]
+    - [384, 384]
+    - [1536, 320]
+    - [3072, 768]
+    - [1024, 1024]
+    - [4096, 1024]
+    - [5333, 71]
+    - [8192, 2048]
+    - [10240, 4096]
+  shape_desc: "M, K"
+
 # NORM shapes can be either 3D or 4D:
 # - 3D shapes are represented as [batch_size, channels, hidden_size]
 # - 4D shapes are represented as [batch_size, channels, height, width]

benchmark/test_blas_perf.py

@@ -135,6 +135,21 @@ def get_tflops(self, op, *args, **kwargs):
         return total_flops
 
 
+class MmSelfTransposeBenchmark(GenericBenchmark2DOnly):
+    """
+    Benchmark for the mm(a, a.t()) fast path.
+    """
+
+    DEFAULT_METRICS = DEFAULT_METRICS[:] + ["tflops"]
+
+    def set_more_shapes(self):
+        return None
+
+    def get_tflops(self, op, *args, **kwargs):
+        m, k = args[0].shape
+        return 2 * m * m * k
+
+
 def addmm_input_fn(b, m, n, k, cur_dtype, device, b_column_major):
     inp1 = torch.randn([m, k], dtype=cur_dtype, device=device)
     bias = torch.randn([m, n], dtype=cur_dtype, device=device)
@@ -184,6 +199,16 @@ def mm_input_fn(b, m, n, k, cur_dtype, device, b_column_major):
         yield inp1, inp2
 
 
+def torch_mm_self_transpose(inp):
+    return torch.mm(inp, inp.t())
+
+
+def mm_self_transpose_input_fn(shape, cur_dtype, device):
+    m, k = shape
+    inp = torch.randn([k, m], dtype=cur_dtype, device=device).t()
+    yield inp,
+
+
 def group_mm_input_fn(groups, N, K, cur_dtype, device):
     assert cur_dtype == torch.bfloat16
     import random
@@ -439,6 +464,17 @@ def test_mv_and_outer_benchmark(op_name, torch_op, input_fn):
     bench.run()
 
 
+@pytest.mark.mm
+def test_mm_self_transpose_benchmark():
+    bench = MmSelfTransposeBenchmark(
+        input_fn=mm_self_transpose_input_fn,
+        op_name="mm_self_transpose",
+        torch_op=torch_mm_self_transpose,
+        dtypes=FLOAT_DTYPES,
+    )
+    bench.run()
+
+
 class AddmvBenchmark(GenericBenchmark2DOnly):
     """
     Benchmark for addmv

benchmark/test_clip_perf.py

@@ -0,0 +1,37 @@
+import pytest
+import torch
+
+from benchmark.attri_util import FLOAT_DTYPES
+from benchmark.conftest import BenchLevel, Config
+from benchmark.performance_utils import GenericBenchmark, generate_tensor_input
+
+
+def clip_input_fn(shape, cur_dtype, device):
+    inp = generate_tensor_input(shape, cur_dtype, device)
+    yield inp, -0.5, 0.5
+    if Config.bench_level == BenchLevel.COMPREHENSIVE:
+        yield inp, None, 0.5
+        yield inp, -0.5, None
+
+
+@pytest.mark.clip
+def test_clip():
+    bench = GenericBenchmark(
+        input_fn=clip_input_fn,
+        op_name="clip",
+        torch_op=torch.clip,
+        dtypes=FLOAT_DTYPES,
+    )
+    bench.run()
+
+
+@pytest.mark.clip_
+def test_clip_inplace():
+    bench = GenericBenchmark(
+        input_fn=clip_input_fn,
+        op_name="clip_",
+        torch_op=torch.clip_,
+        dtypes=FLOAT_DTYPES,
+        is_inplace=True,
+    )
+    bench.run()

src/flag_gems/__init__.py

@@ -129,6 +129,8 @@ def torch_ge(v):
     ("clamp_", clamp_),
     ("clamp_.Tensor", clamp_tensor_),
     ("clamp_min_", clamp_min_),
+    ("clip", clip),
+    ("clip_", clip_),
     ("conj_physical", conj_physical),
     ("constant_pad_nd", constant_pad_nd),
     # ("contiguous", contiguous),

src/flag_gems/ops/__init__.py

@@ -71,6 +71,7 @@
     clamp_tensor,
     clamp_tensor_,
 )
+from flag_gems.ops.clip import clip, clip_
 from flag_gems.ops.conj_physical import conj_physical
 from flag_gems.ops.contiguous import contiguous
 from flag_gems.ops.conv1d import conv1d
@@ -410,6 +411,8 @@
     "clamp_min_",
     "clamp_tensor",
     "clamp_tensor_",
+    "clip",
+    "clip_",
     "constant_pad_nd",
     "contiguous",
     "conv1d",

src/flag_gems/ops/clip.py

@@ -0,0 +1,52 @@
+import logging
+
+import triton
+import triton.language as tl
+
+from flag_gems.utils import pointwise_dynamic
+
+logger = logging.getLogger(__name__)
+
+
+@pointwise_dynamic(
+    is_tensor=[True, False, False], promotion_methods=[(0, 1, 2, "DEFAULT")]
+)
+@triton.jit
+def clip_func(x, mini, maxi):
+    return tl.minimum(maxi, tl.maximum(mini, x.to(tl.float32)))
+
+
+@pointwise_dynamic(is_tensor=[True, False], promotion_methods=[(0, 1, "DEFAULT")])
+@triton.jit
+def clip_func_min(x, mini):
+    return tl.maximum(mini, x.to(tl.float32))
+
+
+@pointwise_dynamic(is_tensor=[True, False], promotion_methods=[(0, 1, "DEFAULT")])
+@triton.jit
+def clip_func_max(x, maxi):
+    return tl.minimum(maxi, x.to(tl.float32))
+
+
+def clip(A, mini=None, maxi=None):
+    logger.debug("GEMS CLIP")
+    if mini is None and maxi is None:
+        raise ValueError("At least one of mini or maxi must not be None")
+    elif mini is None:
+        return clip_func_max(A, maxi)
+    elif maxi is None:
+        return clip_func_min(A, mini)
+    else:
+        return clip_func(A, mini, maxi)
+
+
+def clip_(A, mini=None, maxi=None):
+    logger.debug("GEMS CLIP_")
+    if mini is None and maxi is None:
+        raise ValueError("At least one of mini or maxi must not be None")
+    elif mini is None:
+        return clip_func_max(A, maxi, out0=A)
+    elif maxi is None:
+        return clip_func_min(A, mini, out0=A)
+    else:
+        return clip_func(A, mini, maxi, out0=A)

src/flag_gems/ops/mm.py

@@ -167,6 +167,117 @@ def general_mm(a, b, c, M, N, K):
     return c
 
 
+@libentry()
+@libtuner(
+    configs=runtime.get_tuned_config("mm_self_transpose"),
+    key=["M", "K", "stride_am", "stride_ak"],
+    strategy=["align32", "align32", "align32", "align32"],
+    warmup=2,
+    rep=4,
+)
+@triton.jit
+def mm_kernel_syrk(
+    A,
+    C,
+    M,
+    K,
+    stride_am,
+    stride_ak,
+    stride_cm,
+    stride_cn,
+    BLOCK_M: tl.constexpr,
+    BLOCK_K: tl.constexpr,
+):
+    pid = tl.program_id(0)
+
+    # Packed lower-triangular launch domain:
+    #   pid = row * (row + 1) / 2 + col, where 0 <= col <= row.
+    #
+    # Invert the triangular-number indexing by solving:
+    #   row^2 + row - 2 * pid = 0
+    # => row = (-1 + sqrt(1 + 8 * pid)) / 2
+    #
+    # We take floor(...) as the candidate row, then apply an integer +/-1 correction
+    # because fp32 sqrt can be off near triangular-number boundaries.
+    pid_f = pid.to(tl.float32)
+    pid_m = tl.floor((tl.sqrt(8.0 * pid_f + 1.0) - 1.0) / 2.0).to(tl.int32)
+    tri_start = pid_m * (pid_m + 1) // 2
+    pid_m = tl.where(tri_start > pid, pid_m - 1, pid_m)
+    next_tri_start = (pid_m + 1) * (pid_m + 2) // 2
+    pid_m = tl.where(next_tri_start <= pid, pid_m + 1, pid_m)
+    tri_start = pid_m * (pid_m + 1) // 2
+    pid_n = pid - tri_start
+
+    rm = pid_m * BLOCK_M + tl.arange(0, BLOCK_M)
+    rn = pid_n * BLOCK_M + tl.arange(0, BLOCK_M)
+    ram = tl.max_contiguous(tl.multiple_of(rm % M, BLOCK_M), BLOCK_M).to(tl.int64)
+    ran = tl.max_contiguous(tl.multiple_of(rn % M, BLOCK_M), BLOCK_M).to(tl.int64)
+    rm = rm.to(tl.int64)
+    rn = rn.to(tl.int64)
+    acc = tl.zeros((BLOCK_M, BLOCK_M), dtype=tl.float32)
+
+    for start_k in range(0, K, BLOCK_K):
+        rk = (start_k + tl.arange(0, BLOCK_K)).to(tl.int64)
+        mask_k = rk < K
+        a = tl.load(
+            A + (ram[:, None] * stride_am + rk[None, :] * stride_ak),
+            mask=mask_k[None, :],
+            other=0.0,
+        )
+        b = tl.load(
+            A + (rk[:, None] * stride_ak + ran[None, :] * stride_am),
+            mask=mask_k[:, None],
+            other=0.0,
+        )
+        acc += tl.dot(a, b, out_dtype=tl.float32, allow_tf32=False)
+
+    out = acc.to(C.dtype.element_ty)
+    c_ptr = C + (rm[:, None] * stride_cm + rn[None, :] * stride_cn)
+    mask = (rm < M)[:, None] & (rn < M)[None, :]
+    tl.store(c_ptr, out, mask=mask)
+
+    if pid_m > pid_n:
+        c_t_ptr = C + (rn[:, None] * stride_cm + rm[None, :] * stride_cn)
+        mask_t = (rn < M)[:, None] & (rm < M)[None, :]
+        tl.store(c_t_ptr, tl.trans(out), mask=mask_t)
+
+
+def is_syrk_transpose_pair(a, b):
+    return (
+        a.ndim == 2
+        and b.ndim == 2
+        and a.shape[0] == b.shape[1]
+        and a.shape[1] == b.shape[0]
+        and a.stride(0) == b.stride(1)
+        and a.stride(1) == b.stride(0)
+        and a.storage_offset() == b.storage_offset()
+        and a.data_ptr() == b.data_ptr()
+    )
+
+
+def syrk_mm(a, c, M, K):
+    grid = lambda META: (
+        # Number of tile rows is tiles = ceil(M / BLOCK_M).
+        # Packed lower triangle contains:
+        #   1 + 2 + ... + tiles = tiles * (tiles + 1) / 2
+        triton.cdiv(M, META["BLOCK_M"])
+        * (triton.cdiv(M, META["BLOCK_M"]) + 1)
+        // 2,
+    )
+    with torch_device_fn.device(a.device):
+        mm_kernel_syrk[grid](
+            a,
+            c,
+            M,
+            K,
+            a.stride(0),
+            a.stride(1),
+            c.stride(0),
+            c.stride(1),
+        )
+    return c
+
+
 def streamk_scenario(a, b, M, N, K):
     # TODO: this my change sometime according to the realbenchmark result
     # Currently, the best configuration for streamk has only been tested on A100(capability[0] == 8).
@@ -185,6 +296,10 @@ def streamk_scenario(a, b, M, N, K):
 
 def mm(a, b):
     device = a.device
+    if is_syrk_transpose_pair(a, b):
+        M, K = a.shape
+        c = torch.empty((M, M), device=device, dtype=a.dtype)
+        return syrk_mm(a, c, M, K)
     # handle non-contiguous inputs if necessary
     if a.stride(0) > 1 and a.stride(1) > 1:
         a = a.contiguous()
@@ -206,6 +321,9 @@ def mm(a, b):
 
 
 def mm_out(a, b, *, out):
+    if is_syrk_transpose_pair(a, b):
+        M, K = a.shape
+        return syrk_mm(a, out, M, K)
     # handle non-contiguous inputs if necessary
     if a.stride(0) > 1 and a.stride(1) > 1:
         a = a.contiguous()