# Copyright (c) 2025 PaddlePaddle Authors. All Rights Reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

"""Test PyLayer tensor_hold_helper for _clear_dataptr protection.

Pipeline-parallel pattern:
  1. outputs = Layer.apply(inputs)    # forward: data is valid
  2. loss    = f(outputs)             # build loss graph BEFORE clearing
  3. outputs._clear_dataptr()        # free activation memory
  4. loss.backward()                  # backward via tensor_hold_helper recovery

tensor_hold_helper is a vector<shared_ptr<DenseTensor>> on PyLayerObject that
holds strong references to every DenseTensor impl saved via save_for_backward.
It is born with set_container (save_for_backward) and destroyed with the
PyLayerObject itself, preventing _clear_dataptr from freeing the underlying
allocation before backward runs.
"""

import gc
import unittest

import numpy as np

import paddle
from paddle.autograd import PyLayer


def _clear(tensors):
    """Call _clear_dataptr on a single tensor or iterable of tensors."""
    if isinstance(tensors, (list, tuple)):
        for t in tensors:
            if hasattr(t, '_clear_dataptr'):
                t._clear_dataptr()
    elif hasattr(tensors, '_clear_dataptr'):
        tensors._clear_dataptr()


class TestPyLayerClearDataptr(unittest.TestCase):
    """Core tests: _clear_dataptr on outputs does not break backward."""

    def test_basic_clear_dataptr(self):
        """Single output, single saved tensor."""

        class TanhLayer(PyLayer):
            @staticmethod
            def forward(ctx, x):
                y = paddle.tanh(x)
                ctx.save_for_backward(y)
                return y

            @staticmethod
            def backward(ctx, dy):
                (y,) = ctx.saved_tensor()
                return dy * (1 - paddle.square(y))

        x = paddle.randn([2, 3]).astype('float64')
        x.stop_gradient = False
        out = TanhLayer.apply(x)
        loss = out.mean()  # build graph first
        _clear(out)  # then free activation
        loss.backward()
        self.assertIsNotNone(x.grad)

    def test_multiple_saved_tensors(self):
        """Multiple tensors passed to save_for_backward."""

        class AddLayer(PyLayer):
            @staticmethod
            def forward(ctx, x, y):
                ctx.save_for_backward(x, y)
                return x + y

            @staticmethod
            def backward(ctx, dy):
                x, y = ctx.saved_tensor()
                return dy, dy

        x = paddle.randn([2, 3]).astype('float64')
        y = paddle.randn([2, 3]).astype('float64')
        x.stop_gradient = False
        y.stop_gradient = False
        out = AddLayer.apply(x, y)
        loss = out.mean()
        _clear(out)
        loss.backward()
        self.assertIsNotNone(x.grad)
        self.assertIsNotNone(y.grad)

    def test_multiple_outputs(self):
        """Tuple output: both outputs are cleared."""

        class MultiOutLayer(PyLayer):
            @staticmethod
            def forward(ctx, x):
                y1 = paddle.tanh(x)
                y2 = paddle.sin(x)
                ctx.save_for_backward(y1, y2)
                return y1, y2

            @staticmethod
            def backward(ctx, dy1, dy2):
                y1, y2 = ctx.saved_tensor()
                return dy1 * (1 - paddle.square(y1)) + dy2 * paddle.cos(y2)

        x = paddle.randn([2, 3]).astype('float64')
        x.stop_gradient = False
        y1, y2 = MultiOutLayer.apply(x)
        loss = (y1 + y2).mean()  # build graph while data is valid
        _clear([y1, y2])
        loss.backward()
        self.assertIsNotNone(x.grad)

    def test_chained_computation(self):
        """Final output of a chain is cleared; intermediate kept for input."""

        class TanhLayer(PyLayer):
            @staticmethod
            def forward(ctx, x):
                y = paddle.tanh(x)
                ctx.save_for_backward(y)
                return y

            @staticmethod
            def backward(ctx, dy):
                (y,) = ctx.saved_tensor()
                return dy * (1 - paddle.square(y))

        x = paddle.randn([2, 3]).astype('float64')
        x.stop_gradient = False
        y = TanhLayer.apply(x)  # intermediate – not cleared
        z = TanhLayer.apply(y)  # final output
        loss = z.mean()
        _clear(z)  # only clear final activation
        loss.backward()
        self.assertIsNotNone(x.grad)

    def test_different_dtypes(self):
        """float32 / float64 (and float16 on GPU) all work after _clear_dataptr."""

        class TanhLayer(PyLayer):
            @staticmethod
            def forward(ctx, x):
                y = paddle.tanh(x)
                ctx.save_for_backward(y)
                return y

            @staticmethod
            def backward(ctx, dy):
                (y,) = ctx.saved_tensor()
                return dy * (1 - paddle.square(y))

        dtypes = ['float32', 'float64']
        if paddle.is_compiled_with_cuda():
            dtypes.append('float16')
        for dtype in dtypes:
            x = paddle.randn([2, 3]).astype(dtype)
            x.stop_gradient = False
            out = TanhLayer.apply(x)
            loss = out.mean()
            _clear(out)
            loss.backward()
            self.assertIsNotNone(x.grad)

    def test_memory_cleanup(self):
        """Multiple iterations: per-iteration objects are collectible."""
        import weakref

        class TanhLayer(PyLayer):
            @staticmethod
            def forward(ctx, x):
                y = paddle.tanh(x)
                ctx.save_for_backward(y)
                return y

            @staticmethod
            def backward(ctx, dy):
                (y,) = ctx.saved_tensor()
                return dy * (1 - paddle.square(y))

        # Track the first iteration's `out` via weakref; after the loop ends
        # and gc runs, it must be collected.  Catches holder leaks where
        # tensor_hold_helper accidentally retains a strong reference across
        # ctx lifetimes.
        first_out_ref = None
        for i in range(10):
            x = paddle.randn([64, 64]).astype('float32')
            x.stop_gradient = False
            out = TanhLayer.apply(x)
            if i == 0:
                first_out_ref = weakref.ref(out)
            loss = out.mean()
            _clear(out)
            loss.backward()
            del x, out, loss
            gc.collect()
        self.assertIsNone(first_out_ref())


