Reuse assert_close for benchmarks from triton.testing (#3161)

Part of #3160

Signed-off-by: Anatoly Myachev <[email protected]>

Author: anmyachev

benchmarks/triton_kernels_benchmark/__init__.py

@@ -1,6 +1,16 @@
 import os
 
-from .benchmark_testing import do_bench, assert_close, perf_report, Benchmark, BENCHMARKING_METHOD  # type: ignore # noqa: F401
+from triton.testing import assert_close
+
+from .benchmark_testing import do_bench, perf_report, Benchmark, BENCHMARKING_METHOD
 
 if BENCHMARKING_METHOD == "UPSTREAM_PYTORCH_PROFILER":
     os.environ["INJECT_PYTORCH"] = "True"
+
+__all__ = [
+    "assert_close",
+    "do_bench",
+    "perf_report",
+    "Benchmark",
+    "BENCHMARKING_METHOD",
+]

benchmarks/triton_kernels_benchmark/benchmark_testing.py

@@ -160,43 +160,6 @@ def extract_kernels(funcs):
     raise NotImplementedError(f"BENCHMARKING_METHOD: {BENCHMARKING_METHOD} isn't implemented")
 
 
-def assert_close(x, y, atol=None, rtol=None, err_msg=""):
-    import numpy as np
-    import torch
-
-    # canonicalize arguments to be tensors
-    if not isinstance(x, torch.Tensor):
-        x = torch.tensor(x)
-    if not isinstance(y, torch.Tensor):
-        y = torch.tensor(y)
-    # absolute tolerance
-    if atol is None:
-        atol = 1e-2
-    atol = atol(x.dtype) if callable(atol) else atol
-    # relative tolerance hook
-    if rtol is None:
-        rtol = 0.
-    rtol = rtol(x.dtype) if callable(rtol) else rtol
-    # we use numpy instead of pytorch
-    # as it seems more memory efficient
-    # pytorch tends to oom on large tensors
-    if isinstance(x, torch.Tensor):
-        if x.dtype == torch.bfloat16:
-            x = x.float()
-        x = x.cpu().detach().numpy()
-    if isinstance(y, torch.Tensor):
-        if y.dtype == torch.bfloat16:
-            y = y.float()
-        y = y.cpu().detach().numpy()
-    # we handle size==1 case separately as we can
-    # provide better error message there
-    if x.size > 1 or y.size > 1:
-        np.testing.assert_allclose(x, y, atol=atol, rtol=rtol, equal_nan=True)
-        return
-    if not np.allclose(x, y, atol=atol, rtol=rtol):
-        raise AssertionError(f"{err_msg} {x} is not close to {y} (atol={atol}, rtol={rtol})")
-
-
 def perf_report(benchmarks):
     """
     Mark a function for benchmarking. The benchmark can then be executed by using the :code:`.run` method on the return value.