Add clip operator implementation, tests and benchmark

Author: factnn

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()

benchmark/test_generic_pointwise_perf.py

@@ -161,25 +161,31 @@ def test_generic_pointwise_benchmark(op_name, torch_op, input_fn, dtypes):
     bench.run()
 
 
-@pytest.mark.clamp_
-def test_clamp_inplace():
-    bench = GenericBenchmark(
-        input_fn=clamp_input_fn,
-        op_name="clamp_",
-        torch_op=torch.clamp_,
-        dtypes=FLOAT_DTYPES,
-        is_inplace=True,
-    )
-    bench.run()
-
-
-@pytest.mark.clamp_min_
-def test_clamp_min_inplace():
+@pytest.mark.parametrize(
+    "op_name, torch_op, input_fn, dtypes",
+    [
+        pytest.param(
+            "clamp_",
+            torch.clamp_,
+            clamp_input_fn,
+            FLOAT_DTYPES,
+            marks=pytest.mark.clamp_,
+        ),
+        pytest.param(
+            "clamp_min_",
+            torch.clamp_min_,
+            clamp_min_input_fn,
+            FLOAT_DTYPES,
+            marks=pytest.mark.clamp_min_,
+        ),
+    ],
+)
+def test_generic_inplace_pointwise_benchmark(op_name, torch_op, input_fn, dtypes):
     bench = GenericBenchmark(
-        input_fn=clamp_min_input_fn,
-        op_name="clamp_min_",
-        torch_op=torch.clamp_min_,
-        dtypes=FLOAT_DTYPES,
+        input_fn=input_fn,
+        op_name=op_name,
+        torch_op=torch_op,
+        dtypes=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)

tests/test_clip.py

@@ -0,0 +1,55 @@
+import pytest
+import torch
+
+import flag_gems
+
+from .accuracy_utils import (
+    FLOAT_DTYPES,
+    POINTWISE_SHAPES,
+    SCALARS,
+    gems_assert_equal,
+    to_reference,
+)
+
+
+@pytest.mark.clip
+@pytest.mark.parametrize("shape", POINTWISE_SHAPES)
+@pytest.mark.parametrize("maxi", SCALARS)
+@pytest.mark.parametrize("mini", SCALARS)
+@pytest.mark.parametrize("isnone", [None, "max", "min"])
+@pytest.mark.parametrize("dtype", FLOAT_DTYPES)
+def test_accuracy_clip(shape, maxi, mini, isnone, dtype):
+    inp = torch.randn(shape, dtype=dtype, device=flag_gems.device)
+    if isnone == "min":
+        mini = None
+    elif isnone == "max":
+        maxi = None
+    ref_inp = to_reference(inp)
+
+    ref_out = torch.clip(ref_inp, min=mini, max=maxi)
+    with flag_gems.use_gems():
+        res_out = torch.clip(inp, min=mini, max=maxi)
+
+    gems_assert_equal(res_out, ref_out)
+
+
+@pytest.mark.inplace
+@pytest.mark.clip_
+@pytest.mark.parametrize("shape", POINTWISE_SHAPES)
+@pytest.mark.parametrize("maxi", SCALARS)
+@pytest.mark.parametrize("mini", SCALARS)
+@pytest.mark.parametrize("isnone", [None, "max", "min"])
+@pytest.mark.parametrize("dtype", FLOAT_DTYPES)
+def test_accuracy_clip_(shape, maxi, mini, isnone, dtype):
+    inp = torch.randn(shape, dtype=dtype, device=flag_gems.device)
+    if isnone == "min":
+        mini = None
+    elif isnone == "max":
+        maxi = None
+    ref_inp = to_reference(inp.clone())
+
+    ref_out = torch.clip_(ref_inp, min=mini, max=maxi)
+    with flag_gems.use_gems():
+        res_out = torch.clip_(inp, min=mini, max=maxi)
+
+    gems_assert_equal(res_out, ref_out)