class TestCtxDirect(unittest.TestCase):
    """Unit tests for ctx API without going through PyLayer.apply().

    These tests create a ctx object directly via cls._backward_function() and
    exercise save_for_backward / saved_tensor / pop_saved_impl in isolation,
    independently of the forward/backward dispatch machinery.

    Key design: cls._backward_function is a subclass of PyLayerBackward which
    inherits core.eager.PyLayer (C++ PyLayerObject).  Instantiating it calls
    PyLayerNew, giving a fully-initialized ctx with an empty tensor_hold_helper.
    """

    def _make_ctx(self):
        """Create a bare ctx (PyLayerObject) without running forward."""

        class _Stub(PyLayer):
            @staticmethod
            def forward(ctx, x):
                return x

            @staticmethod
            def backward(ctx, dy):
                return dy

        return _Stub._backward_function()

    # ------------------------------------------------------------------
    # Basic save / recover
    # ------------------------------------------------------------------

    def test_direct_single_tensor_recover(self):
        """save_for_backward + _clear_dataptr + saved_tensor, no apply."""
        ctx = self._make_ctx()
        t = paddle.randn([2, 3]).astype('float64')
        ctx.save_for_backward(t)

        _clear(t)

        (recovered,) = ctx.saved_tensor()
        self.assertIsNotNone(recovered)
        self.assertEqual(list(recovered.shape), [2, 3])

    def test_direct_multiple_tensors_recover(self):
        """All tensors are recovered after _clear_dataptr, no apply."""
        ctx = self._make_ctx()
        a = paddle.randn([3]).astype('float32')
        b = paddle.ones([4, 2]).astype('float64')
        ctx.save_for_backward(a, b)

        _clear(a)
        _clear(b)

        recovered = ctx.saved_tensor()
        self.assertEqual(len(recovered), 2)
        self.assertEqual(list(recovered[0].shape), [3])
        self.assertEqual(list(recovered[1].shape), [4, 2])

    def test_direct_no_clear(self):
        """saved_tensor returns correct values when _clear_dataptr was not called."""
        ctx = self._make_ctx()
        t = paddle.randn([2, 3]).astype('float32')
        expected = t.numpy().copy()
        ctx.save_for_backward(t)

        (recovered,) = ctx.saved_tensor()
        np.testing.assert_allclose(recovered.numpy(), expected, rtol=1e-6)

    # ------------------------------------------------------------------
    # pop_saved_impl
    # ------------------------------------------------------------------

    def test_pop_saved_impl_single(self):
        """pop_saved_impl removes the holder entry; recovered tensor stays valid."""
        ctx = self._make_ctx()
        t = paddle.randn([2, 3]).astype('float32')
        orig = t.numpy().copy()
        ctx.save_for_backward(t)

        _clear(t)
        (recovered,) = ctx.saved_tensor()
        # Verify the recovered tensor carries the correct data (not just non-None).
        np.testing.assert_allclose(recovered.numpy(), orig, rtol=1e-6)

        # Pop removes the holder entry; recovered's own shared_ptr keeps data alive.
        ctx._pop_saved_impl(recovered)
        self.assertEqual(list(recovered.shape), [2, 3])

    def test_pop_saved_impl_partial(self):
        """Pop both saved tensors one by one; proves each entry is stored independently."""
        ctx = self._make_ctx()
        a = paddle.randn([2]).astype('float32')
        b = paddle.randn([3]).astype('float32')
        ctx.save_for_backward(a, b)

        _clear(a)
        _clear(b)
        recovered = ctx.saved_tensor()
        self.assertEqual(len(recovered), 2)

        # Pop the first entry; if holder only had one entry this would erase it
        # and the second pop below would be a no-op instead of finding b's entry.
        ctx._pop_saved_impl(recovered[0])
        # Pop the second entry; succeeds only if b's entry is still in holder
        # (i.e. the two entries are stored independently).
        ctx._pop_saved_impl(recovered[1])
        # Both recovered handles remain valid via their own shared_ptr copies.
        self.assertEqual(list(recovered[0].shape), [2])
        self.assertEqual(list(recovered[1].shape), [3])

    def test_pop_saved_impl_no_clear(self):
        """pop_saved_impl does not crash when tensor was never cleared.

        Also verifies the pop targets a specific entry: after popping t's
        holder entry, a subsequent saved_tensor() call still succeeds and
        returns t with its original data (pop did not corrupt container).
        """
        ctx = self._make_ctx()
        t = paddle.randn([5]).astype('float32')
        orig = t.numpy().copy()
        ctx.save_for_backward(t)

        # No _clear_dataptr; pop should still succeed silently
        ctx._pop_saved_impl(t)
        # saved_tensor() must still return the tensor correctly.
        (recovered,) = ctx.saved_tensor()
        np.testing.assert_allclose(recovered.numpy(), orig, rtol=1e-6)

    # ------------------------------------------------------------------
    # Deep-traversal via nested list in container
    # ------------------------------------------------------------------

    def test_nested_list_holder_populated(self):
        """Container with a nested list: CollectDenseTensors populates holder for all tensors.

        save_for_backward packs args as a tuple, so the container at the top
        level is always a tuple.  But tuple *elements* may themselves be lists
        (e.g. when a list is passed as one argument).  CollectDenseTensors
        recurses into them; verify via pop_saved_impl that both were collected.
        """
        ctx = self._make_ctx()
        t1 = paddle.randn([2]).astype('float32')
        t2 = paddle.randn([3]).astype('float32')

        # Directly assign a tuple whose sole element is a list of tensors.
        # This bypasses save_for_backward's *args flattening so we can test
        # the deep-traversal branch.
        ctx.container = ([t1, t2],)

        # Each pop finds and removes its entry; if CollectDenseTensors missed
        # an entry, the corresponding pop is a silent no-op — so we follow
        # each pair of pops with a redundant third pop that must also not crash,
        # confirming the erase path is robust against missing entries.
        ctx._pop_saved_impl(t1)
        ctx._pop_saved_impl(t2)
        ctx._pop_saved_impl(t1)  # already removed — must be a silent no-op

    def test_nested_tuple_holder_populated(self):
        """Container with a nested tuple: all inner tensors are held."""
        ctx = self._make_ctx()
        t1 = paddle.randn([2]).astype('float32')
        t2 = paddle.randn([3]).astype('float32')

        ctx.container = ((t1, t2),)

        ctx._pop_saved_impl(t1)
        ctx._pop_saved_impl(t2)
        ctx._pop_saved_impl(t1)  # already removed — must be a silent no-op


