MagiCompiler/tests/feature_tests/test_cuda_graph.py

# Copyright (c) 2025 SandAI. 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.

"""
Tests for CUDA Graph compilation modes.

This module tests:
- FULL mode: entire forward pass captured as a single graph
- PIECEWISE mode: forward pass split at custom ops into multiple graphs
- Parameter handling: nn.Parameter should be ignored in graph capture
"""

import os
from typing import Optional
from unittest.mock import patch

import pytest
import torch
import torch.nn as nn
from torch.testing import assert_close

from magi_compiler.api import magi_compile, magi_register_custom_op
from magi_compiler.config import CompileConfig, CompileMode, CudaGraphMode
from magi_compiler.magi_backend.cuda_graph_mgr import CudaGraphMgr

pytestmark = pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA not available")


@magi_register_custom_op("athena::my_split_op", infer_output_meta_fn=["x"], is_subgraph_boundary=True)
def my_split_op(x: torch.Tensor) -> torch.Tensor:
    return x.clone()


@pytest.fixture(autouse=True)
def cuda_graph_test_env(tmp_path):
    """Fixture to set up a clean CUDA graph test environment."""

    # Create an isolated CudaGraphMgr instance for testing
    test_mgr = CudaGraphMgr()
    test_mgr.cache = dict()

    with patch("magi_compiler.magi_backend.cuda_graph_mgr.cuda_graph_mgr", return_value=test_mgr):
        yield tmp_path, test_mgr


def _create_compile_config(cache_dir: str, cudagraph_mode: CudaGraphMode) -> CompileConfig:
    """Create a compile configuration for CUDA graph testing."""
    from magi_compiler.config import get_compile_config

    return CompileConfig(
        compile_mode=CompileMode.MAGI_COMPILE,
        backend="inductor",
        cudagraph_mode=cudagraph_mode,
        splitting_ops=get_compile_config().splitting_ops,
        cache_root_dir=cache_dir,
    )


class TestCudaGraphFullMode:
    """Tests for CudaGraphMode.FULL - entire forward as single graph."""

    def test_full_mode_basic(self, cuda_graph_test_env):
        """Test basic FULL mode functionality with graph reuse."""
        tmp_path, test_mgr = cuda_graph_test_env
        cache_dir = os.path.join(str(tmp_path), "cache_full")
        os.makedirs(cache_dir, exist_ok=True)

        compile_config = _create_compile_config(cache_dir, CudaGraphMode.FULL)

        with patch("magi_compiler._api.get_compile_config", return_value=compile_config), patch(
            "torch.distributed.get_rank", return_value=0
        ):

            @magi_compile(dynamic_arg_dims={"x": 0})
            class FullModeModel(nn.Module):
                def __init__(self, model_config: Optional[dict] = None):
                    super().__init__()
                    self.linear1 = nn.Linear(10, 20).cuda()
                    self.linear2 = nn.Linear(20, 5).cuda()

                def forward(self, x: torch.Tensor) -> torch.Tensor:
                    x = torch.relu(self.linear1(x))
                    return self.linear2(x)

            model = FullModeModel(model_config=None).cuda()

            with torch.no_grad():
                # Prepare test inputs
                x1 = torch.randn(4, 10).cuda()
                x2 = x1.clone()
                x3 = torch.randn(2, 10).cuda()
                x4 = torch.randn(6, 10).cuda()
                x5 = x1.clone()

                import magi_compiler.magi_backend.cuda_graph_mgr as cgm

                active_mgr = cgm.cuda_graph_mgr()

                # First run: capture graph
                output1 = model(x1)
                assert output1.shape == (4, 5)
                assert active_mgr.graph_count == 1
                assert active_mgr.tensor_entry_count == 1

                # Same shape input: reuse tensor and graph
                output2 = model(x2)
                assert_close(output1, output2, rtol=1e-4, atol=1e-4)
                assert active_mgr.tensor_entry_count == 1
                assert active_mgr.graph_count == 1

                # Smaller batch: reuse tensor, new graph
                output3 = model(x3)
                assert output3.shape == (2, 5)
                assert active_mgr.tensor_entry_count == 1
                assert active_mgr.graph_count == 2

                # Larger batch: expand tensor, invalidate previous graphs
                output4 = model(x4)
                assert output4.shape == (6, 5)
                assert active_mgr.tensor_entry_count == 1
                assert active_mgr.graph_count == 1

                # Return to original batch size: recapture graph
                output5 = model(x5)
                assert_close(output1, output5, rtol=1e-4, atol=1e-4)
                assert active_mgr.tensor_entry_count == 1
                assert active_mgr.graph_count == 2


