Repository: nv-legate/cunumeric

Issue:
Handle ufunc coverage wrappers more robustly
ref: https://github.com/nv-legate/cunumeric/pull/268/files#r846513290

As noted in the above conversation, the generic callable wrapping for adding coverage reporting is not sufficient. 

Numpy (and thus cunumeric) `ufunc` are objects with their own API (https://numpy.org/doc/stable/reference/ufuncs.html) and just using a plain function wrapper makes those methods invisible. Some requirements to decide first:

* Do all the methods of a `ufunc` need to be included in coverage reporting? Or just its `__call__`

If yes, we will need to resort to a wrapping object (and then: is it sufficient to just create a purpose-built `ufunc_wrapper` or do we need a generic forwarding wrapper?)

If not, we may be able to just wrap and replace `__call__` using the function wrappers similar to the existing ones.


Discussion:
Though in theory we should report coverage on all those methods in `ufunc`, I doubt that they are actually used that much. Plus, many of those functions have aliases in the top-level module that I believe are more likely being used. So, unless we find someone who really wants to see those methods counted, let's track only two methods `__call__` and `reduce` for now so that our own reporting tool stops breaking the tests.
@magnatelee do we actually need to track `reduce` or just make sure it is still available and callable?
@bryevdv I guess we don't want to track it, as many of the ufuncs are missing that function and we don't intend to implement those.

Resolving patch:
diff --git a/cunumeric/coverage.py b/cunumeric/coverage.py
--- a/cunumeric/coverage.py
+++ b/cunumeric/coverage.py
@@ -24,7 +24,7 @@
 from .runtime import runtime
 from .utils import find_last_user_frames, find_last_user_stacklevel
 
-__all__ = ("clone_module",)
+__all__ = ("clone_class", "clone_module")
 
 FALLBACK_WARNING = (
     "cuNumeric has not implemented {name} "
@@ -66,12 +66,9 @@ def __call__(self, *args: Any, **kwargs: Any) -> Any:
         ...
 
 
-class CuWrapped(Protocol):
+class CuWrapped(AnyCallable, Protocol):
     _cunumeric_implemented: bool
 
-    def __call__(self, *args: Any, **kwargs: Any) -> Any:
-        ...
-
 
 def implemented(
     func: AnyCallable, prefix: str, name: str, *, reporting: bool = True
@@ -92,7 +89,9 @@ def wrapper(*args: Any, **kwargs: Any) -> Any:
 
     else:
 
-        wrapper = cast(CuWrapped, func)
+        @wraps(func)
+        def wrapper(*args: Any, **kwargs: Any) -> Any:
+            return func(*args, **kwargs)
 
     wrapper._cunumeric_implemented = True
 
@@ -162,12 +161,10 @@ def clone_module(
         omit_types=(ModuleType,),
     )
 
-    from numpy import ufunc as npufunc
+    reporting = runtime.report_coverage
 
     from ._ufunc.ufunc import ufunc as lgufunc
 
-    reporting = runtime.report_coverage
-
     for attr, value in new_globals.items():
         if isinstance(value, (FunctionType, lgufunc)):
             wrapped = implemented(
@@ -175,6 +172,8 @@ def clone_module(
             )
             new_globals[attr] = wrapped
 
+    from numpy import ufunc as npufunc
+
     for attr, value in missing.items():
         if isinstance(value, (FunctionType, npufunc)):
             wrapped = unimplemented(value, mod_name, attr, reporting=reporting)


Test patch:
diff --git a/tests/reduction.py b/tests/reduction.py
--- a/tests/reduction.py
+++ b/tests/reduction.py
@@ -36,49 +36,36 @@ def test():
 
     x = num.array([1, 2, 3, 4.0])
     r = num.sum(x)
-    r2 = num.add.reduce(x)
 
-    assert r == r2 == 10
+    assert r == 10
 
     x = num.array([1, 2, 3, 4.0, 5.0])
     r = num.prod(x)
-    r2 = num.multiply.reduce(x)
-    assert r == r2 == 120
+    assert r == 120
 
     asserts.assert_equal(num.sum([]), np.sum([]))
-    asserts.assert_equal(num.add.reduce([]), np.add.reduce([]))
 
     asserts.assert_equal(num.sum([[], []]), np.sum([[], []]))
-    asserts.assert_equal(num.add.reduce([[], []]), np.add.reduce([[], []]))
 
     asserts.assert_equal(num.sum(num.array([0])), np.sum(np.array([0])))
-    asserts.assert_equal(
-        num.add.reduce(num.array([0])), np.add.reduce(np.array([0]))
-    )
 
     asserts.assert_equal(num.sum([1]), np.sum([1]))
-    asserts.assert_equal(num.add.reduce([1]), np.add.reduce([1]))
 
     asserts.assert_equal(num.sum(0), np.sum(0))
-    asserts.assert_equal(num.add.reduce(0), np.add.reduce(0))
 
     asserts.assert_equal(num.sum(1), np.sum(1))
-    asserts.assert_equal(num.add.reduce(1), np.add.reduce(1))
 
     x = num.array([1, 0, 2, -1, 0, 0, 8])
     x_np = np.array([1, 0, 2, -1, 0, 0, 8])
     asserts.assert_equal(num.sum(x), np.sum(x_np))
-    asserts.assert_equal(num.add.reduce(x), np.add.reduce(x_np))
 
     x = num.array([[0, 1, 0], [2, 0, 3]])
     x_np = np.array([[0, 1, 0], [2, 0, 3]])
     asserts.assert_equal(num.sum(x), np.sum(x_np))
-    asserts.assert_equal(num.add.reduce(x), np.add.reduce(x_np))
 
     x = num.eye(3)
     x_np = np.eye(3)
     asserts.assert_equal(num.sum(x), np.sum(x_np))
-    asserts.assert_equal(num.add.reduce(x), np.add.reduce(x_np))
 
     x = num.array(
         [
@@ -95,17 +82,7 @@ def test():
     asserts.assert_equal(num.sum(x, axis=0), np.sum(x_np, axis=0))
     asserts.assert_equal(num.sum(x, axis=1), np.sum(x_np, axis=1))
     asserts.assert_equal(num.sum(x, axis=2), np.sum(x_np, axis=2))
-    asserts.assert_equal(
-        num.add.reduce(x, axis=0), np.add.reduce(x_np, axis=0)
-    )
-    asserts.assert_equal(
-        num.add.reduce(x, axis=1), np.add.reduce(x_np, axis=1)
-    )
-    asserts.assert_equal(
-        num.add.reduce(x, axis=2), np.add.reduce(x_np, axis=2)
-    )
     asserts.assert_equal(num.sum(x), np.sum(x_np))
-    asserts.assert_equal(num.add.reduce(x), np.add.reduce(x_np))
 
     x_np = np.concatenate((x_np,) * 2000, axis=1)
     x = num.array(x_np)
@@ -113,16 +90,6 @@ def test():
     asserts.assert_equal(num.sum(x, axis=1), np.sum(x_np, axis=1))
     asserts.assert_equal(num.sum(x, axis=2), np.sum(x_np, axis=2))
     asserts.assert_equal(num.sum(x), np.sum(x_np))
-    asserts.assert_equal(
-        num.add.reduce(x, axis=0), np.add.reduce(x_np, axis=0)
-    )
-    asserts.assert_equal(
-        num.add.reduce(x, axis=1), np.add.reduce(x_np, axis=1)
-    )
-    asserts.assert_equal(
-        num.add.reduce(x, axis=2), np.add.reduce(x_np, axis=2)
-    )
-    asserts.assert_equal(num.add.reduce(x), np.add.reduce(x_np))
 
     x_np = np.random.randn(100)
     indices = np.random.choice(
@@ -131,14 +98,10 @@ def test():
     x_np[indices] = 0
     x = num.array(x_np)
     asserts.assert_allclose(num.sum(x), np.sum(x_np))
-    asserts.assert_allclose(num.add.reduce(x), np.add.reduce(x_np))
 
     x_np = x_np.reshape(10, 10)
     x = num.array(x_np)
     asserts.assert_allclose(num.sum(x), np.sum(x_np))
-    asserts.assert_allclose(num.add.reduce(x), np.add.reduce(x_np))
-
-    return
 
 
 if __name__ == "__main__":
diff --git a/tests/unit/cunumeric/test_coverage.py b/tests/unit/cunumeric/test_coverage.py
--- a/tests/unit/cunumeric/test_coverage.py
+++ b/tests/unit/cunumeric/test_coverage.py
@@ -115,27 +115,95 @@ def _test_func(a: int, b: int) -> int:
     return a + b
 
 
-@patch("cunumeric.runtime.record_api_call")
-def test_implemented(mock_record_api_call: MagicMock) -> None:
-    wrapped = m.implemented(_test_func, "foo", "_test_func")
+class _Test_ufunc(cunumeric._ufunc.ufunc):
+    """docstring"""
+
+    def __call__(self, a: int, b: int) -> int:
+        return a + b
+
+
+_test_ufunc = _Test_ufunc()
+
+
+class Test_implemented:
+    @patch("cunumeric.runtime.record_api_call")
+    def test_reporting_True_func(
+        self, mock_record_api_call: MagicMock
+    ) -> None:
+        wrapped = m.implemented(_test_func, "foo", "_test_func")
+
+        assert wrapped.__doc__ == _test_func.__doc__
+        assert wrapped.__wrapped__ is _test_func
+
+        assert wrapped(10, 20) == 30
+
+        mock_record_api_call.assert_called_once()
+        assert mock_record_api_call.call_args[0] == ()
+        assert mock_record_api_call.call_args[1]["name"] == "foo._test_func"
+        assert mock_record_api_call.call_args[1]["implemented"]
+        filename, lineno = mock_record_api_call.call_args[1]["location"].split(
+            ":"
+        )
+        assert filename == __file__
+        assert int(lineno)
+
+    @patch("cunumeric.runtime.record_api_call")
+    def test_reporting_False_func(
+        self, mock_record_api_call: MagicMock
+    ) -> None:
+        wrapped = m.implemented(
+            _test_func, "foo", "_test_func", reporting=False
+        )
+
+        assert wrapped.__doc__ == _test_func.__doc__
+        assert wrapped.__wrapped__ is _test_func
+
+        assert wrapped(10, 20) == 30
+
+        mock_record_api_call.assert_not_called()
+
+    @patch("cunumeric.runtime.record_api_call")
+    def test_reporting_True_ufunc(
+        self, mock_record_api_call: MagicMock
+    ) -> None:
+        wrapped = m.implemented(_test_ufunc, "foo", "_test_ufunc")
+
+        assert wrapped.__doc__ == _test_ufunc.__doc__
+        assert wrapped.__wrapped__ is _test_ufunc
+
+        assert wrapped(10, 20) == 30
+
+        mock_record_api_call.assert_called_once()
+        assert mock_record_api_call.call_args[0] == ()
+        assert mock_record_api_call.call_args[1]["name"] == "foo._test_ufunc"
+        assert mock_record_api_call.call_args[1]["implemented"]
+        filename, lineno = mock_record_api_call.call_args[1]["location"].split(
+            ":"
+        )
+        assert filename == __file__
+        assert int(lineno)
 
-    assert wrapped.__doc__ == _test_func.__doc__
-    assert wrapped.__wrapped__ is _test_func
+    @patch("cunumeric.runtime.record_api_call")
+    def test_reporting_False_ufunc(
+        self, mock_record_api_call: MagicMock
+    ) -> None:
+        wrapped = m.implemented(
+            _test_ufunc, "foo", "_test_func", reporting=False
+        )
 
-    assert wrapped(10, 20) == 30
+        assert wrapped.__doc__ == _test_ufunc.__doc__
+        assert wrapped.__wrapped__ is _test_ufunc
+
+        assert wrapped(10, 20) == 30
 
-    mock_record_api_call.assert_called_once()
-    assert mock_record_api_call.call_args[0] == ()
-    assert mock_record_api_call.call_args[1]["name"] == "foo._test_func"
-    assert mock_record_api_call.call_args[1]["implemented"]
-    filename, lineno = mock_record_api_call.call_args[1]["location"].split(":")
-    assert filename == __file__
-    assert int(lineno)
+        mock_record_api_call.assert_not_called()
 
 
 class Test_unimplemented:
     @patch("cunumeric.runtime.record_api_call")
-    def test_reporting_True(self, mock_record_api_call: MagicMock) -> None:
+    def test_reporting_True_func(
+        self, mock_record_api_call: MagicMock
+    ) -> None:
         wrapped = m.unimplemented(_test_func, "foo", "_test_func")
 
         assert wrapped.__doc__ == _test_func.__doc__
@@ -154,7 +222,9 @@ def test_reporting_True(self, mock_record_api_call: MagicMock) -> None:
         assert int(lineno)
 
     @patch("cunumeric.runtime.record_api_call")
-    def test_reporting_False(self, mock_record_api_call: MagicMock) -> None:
+    def test_reporting_False_func(
+        self, mock_record_api_call: MagicMock
+    ) -> None:
         wrapped = m.unimplemented(
             _test_func, "foo", "_test_func", reporting=False
         )
@@ -172,6 +242,48 @@ def test_reporting_False(self, mock_record_api_call: MagicMock) -> None:
 
         mock_record_api_call.assert_not_called()
 
+    @patch("cunumeric.runtime.record_api_call")
+    def test_reporting_True_ufunc(
+        self, mock_record_api_call: MagicMock
+    ) -> None:
+        wrapped = m.unimplemented(_test_ufunc, "foo", "_test_ufunc")
+
+        assert wrapped.__doc__ == _test_ufunc.__doc__
+        assert wrapped.__wrapped__ is _test_ufunc
+
+        assert wrapped(10, 20) == 30
+
+        mock_record_api_call.assert_called_once()
+        assert mock_record_api_call.call_args[0] == ()
+        assert mock_record_api_call.call_args[1]["name"] == "foo._test_ufunc"
+        assert not mock_record_api_call.call_args[1]["implemented"]
+        filename, lineno = mock_record_api_call.call_args[1]["location"].split(
+            ":"
+        )
+        assert filename == __file__
+        assert int(lineno)
+
+    @patch("cunumeric.runtime.record_api_call")
+    def test_reporting_False_ufunc(
+        self, mock_record_api_call: MagicMock
+    ) -> None:
+        wrapped = m.unimplemented(
+            _test_ufunc, "foo", "_test_ufunc", reporting=False
+        )
+
+        assert wrapped.__doc__ == _test_ufunc.__doc__
+        assert wrapped.__wrapped__ is _test_ufunc
+
+        with pytest.warns(RuntimeWarning) as record:
+            assert wrapped(10, 20) == 30
+
+        assert len(record) == 1
+        assert record[0].message.args[0] == m.FALLBACK_WARNING.format(
+            name="foo._test_ufunc"
+        )
+
+        mock_record_api_call.assert_not_called()
+
 
 _OriginMod = ModuleType("origin")
 exec(