class TestCtxHoldRestore(unittest.TestCase):
    """Direct-ctx tests for _hold_tensors / _restore_held_tensors.

    These cover the C++ WalkDenseTensors recursion (Tensor / tuple / list /
    dict), the SavedTensorsHooks short-circuit in pylayer_hold_tensors, and
    the ``impl() != nullptr`` early-return in pylayer_restore_held_tensors.
    """

    def _make_ctx(self):
        class _Stub(PyLayer):
            @staticmethod
            def forward(ctx, x):
                return x

            @staticmethod
            def backward(ctx, dy):
                return dy

        return _Stub._backward_function()

    def test_hold_restore_basic(self):
        """hold(tensor) + _clear_dataptr + restore re-installs impl_."""
        ctx = self._make_ctx()
        t = paddle.randn([2, 3]).astype('float32')
        orig = t.numpy().copy()

        ctx._hold_tensors(t)
        _clear(t)
        self.assertFalse(t._is_initialized())

        ctx._restore_held_tensors()
        self.assertTrue(t._is_initialized())
        np.testing.assert_allclose(t.numpy(), orig, rtol=1e-6)

    def test_hold_nested_containers(self):
        """tuple / list / dict values are all deep-traversed."""
        ctx = self._make_ctx()
        t_tuple = paddle.randn([2]).astype('float32')
        t_list = paddle.randn([3]).astype('float32')
        t_dict = paddle.randn([4]).astype('float32')
        originals = [t.numpy().copy() for t in (t_tuple, t_list, t_dict)]

        # One call with a container mixing all three Python collection types.
        ctx._hold_tensors(((t_tuple,), [t_list], {'k': t_dict}))

        _clear([t_tuple, t_list, t_dict])
        ctx._restore_held_tensors()

        for got, orig in zip((t_tuple, t_list, t_dict), originals):
            self.assertTrue(got._is_initialized())
            np.testing.assert_allclose(got.numpy(), orig, rtol=1e-6)

    def test_hold_none_is_noop(self):
        """_hold_tensors(None) collects nothing; restore is a no-op."""
        ctx = self._make_ctx()
        ctx._hold_tensors(None)
        ctx._restore_held_tensors()  # must not crash

    def test_hold_scalar_top_level_noop(self):
        """_hold_tensors on a bare non-container scalar collects nothing."""
        ctx = self._make_ctx()
        for val in (42, 3.14, "str", b"bytes"):
            ctx._hold_tensors(val)
            ctx._restore_held_tensors()  # must not crash

    def test_restore_skips_valid_impl(self):
        """Restore leaves tensors whose impl is still valid untouched."""
        ctx = self._make_ctx()
        t_cleared = paddle.randn([2]).astype('float32')
        t_kept = paddle.randn([3]).astype('float32')
        orig_cleared = t_cleared.numpy().copy()
        orig_kept = t_kept.numpy().copy()

        ctx._hold_tensors([t_cleared, t_kept])
        _clear(t_cleared)  # only one is cleared
        ctx._restore_held_tensors()

        # cleared tensor resurrected
        np.testing.assert_allclose(t_cleared.numpy(), orig_cleared, rtol=1e-6)
        # kept tensor's impl untouched — covers the ``if (!tensor.impl())``
        # false branch in pylayer_restore_held_tensors.
        self.assertTrue(t_kept._is_initialized())
        np.testing.assert_allclose(t_kept.numpy(), orig_kept, rtol=1e-6)

    def test_hold_non_tensor_leaves_ignored(self):
        """Non-Tensor leaves (int/float/str/None/bytes) are silently skipped."""
        ctx = self._make_ctx()
        t1 = paddle.randn([2]).astype('float32')
        t2 = paddle.randn([3]).astype('float32')
        orig1 = t1.numpy().copy()
        orig2 = t2.numpy().copy()

        # Container mixes Tensors with int / float / str / None / bytes /
        # a dict whose values are non-Tensor; WalkDenseTensors must descend
        # into the containers, collect t1 / t2, and ignore everything else.
        mixed = (
            t1,
            42,
            "hello",
            None,
            [3.14, t2, b"bytes"],
            {'tag': 'x', 'n': 7, 'nested': (None, 'str')},
        )
        ctx._hold_tensors(mixed)

        _clear([t1, t2])
        ctx._restore_held_tensors()

        np.testing.assert_allclose(t1.numpy(), orig1, rtol=1e-6)
        np.testing.assert_allclose(t2.numpy(), orig2, rtol=1e-6)

    def test_hold_skipped_under_saved_tensors_hooks(self):
        """When saved_tensors_hooks is enabled _hold_tensors collects nothing."""
        ctx = self._make_ctx()
        t = paddle.randn([2, 3]).astype('float32')

        with paddle.autograd.saved_tensors_hooks(lambda x: x, lambda x: x):
            ctx._hold_tensors(t)

        _clear(t)
        ctx._restore_held_tensors()
        # holder was not populated, so impl stays empty after _clear_dataptr.
        self.assertFalse(t._is_initialized())