class TestCudaGraphPiecewiseMode:
    """Tests for CudaGraphMode.PIECEWISE - split at custom ops."""

    def test_piecewise_mode_with_split_op(self, cuda_graph_test_env):
        """Test PIECEWISE mode with custom splitting ops."""
        tmp_path, test_mgr = cuda_graph_test_env
        cache_dir = os.path.join(str(tmp_path), "cache_piecewise")
        os.makedirs(cache_dir, exist_ok=True)

        compile_config = _create_compile_config(cache_dir, CudaGraphMode.PIECEWISE)

        with patch("magi_compiler._api.get_compile_config", return_value=compile_config), patch(
            "torch.distributed.get_rank", return_value=0
        ):

            @magi_compile(dynamic_arg_dims={"x": 0})
            class PiecewiseModel(nn.Module):
                def __init__(self, *, model_config):
                    super().__init__()
                    self.linear1 = nn.Linear(10, 20).cuda()
                    self.linear2 = nn.Linear(20, 5).cuda()

                def forward(self, x: torch.Tensor) -> torch.Tensor:
                    x = self.linear1(x)
                    x = torch.ops.athena.my_split_op(x)  # Split point
                    x = self.linear2(x)
                    return x

            model = PiecewiseModel(model_config=None).cuda()

            with torch.no_grad():
                x1 = torch.randn(4, 10).cuda()
                x2 = torch.randn(2, 10).cuda()
                x3 = torch.randn(6, 10).cuda()
                x4 = torch.randn(4, 10).cuda()
                x5 = x1.clone()

                import magi_compiler.magi_backend.cuda_graph_mgr as cgm

                active_mgr = cgm.cuda_graph_mgr()

                # First run: capture 2 sub-graphs (before and after split op)
                output1 = model(x1)
                assert output1.shape == (4, 5)
                assert active_mgr.tensor_entry_count == 2  # Two tensor entries for two sub-graphs
                assert active_mgr.graph_count == 2

                # Same input: reuse all
                output2 = model(x1)
                assert_close(output1, output2, rtol=1e-4, atol=1e-4)
                assert active_mgr.tensor_entry_count == 2
                assert active_mgr.graph_count == 2

                # Smaller batch: reuse tensors, new sub-graphs
                output3 = model(x2)
                assert output3.shape == (2, 5)
                assert active_mgr.tensor_entry_count == 2
                assert active_mgr.graph_count == 4

                # Larger batch: expand tensors, invalidate previous sub-graphs
                output4 = model(x3)
                assert output4.shape == (6, 5)
                assert active_mgr.tensor_entry_count == 2
                assert active_mgr.graph_count == 2

                # Return to batch=4: reuse tensors, recapture sub-graphs
                output5 = model(x4)
                assert output5.shape == (4, 5)
                assert active_mgr.tensor_entry_count == 2
                assert active_mgr.graph_count == 4

                # Same as first input: verify output consistency
                output6 = model(x5)
                assert_close(output1, output6, rtol=1e-4, atol=1e-4)
                assert active_mgr.tensor_entry_count == 2
                assert active_mgr.graph_count == 4


class TestCudaGraphParameterHandling:
    """Tests for nn.Parameter handling in CUDA graph capture."""

    def test_parameters_excluded_from_graph_inputs(self, cuda_graph_test_env):
        """Test that nn.Parameters are not included in graph input tensors."""
        tmp_path, test_mgr = cuda_graph_test_env
        cache_dir = os.path.join(str(tmp_path), "cache_params")
        os.makedirs(cache_dir, exist_ok=True)

        compile_config = _create_compile_config(cache_dir, CudaGraphMode.FULL)

        with patch("magi_compiler._api.get_compile_config", return_value=compile_config), patch(
            "torch.distributed.get_rank", return_value=0
        ):

            @magi_compile(dynamic_arg_dims={"x": 0})
            class ParamModel(nn.Module):
                def __init__(self, model_config: Optional[dict] = None):
                    super().__init__()
                    # Multiple nn.Parameters (simulating model weights)
                    self.weight1 = nn.Parameter(torch.randn(10, 20).cuda())
                    self.bias1 = nn.Parameter(torch.randn(20).cuda())
                    self.weight2 = nn.Parameter(torch.randn(20, 5).cuda())
                    self.bias2 = nn.Parameter(torch.randn(5).cuda())
                    self.nested_params = nn.ParameterList(
                        [nn.Parameter(torch.randn(3, 5).cuda()), nn.Parameter(torch.randn(4, 4).cuda())]
                    )

                def forward(self, x: torch.Tensor) -> torch.Tensor:
                    x = torch.matmul(x, self.weight1) + self.bias1
                    x = torch.relu(x)
                    x = torch.matmul(x, self.weight2) + self.bias2
                    x = x + torch.matmul(torch.ones_like(x[:, :3]), self.nested_params[0])
                    return x

            model = ParamModel(model_config=None).cuda()

            with torch.no_grad():
                x1 = torch.randn(4, 10).cuda()
                output1 = model(x1)
                assert output1.shape == (4, 5)

                import magi_compiler.magi_backend.cuda_graph_mgr as cgm

                active_mgr = cgm.cuda_graph_mgr()

                # Only 1 tensor entry (the input x), not the parameters
                assert active_mgr.tensor_entry_count == 1

                # Verify the static entry contains only the input tensor
                static_entry = next(iter(active_mgr.cache.values()))
                assert len(static_entry.input_tensors) == 1
                assert len(static_entry.output_tensors) == 1
                assert isinstance(static_entry.input_tensors[0], torch.Tensor)
                assert not isinstance(static_entry.input_tensors[0], nn.Parameter)

            # Additional verification: ArgsUtils extracts only input tensor
            from magi_compiler.magi_backend.cuda_graph_mgr import ArgsUtils

            input_obj = {"args": (x1,), "kwargs": {}}
            extracted_tensors, _, _ = ArgsUtils.recursive_extract_core(input_obj)
            assert len(extracted_tensors) == 1
            assert extracted_tensors[0].data_ptr() == x1.data_ptr()


if __name__ == "__main__":
    pytest.main([__file__, "-v"])