class TestRecomputeClosureHold(unittest.TestCase):
    """End-to-end recompute coverage of the Python-side closure helper.

    Covers ``_closure_cell_values`` (plain fn / nn.Layer / no-closure) and the
    ``_has_held_tensors`` True/False branches in RecomputeFunction.
    """

    def setUp(self):
        np.random.seed(1234)
        paddle.seed(1234)

    @staticmethod
    def _clone_leaf(t):
        out = paddle.to_tensor(t.numpy(), dtype=t.dtype)
        out.stop_gradient = False
        return out

    def test_closure_cell_values_empty_cell(self):
        """Empty cell triggers ValueError branch; valid cells still collected."""
        from paddle.distributed.fleet.recompute.recompute import (
            _closure_cell_values,
        )

        def outer():
            x = 1  # will be deleted → empty cell
            y = paddle.randn([2])

            def inner(a):
                return a + x + y  # noqa: F821

            del x
            return inner, y

        fn, y = outer()
        vals = _closure_cell_values(fn)
        # Empty cell dropped by the ValueError branch; only y remains.
        self.assertEqual(vals, (y,))

    def test_recompute_no_closure(self):
        """run_fn has no __closure__: _has_held_tensors=False, restore skipped."""
        from paddle.distributed.fleet.utils import recompute

        def run_fn(a, b):
            return (a * b + a).sum()

        a = paddle.randn([4, 4])
        a.stop_gradient = False
        b = paddle.randn([4, 4])
        b.stop_gradient = False
        a_ref = self._clone_leaf(a)
        b_ref = self._clone_leaf(b)

        loss = recompute(run_fn, a, b)
        _clear([a, b])
        # Sanity: _clear actually nulled impls — otherwise "restore succeeded"
        # would be trivially true and mask regressions.
        self.assertFalse(a._is_initialized())
        self.assertFalse(b._is_initialized())
        loss.backward()
        run_fn(a_ref, b_ref).backward()

        np.testing.assert_allclose(
            a.grad.numpy(), a_ref.grad.numpy(), rtol=1e-4
        )
        np.testing.assert_allclose(
            b.grad.numpy(), b_ref.grad.numpy(), rtol=1e-4
        )

    def test_recompute_closure_tensors(self):
        """Closure captures Tensor / tuple / list / dict: all restored."""
        from paddle.distributed.fleet.utils import recompute

        w_s = paddle.randn([4, 4])
        w_s.stop_gradient = False
        w_a = paddle.randn([4, 4])
        w_a.stop_gradient = False
        w_b = paddle.randn([4, 4])
        w_b.stop_gradient = False
        w_d = paddle.randn([4, 4])
        w_d.stop_gradient = False
        refs = [self._clone_leaf(t) for t in (w_s, w_a, w_b, w_d)]

        def make_fn(s, pair, mapping):
            def fn(x):
                a, b = pair
                return (x @ s + a * x + b * x + mapping['k'] * x).sum()

            return fn

        x = paddle.randn([4, 4])
        x.stop_gradient = False
        x_ref = self._clone_leaf(x)

        run_fn = make_fn(w_s, (w_a, w_b), {'k': w_d})
        ref_fn = make_fn(refs[0], (refs[1], refs[2]), {'k': refs[3]})

        loss = recompute(run_fn, x)
        _clear([x, w_s, w_a, w_b, w_d])
        for t in (x, w_s, w_a, w_b, w_d):
            self.assertFalse(t._is_initialized())
        loss.backward()
        ref_fn(x_ref).backward()

        for got, expect in zip((x, w_s, w_a, w_b, w_d), (x_ref, *refs)):
            self.assertIsNotNone(got.grad)
            np.testing.assert_allclose(
                got.grad.numpy(), expect.grad.numpy(), rtol=1e-4
            )

    def test_recompute_all_grad_from_closure(self):
        """Trainable tensors captured via closure must receive grads.

        Real-world pattern: trainable weights are closure-captured while the
        PyLayer arg is a regular activation.  Verifies that closure-captured
        ``w1`` / ``w2`` tensors are held across ``_clear_dataptr()`` and their
        grads are computed correctly during the recomputed backward.
        """
        from paddle.distributed.fleet.utils import recompute

        w1 = paddle.randn([4, 4])
        w1.stop_gradient = False
        w2 = paddle.randn([4, 4])
        w2.stop_gradient = False
        w1_ref = self._clone_leaf(w1)
        w2_ref = self._clone_leaf(w2)

        def make_fn(a, b):
            def fn(inp):
                return (inp * a * b).sum()

            return fn

        run_fn = make_fn(w1, w2)
        ref_fn = make_fn(w1_ref, w2_ref)

        inp = paddle.ones([4, 4])
        inp.stop_gradient = False
        inp_ref = paddle.ones([4, 4])
        inp_ref.stop_gradient = False

        loss = recompute(run_fn, inp)
        _clear([inp, w1, w2])
        for t in (inp, w1, w2):
            self.assertFalse(t._is_initialized())
        loss.backward()
        ref_fn(inp_ref).backward()

        np.testing.assert_allclose(
            w1.grad.numpy(), w1_ref.grad.numpy(), rtol=1e-4
        )
        np.testing.assert_allclose(
            w2.grad.numpy(), w2_ref.grad.numpy(), rtol=1e-4
        )

    def test_recompute_layer_forward_closure(self):
        """paddle.nn.Layer branch of _closure_cell_values."""
        from paddle.distributed.fleet.utils import recompute

        bias = paddle.randn([4, 4])
        bias.stop_gradient = False
        bias_ref = self._clone_leaf(bias)

        class MyLayer(paddle.nn.Layer):
            def __init__(self, captured):
                super().__init__()

                def forward(x):
                    return (x + captured).sum()

                self.forward = forward

            def forward(self, x):  # pragma: no cover
                raise RuntimeError

        layer = MyLayer(bias)
        layer_ref = MyLayer(bias_ref)
        x = paddle.randn([4, 4])
        x.stop_gradient = False
        x_ref = self._clone_leaf(x)

        loss = recompute(layer, x)
        _clear([x, bias])
        self.assertFalse(x._is_initialized())
        self.assertFalse(bias._is_initialized())
        loss.backward()
        layer_ref(x_ref).backward()

        np.testing.assert_allclose(
            x.grad.numpy(), x_ref.grad.numpy(), rtol=1e-4
        )
        np.testing.assert_allclose(
            bias.grad.numpy(), bias_ref.grad.numpy(), rtol=1e-4
        )


if __name__ == '__main__':
    unittest.main()