diff --git a/.gitattributes b/.gitattributes index 119bef84d585fe684d0fada4c68fa2045db07d63..434cf9e3aedd83b763919659fd354334dc016c20 100644 --- a/.gitattributes +++ b/.gitattributes @@ -4,3 +4,5 @@ *.png filter=lfs diff=lfs merge=lfs -text *.jpg filter=lfs diff=lfs merge=lfs -text *.jpeg filter=lfs diff=lfs merge=lfs -text +*.pdf filter=lfs diff=lfs merge=lfs -text +*.whl filter=lfs diff=lfs merge=lfs -text diff --git a/Dockerfile b/Dockerfile index 023e5ce0fe3d6869d3845bf49b9761172490e049..bf3ea9020beb4d363cdec8c065b7ad89aa94a723 100644 --- a/Dockerfile +++ b/Dockerfile @@ -1,28 +1,46 @@ # ============================================================================= # HF Spaces Docker image for daVinci-MagiHuman -# Hardware: A100-80GB (or H100) +# Hardware: A100-80GB (recommended) # ============================================================================= -# Based on the official MagiCompiler image which includes: -# - CUDA 12.4, cuDNN, Python 3.12, PyTorch 2.9 -# - MagiCompiler (pre-installed) -# - Flash Attention 3 (Hopper) (pre-installed) -# ============================================================================= -FROM sandai/magi-compiler:latest +FROM nvidia/cuda:12.4.1-cudnn-devel-ubuntu22.04 ENV DEBIAN_FRONTEND=noninteractive ENV PYTHONUNBUFFERED=1 ENV GRADIO_SERVER_NAME=0.0.0.0 ENV GRADIO_SERVER_PORT=7860 -# System deps needed for audio/video processing +# System deps RUN apt-get update && apt-get install -y --no-install-recommends \ - ffmpeg libsndfile1 && \ - rm -rf /var/lib/apt/lists/* + python3.12 python3.12-dev python3.12-venv python3-pip \ + git ffmpeg libsndfile1 ninja-build && \ + rm -rf /var/lib/apt/lists/* && \ + ln -sf /usr/bin/python3.12 /usr/bin/python && \ + ln -sf /usr/bin/python3.12 /usr/bin/python3 WORKDIR /app # --------------------------------------------------------------------------- -# Python dependencies +# PyTorch (must be installed first — MagiCompiler build depends on it) +# --------------------------------------------------------------------------- +RUN pip install --no-cache-dir --upgrade pip && \ + pip install --no-cache-dir torch torchvision torchaudio \ + --index-url https://download.pytorch.org/whl/cu124 + +# --------------------------------------------------------------------------- +# Local packages: MagiCompiler + stable-audio whl +# --------------------------------------------------------------------------- +COPY pkgs/ pkgs/ +RUN pip install -e ./pkgs/MagiCompiler \ + --no-build-isolation --config-settings editable_mode=compat && \ + pip install --no-cache-dir pkgs/magife_stable_audio_open-1.0.0+mav.1-py3-none-any.whl + +# --------------------------------------------------------------------------- +# Flash Attention (pre-built wheel for CUDA 12.4 + PyTorch 2.9) +# --------------------------------------------------------------------------- +RUN pip install --no-cache-dir flash-attn --no-build-isolation + +# --------------------------------------------------------------------------- +# Project Python dependencies # --------------------------------------------------------------------------- COPY requirements.txt requirements-nodeps.txt ./ RUN pip install --no-cache-dir -r requirements.txt && \ @@ -36,16 +54,10 @@ COPY inference/ inference/ COPY example/ example/ COPY app.py . -# --------------------------------------------------------------------------- # Model weights are downloaded at runtime from HF Hub. -# Set HF_TOKEN as a Space secret if any repos are gated/private. -# -# Persistent storage (/data) is recommended on HF Spaces so weights survive -# container restarts. Enable it in Space settings → "Persistent storage". -# --------------------------------------------------------------------------- +# Enable "Persistent storage" in Space settings so /data survives restarts. ENV MODEL_ROOT=/data/models -# HF Spaces requires the app to listen on port 7860 EXPOSE 7860 CMD ["python", "app.py"] diff --git a/README.md b/README.md index 5a69647909601d9425104f341801aa71ccf7b46e..62fed289515daa14bcae7063c003b4d2ea411f0e 100644 --- a/README.md +++ b/README.md @@ -3,8 +3,7 @@ title: daVinci-MagiHuman emoji: 🎬 colorFrom: blue colorTo: purple -sdk: gradio -sdk_version: 5.23.0 +sdk: docker app_port: 7860 --- diff --git a/app.py b/app.py index 41d1862a492dd30b6992795a50ca88ef861234af..b9fd806a457742b43db8736d0b887650f399346f 100644 --- a/app.py +++ b/app.py @@ -2,7 +2,7 @@ """ Gradio frontend for daVinci-MagiHuman distilled model. -Designed for Hugging Face Spaces with ZeroGPU (Gradio SDK). +Designed for Hugging Face Spaces (Docker SDK, A100-80GB GPU). Accepts an image + text prompt + duration, generates audio-video output. """ @@ -12,8 +12,6 @@ import sys import tempfile import uuid -import spaces - # --------------------------------------------------------------------------- # 1. Download all model weights from HF Hub (runs on CPU, cached) # --------------------------------------------------------------------------- @@ -132,7 +130,6 @@ print("[app] Pipeline ready.") # 4. Inference wrapper — @spaces.GPU requests a ZeroGPU allocation # duration= sets the max GPU time in seconds (default 60, max 300) # --------------------------------------------------------------------------- -@spaces.GPU(duration=300) def generate_video( image, prompt: str, diff --git a/pkgs/MagiCompiler/.gitignore b/pkgs/MagiCompiler/.gitignore new file mode 100644 index 0000000000000000000000000000000000000000..225e15a88bb9e1fbd38b54c78ccca1d33512b34d --- /dev/null +++ b/pkgs/MagiCompiler/.gitignore @@ -0,0 +1,216 @@ +# magi_compiler +magi_compiler/_version.py +magi_dump_src_dir/ +*.nsys-rep +*.ncu-rep + +# Byte-compiled / optimized / DLL files +__pycache__/ +*.py[codz] +*$py.class + +# C extensions +*.so + +# Distribution / packaging +.Python +build/ +develop-eggs/ +dist/ +downloads/ +eggs/ +.eggs/ +lib/ +lib64/ +parts/ +sdist/ +var/ +wheels/ +share/python-wheels/ +*.egg-info/ +.installed.cfg +*.egg +MANIFEST + +# PyInstaller +# Usually these files are written by a python script from a template +# before PyInstaller builds the exe, so as to inject date/other infos into it. +*.manifest +*.spec + +# Installer logs +pip-log.txt +pip-delete-this-directory.txt + +# Unit test / coverage reports +htmlcov/ +.tox/ +.nox/ +.coverage +.coverage.* +.cache +nosetests.xml +coverage.xml +*.cover +*.py.cover +.hypothesis/ +.pytest_cache/ +cover/ + +# Translations +*.mo +*.pot + +# Vscode stuff: +.vscode + +# Django stuff: +*.log +local_settings.py +db.sqlite3 +db.sqlite3-journal + +# Flask stuff: +instance/ +.webassets-cache + +# Scrapy stuff: +.scrapy + +# Sphinx documentation +docs/_build/ + +# PyBuilder +.pybuilder/ +target/ + +# Jupyter Notebook +.ipynb_checkpoints + +# IPython +profile_default/ +ipython_config.py + +# pyenv +# For a library or package, you might want to ignore these files since the code is +# intended to run in multiple environments; otherwise, check them in: +# .python-version + +# pipenv +# According to pypa/pipenv#598, it is recommended to include Pipfile.lock in version control. +# However, in case of collaboration, if having platform-specific dependencies or dependencies +# having no cross-platform support, pipenv may install dependencies that don't work, or not +# install all needed dependencies. +#Pipfile.lock + +# UV +# Similar to Pipfile.lock, it is generally recommended to include uv.lock in version control. +# This is especially recommended for binary packages to ensure reproducibility, and is more +# commonly ignored for libraries. +#uv.lock + +# poetry +# Similar to Pipfile.lock, it is generally recommended to include poetry.lock in version control. +# This is especially recommended for binary packages to ensure reproducibility, and is more +# commonly ignored for libraries. +# https://python-poetry.org/docs/basic-usage/#commit-your-poetrylock-file-to-version-control +#poetry.lock +#poetry.toml + +# pdm +# Similar to Pipfile.lock, it is generally recommended to include pdm.lock in version control. +# pdm recommends including project-wide configuration in pdm.toml, but excluding .pdm-python. +# https://pdm-project.org/en/latest/usage/project/#working-with-version-control +#pdm.lock +pdm.toml +.pdm-python +.pdm-build/ + +# pixi +# Similar to Pipfile.lock, it is generally recommended to include pixi.lock in version control. +#pixi.lock +# Pixi creates a virtual environment in the .pixi directory, just like venv module creates one +# in the .venv directory. It is recommended not to include this directory in version control. +.pixi + +# PEP 582; used by e.g. github.com/David-OConnor/pyflow and github.com/pdm-project/pdm +__pypackages__/ + +# Celery stuff +celerybeat-schedule +celerybeat.pid + +# SageMath parsed files +*.sage.py + +# Environments +.env +.envrc +.venv +env/ +venv/ +ENV/ +env.bak/ +venv.bak/ + +# Spyder project settings +.spyderproject +.spyproject + +# Rope project settings +.ropeproject + +# mkdocs documentation +/site + +# mypy +.mypy_cache/ +.dmypy.json +dmypy.json + +# Pyre type checker +.pyre/ + +# pytype static type analyzer +.pytype/ + +# Cython debug symbols +cython_debug/ + +# PyCharm +# JetBrains specific template is maintained in a separate JetBrains.gitignore that can +# be found at https://github.com/github/gitignore/blob/main/Global/JetBrains.gitignore +# and can be added to the global gitignore or merged into this file. For a more nuclear +# option (not recommended) you can uncomment the following to ignore the entire idea folder. +#.idea/ + +# Abstra +# Abstra is an AI-powered process automation framework. +# Ignore directories containing user credentials, local state, and settings. +# Learn more at https://abstra.io/docs +.abstra/ + +# Visual Studio Code +# Visual Studio Code specific template is maintained in a separate VisualStudioCode.gitignore +# that can be found at https://github.com/github/gitignore/blob/main/Global/VisualStudioCode.gitignore +# and can be added to the global gitignore or merged into this file. However, if you prefer, +# you could uncomment the following to ignore the entire vscode folder +# .vscode/ + +# Ruff stuff: +.ruff_cache/ + +# PyPI configuration file +.pypirc + +# Cursor +# Cursor is an AI-powered code editor. `.cursorignore` specifies files/directories to +# exclude from AI features like autocomplete and code analysis. Recommended for sensitive data +# refer to https://docs.cursor.com/context/ignore-files +.cursorignore +.cursorindexingignore + +# Marimo +marimo/_static/ +marimo/_lsp/ +__marimo__/ diff --git a/pkgs/MagiCompiler/.pre-commit-config.yaml b/pkgs/MagiCompiler/.pre-commit-config.yaml new file mode 100644 index 0000000000000000000000000000000000000000..ff7f71d493a57427cc502b3eeb6778548cee80b9 --- /dev/null +++ b/pkgs/MagiCompiler/.pre-commit-config.yaml @@ -0,0 +1,60 @@ +exclude: \.patch$ +repos: +- repo: local + hooks: + - id: copyright_checker + name: copyright_checker + entry: python3 ./.github/.codestyle/copyright.hook + language: system + files: \.(c|cc|cxx|cpp|cu|h|hpp|hxx|proto|py|sh)$ +- repo: https://github.com/pre-commit/pre-commit-hooks + rev: v4.4.0 + hooks: + - id: check-added-large-files + args: + - --maxkb=30720 + - id: check-merge-conflict + - id: check-symlinks + - id: detect-private-key + files: (?!.*third_party)^.*$ | (?!.*book)^.*$ + - id: end-of-file-fixer + - id: trailing-whitespace + - id: requirements-txt-fixer + - id: sort-simple-yaml +- repo: https://github.com/Lucas-C/pre-commit-hooks.git + rev: v1.5.1 + hooks: + - id: remove-crlf + files: (?!.*third_party)^.*$ | (?!.*book)^.*$ + - id: remove-tabs + name: Tabs remover (C++) + files: \.(c|cc|cxx|cpp|cu|h|hpp|hxx|xpu|kps)$ + args: [--whitespaces-count, '2'] + - id: remove-tabs + name: Tabs remover (Python) + files: (.*\.(py|bzl)|BUILD|.*\.BUILD|WORKSPACE)$ + args: [--whitespaces-count, '4'] +- repo: https://github.com/psf/black.git + rev: 23.3.0 + hooks: + - id: black + args: [--line-length=127, --skip-string-normalization, --skip-magic-trailing-comma] + files: (.*\.(py|pyi|bzl)|BUILD|.*\.BUILD|WORKSPACE)$ +- repo: https://github.com/pre-commit/mirrors-isort + rev: v5.10.1 + hooks: + - id: isort + args: [--profile=black, --line-length=127, --multi-line=3, --force-grid-wrap=0] + files: \.py$ +- repo: https://github.com/PyCQA/autoflake + rev: v2.3.1 + hooks: + - id: autoflake + args: [--remove-all-unused-imports, --remove-unused-variables, --in-place, --ignore-init-module-imports, --ignore-pass-after-docstring] + files: \.py$ +- repo: https://github.com/macisamuele/language-formatters-pre-commit-hooks.git + rev: v2.9.0 + hooks: + - id: pretty-format-yaml + args: [--autofix, --indent, '4'] + additional_dependencies: [setuptools] diff --git a/pkgs/MagiCompiler/LICENSE b/pkgs/MagiCompiler/LICENSE new file mode 100644 index 0000000000000000000000000000000000000000..261eeb9e9f8b2b4b0d119366dda99c6fd7d35c64 --- /dev/null +++ b/pkgs/MagiCompiler/LICENSE @@ -0,0 +1,201 @@ + Apache License + Version 2.0, January 2004 + http://www.apache.org/licenses/ + + TERMS AND CONDITIONS FOR USE, REPRODUCTION, AND DISTRIBUTION + + 1. Definitions. + + "License" shall mean the terms and conditions for use, reproduction, + and distribution as defined by Sections 1 through 9 of this document. + + "Licensor" shall mean the copyright owner or entity authorized by + the copyright owner that is granting the License. + + "Legal Entity" shall mean the union of the acting entity and all + other entities that control, are controlled by, or are under common + control with that entity. For the purposes of this definition, + "control" means (i) the power, direct or indirect, to cause the + direction or management of such entity, whether by contract or + otherwise, or (ii) ownership of fifty percent (50%) or more of the + outstanding shares, or (iii) beneficial ownership of such entity. + + "You" (or "Your") shall mean an individual or Legal Entity + exercising permissions granted by this License. + + "Source" form shall mean the preferred form for making modifications, + including but not limited to software source code, documentation + source, and configuration files. + + "Object" form shall mean any form resulting from mechanical + transformation or translation of a Source form, including but + not limited to compiled object code, generated documentation, + and conversions to other media types. + + "Work" shall mean the work of authorship, whether in Source or + Object form, made available under the License, as indicated by a + copyright notice that is included in or attached to the work + (an example is provided in the Appendix below). + + "Derivative Works" shall mean any work, whether in Source or Object + form, that is based on (or derived from) the Work and for which the + editorial revisions, annotations, elaborations, or other modifications + represent, as a whole, an original work of authorship. For the purposes + of this License, Derivative Works shall not include works that remain + separable from, or merely link (or bind by name) to the interfaces of, + the Work and Derivative Works thereof. + + "Contribution" shall mean any work of authorship, including + the original version of the Work and any modifications or additions + to that Work or Derivative Works thereof, that is intentionally + submitted to Licensor for inclusion in the Work by the copyright owner + or by an individual or Legal Entity authorized to submit on behalf of + the copyright owner. For the purposes of this definition, "submitted" + means any form of electronic, verbal, or written communication sent + to the Licensor or its representatives, including but not limited to + communication on electronic mailing lists, source code control systems, + and issue tracking systems that are managed by, or on behalf of, the + Licensor for the purpose of discussing and improving the Work, but + excluding communication that is conspicuously marked or otherwise + designated in writing by the copyright owner as "Not a Contribution." + + "Contributor" shall mean Licensor and any individual or Legal Entity + on behalf of whom a Contribution has been received by Licensor and + subsequently incorporated within the Work. + + 2. Grant of Copyright License. Subject to the terms and conditions of + this License, each Contributor hereby grants to You a perpetual, + worldwide, non-exclusive, no-charge, royalty-free, irrevocable + copyright license to reproduce, prepare Derivative Works of, + publicly display, publicly perform, sublicense, and distribute the + Work and such Derivative Works in Source or Object form. + + 3. Grant of Patent License. Subject to the terms and conditions of + this License, each Contributor hereby grants to You a perpetual, + worldwide, non-exclusive, no-charge, royalty-free, irrevocable + (except as stated in this section) patent license to make, have made, + use, offer to sell, sell, import, and otherwise transfer the Work, + where such license applies only to those patent claims licensable + by such Contributor that are necessarily infringed by their + Contribution(s) alone or by combination of their Contribution(s) + with the Work to which such Contribution(s) was submitted. If You + institute patent litigation against any entity (including a + cross-claim or counterclaim in a lawsuit) alleging that the Work + or a Contribution incorporated within the Work constitutes direct + or contributory patent infringement, then any patent licenses + granted to You under this License for that Work shall terminate + as of the date such litigation is filed. + + 4. Redistribution. You may reproduce and distribute copies of the + Work or Derivative Works thereof in any medium, with or without + modifications, and in Source or Object form, provided that You + meet the following conditions: + + (a) You must give any other recipients of the Work or + Derivative Works a copy of this License; and + + (b) You must cause any modified files to carry prominent notices + stating that You changed the files; and + + (c) You must retain, in the Source form of any Derivative Works + that You distribute, all copyright, patent, trademark, and + attribution notices from the Source form of the Work, + excluding those notices that do not pertain to any part of + the Derivative Works; and + + (d) If the Work includes a "NOTICE" text file as part of its + distribution, then any Derivative Works that You distribute must + include a readable copy of the attribution notices contained + within such NOTICE file, excluding those notices that do not + pertain to any part of the Derivative Works, in at least one + of the following places: within a NOTICE text file distributed + as part of the Derivative Works; within the Source form or + documentation, if provided along with the Derivative Works; or, + within a display generated by the Derivative Works, if and + wherever such third-party notices normally appear. The contents + of the NOTICE file are for informational purposes only and + do not modify the License. You may add Your own attribution + notices within Derivative Works that You distribute, alongside + or as an addendum to the NOTICE text from the Work, provided + that such additional attribution notices cannot be construed + as modifying the License. + + You may add Your own copyright statement to Your modifications and + may provide additional or different license terms and conditions + for use, reproduction, or distribution of Your modifications, or + for any such Derivative Works as a whole, provided Your use, + reproduction, and distribution of the Work otherwise complies with + the conditions stated in this License. + + 5. Submission of Contributions. Unless You explicitly state otherwise, + any Contribution intentionally submitted for inclusion in the Work + by You to the Licensor shall be under the terms and conditions of + this License, without any additional terms or conditions. + Notwithstanding the above, nothing herein shall supersede or modify + the terms of any separate license agreement you may have executed + with Licensor regarding such Contributions. + + 6. Trademarks. This License does not grant permission to use the trade + names, trademarks, service marks, or product names of the Licensor, + except as required for reasonable and customary use in describing the + origin of the Work and reproducing the content of the NOTICE file. + + 7. Disclaimer of Warranty. Unless required by applicable law or + agreed to in writing, Licensor provides the Work (and each + Contributor provides its Contributions) on an "AS IS" BASIS, + WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or + implied, including, without limitation, any warranties or conditions + of TITLE, NON-INFRINGEMENT, MERCHANTABILITY, or FITNESS FOR A + PARTICULAR PURPOSE. You are solely responsible for determining the + appropriateness of using or redistributing the Work and assume any + risks associated with Your exercise of permissions under this License. + + 8. Limitation of Liability. In no event and under no legal theory, + whether in tort (including negligence), contract, or otherwise, + unless required by applicable law (such as deliberate and grossly + negligent acts) or agreed to in writing, shall any Contributor be + liable to You for damages, including any direct, indirect, special, + incidental, or consequential damages of any character arising as a + result of this License or out of the use or inability to use the + Work (including but not limited to damages for loss of goodwill, + work stoppage, computer failure or malfunction, or any and all + other commercial damages or losses), even if such Contributor + has been advised of the possibility of such damages. + + 9. Accepting Warranty or Additional Liability. While redistributing + the Work or Derivative Works thereof, You may choose to offer, + and charge a fee for, acceptance of support, warranty, indemnity, + or other liability obligations and/or rights consistent with this + License. However, in accepting such obligations, You may act only + on Your own behalf and on Your sole responsibility, not on behalf + of any other Contributor, and only if You agree to indemnify, + defend, and hold each Contributor harmless for any liability + incurred by, or claims asserted against, such Contributor by reason + of your accepting any such warranty or additional liability. + + END OF TERMS AND CONDITIONS + + APPENDIX: How to apply the Apache License to your work. + + To apply the Apache License to your work, attach the following + boilerplate notice, with the fields enclosed by brackets "[]" + replaced with your own identifying information. (Don't include + the brackets!) The text should be enclosed in the appropriate + comment syntax for the file format. We also recommend that a + file or class name and description of purpose be included on the + same "printed page" as the copyright notice for easier + identification within third-party archives. + + Copyright [yyyy] [name of copyright owner] + + 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. diff --git a/pkgs/MagiCompiler/README.md b/pkgs/MagiCompiler/README.md new file mode 100644 index 0000000000000000000000000000000000000000..40cf60eba614fb6e4c21ec042f02e7eac3931f59 --- /dev/null +++ b/pkgs/MagiCompiler/README.md @@ -0,0 +1,186 @@ +## MagiCompiler + +An engineering-oriented compiler and execution augmentation library for PyTorch 2.8+, providing module-level compilation decorators, backend adapters, graph partitioning strategies, readable and reusable compile artifacts, and tightly integrated runtime scheduling for any inference engine. The design goal is to systematically expose capabilities of PyTorch Dynamo / AOTAutograd / Inductor / Triton while prioritizing correctness, stability, observability, and maintainability. + +### Design Overview + +- Compilation entrypoint: the `@magi_compile` decorator augments `nn.Module.forward` with compilation, including dynamic-shape annotation and argument validation. +- Partitioning and passes: configurable graph partitioning and pass management (e.g., `InductorPass`, `PostGradPassManager`) for fusion, kernel generation, and tuning. +- Artifact system: persists compile artifacts using a Python file/directory layout for readability, auditability, and portability (see “Compile Cache Overview”). +- Configuration: `CompileConfig` centralizes backend, partition rules, cache root, runtime shapes, and other key parameters. + +### Key Features + +- Dynamic-shape annotations: + - Automatic inference: when a `forward` parameter is annotated as `torch.Tensor` or `torch.Tensor | None`, dimension 0 is treated as dynamic by default. + - Explicit specification: use `@magi_compile(dynamic_arg_dims={...})` to mark dimensions (negative indices supported). + - Consistency constraints: parameters that alternately appear as `None` and non-`None` across the model lifetime cannot be captured into the same computation graph. +- Backend selection and standalone compilation: + - `inductor` mode defaults to PyTorch 2.8+ `standalone_compile`, producing reusable artifacts. + - `eager` mode is available for debugging or fallback paths. +- Partitioning and passes: operator-set-driven partition rules and pass contexts that stabilize subgraph boundaries and kernel generation across runtime shapes. +- Readable, portable artifacts: structured directories with Python files for quick triage and cross-environment debugging. +- Engine integration: the decorator reads engine-level `CompileConfig` to stay aligned with distributed/scheduling components. + +## Installation and Requirements + +- Python ≥ 3.10 +- PyTorch ≥ 2.8 (with `torch._inductor.standalone_compile` available) +- Recommended to be used within the Athena environment, together with its dependencies and distributed components (e.g., CUDA Graph manager). + +For local development, install in editable mode: + +```bash +pip install -e . --no-build-isolation --config-settings editable_mode=compat +``` + +## Quick Start + +### Minimal Example (automatic dynamic-dim inference) + +```python +import torch +from torch import nn +from magi_compiler.decorator import magi_compile + +@magi_compile +class MyModel(nn.Module): + def __init__(self, *, model_config): + super().__init__() + self.linear = nn.Linear(10, 5) + + def forward(self, x: torch.Tensor, y: torch.Tensor | None) -> torch.Tensor: + if y is not None: + return self.linear(x + y) + return self.linear(x) + +# In Athena, model_config is typically provided by the engine +model = MyModel(model_config=...) +out1 = model(torch.randn(4, 10), torch.randn(4, 10)) +out2 = model(torch.randn(8, 10), None) # dynamic batch dimension +``` + +### Explicit Dynamic-Dim Specification + +```python +@magi_compile(dynamic_arg_dims={"x": -1}) # mark the last dimension as dynamic +class DynamicDimModel(nn.Module): + def __init__(self, *, model_config): + super().__init__() + self.proj = nn.Linear(16, 16, bias=False) + + def forward(self, x: torch.Tensor) -> torch.Tensor: + return self.proj(x) + +m = DynamicDimModel(model_config=...) +_ = m(torch.randn(2, 16)) +_ = m(torch.randn(2, 32)) # allow the last dimension to vary +``` + +## Configuration and Modes + +- `CompileConfig`: centralizes compile parameters (backend, cache paths, partition strategy, dynamic shapes, traced files, etc.). +- `CompileMode`: typical setting is `CompileMode.TORCH_COMPILE`. +- Backends: + - `inductor`: uses `standalone_compile` to produce reusable artifacts, ideal for production deployments. + - `eager`: convenient for rapid debugging or as a fallback. + +## Architecture and Execution Flow (Brief) + +1. `@magi_compile` wraps `nn.Module`: + - infers/validates `dynamic_arg_dims`; + - extends MRO by injecting `MagiCompilerBase`; + - reads engine-level `CompileConfig` in MagiCompiler. +2. `CompilerManager`: + - defines cache keys using `(runtime_shape, graph_index, backend)`; + - dispatches to backends via `CompilerInterface` (`InductorStandaloneAdaptor` or `EagerAdaptor`); + - applies partition rules and pass contexts within `compile_context(...)`; + - serializes compile artifacts into a human-readable directory structure. +3. Monitoring and statistics: + - counters and timestamps report per-shape/per-subgraph latencies and milestones. + +## Compile Cache Overview + +This document summarizes the cache files generated by `torch.compile` (TorchDynamo + TorchInductor + AOTAutograd + Triton). Reference path: `cache/`. + +### Directory Layout (Tree) + +```text +cache/ +├─ depyf/ +│ └─ rank_0/ +│ ├─ __transformed_code_0_for_forward.py +│ ├─ decompiled_code.py +│ ├─ full_code_for_forward_0.py +│ ├─ __compiled_fn_1.BEFORE_PRE_GRAD.{0..N}.py +│ ├─ __compiled_fn_1.kernel_{0..K}.py +│ ├─ __compiled_fn_1.__compiled_fn_1_.0.py +│ ├─ __compiled_fn_1.Before_split.0.py +│ ├─ __compiled_fn_1.After_split.0.py +│ ├─ __compiled_fn_1.pre_split_module.0.py +│ ├─ __compiled_fn_1.post_split_module.0.py +│ └─ __compiled_fn_1.pre_insert_deferred_runtime_asserts__.0.py +│ +└─ torch_compile_cache/ + └─ bfa0df33ea/ # Hash for graph + compile options + device, etc. + └─ rank_0/ # Rank id in distributed/multi-GPU runs + └─ backbone/ + ├─ computation_graph.py + ├─ magi_compile_cache.py + ├─ artifact_shape_None_subgraph_0/ + ├─ artifact_shape_None_subgraph_1/ + ├─ ... + └─ artifact_shape_None_subgraph_30/ + ├─ ir/ + │ └─ *.py # Python/Triton kernels generated by Inductor for this subgraph + ├─ fxgraph/ + │ └─ */*/ # Binary FX IR/metadata snapshots (not human-readable) + ├─ aotautograd/ + │ └─ */*/ # AOTAutograd partition/capture metadata and artifacts + ├─ 44/ + │ └─ *.py # Other sharded/generated code buckets + └─ ... # Structure may vary slightly across subgraphs +``` + +### What Each File/Dir Is For + +- `cache/depyf/` (TorchDynamo/Depyf debug exports) + - `__transformed_code_0_for_forward.py`: The Dynamo-transformed `forward` code (diff-friendly view of pre/post transformation). + - `decompiled_code.py`: Decompiled snapshot to help map traced graphs back to original Python. + - `full_code_for_forward_0.py`: A more complete expanded `forward` for inspection. + - `__compiled_fn_1.BEFORE_PRE_GRAD.{i}.py`: Intermediate wrapper snapshots at specific compile stages (e.g., before autodiff). + - `__compiled_fn_1.kernel_{k}.py`: Entrypoints/wrappers for kernels generated at various stages. + - `Before_split` / `After_split` / `pre_split_module` / `post_split_module`: Intermediate forms around graph partitioning. + - `pre_insert_deferred_runtime_asserts__*.py`: Snapshot before inserting deferred runtime assertions (dynamic shapes/guards). + +- `cache/torch_compile_cache/` (TorchInductor artifacts) + - `bfa0df33ea/`: Namespace keyed by a hash of model structure, compile settings, and device info. + - `rank_0/`: Bucket per process rank for distributed runs. + - `backbone/`: + - `computation_graph.py`: Full model FX GraphModule with symbolic dims; shared across subgraph kernels. + - `magi_compile_cache.py`: + - Maps subgraph indices to artifact directories, e.g. `(None, i, 'inductor_standalone') -> artifact_shape_None_subgraph_i/`. + - Registers and asynchronously compiles Triton kernels via `AsyncCompile.triton(...)`, including autotune metadata, device properties, scheduling hints, etc. + - `artifact_shape_None_subgraph_{N}/`: + - `ir/*.py`: Inductor-generated Python/Triton kernels and scheduling code for this subgraph (readable). + - `fxgraph/*/*/`: FX IR/metadata snapshots for fast graph reconstruction (binary; do not edit). + - `aotautograd/*/*/`: AOTAutograd partitions/captures and replay requirements. + - Additional hashed/prefixed buckets (e.g., `44/`, `o5/`, `55/`, `br/`) containing generated operator/subtask code. + +### FAQ + +- How are these caches produced? + - At runtime by `torch.compile(...)`, after TorchDynamo tracing, AOTAutograd partitioning, TorchInductor lowering/fusion, and Triton codegen. +- Will they change across runs? + - Yes. Different input shapes, env vars, device info, or compile options can produce different hash namespaces (e.g., a new `bfa0df33ea`). +- Is it safe to delete them? + - Yes. You can delete `cache/`. It will be rebuilt on demand; the next run will be slower due to recompilation. + +## Compatibility and Recommendations + +- Prefer official PyTorch ≥ 2.8 builds to ensure `standalone_compile` availability. +- For highly dynamic models, explicitly mark key dynamic dimensions to improve graph capture and cache reuse. + +## Acknowledgments + +This library builds upon capabilities of PyTorch Dynamo, AOTAutograd, Inductor, and Triton, and incorporates engineering practices and interface designs inspired by the vLLM community. We thank the relevant open-source communities and contributors. diff --git a/pkgs/MagiCompiler/docs/AutoCudaGraphDesign.md b/pkgs/MagiCompiler/docs/AutoCudaGraphDesign.md new file mode 100644 index 0000000000000000000000000000000000000000..cb8e7cfc32fe88e9ab913f1657b01776f81802ee --- /dev/null +++ b/pkgs/MagiCompiler/docs/AutoCudaGraphDesign.md @@ -0,0 +1,174 @@ +## AutoCudaGraph Design + +Author: ZhiyaoCen + +## Overview +AutoCudaGraph is a CUDA Graph optimization module integrated into the MagiCompiler framework, designed to automate CUDA Graph capture, caching, replay, and tensor memory management for PyTorch-based neural network inference. It targets Transformer architectures with dynamic sequence lengths, optimizing kernel execution by reusing pre-captured computation graphs and static tensor buffers. Core Goals: +* Automate CUDA Graph lifecycle (capture/replay/cache) with minimal code intrusion +* Support dynamic shape adaptation (sequence length expansion) +* Optimize memory efficiency via global memory pool and static tensor reuse +* Ensure consistency between cached graphs and runtime inputs/outputs +## Key Components + + +### CudaGraphMgr (Core Manager) +Singleton class managing all CUDA Graph operations: +```python +class CudaGraphMgr: + def __init__(self): + self.cache: Dict[StaticSignature, StaticTensorEntry] = dict() + self.graph_mem_pool: Optional[torch.cuda.graph_pool_handle] = None +``` + +**Core Methods** +| Method | Purpose | +|---------------------------------|----------------------------------------| +| run() | Main entry: Replay cached graph or warm up & capture new graph| +| wrapped_graph_capture() | Capture CUDA Graph with sliced static input/output tensors | +| wrapped_graph_replay() | Replay cached CUDA Graph with sliced static tensors and output template wrapping +| get_expanded_static_tensors() | Expand static tensors, reuse buffers if dimensionally compatible| + + +### Signature System + +StaticSignature +```python +@dataclass(unsafe_hash=True) +class StaticSignature(HashableDataclass): + func_name: str = "" + tensor_static_infos: Tuple[TensorStaticInfo, ...] = tuple() +``` +* Encodes fixed properties of input tensors (dtype, static dimensions) +* Used as primary key for static tensor buffer caching + +DynamicSignature +```python +@dataclass(unsafe_hash=True) +class DynamicSignature(HashableDataclass): + tensor_dynamic_infos: Tuple[TensorDynamicInfo, ...] = tuple() + literals_info: LiteralsInfo = None +``` +* Tracks dynamic dimensions (sequence length) and literal parameters +* Secondary key for graph entry lookup + +### Tensor Management +```python +@dataclass +class StaticTensorEntry: + input_tensors: Optional[List[torch.Tensor]] = None + output_tensors: Optional[List[torch.Tensor]] = None + template_entry_dict: Dict[DynamicSignature, OutputTemplateEntry] = None +``` +* Memory Reuse: Reuse existing tensor buffers when possible to avoid reallocation +* Dynamic Expansion: Only expand static tensors when new input dimensions exceed current buffer size +* Shape Validation: Ensure static dimensions (non-sequence) match between cached and new tensors + + +### Graph Management +```python +@dataclass +class GraphEntry: + graph: Optional[torch.cuda.CUDAGraph] = None + inconsistent: bool = False + invalid: bool = False + +@dataclass +class OutputTemplateEntry: + graph_entry_dict: Dict[int, GraphEntry] = None + output_template: Any = None +``` +* Graph State Tracking: GraphEntry tracks CUDA Graph instances and validity states to control replay eligibility. +* Layer-wise Organization: OutputTemplateEntry maps dynamic signatures to per-layer GraphEntry for layer-specific graph reuse. +* Output Consistency: output_template preserves output object structure to ensure consistent result wrapping during replay. + +## Execution Flow +### Inline Replay (Fast Path) +* Extract input signatures from runtime arguments +* Look up cached CUDA Graph via StaticSignature + DynamicSignature + layer number +* Validate graph consistency (not inconsistent/invalid) +* Reuse static tensors with dynamic slicing +* Replay graph and return sliced output +### Graph Capture (Slow Path) +Triggered when no valid cached graph exists or tensor expansion is needed: +* Execute function to get output tensors +* Ensure input signatures match post-warmup +* Expand static buffers if new shapes require it +* Capture new CUDA Graph with static tensors +* Store new graph and update tensor entries +* Return warmup execution output as final result +### Sequence Length Handling +* Only last dimension is static for ND tensors (ND > 1) +* All dimension is dynamic for 1D tensors (ND=1) +* Automatic buffer expansion for increasing sequence lengths +* Invalidates old graphs when tensors are expanded + + +## Examples +```python +import torch +import torch.nn as nn +from magi_compiler.cuda_graph_mgr import cuda_graph_mgr, cuda_graph_enable_if + +class SimpleTransformerLayer(nn.Module): + def __init__(self, hidden_dim: int = 1024, num_heads: int = 8): + super().__init__() + self.self_attn = nn.MultiheadAttention(hidden_dim, num_heads, batch_first=True) + self.linear = nn.Linear(hidden_dim, hidden_dim) + self.layer_norm = nn.LayerNorm(hidden_dim) + self.layer_number = 0 + + @cuda_graph_enable_if(lambda: torch.cuda.is_available()) + def forward(self, x: torch.Tensor): + attn_out, _ = self.self_attn(x, x, x) + out = self.linear(self.layer_norm(x + attn_out)) + return out + +device = torch.device("cuda" if torch.cuda.is_available() else "cpu") +model = SimpleTransformerLayer(hidden_dim=1024, num_heads=8).to(device).eval() +graph_mgr = cuda_graph_mgr() + +with torch.no_grad(): + input_1 = torch.randn(2, 512, 1024, device=device) + output_1 = model(input_1) + print(f"First run (graph capture): Output shape = {output_1.shape}") + print(f"Cached graphs count: {graph_mgr.graph_count}") + + input_2 = torch.randn(2, 512, 1024, device=device) + output_2 = model(input_2) + print(f"Second run (graph replay): Output shape = {output_2.shape}") + print(f"Cached graphs count: {graph_mgr.graph_count}") + + input_3 = torch.randn(2, 1024, 1024, device=device) + output_3 = model(input_3) + print(f"Third run (tensor expansion): Output shape = {output_3.shape}") + print(f"Cached graphs count: {graph_mgr.graph_count}") + print(f"Static tensor memory usage: {graph_mgr.tensor_mem_size:.2f} MB") + + print("\nCUDA Graph Cache Details:") + print(graph_mgr.formatted_cache_str()) + + # StaticSignature: StaticSignature(_cached_hash=None, func_name='SimpleTransformerLayer.forward', tensor_static_infos=(TensorStaticInfo(_cached_hash=None, name='', shapes=(-1, -1, 1024), dtype='torch.float32'),)) + # Input Static Tensors: [shape=[2, 1024, 1024],dtype=torch.float32] + # Output Static Tensors: [shape=[2, 1024, 1024],dtype=torch.float32] + # DynamicSignature: DynamicSignature(_cached_hash=None, tensor_dynamic_infos=(TensorDynamicInfo(_cached_hash=None, name='', shapes=(2, 512, -1)),), literals_info=LiteralsInfo(_cached_hash=None, literals=())) + # Output Template: FakeTensor(shape=[2, 512, 1024], dtype='torch.float32', device='cuda:0') + # Layer 0: Graph Status: Invalid + # DynamicSignature: DynamicSignature(_cached_hash=None, tensor_dynamic_infos=(TensorDynamicInfo(_cached_hash=None, name='', shapes=(2, 1024, -1)),), literals_info=LiteralsInfo(_cached_hash=None, literals=())) + # Output Template: FakeTensor(shape=[2, 1024, 1024], dtype='torch.float32', device='cuda:0') + # Layer 0: Graph Status: Valid +``` + +## Limitations and Constraints +* No support for data-dependent control flow in captured functions +* Graph capture fails if function contains CPU/GPU synchronization +* Only supports CUDA tensors (CPU tensors trigger fallback) +* Custom input classes must inherit from InplaceSubstituteFakeClass +* Assumes input tensors of captured graphs are not reused externally (risk of cross-scenario static tensor reuse) +* Relies on identical function, input tensors shapes, and constants for valid graph reuse +* No support for multi-stream execution scenarios + +## Best Practices +* Dynamic Dimensions: Tensor use sequence length as dimension 0 where possible +* Monitor Memory Usage: Track graph_mem_pool_size and tensor_mem_size to avoid OOM +* Specify Layer IDs: Use layer_number to distinguish graphs across different models/layers +* LRU Cache (Future): Implement cache eviction to limit total graph/tensor count diff --git a/pkgs/MagiCompiler/docs/Hunyuan15Benchmark.md b/pkgs/MagiCompiler/docs/Hunyuan15Benchmark.md new file mode 100644 index 0000000000000000000000000000000000000000..9700e9696f5995c36298093b1f1a00291219a584 --- /dev/null +++ b/pkgs/MagiCompiler/docs/Hunyuan15Benchmark.md @@ -0,0 +1,79 @@ +## Hunyuan1.5 Benchmark + +### Executive Summary +This report presents a comprehensive performance evaluation of the **[Athena](https://github.com/world-sim-dev/athena)** framework compared to the baseline **[LightX2V](https://github.com/ModelTC/LightX2V)** framework. The benchmarks were conducted using the **[Hunyuan-1.5](https://github.com/Tencent-Hunyuan/HunyuanVideo-1.5)** model on NVIDIA H100 hardware. + +--- +### 🎯Test Environment & Versioning +#### Hardware & Settings + +| Parameter | Value | +| ------------------- | -------------- | +| Hardware | NVIDIA H100 | +| Model | Hunyuan-1.5 480p_t2v_distilled | +| Precision | torch.bfloat16 | +| Inference Steps | 20 | +| Resolution | 480p | +| FPS | 24 | +| CFG | Disable | +#### Software Versioning +To ensure reproducibility, the following specific commits were used for this benchmark: +| Framework | Branch / Tag | Commit | +| --------- | ------------ | ------ | +| Athena | main|[5e6086b](https://github.com/world-sim-dev/athena/commit/5e6086b4dc2ab60bc4d44dbe39745b4354075121) | +| LightX2V | main | [5573905](https://github.com/ModelTC/LightX2V/commit/5573905f3f38d876d468b815f86d417a608975b6) | + +### 🏆 Performance Benchmarks +📊 We compared the iteration speed (seconds per iteration) between Athena and LightX2V across three distinct Context Parallel (CP) configurations. +| Configuration | Frames | LightX2V (s/it) | Athena (s/it) | Speedup | +| ------------- | ------ | -------------- | -------------- | ------- | +| CP1 | 121 | 2.42 | **2.06** | **1.17x** 🚀| +| CP2 | 121 | 1.38 | **1.13** | **1.22x** 🚀| +| CP4 | 241 | 2.25 | **1.85** | **1.22x** 🚀| +| CP8 | 241 | 1.28 | **1.01** | **1.27x** 🚀| + +--- +### 📹 Output Comparison +| Framework | Video Result | +| --------- | ---------------------------- | +| Athena | | +| LightX2V | | + + +### 💡 Reproduction Guide +To reproduce the results presented in this report, follow the steps below using the specified commit hashes. +#### Setup +```bash +git clone https://github.com/world-sim-dev/athena +cd athena +git checkout 5e6086b +pip install -r requirements.txt +pip install -r requirements-nodeps.txt +pip install -e ./pkgs/MagiCompiler --no-build-isolation --config-settings editable_mode=compat + + +# Clone and install LightX2V (for baseline comparison) +git clone https://github.com/ModelTC/LightX2V +cd lightx2v +git checkout 5573905 +pip install -v . +``` + +#### Running Benchmarks +For Athena, run: +``` +RESOLUTION=480p CFG_DISTILLED=true TASK=t2v CHECKPOINT_PATH=path/to/480p_t2v_distilled bash ./scripts/run_hunyuan.sh +``` +For LightX2V: +Clone the scripts from [Benchmark for LightX2V](https://gist.github.com/wtr0504/d80bbebb7da1ef7b58f3e6faf1c68880) and run: +``` +git clone https://gist.github.com/wtr0504/d80bbebb7da1ef7b58f3e6faf1c68880 +MODEL_PATH=path/to/HunyuanVideo-1.5 DISTILL_CKPT=path/to/480p_t2v_distilled/diffusion_pytorch_model.safetensors bash run_hunyuan.sh +``` + +### 🔎 MagiCompiler Optimization Methodology +**Whole Graph Compilation** +Constant Folding & Dead Code Elimination: Streamlining the computation graph prior to execution. + +**Coarse-grained Kernel Fusion** +MagiCompiler aggregates multiple smaller operators into larger, fused kernels. This optimization is critical for efficient execution on the GPU. diff --git a/pkgs/MagiCompiler/docs/Wan2.2Benchmark.md b/pkgs/MagiCompiler/docs/Wan2.2Benchmark.md new file mode 100644 index 0000000000000000000000000000000000000000..ca4233d8491aeb484338659e1336afb1bf8a9eb3 --- /dev/null +++ b/pkgs/MagiCompiler/docs/Wan2.2Benchmark.md @@ -0,0 +1,72 @@ +## Wan2.2 Benchmark + +### Executive Summary +This report presents a comprehensive performance evaluation of the **[Athena](https://github.com/world-sim-dev/athena)** framework compared to the baseline **[LightX2V](https://github.com/ModelTC/LightX2V)** framework. The benchmarks were conducted using the **[Wan2.2-TI2V-5B](https://huggingface.co/Wan-AI)** model on NVIDIA H100 hardware. + +--- +### 🎯Test Environment & Versioning +#### Hardware & Settings + +| Parameter | Value | +| ------------------- | -------------- | +| Hardware | NVIDIA H100 | +| Model | Wan2.2-TI2V-5B | +| Precision | torch.bfloat16 | +| Inference Steps | 50 | +| Resolution | 704 × 1280(720p)| +| FPS | 24 | +| CFG | Enabled | +#### Software Versioning +To ensure reproducibility, the following specific commits were used for this benchmark: +| Framework | Branch / Tag | Commit | +| --------- | ------------ | ------ | +| Athena | main|[f676ae6](https://github.com/world-sim-dev/athena/commit/f676ae64ad2fc581289d1c3ae5eb51c15ce76f1d) | +| LightX2V | main | [33f0f67](https://github.com/ModelTC/LightX2V/commit/33f0f67f4ecdff86b1db676d3e0786628cc31c7b) | + +### 🏆 Performance Benchmarks +📊 We compared the iteration speed (seconds per iteration) between Athena and LightX2V across three distinct Context Parallel (CP) configurations. +| Configuration | Frames | LightX2V (s/it) | Athena (s/it) | Speedup | +| ------------- | ------ | -------------- | -------------- | ------- | +| CP1 | 121 | 1.928 | **1.69** | **1.14x** 🚀| +| CP2 | 121 | 1.197 | **1.06** | **1.13x** 🚀| +| CP4 | 241 | 1.767 | **1.32** | **1.34x** 🚀| +| CP8 | 241 | 1.507 | **1.35** | **1.12x** 🚀| + +--- + +### 💡 Reproduction Guide +To reproduce the results presented in this report, follow the steps below using the specified commit hashes. +#### Setup +```bash +git clone https://github.com/world-sim-dev/athena +cd athena +git checkout f676ae6 +pip install -r requirements.txt + +# Clone and install LightX2V (for baseline comparison) +git clone https://github.com/ModelTC/LightX2V +cd lightx2v +git checkout 33f0f67 +pip install -r requirements.txt + +``` + +#### Running Benchmarks +For Athena, run: +``` +bash ./scripts/run_wan2_2_ti2v_i2v.sh +``` +For LightX2V: +Clone the scripts from [Benchmark for LightX2V](https://gist.github.com/wtr0504/629388f17ed38d1c12d5ef5c25a15197) and run: +``` +git clone https://gist.github.com/wtr0504/629388f17ed38d1c12d5ef5c25a15197 +bash run_wan.sh +``` + +### 🔎 MagiCompiler Optimization Methodology +**Whole Graph Compilation** +Constant Folding & Dead Code Elimination: Streamlining the computation graph prior to execution. +**Coarse-grained Kernel Fusion** +MagiCompiler aggregates multiple smaller operators into larger, fused kernels. This optimization is critical for efficient execution on the GPU. +**All to All Communication** +MagiCompiler Uses ``all_to_all_single`` (1 communication op per sync point) while LightX2V Uses all_to_all x 3 (3 separate communication ops). diff --git a/pkgs/MagiCompiler/docs/WhyMagiCompiler.md b/pkgs/MagiCompiler/docs/WhyMagiCompiler.md new file mode 100644 index 0000000000000000000000000000000000000000..21d4194f4070a7916d26dc7831ced834994196cb --- /dev/null +++ b/pkgs/MagiCompiler/docs/WhyMagiCompiler.md @@ -0,0 +1,246 @@ +# Why MagiCompiler? + +## 1. Compiler Overview + +### 1.1 Background + +We have long encountered several significant challenges in model optimization: + +1. **Blurred Acceleration Boundaries:** There is ambiguity regarding the extent of optimization required to achieve "extreme" performance. +2. **Complex Performance Tuning:** Optimization strategies are often tightly coupled with model architectures, necessitating extensive and repetitive manual intervention. +3. **Deficiency in Optimization Tools:** The infrastructure lacks sufficient mechanisms for computational graph-level optimizations, such as operator substitution and communication overlap. + +MagiCompiler addresses these issues through the following approaches: + +* **Addressing Challenge 1:** It adopts **whole-graph compilation**, thoroughly transcending the boundaries of `TransformerLayer` to maximize the scope of kernel fusion. +* **Addressing Challenge 2:** It integrates infrastructure optimizations directly into MagiCompiler, implementing features such as `AutoCudaGraph` and `AutoCheckpointing(WIP)`. +* **Addressing Challenge 3:** It leverages the dynamic-to-static capabilities provided by **Dynamo**, capturing `fx.graph` IR in eager mode to perform pass optimizations at the IR level. + +#### Illustrative Example + +```python +from magi_compiler import magi_compile + +@magi_compile() +class TinyModel(nn.Module): + def __init__(self): + super().__init__() + self.linear = nn.Linear(1024, 1024, device="cuda") + + @no_grad() + def forward(self, x: torch.Tensor, y: torch.Tensor, z: torch.Tensor) -> torch.Tensor: + return self.linear(x + y - z + 1) + + +def magi_compiler_demo(): + model = TinyModel() + x = torch.randn(1024, 1024, device="cuda") + y = torch.randn(1024, 1024, device="cuda") + z = torch.randn(1024, 1024, device="cuda") + model(x, y, z) +``` + +**Optimized Code (Triton Kernel):** + +```python +triton_poi_fused_add_sub_0 = async_compile.triton('triton_poi_fused_add_sub_0', ''' +import triton +import triton.language as tl + +from torch._inductor.runtime import triton_helpers, triton_heuristics +from torch._inductor.runtime.triton_helpers import libdevice, math as tl_math +from torch._inductor.runtime.hints import AutotuneHint, ReductionHint, TileHint, DeviceProperties +triton_helpers.set_driver_to_gpu() + +@triton_heuristics.pointwise( + size_hints={'x': 1048576}, + filename=__file__, + triton_meta={'signature': {'in_ptr0': '*fp32', 'in_ptr1': '*fp32', 'in_ptr2': '*fp32', 'out_ptr0': '*fp32', 'xnumel': 'i32', 'XBLOCK': 'constexpr'}, 'device': DeviceProperties(type='cuda', index=0, multi_processor_count=132, cc=90, major=9, regs_per_multiprocessor=65536, max_threads_per_multi_processor=2048, warp_size=32), 'constants': {}, 'configs': [{(0,): [['tt.divisibility', 16]], (1,): [['tt.divisibility', 16]], (2,): [['tt.divisibility', 16]], (3,): [['tt.divisibility', 16]], (4,): [['tt.divisibility', 16]]}]}, + inductor_meta={'grid_type': 'Grid1D', 'autotune_hints': set(), 'kernel_name': 'triton_poi_fused_add_sub_0', 'mutated_arg_names': [], 'optimize_mem': True, 'no_x_dim': False, 'num_load': 3, 'num_reduction': 0, 'backend_hash': 'B8F4209CBFC2377D6AF9CF3C65D610CA2B56C138A443862350DE1E56F5BF54C3', 'are_deterministic_algorithms_enabled': False, 'assert_indirect_indexing': True, 'autotune_local_cache': True, 'autotune_pointwise': True, 'autotune_remote_cache': None, 'force_disable_caches': False, 'dynamic_scale_rblock': True, 'max_autotune': False, 'max_autotune_pointwise': False, 'min_split_scan_rblock': 256, 'spill_threshold': 16, 'store_cubin': False}, + min_elem_per_thread=0 +) +@triton.jit +def triton_poi_fused_add_sub_0(in_ptr0, in_ptr1, in_ptr2, out_ptr0, xnumel, XBLOCK : tl.constexpr): + xoffset = tl.program_id(0) * XBLOCK + xindex = xoffset + tl.arange(0, XBLOCK)[:] + xmask = xindex < xnumel + x0 = xindex + tmp0 = tl.load(in_ptr0 + (x0), xmask) + tmp1 = tl.load(in_ptr1 + (x0), xmask) + tmp3 = tl.load(in_ptr2 + (x0), xmask) + tmp2 = tmp0 + tmp1 + tmp4 = tmp2 - tmp3 + tmp5 = 1.0 + tmp6 = tmp4 + tmp5 + tl.store(out_ptr0 + (x0), tmp6, xmask) +''', device_str='cuda') +``` + +### 1.2 Frontend (Dynamo) + +![Dynamo](./assets/why_magicompiler_1_dynamo.jpeg) + +* **PyFrameObject (Dynamic Call Stack):** + * Represents the context environment during function execution. Python creates a new `PyFrameObject` for each function call. +* **PyCodeObject (Static Bytecode):** + * The compiled product of Python code, which is static and read-only. A single `PyCodeObject` exists regardless of how many times the function is invoked. + +```python +def f(x, mod): + for guard, transformed_code in f.compiled_entries: + if guard(x, mod): + return transformed_code(x, mod) + try: + guard, transformed_code = compile_and_optimize(x, mod) + f.compiled_entries.append([guard, transformed_code]) + return transformed_code(x, mod) + except FailToCompileError: + y = mod(x) + z = torch.log(y) + return z +``` + +#### Symbolic Shape + +MagiCompiler specifically targets the Transformer architecture and supports custom `dynamic_arg_dims` (typically for `seq_len`). + +**Example:** + +```python +@magi_compile(dynamic_arg_dims={"x": 0, "y": 0, "z": 0}) +class TinyModel(nn.Module): + def __init__(self): + super().__init__() + self.linear = nn.Linear(1024, 1024, device="cuda") + + @no_grad() + def forward(self, x: torch.Tensor, y: torch.Tensor, z: torch.Tensor) -> torch.Tensor: + return self.linear(x + y - z + 1) +``` + +**Guard Mechanism and Elimination in Symbolic Shape Deduction:** + +```log +I1204 16:31:35.745000 1859360 torch/_dynamo/symbolic_convert.py:3842] [0/0] Step 1: torchdynamo start tracing inner /usr/local/lib/python3.12/dist-packages/torch/_dynamo/external_utils.py:66 +I1204 16:31:35.746000 1859360 torch/fx/experimental/symbolic_shapes.py:3775] [0/0] create_env +I1204 16:31:35.781000 1859360 torch/fx/experimental/symbolic_shapes.py:5120] [0/0] create_symbol s33 = 1024 for L['args'][0].size()[0] [2, int_oo] return self.linear(x + y - z + 1) # ome/niubility2/hongyu/athena/integration_test/scripts/linear_demo.py:50 in forward (_dynamo/variables/builder.py:3501 in ), for more info run with TORCHDYNAMO_EXTENDED_DEBUG_CREATE_SYMBOL="s33" or to suppress this message run with TORCHDYNAMO_EXTENDED_ADVICE="0" +I1204 16:31:35.785000 1859360 torch/fx/experimental/symbolic_shapes.py:5120] [0/0] create_symbol s6 = 1024 for L['args'][1].size()[0] [2, int_oo] return self.linear(x + y - z + 1) # ome/niubility2/hongyu/athena/integration_test/scripts/linear_demo.py:50 in forward (_dynamo/variables/builder.py:3501 in ), for more info run with TORCHDYNAMO_EXTENDED_DEBUG_CREATE_SYMBOL="s6" or to suppress this message run with TORCHDYNAMO_EXTENDED_ADVICE="0" +I1204 16:31:35.794000 1859360 torch/fx/experimental/symbolic_shapes.py:7213] [0/0] eval Eq(s33, s6) [guard added] return self.linear(x + y - z + 1) # ome/niubility2/hongyu/athena/integration_test/scripts/linear_demo.py:50 in forward (_subclasses/fake_impls.py:1148 in infer_size), for more info run with TORCHDYNAMO_EXTENDED_DEBUG_GUARD_ADDED="Eq(s33, s6)" +I1204 16:31:35.795000 1859360 torch/fx/experimental/symbolic_shapes.py:6792] [0/0] set_replacement s6 = s33 (solve) VR[2, int_oo] +I1204 16:31:35.800000 1859360 torch/fx/experimental/symbolic_shapes.py:5120] [0/0] create_symbol s21 = 1024 for L['args'][2].size()[0] [2, int_oo] return self.linear(x + y - z + 1) # ome/niubility2/hongyu/athena/integration_test/scripts/linear_demo.py:50 in forward (_dynamo/variables/builder.py:3501 in ), for more info run with TORCHDYNAMO_EXTENDED_DEBUG_CREATE_SYMBOL="s21" or to suppress this message run with TORCHDYNAMO_EXTENDED_ADVICE="0" +I1204 16:31:35.806000 1859360 torch/fx/experimental/symbolic_shapes.py:7213] [0/0] eval Eq(s33, s21) [guard added] return self.linear(x + y - z + 1) # ome/niubility2/hongyu/athena/integration_test/scripts/linear_demo.py:50 in forward (_subclasses/fake_impls.py:1148 in infer_size), for more info run with TORCHDYNAMO_EXTENDED_DEBUG_GUARD_ADDED="Eq(s33, s21)" +I1204 16:31:35.807000 1859360 torch/fx/experimental/symbolic_shapes.py:6792] [0/0] set_replacement s33 = s21 (solve) VR[2, int_oo] +I1204 16:31:35.828000 1859360 torch/_dynamo/symbolic_convert.py:4059] [0/0] Step 1: torchdynamo done tracing inner (RETURN_VALUE) +I1204 16:31:35.837000 1859360 torch/fx/experimental/symbolic_shapes.py:6792] [0/0] set_replacement s6 = s21 (find) VR[2, int_oo] +``` + +### 1.3 Backend (Inductor, MagiBackend, etc.) + +![Backend Architecture](./assets/why_magicompiler_2_arch.png) + +MagiCompiler hijack the `torch.compile` logic through the following components: + +* **`custom_partitioner_fn`:** Segments the forward and backward computational graphs and determines which intermediate results are transmitted to the backward pass. +* **`post_grad_custom_pre_pass`:** Performs pass optimizations at the whole-graph level (computational graph matching and rewriting). +* **`PartitionFunc`:** Implements custom subgraph partitioning logic, utilizing attention mechanisms as splitting points. + +![Partition](./assets/why_magicompiler_3_partition.png) + +* **`post_grad_custom_post_pass`:** Executes pass optimizations at the subgraph level (computation/communication overlap). + +--- + +## 2. Best Practices + +### 2.1 Model Adaptation + +MagiCompiler has certain limitations, such as mandatory whole-graph capture and the inability to support implicit subgraph interruptions. Consequently, manual adaptation is required in specific scenarios: + +**1. Computational Graph Dependencies or CPU/GPU Synchronization** + +```python +@magi_compile +class MeanModule(torch.nn.Module): + def __init__(self): + super().__init__() + + def forward(self, x: torch.Tensor, y: torch.Tensor): + x = x.cos().sin() + if x.mean() > 0.5: + x = x - 1 + return x * y +``` + +> **Note:** In typical Transformer models, certain pre/post-processing operations are unavoidable. Therefore, the recommended practice for `magi_compiler` is to perform **whole-graph capture at the `TransformerBlock` level**, as `TransformerBlock` computations constitute over 95% of the total workload. + +**2. Custom Operators (e.g., FlashAttention, FlexFlashAttention, MoE kernels)** + +* **Operator Registration:** A logic for operator registration is provided. Commonly used operators like FlashAttention (FA) and FlexFlashAttention (FFA) are already registered. + +```python +# Operator Registration +@torch.library.custom_op("athena::flash_attn_func", mutates_args=()) +def flash_attn_func(query: torch.Tensor, key: torch.Tensor, value: torch.Tensor) -> torch.Tensor: + ... + +# Operator Deduce Function +@flash_attn_func.register_fake +def _(query: torch.Tensor, key: torch.Tensor, value: torch.Tensor) -> torch.Tensor: + return torch.empty_like(query) + +# Call flash_attn_func +self_attn_out = torch.ops.athena.flash_attn_func(q, k, v) +out, _ = torch.ops.athena.flex_flash_attn_func(q, k, v, q_ranges=ffa_handler.q_ranges, k_ranges=ffa_handler.k_ranges) +``` + +* **Unit Testing:** Independent unit tests for operators should be provided in the production environment. + +```python +@pytest.mark.parametrize("batch_size", [1]) +@pytest.mark.parametrize("seq_len", [1024, 2048, 4096]) +@pytest.mark.parametrize("query_head", [48]) +@pytest.mark.parametrize("kv_head", [4, 8]) +@pytest.mark.parametrize("head_dim", [128, 256]) +def test_fake_fa3(batch_size, seq_len, query_head, kv_head, head_dim): + q = torch.randn((batch_size, seq_len, query_head, head_dim), device="cuda", dtype=torch.bfloat16) + k = torch.randn((batch_size, seq_len, kv_head, head_dim), device="cuda", dtype=torch.bfloat16) + v = torch.randn((batch_size, seq_len, kv_head, head_dim), device="cuda", dtype=torch.bfloat16) + torch.library.opcheck(torch.ops.athena.flash_attn_func, (q, k, v)) +``` + +### 2.2 Debugging Methods + +Key questions for debugging: +* Is the bug originating from the compiler? +* Which specific component of the compiler is causing the bug? + +![Debugging](./assets/why_magicompiler_4_debug.png) + +```python +class CompileConfig(BaseModel): + # Basic configs + backend: str = Field("inductor", description="Compilation backend.") + compile_mode: CompileMode = Field(CompileMode.MAGI_COMPILE, description="Compilation mode.") + ... + + # Cudagraph configs + cudagraph_mode: CudaGraphMode = Field(CudaGraphMode.NONE, description="Cudagraph mode.") + ... + + # Pass configs + pass_config: PassConfig = Field(PassConfig(), description="Pass configuration.") + ... +``` + +### 2.3 Profiling Results + +For further details, please refer to the [**Wan2.2 Benchmark**](Wan2.2Benchmark.md). + +--- + +## References + +1. [PyTorch 2.0 Overview](https://docs.pytorch.org/assets/pytorch2-2.pdf) +2. [TorchDynamo: An Experiment in Dynamic Python Bytecode Transformation](https://dev-discuss.pytorch.org/t/torchdynamo-an-experiment-in-dynamic-python-bytecode-transformation/361) +3. [Depyf Walkthrough](https://depyf.readthedocs.io/en/latest/walk_through.html) +4. [Getting Started with PyTorch Compiler](https://docs.pytorch.org/docs/main/torch.compiler_get_started.html) diff --git a/pkgs/MagiCompiler/docs/WhyMagiDepyf.md b/pkgs/MagiCompiler/docs/WhyMagiDepyf.md new file mode 100644 index 0000000000000000000000000000000000000000..616d9a4a4ca9d078c2bb7b06521ad7033d2ddfd0 --- /dev/null +++ b/pkgs/MagiCompiler/docs/WhyMagiDepyf.md @@ -0,0 +1,175 @@ +# magi_depyf + +A structured inspector for `torch.compile` and MagiCompiler compilation +artifacts — decompiled source, Inductor kernels, guard conditions, graph break +chains, and more — all organized in a navigable directory tree. + +## Why + +### The problem: compilation is a black box + +`torch.compile` and MagiCompiler accelerate models by transforming Python +functions through a deep pipeline: Dynamo captures bytecode into FX graphs, +a backend (Inductor, etc.) compiles them into optimized kernels, and the +runtime dispatches through a chain of cache entries, compiled functions, and +resume functions. The result is fast — but opaque. + +When something goes wrong — a correctness bug, an unexpected graph break, a +performance cliff — you need to see what the compiler actually produced. +What bytecode did Dynamo generate? Which subgraphs went to Inductor vs. +eager fallback? What do the kernels look like? How do resume functions +chain together? MagiCompiler adds further layers: CUDA graph capture +regions, piecewise subgraph splits, and its own dispatch logic. + +None of this is easily accessible. + +### depyf: a pioneering effort + +[depyf](https://github.com/thuml/depyf) was the first tool to address this, +hooking into `torch._dynamo` to dump decompiled source, FX graphs, and +Inductor output. It made `torch.compile` significantly more transparent. + +### Why a new tool? + +magi_depyf is purpose-built for MagiCompiler's compilation stack, and takes +a fundamentally different approach from depyf: + +| | depyf | magi_depyf | +|-|-------|------------| +| **When artifacts are collected** | During compilation, via monkey-patching internal hooks | **After** compilation completes, by walking the final CacheEntry chain — a single, clean post-hoc pass | +| **Output structure** | Flat files (`full_code_0.py`, `__transformed_code_0_for_xxx.py`, …) — hard to navigate for complex models | **Hierarchical directory tree** mirroring the compilation structure: function → entries → compiled\_fns / resume\_fns | +| **MagiCompiler support** | None | First-class: per-subgraph Inductor source, CUDA graph mode, piecewise split metadata | +| **Decompiler** | Monolithic class supporting Python 3.8–3.12 | Modular handler registry; focused on 3.10+ | + +### Key features + +**See everything the compiler produced, in one structured tree.** +One context manager call gives you a complete, navigable dump: decompiled +bytecode (before and after Dynamo), Inductor kernel source for every compiled +function, guard conditions, bytecode metadata (`co_flags`, `co_consts`, +`dis` output), and the full resume function chain — recursively. + +**MagiCompiler-native.** +Understands MagiCompiler's backend, extracting per-subgraph Inductor source, +CUDA graph capture mode (full / piecewise), and split metadata that +`torch.compile`-only tools cannot see. + +**Post-hoc introspection.** +Artifacts are collected after compilation finishes, by walking the CacheEntry +linked list and extracting what Dynamo and the backend actually produced. +No monkey-patching of internal compilation hooks, no interference with the +compilation process itself. + +## Usage + +### `dump_src` — the main entry point + +```python +import torch +from magi_compiler.magi_depyf.inspect import dump_src + +@torch.compile +def toy_example(a, b): + x = a / (torch.abs(a) + 1) + if b.sum() < 0: + b = b * -1 + return x * b + +with dump_src("./output"): + for _ in range(100): + toy_example(torch.randn(10), torch.randn(10)) +``` + +This produces: + +``` +output/ + toy_example/ + overview.md # Navigable index with links to everything + decompiled_code.py # Original function source + bytecode_info.txt # CodeType metadata + dis output + entry_0/ + decompiled_code.py # Dynamo-transformed bytecode → Python + bytecode_info.txt # Transformed code metadata + guards.txt # Guard conditions for this cache entry + compiled_fns/ + __compiled_fn_1_xxx.py # FX graph (readable) + __compiled_fn_1_xxx_post_grad.py # Post-grad graph + __compiled_fn_1_xxx_runnable.py # Inductor kernel source + resume_fns/ + __resume_at_94_2/ # Resume function after graph break + overview.md + decompiled_code.py # Resume function source + bytecode_info.txt + entry_0/ # Dynamo compiles resume fns too + decompiled_code.py + guards.txt + compiled_fns/ + ... + __resume_at_104_3/ + ... +``` + +### Programmatic API + +```python +from magi_compiler.magi_depyf import decompile + +# Decompile a code object to Python source +source = decompile(my_function.__code__) + +# Introspect a compiled function +from magi_compiler.magi_depyf.inspect import Introspector +info = Introspector.build_function_info(fn, fn_globals=fn.__globals__) +# info.entries[0].decompiled_src — decompiled transformed code +# info.entries[0].compiled_fns — backend-compiled functions +# info.entries[0].resume_fns — resume functions after graph breaks +``` + +### Tested model architectures + +The test suite verifies the decompile → recompile round-trip on real model +structures, ensuring the decompiler produces correct source for Dynamo output: + +| Category | Models | +|----------|--------| +| **PyTorch core** | MLP, Conv-BN-ReLU, MultiheadAttention, TransformerEncoderLayer, Embedding, residual blocks, depthwise separable conv | +| **Diffusion blocks** | GEGLU, RMSNorm, sinusoidal embeddings, cross-attention, AdaLayerNorm, DiT blocks, timestep MLP | +| **HuggingFace transformers** | BERT, GPT-2, T5 encoder (tiny configs) | +| **HuggingFace diffusers** | Attention (self / cross), BasicTransformerBlock | +| **timm** | ResNet-18, MobileNetV3, EfficientNet-B0, ViT, ConvNeXt, Swin, DeiT | +| **Graph breaks** | `print()` breaks, explicit `graph_break()`, multi-break chains — with recursive resume function round-tripping | + +## Code structure + +``` +magi_depyf/ +├── __init__.py # Public API: decompile, safe_decompile +│ +├── decompile/ # Bytecode → Python source (no torch dependency) +│ ├── decompiler.py # Decompiler: orchestrates the pipeline +│ ├── recompiler.py # CodeRecompiler: decompile → compile() → CodeType +│ ├── bytecode/ +│ │ ├── instruction.py # Mutable wrapper over dis.Instruction +│ │ ├── source_emitter.py # Stack machine + source accumulator +│ │ ├── decompile_context.py # Read-only context for handlers +│ │ ├── handler_registry.py # Opcode → handler dispatch table +│ │ └── handlers/ # One module per opcode category +│ └── postprocess/ # Source-level cleanup passes +│ +└── inspect/ # torch.compile introspection (requires torch) + ├── dump_src.py # dump_src(): the main entry point + ├── introspect.py # Introspector: walk CacheEntry chain + ├── model.py # Data model (FunctionInfo, EntryInfo, ...) + ├── writer.py # Serialize to directory tree + ├── session.py # CaptureSession: bytecode hook lifecycle + └── result.py # CaptureResult: one compilation event +``` + +## Compatibility + +| Requirement | Version | +|-------------|---------| +| **Python** | >= 3.10 | +| **PyTorch** | >= 2.0 (requires `torch._dynamo` internals) | +| **depyf** | Optional; used as fallback by `safe_decompile` | diff --git a/pkgs/MagiCompiler/docs/assets/submod_0_rank_0.pdf b/pkgs/MagiCompiler/docs/assets/submod_0_rank_0.pdf new file mode 100644 index 0000000000000000000000000000000000000000..10137762a1ce3849ed200776914aa47fc0c11e70 --- /dev/null +++ b/pkgs/MagiCompiler/docs/assets/submod_0_rank_0.pdf @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:583d97460eb7ebf48efbdeb7a6ae424f640de9ff99e6f33bbadef432583f40d3 +size 16122 diff --git a/pkgs/MagiCompiler/magi_compiler/__init__.py b/pkgs/MagiCompiler/magi_compiler/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..4ccf87cd7ab50b4c688875dec348e45a46fedf40 --- /dev/null +++ b/pkgs/MagiCompiler/magi_compiler/__init__.py @@ -0,0 +1,17 @@ +# 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. + +from .api import magi_compile + +__all__ = ["magi_compile"] diff --git a/pkgs/MagiCompiler/magi_compiler/_cache_data_cls.py b/pkgs/MagiCompiler/magi_compiler/_cache_data_cls.py new file mode 100644 index 0000000000000000000000000000000000000000..cefcabf78a68b2d50a97d0eb2cb3b44ed874de22 --- /dev/null +++ b/pkgs/MagiCompiler/magi_compiler/_cache_data_cls.py @@ -0,0 +1,28 @@ +# 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. + +import dataclasses + + +@dataclasses.dataclass(frozen=True) +class CacheHandle: + key: str | None + path: str + + +@dataclasses.dataclass(frozen=True) +class CacheEntry: + runtime_shape: int | None + graph_index: int + backend_name: str diff --git a/pkgs/MagiCompiler/magi_compiler/api.py b/pkgs/MagiCompiler/magi_compiler/api.py new file mode 100644 index 0000000000000000000000000000000000000000..18615f02c159e02ed61ac4ac506105a6e5d37d24 --- /dev/null +++ b/pkgs/MagiCompiler/magi_compiler/api.py @@ -0,0 +1,666 @@ +# 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. + +import copy +import gc +import inspect +import os +from contextlib import contextmanager +from typing import Callable, TypeVar, get_args, get_origin, overload +from unittest.mock import patch + +import magi_compiler.utils.envs as envs +import torch +from magi_compiler.cuda.cudart import pin_memory_in_place +from magi_compiler.magi_compiler_base import MagiCompilerBase +from magi_compiler.utils import compilation_counter, magi_logger +from magi_compiler.utils.compile_time_monitor import CompileMonitor +from torch import distributed as dist +from torch import nn +from torch._dynamo.symbolic_convert import InliningInstructionTranslator + +from .config import CompileConfig, CompileMode, get_compile_config + + +# ============================================================================= +# Workaround: TorchInductor autotune get_raw_stream +# ============================================================================= +# TorchInductor autotune code blocks may reference get_raw_stream() without +# defining it, causing "name 'get_raw_stream' is not defined" at runtime. +# Register it as a builtin so the exec'd autotune snippets can always find it. +def _patch_get_raw_stream(): + try: + import builtins + + from torch._C import _cuda_getCurrentRawStream as _get_raw_stream + except Exception: + return + if not hasattr(builtins, "get_raw_stream"): + builtins.get_raw_stream = _get_raw_stream + + +_patch_get_raw_stream() + +# ============================================================================= +# Dynamo Config Isolation +# ============================================================================= +# Capture the default dynamo config at module load time (before any torch.compile). +# This ensures we have a "clean" baseline config that hasn't been modified by +# external torch.compile calls (e.g., with dynamic=True). +_DEFAULT_DYNAMO_CONFIG: dict = torch._dynamo.config.get_config_copy() + + +@contextmanager +def _isolated_dynamo_config(): + """ + Context manager that provides an isolated dynamo config environment. + """ + with torch._dynamo.config.patch(**_DEFAULT_DYNAMO_CONFIG): + yield + + +_T = TypeVar("_T", bound=type[nn.Module]) +_W = TypeVar("_W", bound="MagiCompilerBase") + + +@overload +def magi_compile(*, enable_if: Callable[None, bool] | None = None) -> Callable[[_T], _T]: + ... + + +@overload +def magi_compile(*, dynamic_arg_dims: dict[str, int | list[int]] | None) -> Callable[[_T], _T]: + ... + + +@overload +def magi_compile(*, config_patch: Callable[[CompileConfig], CompileConfig] | None = None) -> Callable[[_T], _T]: + ... + + +@overload +def magi_compile(cls: _T) -> _T: + ... + + +def magi_compile( + cls: _T | None = None, + *, + model_tag: str | None = None, + dynamic_arg_dims: dict[str, int | list[int]] | None = None, + enable_if: Callable[None, bool] | None = None, + config_patch: Callable[[CompileConfig], CompileConfig] | None = None, +) -> Callable[[_T], _T] | _T: + """ + A decorator to add support for compiling the forward method of a class. + + Usage: + 1. use directly as a decorator without arguments: + ```python + @magi_compile + class MyModel(nn.Module): + def forward(self, x: torch.Tensor, y: Optional[torch.Tensor]): ... + ``` + + 2. use as a decorator with arguments: + ```python + @magi_compile(dynamic_arg_dims={"x": 0, "y": 0}) + class MyModel(nn.Module): + def forward(self, x: torch.Tensor, y: Optional[torch.Tensor]): ... + ``` + + Arguments: + - model_tag: optional tag in cache path (e.g. "wan_ti2v"). If not set, class name is used. + Path segment: model_{idx}_{model_tag}_rank_{rank}. + - dynamic_arg_dims: a dictionary that maps argument names to the dynamic + dimensions of the argument. The dynamic dimensions can be either a single + integer or a list of integers. + - enable_if: a function that returns a boolean value indicating whether to compile the model or not. + This is useful if you want to compile the model only when certain conditions are met. + + Notes: + - dynamic_arg_dims will be inferred from the type annotation of the forward method if not provided, + if the argument is annotated as `torch.Tensor` or `Optional[torch.Tensor]`, + the first dimension will be marked as dynamic. + + - if an argument is `None`, it should always be passed as `None` during + the lifetime of the model, otherwise, it cannot be captured as a single + computation graph. + + """ + + def cls_decorator_helper(cls: _T) -> _T: + nonlocal dynamic_arg_dims + dynamic_arg_dims = dynamic_arg_dims or _infer_dynamic_arg_dims(cls) + + # Accuracy check + assert hasattr(cls, "forward"), "decorated class should have a forward method." + assert len(dynamic_arg_dims) > 0, ( + "No dynamic dimensions found in the forward method of " f"{cls}. Please provide dynamic_arg_dims explicitly." + ) + for k in dynamic_arg_dims: + assert k in inspect.signature(cls.forward).parameters, f"Argument {k} not found in the forward method of {cls}" + + return _magi_compile(cls, dynamic_arg_dims, enable_if, config_patch, model_tag=model_tag) + + if cls is not None: + # use `magi_compile` as a decorator without arguments, cls is the class to be decorated + assert isinstance(cls, type) + return cls_decorator_helper(cls) + + return cls_decorator_helper + + +def offload(obj): + if isinstance(obj, torch.Tensor): + return obj.cpu() + elif isinstance(obj, dict): + return {k: offload(v) for k, v in obj.items()} + elif isinstance(obj, (list, tuple)): + return type(obj)(offload(item) for item in obj) + return obj + + +def _magi_compile( + cls: _T, + dynamic_arg_dims: dict[str, int | list[int]], + enable_if: Callable[None, bool] | None = None, + config_patch: Callable[[CompileConfig], CompileConfig] | None = None, + model_tag: str | None = None, +) -> _T: + """ + A decorator to add support for compiling the forward method of a class. + """ + if MagiCompilerBase in cls.__bases__: + return cls + + # take care of method resolution order, make sure super().__init__ is called on the base class + # other than MagiCompilerBase + cls.__bases__ = cls.__bases__ + (MagiCompilerBase,) + + if get_compile_config().offload_config.model_cpu_offload: + magi_logger.info(f"Enabling CPU offload for {cls}") + _orig_apply = cls._apply + + def _cpu_apply(self, fn): + if getattr(self, "_magi_offloaded_once", False): + return _orig_apply(self, fn) + + # First, move all parameters/buffers to CPU + def _force_cpu(t): + return fn(t).cpu() + + _orig_apply(self, _force_cpu) + + # create shared memory tensors for all parameters/buffers on CPU + if dist.is_initialized(): + local_rank = int(os.environ.get("LOCAL_RANK", 0)) + full_state_dict = self.state_dict() + + grouped_params = {} # {dtype: [(name, tensor), ...]} + for name, tensor in full_state_dict.items(): + if tensor.device.type == 'cpu': + dt = tensor.dtype + if dt not in grouped_params: + grouped_params[dt] = [] + grouped_params[dt].append((name, tensor)) + + shared_state_dict = {} + self._magi_giant_buffers = [] + + dist.barrier() + + for dtype, param_list in grouped_params.items(): + dtype_str = str(dtype).split('.')[-1] + shared_bin_path = ( + f"{envs.MAGI_SHARED_BIN_PATH}/magi_model_shared_{dtype_str}_{self.__class__.__name__}.bin" + ) + + total_numel = sum(t.numel() for _, t in param_list) + + if local_rank == 0: + flat_buffer = torch.zeros(total_numel, dtype=dtype) + offset = 0 + for _, tensor in param_list: + numel = tensor.numel() + flat_buffer[offset : offset + numel].copy_(tensor.view(-1)) + offset += numel + + if dtype == torch.bfloat16: + flat_buffer.view(torch.int16).numpy().tofile(shared_bin_path) + elif dtype.itemsize == 1 and dtype.is_floating_point: + # fp8 + flat_buffer.view(torch.uint8).numpy().tofile(shared_bin_path) + else: + flat_buffer.numpy().tofile(shared_bin_path) + + del flat_buffer + gc.collect() + + dist.barrier() + + giant_shared_tensor = torch.from_file( + shared_bin_path, shared=True, size=total_numel, dtype=dtype, device="cpu" + ) + self._magi_giant_buffers.append(giant_shared_tensor) + + pin_memory_in_place(giant_shared_tensor) + + offset = 0 + for name, original_tensor in param_list: + numel = original_tensor.numel() + shared_param = giant_shared_tensor[offset : offset + numel].view(original_tensor.shape) + + if original_tensor.requires_grad: + shared_param.requires_grad_(True) + + shared_state_dict[name] = shared_param + offset += numel + + dist.barrier() + if local_rank == 0: + if os.path.exists(shared_bin_path): + os.remove(shared_bin_path) + + self.load_state_dict(shared_state_dict, assign=True) + + else: + + def _pinner(t): + return t.pin_memory() + + _orig_apply(self, _pinner) + + self._magi_offloaded_once = True + return self + + cls._apply = _cpu_apply + + old_init = cls.__init__ + + def __init__(self: _W, *args, **kwargs): + old_init(self, *args, **kwargs) + compile_config = get_compile_config() + if config_patch is not None: + compile_config = config_patch(compile_config) + # deepcopy the compile config to avoid modifying the original compile config + self.compile_config = copy.deepcopy(compile_config) + + enable_compile = enable_if is None or enable_if() + self.enable_compile = self.compile_config.compile_mode != CompileMode.NONE and enable_compile + if not self.enable_compile: + return + + compilation_counter.num_models_seen += 1 + self.compile_config.model_idx = compilation_counter.num_models_seen + self.compile_config.model_tag = model_tag if model_tag is not None else self.__class__.__name__ + MagiCompilerBase.__init__(self, compile_config=self.compile_config) + + cls.__init__ = __init__ + + old_call = cls.__call__ + + def __call__(self: _W, *args, **kwargs): + ### Step1: Run compiled module directly if disable compile or captured before ### + if self.compile_config.offload_config.model_cpu_offload and self.compiled_code is None: + args = offload(args) + kwargs = offload(kwargs) + + if not self.enable_compile or torch.compiler.is_compiling(): + # Skip compiling the model if inside the compilation process. + return old_call(self, *args, **kwargs) + + if self.compiled_code is not None: + # Run the compiled function if compiled code is available. + with self.dispatch_to_compiled_fwd(mode="jit"): + return old_call(self, *args, **kwargs) + + if envs.MAGI_AOT_COMPILE: + # Try load AOT artifacts from cache and run directly. + self.aot_compiled_fn = self.try_load_aot_compile_artifacts() + if self.aot_compiled_fn is not None: + with self.dispatch_to_compiled_fwd(mode="aot"): + return old_call(self, *args, **kwargs) + + ### Step2: Mark dynamic shapes for the first compilation ### + bound_args = inspect.signature(self.__class__.forward).bind(self, *args, **kwargs) + bound_args.apply_defaults() + for k, dims in dynamic_arg_dims.items(): + arg = bound_args.arguments.get(k) + if arg is None: + continue + dims = [dims] if isinstance(dims, int) else dims + assert isinstance(arg, torch.Tensor), f"Unsupported dynamic dim {dims} for argument {k} with type {type(arg)}." + dims = [arg.ndim + dim if dim < 0 else dim for dim in dims] + torch._dynamo.mark_dynamic(arg, dims) + + ### Step3: Start compiling the model ### + magi_logger.info(f"Start compiling function {self.original_code_object}") + + CompileMonitor().start( + self.compile_config.compile_mode == CompileMode.MAGI_COMPILE, self.compile_config.debug_dump_path() + ) + # Dynamo reuse the compilation across instances, but we need to make sure the compiled code is not reused. + torch._dynamo.eval_frame.remove_from_cache(self.original_code_object) + + with ( + _hijack_inline_call_to_collect_traced_files(self), + patch.object(torch.compiler.config, "dynamic_sources", self.compile_config.dynamic_sources), + patch.object(torch._dynamo.config, "enable_cpp_symbolic_shape_guards", False), + # 允许 mark_dynamic 在 module 属性链上的 tensor 生效 + # (默认 True 会强制 module property tensor 为 static shape，忽略 mark_dynamic) + patch.object(torch._dynamo.config, "force_nn_module_property_static_shapes", False), + patch.dict( + os.environ, {"TORCHINDUCTOR_CACHE_DIR": (self.compile_config.cache_dump_path() / "inductor_cache").as_posix()} + ), + ): + if envs.MAGI_AOT_COMPILE: + self.aot_compiled_fn = self.aot_compile(*args, **kwargs) + self.aot_compiled_fn.save_compiled_function(self.aot_compilation_path) + with self.dispatch_to_compiled_fwd(mode="aot"): + output = old_call(self, *args, **kwargs) + else: + with patch.object(self, "forward", self.jit_compile): + output = old_call(self, *args, **kwargs) + + return output + + # 使用 @torch.compiler.disable 和 _isolated_dynamo_config 包裹整个 __call__ + # 确保 magi compile 在外部嵌套 torch.compile 时也能独立工作不受影响 + isolated_call = _isolated_dynamo_config()(__call__) + cls.__call__ = torch.compiler.disable(isolated_call) + return cls + + +# Collect all relevant files traced by Dynamo, re-compile the model when any of these files change. +# 1. the file containing the top-level forward function +# 2. hijack function to know all the functions called during Dynamo tracing, every time Dynamo sees a function call, it will inline +# the function by calling InliningInstructionTranslator.inline_call_ +def _hijack_inline_call_to_collect_traced_files(owner: _W): + owner.compile_config.traced_files.add(owner.original_code_object.co_filename) + inline_call = InliningInstructionTranslator.inline_call_ + + def patched_inline_call(self_): + code = self_.f_code + owner.compile_config.traced_files.add(code.co_filename) + return inline_call(self_) + + return patch.object(InliningInstructionTranslator, "inline_call_", patched_inline_call) + + +def _infer_dynamic_arg_dims(cls: _T) -> dict[str, int | list[int]]: + sig = inspect.signature(cls.forward) + inferred_dynamic_arg_dims = {} + for k, v in sig.parameters.items(): + if v.annotation in [torch.Tensor, torch.Tensor | None]: + inferred_dynamic_arg_dims[k] = 0 + + magi_logger.info(f"Inferred dynamic dimensions for forward method of {cls}: {list(inferred_dynamic_arg_dims.keys())}") + return inferred_dynamic_arg_dims + + +def _get_num_outputs_from_return_annotation(fn: Callable) -> int: + """ + Get the number of outputs from the function's return type annotation. + + Returns: + - 1 if the return type is a single Tensor + - N if the return type is tuple[Tensor, Tensor, ...] with N elements + - 1 if no annotation or unrecognized annotation (default to single output) + """ + sig = inspect.signature(fn) + return_annotation = sig.return_annotation + + if return_annotation is inspect.Parameter.empty: + return 1 + + # Check if it's a tuple type (e.g., tuple[Tensor, Tensor]) + origin = get_origin(return_annotation) + if origin is tuple: + args = get_args(return_annotation) + # Filter out ellipsis (for variable-length tuples like tuple[Tensor, ...]) + if args and args[-1] is not ...: + return len(args) + return 1 + + return 1 + + +def _generate_op_name(fn: Callable) -> str: + """ + Generate a unique operator name from function's name and source file. + + The generated name follows the format: namespace::op_name + - namespace: derived from the source file path (module-like structure) + - op_name: the function name + + Example: + Function `_my_custom_op` in file `/path/to/my_module.py` + -> "my_module::_my_custom_op" + """ + import re + from pathlib import Path + + func_name = fn.__name__ + + # Get the source file path + try: + source_file = inspect.getfile(fn) + # Extract the file stem (without extension) as namespace + namespace = Path(source_file).stem + # Clean up namespace: replace invalid characters with underscores + namespace = re.sub(r"[^a-zA-Z0-9_]", "_", namespace) + except (TypeError, OSError): + # If we can't get the source file, use a default namespace + namespace = "magi_custom" + + return f"{namespace}::{func_name}" + + +def _create_identity_meta_fn(fn: Callable) -> Callable: + """ + Create a default identity meta function for the given function. + + This identity meta function assumes that: + - The number of outputs is determined by the function's return type annotation + - Each output's metadata (shape, dtype, device) matches the corresponding input tensor + + For example, if a function has signature: + def my_op(a: Tensor, b: Tensor, scale: float) -> tuple[Tensor, Tensor]: + The identity meta function will return: + (torch.empty_like(a), torch.empty_like(b)) + """ + num_outputs = _get_num_outputs_from_return_annotation(fn) + sig = inspect.signature(fn) + # Get parameter names, excluding 'self' if present + param_names = [name for name in sig.parameters.keys() if name != "self"] + + def identity_meta_fn(*args, **kwargs): + # Bind arguments to get a mapping of param_name -> value + bound = sig.bind(*args, **kwargs) + bound.apply_defaults() + + # Collect the first `num_outputs` tensor arguments + tensor_args = [] + for name in param_names: + arg = bound.arguments.get(name) + if isinstance(arg, torch.Tensor): + tensor_args.append(arg) + if len(tensor_args) >= num_outputs: + break + + if len(tensor_args) < num_outputs: + raise ValueError( + f"identity_meta_fn requires at least {num_outputs} tensor inputs to match " + f"{num_outputs} outputs, but only found {len(tensor_args)} tensor inputs. " + f"Please provide a custom infer_output_meta_fn." + ) + + # Return outputs with same metadata as the first N inputs + if num_outputs == 1: + return torch.empty_like(tensor_args[0]) + return tuple(torch.empty_like(t) for t in tensor_args[:num_outputs]) + + return identity_meta_fn + + +def _create_meta_fn_from_param_names(fn: Callable, param_names: list[str]) -> Callable: + """ + Create a meta function that returns torch.empty_like() for each specified parameter. + + This is useful when output tensors have the same shape/dtype/device as specific input + parameters, but not necessarily in positional order. + + Example: + param_names = ["weight", "bias"] + def my_op(grad: Tensor, weight: Tensor, bias: Tensor) -> tuple[Tensor, Tensor]: + ... + + Generated meta function returns: + (torch.empty_like(weight), torch.empty_like(bias)) + """ + sig = inspect.signature(fn) + + def meta_fn(*args, **kwargs): + # Bind arguments to get a mapping of param_name -> value + bound = sig.bind(*args, **kwargs) + bound.apply_defaults() + + # Collect tensors for each specified parameter name + tensor_outputs = [] + for name in param_names: + if name not in bound.arguments: + raise ValueError( + f"Parameter '{name}' not found in function signature. " + f"Available parameters: {list(bound.arguments.keys())}" + ) + arg = bound.arguments[name] + if not isinstance(arg, torch.Tensor): + raise ValueError( + f"Parameter '{name}' is not a Tensor (got {type(arg).__name__}). " + f"infer_output_meta_fn list should only contain tensor parameter names." + ) + tensor_outputs.append(torch.empty_like(arg)) + + # Return single tensor or tuple based on number of outputs + if len(tensor_outputs) == 1: + return tensor_outputs[0] + return tuple(tensor_outputs) + + return meta_fn + + +def magi_register_custom_op( + name: str | None = None, + mutates_args: tuple[str, ...] = (), + infer_output_meta_fn: Callable | list[str] | None = None, + setup_context_fn: Callable | None = None, + backward_fn: Callable | None = None, +): + """ + A unified decorator to register a custom operator with PyTorch's library. + + This decorator combines the functionality of: + - @torch.library.custom_op + - @torch.library.register_fake + - fn.register_autograd + + Arguments: + name: The fully qualified name of the operator (e.g., "namespace::op_name"). + If None, auto-generated from the function name and source file. + mutates_args: Tuple of argument names that are mutated by the operator. + infer_output_meta_fn: Specifies output tensor metadata (shape, dtype, device) for tracing. + - None (default): Assumes each output has the same metadata as the corresponding + input tensor (1st output matches 1st tensor input, 2nd matches 2nd, etc.). + - list[str]: Parameter names whose metadata to use for outputs. + E.g., ["weight", "bias"] means output[0] has same shape as `weight`, + output[1] has same shape as `bias`. + - Callable: Custom function with same signature as the op, returns torch.empty_like() + tensors matching the expected output shapes. + setup_context_fn: Function to save tensors/values for backward. + Signature: setup_context_fn(ctx, inputs, output) + backward_fn: Function to compute gradients. + Signature: backward_fn(ctx, *grad_outputs) -> tuple of gradients + + Returns: + The registered custom operator function. + + Examples: + 1. Basic usage (forward only, auto-generated name and meta function): + + >>> @magi_register_custom_op() + ... def my_relu(x: torch.Tensor) -> torch.Tensor: + ... return torch.maximum(x, torch.zeros_like(x)) + + 2. Multiple outputs with explicit output metadata via parameter names: + + >>> @magi_register_custom_op( + ... infer_output_meta_fn=["weight", "bias"], # output shapes match weight and bias + ... ) + ... def compute_gradients( + ... grad_output: torch.Tensor, + ... weight: torch.Tensor, + ... bias: torch.Tensor, + ... ) -> tuple[torch.Tensor, torch.Tensor]: + ... grad_weight = grad_output.sum(dim=0).view_as(weight) + ... grad_bias = grad_output.sum(dim=0).view_as(bias) + ... return grad_weight, grad_bias + + 3. Full custom op with autograd support: + + >>> def _square_meta(x: torch.Tensor) -> torch.Tensor: + ... return torch.empty_like(x) + ... + >>> def _square_setup_context(ctx, inputs, output): + ... (x,) = inputs + ... ctx.save_for_backward(x) + ... + >>> def _square_backward(ctx, grad_output): + ... (x,) = ctx.saved_tensors + ... return grad_output * 2 * x + ... + >>> @magi_register_custom_op( + ... name="my_ops::square", + ... infer_output_meta_fn=_square_meta, + ... setup_context_fn=_square_setup_context, + ... backward_fn=_square_backward, + ... ) + ... def square(x: torch.Tensor) -> torch.Tensor: + ... return x * x + """ + + def decorator(fn: Callable) -> Callable: + # Auto-generate name if not provided + op_name = name if name is not None else _generate_op_name(fn) + + # Step 1: Register the custom op with torch.library.custom_op + registered_op = torch.library.custom_op(op_name, mutates_args=mutates_args)(fn) + + # Step 2: Register the output meta inference function + # Determine meta_fn based on the type of infer_output_meta_fn + if infer_output_meta_fn is None: + meta_fn = _create_identity_meta_fn(fn) + elif isinstance(infer_output_meta_fn, list): + meta_fn = _create_meta_fn_from_param_names(fn, infer_output_meta_fn) + else: + meta_fn = infer_output_meta_fn + torch.library.register_fake(op_name)(meta_fn) + + # Step 3: Register autograd if backward_fn is provided + if backward_fn is not None: + registered_op.register_autograd(backward_fn, setup_context=setup_context_fn) + + return registered_op + + return decorator diff --git a/pkgs/MagiCompiler/magi_compiler/compile_artifacts.py b/pkgs/MagiCompiler/magi_compiler/compile_artifacts.py new file mode 100644 index 0000000000000000000000000000000000000000..ace4b1b93b7f90cfe31c8933f5d848a3d728f559 --- /dev/null +++ b/pkgs/MagiCompiler/magi_compiler/compile_artifacts.py @@ -0,0 +1,125 @@ +# 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. + +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +import inspect +import pickle +from unittest.mock import patch + +import torch +from torch.utils._pytree import tree_map_only + +try: + from torch._dynamo.aot_compile import SerializableCallable +except ImportError: + SerializableCallable = object + +assert isinstance(SerializableCallable, type) + + +class MagiSerializableFunction(SerializableCallable): + """ + A wrapper around a compiled function by vllm. It will forward the tensor + inputs to the compiled function and return the result. + It also implements a serialization interface to support PyTorch's precompile + with custom backend, so that we can save and load the compiled function on + disk. There's no need to wrap around the compiled function if we don't want + to serialize them in particular cases. + Right now serialization for the custom backend is done via + serializing the Dynamo fx graph plus example inputs. + """ + + def __init__(self, graph_module, example_inputs, model_tag, optimized_call): + assert isinstance(graph_module, torch.fx.GraphModule) + self.graph_module = graph_module + self.example_inputs = example_inputs + self.model_tag = model_tag + self.optimized_call = optimized_call + self.shape_env = None + sym_input = next((i for i in self.example_inputs if isinstance(i, torch.SymInt)), None) + if sym_input is not None: + self.shape_env = sym_input.node.shape_env + + def __call__(self, *args, **kwargs): + return self.optimized_call(*args, **kwargs) + + @classmethod + def serialize_compile_artifacts(cls, compiled_fn: "MagiSerializableFunction") -> bytes: + import sympy + from torch._subclasses import FakeTensorMode + from torch.fx._graph_pickler import GraphPickler, Options + + state = compiled_fn.__dict__.copy() + state.pop("optimized_call") + state.pop("shape_env") + for node in state["graph_module"].graph.nodes: + node.meta.pop("source_fn_stack", None) + node.meta.pop("nn_module_stack", None) + + graph_reducer_override = GraphPickler.reducer_override + + def _graph_reducer_override(self, obj): + if inspect.isclass(obj) and issubclass(obj, sympy.Function) and hasattr(obj, "_torch_unpickler"): + return obj._torch_unpickler, (obj._torch_handler_name,) + if isinstance(obj, FakeTensorMode): + return type(None), () + return graph_reducer_override(self, obj) + + # Mask off tensor inputs since they are large and not needed. + state["example_inputs"] = tree_map_only(torch.Tensor, lambda _: None, state["example_inputs"]) + with patch.object(GraphPickler, "reducer_override", _graph_reducer_override): + state["graph_module"] = GraphPickler.dumps(state["graph_module"], Options(ops_filter=None)) + state["example_inputs"] = GraphPickler.dumps(state["example_inputs"]) + return pickle.dumps(state) + + @classmethod + def deserialize_compile_artifacts(cls, data: bytes) -> "MagiSerializableFunction": + from torch._guards import TracingContext, tracing + from torch._subclasses import FakeTensorMode + from torch.fx._graph_pickler import GraphPickler + from torch.fx.experimental.symbolic_shapes import ShapeEnv + + from .config import get_compile_config + from .magi_backend import MagiBackend + + state = pickle.loads(data) + fake_mode = FakeTensorMode(shape_env=ShapeEnv()) + state["graph_module"] = GraphPickler.loads(state["graph_module"], fake_mode) + state["example_inputs"] = GraphPickler.loads(state["example_inputs"], fake_mode) + magi_backend = MagiBackend(get_compile_config(), state["model_tag"]) + + def optimized_call(*example_inputs): + """ + On the first run of the optimized call, we rerun the compiler + backend which should result in a cache hit. After the backend + call returns, we just do a one-time replacement of the optimized + call with the compiled function, so that subsequent calls are on + the AOT compiled path. + """ + compile_inputs = [inp or example_inputs[i] for i, inp in enumerate(fn.example_inputs)] + with tracing(TracingContext(fake_mode)): + fn.optimized_call = magi_backend(state["graph_module"], compile_inputs).optimized_call + return fn.optimized_call(*example_inputs) + + fn = cls(**state, optimized_call=optimized_call) + return fn + + @property + def co_name(self): + """ + Used for depyf debugging. + """ + return "MagiSerializableFunction" diff --git a/pkgs/MagiCompiler/magi_compiler/config.py b/pkgs/MagiCompiler/magi_compiler/config.py new file mode 100644 index 0000000000000000000000000000000000000000..2371c4afd5b1da92f496e52d24c8bc99ebed93b0 --- /dev/null +++ b/pkgs/MagiCompiler/magi_compiler/config.py @@ -0,0 +1,282 @@ +# 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. + +import json +import os +from enum import Enum, unique +from pathlib import Path +from typing import Any, Literal + +import torch +from pydantic import BaseModel, Field +from pydantic_settings import BaseSettings, SettingsConfigDict + +from .utils import OrderedSet, compute_hash, magi_logger + + +@unique +class CompileMode(Enum): + """ + The compilation approach used for torch.compile-based compilation of the model. + + NONE: No torch.compile compilation is applied, model runs in fully eager pytorch mode. The model runs as-is. + TORCH_COMPILE: The standard `torch.compile` compilation pipeline. + MAGI_COMPILE: Custom Inductor-based backend with caching, piecewise compilation, shape specialization, and custom passes. + """ + + NONE = 'NONE' + TORCH_COMPILE = 'TORCH_COMPILE' + MAGI_COMPILE = 'MAGI_COMPILE' + + +@unique +class CudaGraphMode(Enum): + """ + Constants for the cudagraph mode in CompileConfig. + Different from the CUDAGraphMode for llm, PIECEWISE and FULL modes are enough for diffusion models. + + NONE: No cudagraph is used. + PIECEWISE: Cudagraph is used for piecewise compilation. + FULL: Cudagraph is used for full compilation. + """ + + NONE = 'NONE' + PIECEWISE = 'PIECEWISE' + FULL = 'FULL' + + +class PassConfig(BaseModel): + """Configuration for custom Inductor passes""" + + enable_fusion: bool = Field(False, description="Whether to enable the custom fusion (RMSNorm/SiluMul+quant) pass.") + enable_attn_fusion: bool = Field(False, description="Whether to enable the custom attention+quant fusion pass.") + enable_noop: bool = Field(False, description="Whether to enable the custom no-op elimination pass.") + enable_sequence_parallelism: bool = Field(False, description="Whether to enable sequence parallelism.") + enable_async_tp: bool = Field(False, description="Whether to enable async TP.") + enable_fi_allreduce_fusion: bool = Field(False, description="Whether to enable flashinfer allreduce fusion.") + enable_sage_attn: bool = Field(False, description="Whether to replace flash attention with sage attention.") + fi_allreduce_fusion_max_token_num: int = Field( + 16384, description="Max number of tokens to used in flashinfer allreduce fusion." + ) + + def __post_init__(self) -> None: + if not self.enable_noop: + if self.enable_fusion: + magi_logger.warning( + "Fusion enabled but reshape elimination disabled. " "RMSNorm/SiluMul + quant (fp8) fusion might not work" + ) + if self.enable_attn_fusion: + magi_logger.warning( + "Fusion enabled but reshape elimination disabled. " "Attention + quant (fp8) fusion might not work" + ) + + @property + def hash(self) -> str: + return compute_hash(self.model_dump(mode="json")) + + # Compatible with torch pass + def uuid(self) -> str: + return self.hash + + +@unique +class RecomputePolicy(Enum): + """ + Defines the strategy for activation recomputation (rematerialization) to trade off + memory usage against computational overhead. + + HANDCRAFT: + A manual strategy where the user controls the trade-off via a `memory_budget` + parameter. This parameter acts as a threshold (0.0 to 1.0) determining the + target percentage of activations to save. + + HEURISTIC: + A rule-based strategy that selectively saves activations from compute-bound + operators (e.g., MatMul, Attention). Conversely, outputs from memory-bound + or element-wise operators are prioritized for recomputation to save memory. + + AUTOSEARCH: + An automated strategy that searches for the optimal set of saved tensors based + on available device memory. It prioritizes saving tensors with high computational + cost relative to their memory footprint. + + .. note:: + Currently, a `repeat_number` argument is required to stabilize the profiling/search + phase. This requirement is temporary and will be deprecated once full-graph + capture is natively supported. + """ + + HANDCRAFT = "HANDCRAFT" + HEURISTIC = "HEURISTIC" + AUTOSEARCH = "AUTOSEARCH" + + +class RecomputeConfig(BaseModel): + recompute_policy: RecomputePolicy = Field(RecomputePolicy.HEURISTIC, description="Recompute policy.") + memory_budget: float = Field(0.5, description="Activation memory budget for recomputation, only used for handcraft.") + repeat_number: int = Field(default=1, description="Repeat number for recomputation, only used for autosearch.") + + +@unique +class OffloadPolicy(Enum): + """ + The policy for offloading the model to CPU. + + BASE: + The base policy for offloading the model to CPU. + Offload all the submodules to CPU. + COST_EFFECTIVE: + The cost effective policy for offloading the model to CPU. + Offload the submodules to CPU based on the cost effective policy. + HEURISTIC: + The heuristic policy for offloading the model to CPU. + Offload the submodules to CPU based on the heuristic policy. + """ + + BASE = "BASE" + COST_EFFECTIVE = "COST_EFFECTIVE" + HEURISTIC = "HEURISTIC" + + +class OffloadConfig(BaseModel): + model_cpu_offload: bool = Field(False, description="Whether to offload the model to CPU.") + gpu_resident_weight_ratio: float = Field( + 0.3, description="The ratio of GPU memory to keep when offloading the model to CPU." + ) + offload_policy: OffloadPolicy = Field( + OffloadPolicy.COST_EFFECTIVE, description="The policy for offloading the model to CPU." + ) + bandwidth_safety_factor: float = Field(0.9, description="The safety factor for the H2D bandwidth.") + + +class CompileConfig(BaseSettings): + model_config = SettingsConfigDict(cli_parse_args=True, cli_ignore_unknown_args=True, cli_implicit_flags=True) + + # Basic configs + backend: Literal["inductor", "eager"] = Field("inductor", description="Compilation backend.") + compile_mode: CompileMode = Field(CompileMode.MAGI_COMPILE, description="Compilation mode.") + cache_root_dir: str = Field( + default=os.path.expanduser("~/.cache/magi_compiler"), description="Directory to cache the compiled model." + ) + dynamic_sources: str = Field( + default=os.environ.get("TORCH_COMPILE_DYNAMIC_SOURCES", ""), + description="Comma delimited list of sources that should be marked as dynamic.", + ) + + # CPU Offload + offload_config: OffloadConfig = Field(OffloadConfig(), description="Offload configuration.") + + # Inductor configs + # TODO(hongyu): Add unittest for compile_sizes + compile_sizes: list[int] = Field(default_factory=list, description="Sizes to compile the model for.") + use_inductor_graph_partition: bool = Field( + False, description="Whether to use inductor graph partition. Not fully supported yet." + ) + # TODO(hongyu): Find a better way to specify the splitting ops. + splitting_ops: list[str] = Field( + default_factory=lambda: [ + "athena::flash_attn_func", + "athena::flex_flash_attn_func", + "athena::sage_attn_func", + "athena::flash_attn_with_cp", + "athena::flex_flash_attn_with_cp", + ], + description="Operators to split the graph into piecewise graphs.", + ) + + # Pass configs + pass_config: PassConfig = Field(PassConfig(), description="Pass configuration.") + + # Recompute configs + recompute_config: RecomputeConfig = Field(RecomputeConfig(), description="Recompute configuration.") + + # Cudagraph configs + cudagraph_mode: CudaGraphMode = Field(CudaGraphMode.NONE, description="Cudagraph mode.") + cudagraph_copy_inputs: bool = Field(True, description="Whether to copy inputs for cudagraph.") + + # Runtime configs, maybe changed at runtime + model_idx: int = Field(0, description="Index of the model.") + model_tag: str | None = Field( + default=None, description="Tag in cache path: model_{idx}_{model_tag}_rank_{rank}. Class name if unset." + ) + inductor_compile_config: dict[str, Any] = Field(default_factory=dict, description="Inductor compilation configuration.") + traced_files: OrderedSet[str] = Field(default_factory=OrderedSet, description="Files traced by Dynamo.") + + def _model_rank_dir_name(self) -> str: + """Directory name for this model instance: model_{idx}[_{model_tag}]_rank_{rank}.""" + rank = torch.distributed.get_rank() if torch.distributed.is_initialized() else 0 + if self.model_tag: + return f"model_{self.model_idx}_{self.model_tag}_rank_{rank}" + return f"model_{self.model_idx}_rank_{rank}" + + def debug_dump_path(self) -> Path: + return Path(self.cache_root_dir) / "magi_depyf" / self._model_rank_dir_name() + + def cache_dump_path(self) -> Path: + return Path(self.cache_root_dir) / "torch_compile_cache" / self._model_rank_dir_name() + + @property + def hash(self) -> str: + # Create a copy of the config data for serialization + data = self.model_dump(mode="json", exclude={"inductor_compile_config"}) + + # Handle inductor_compile_config separately to serialize objects with uuid() method + # This is a workaround to support serialization of PostGradPassManager in Pydantic models. + if self.inductor_compile_config: + serialized_inductor_config = {} + for key, value in self.inductor_compile_config.items(): + # If the value has a uuid() method (like PostGradPassManager), use it + if hasattr(value, "uuid") and callable(getattr(value, "uuid", None)): + try: + serialized_inductor_config[key] = value.uuid() + except (AttributeError, RuntimeError): + # Fallback to string representation if uuid() fails + serialized_inductor_config[key] = str(value) + else: + # For other types, try to serialize normally + try: + # Try to serialize as JSON-serializable + json.dumps(value) + serialized_inductor_config[key] = value + except (TypeError, ValueError): + # If not JSON-serializable, use string representation + serialized_inductor_config[key] = str(value) + data["inductor_compile_config"] = serialized_inductor_config + + return compute_hash(data) + + def __str__(self, indent: int = 4): + data = self.model_dump(mode="json") + formatted = json.dumps(data, indent=indent, ensure_ascii=False, sort_keys=False) + + # add configuration class name as title + class_name = self.__class__.__name__ + return f"{class_name}:\n{formatted}".replace('"', "") + + def __repr__(self, indent: int = 4): + return self.__str__(indent=indent) + + +_GLOBAL_COMPILE_CONFIG = None + + +def get_compile_config() -> CompileConfig: + global _GLOBAL_COMPILE_CONFIG + if _GLOBAL_COMPILE_CONFIG is None: + _GLOBAL_COMPILE_CONFIG = CompileConfig() + if not torch.distributed.is_initialized() or torch.distributed.get_rank() == 0: + # 仅在首次初始化时打印一次编译配置，默认 WARNING 级别不会输出 + magi_logger.info("compile config: %s", _GLOBAL_COMPILE_CONFIG) + assert _GLOBAL_COMPILE_CONFIG is not None, "compile config is not initialized" + return _GLOBAL_COMPILE_CONFIG diff --git a/pkgs/MagiCompiler/magi_compiler/cuda/cudart.py b/pkgs/MagiCompiler/magi_compiler/cuda/cudart.py new file mode 100644 index 0000000000000000000000000000000000000000..46ff1a77463edc3c1d0339ffe9d0135f8e6f9cd9 --- /dev/null +++ b/pkgs/MagiCompiler/magi_compiler/cuda/cudart.py @@ -0,0 +1,60 @@ +# Copyright (c) 2026 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. + + +import ctypes +import os + +import torch + +_cudart = None + + +def init_cudart(): + global _cudart + if _cudart is not None: + return _cudart + candidates = ["libcudart.so", "libcudart.so.11.0", "libcudart.so.12.0, libcudart.so.13"] + try: + cuda_path = os.path.dirname(torch.utils.cpp_extension._find_cuda_home()) + candidates.append(os.path.join(cuda_path, "lib64", "libcudart.so")) + except: + pass + for lib in candidates: + try: + _cudart = ctypes.CDLL(lib) + return _cudart + except OSError: + continue + return None + + +def pin_memory_in_place(tensor: torch.Tensor): + """ + Pin memory in-place using cudaHostRegister. + """ + if tensor.is_cuda: + return tensor + cudart = init_cudart() + if cudart is None: + return tensor + + ptr = tensor.data_ptr() + size = tensor.numel() * tensor.element_size() + res = cudart.cudaHostRegister(ctypes.c_void_p(ptr), ctypes.c_size_t(size), 0) + + if res == 0: + return tensor + else: + raise RuntimeError(f"cudaHostRegister failed with error code {res}") diff --git a/pkgs/MagiCompiler/magi_compiler/cuda_graph_mgr.py b/pkgs/MagiCompiler/magi_compiler/cuda_graph_mgr.py new file mode 100644 index 0000000000000000000000000000000000000000..724efbd03b6b4d0944891dbf0ac94f1888244461 --- /dev/null +++ b/pkgs/MagiCompiler/magi_compiler/cuda_graph_mgr.py @@ -0,0 +1,931 @@ +# 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. + +from dataclasses import dataclass, fields, is_dataclass +from functools import wraps +from typing import Any, Callable, Dict, List, Optional, Tuple, Union + +import torch + +from .utils import magi_logger, nvtx + + +class InplaceSubstituteFakeClass: + """ + The class which inherits from this class will not be replaced with a new instance, + but the attributes will be updated in-place. + For example, InferenceParams. + """ + + pass + + +@dataclass +class FakeTensor: + shape: Tuple[int, ...] = None + dtype: str = None + device: str = None + + +@dataclass +class HashableDataclass: + _cached_hash: Optional[int] = None + + @nvtx.instrument_nvtx + def _get_hashable_fields(self) -> Tuple[Any, ...]: + hashable_values = [] + for f in fields(self): + if f.name == "_cached_hash": + continue + value = getattr(self, f.name) + if value is None: + continue + if isinstance(value, HashableDataclass): + hashable_values.append(value._get_cached_hash()) + elif isinstance(value, tuple): + tuple_vals = [] + for item in value: + if isinstance(item, (HashableDataclass, str, int, float, bool)): + if isinstance(item, HashableDataclass): + tuple_vals.append(item._get_cached_hash()) + else: + tuple_vals.append(item) + if tuple_vals: + hashable_values.append(tuple(tuple_vals)) + elif isinstance(value, (str, int, float, bool)): + hashable_values.append(value) + return tuple(hashable_values) + + @nvtx.instrument_nvtx + def _compute_hash(self) -> int: + """Computes a hash value based on the dataclass's hashable fields.""" + hashable_fields = self._get_hashable_fields() + return hash(hashable_fields) % (1 << 64) # 限制为 64 位 + + @nvtx.instrument_nvtx + def _get_cached_hash(self) -> int: + if self._cached_hash is None: + self._cached_hash = self._compute_hash() + return self._cached_hash + + @nvtx.instrument_nvtx + def __hash__(self) -> int: + return self._get_cached_hash() + + @nvtx.instrument_nvtx + def __eq__(self, other: Any) -> bool: + if not isinstance(other, self.__class__): + return False + if self._get_cached_hash() != other._get_cached_hash(): + return False + return True + + +@dataclass(unsafe_hash=True) +class LiteralsInfo(HashableDataclass): + literals: Tuple[Any, ...] = tuple() + + +@dataclass(unsafe_hash=True) +class TensorStaticInfo(HashableDataclass): + name: str = "" + shapes: Tuple[int, ...] = tuple() + dtype: str = "" + + +@dataclass(unsafe_hash=True) +class TensorDynamicInfo(HashableDataclass): + name: str = "" + shapes: Tuple[int, ...] = tuple() + + +@dataclass(unsafe_hash=True) +class StaticSignature(HashableDataclass): + func_name: str = "" + tensor_static_infos: Tuple[TensorStaticInfo, ...] = tuple() + + +@dataclass(unsafe_hash=True) +class DynamicSignature(HashableDataclass): + tensor_dynamic_infos: Tuple[TensorDynamicInfo, ...] = tuple() + literals_info: LiteralsInfo = None + + +@dataclass +class GraphEntry: + graph: Optional[torch.cuda.CUDAGraph] = None + inconsistent: bool = False + invalid: bool = False + + +@dataclass +class OutputTemplateEntry: + graph_entry_dict: Dict[int, GraphEntry] = None # key = layer_number + output_template: Any = None # 用于存储输出对象literals的结构模板 + + +@dataclass +class StaticTensorEntry: + input_tensors: Optional[List[torch.Tensor]] = None + output_tensors: Optional[List[torch.Tensor]] = None + template_entry_dict: Dict[DynamicSignature, OutputTemplateEntry] = None + + +class ArgsUtils: + @staticmethod + @nvtx.instrument_nvtx + def generate_both_signatures_from_tensors( + func_name: str, tensors: List[torch.Tensor], names: List[str], literals: List[Any] + ) -> Tuple[StaticSignature, DynamicSignature]: + num_tensors = len(tensors) + tensor_static_infos = [TensorStaticInfo() for _ in range(num_tensors)] + tensor_dynamic_infos = [TensorDynamicInfo() for _ in range(num_tensors)] + + # Local references for performance + TensorStaticInfo_setattr = TensorStaticInfo.__setattr__ + TensorDynamicInfo_setattr = TensorDynamicInfo.__setattr__ + _tuple = tuple + + for i in range(num_tensors): + t = tensors[i] + t_dim = t.dim() + t_shape = t.shape + t_dtype_str = str(t.dtype) + # Last dimension is static, others are dynamic (except for 1D tensor) + static_shapes = ( + _tuple(-1 if idx != t_dim - 1 else dim_size for idx, dim_size in enumerate(t_shape)) if t_dim > 1 else (-1,) + ) + static_info = tensor_static_infos[i] + TensorStaticInfo_setattr(static_info, "shapes", static_shapes) + TensorStaticInfo_setattr(static_info, "dtype", t_dtype_str) + + dynamic_shapes = static_shapes = ( + _tuple(-1 if idx == t_dim - 1 else dim_size for idx, dim_size in enumerate(t_shape)) + if t_dim > 1 + else _tuple(t_shape) + ) + dynamic_info = tensor_dynamic_infos[i] + TensorDynamicInfo_setattr(dynamic_info, "shapes", dynamic_shapes) + + literals_info = LiteralsInfo(literals=_tuple(literals)) + static_sig = StaticSignature(func_name=func_name, tensor_static_infos=_tuple(tensor_static_infos)) + dynamic_sig = DynamicSignature(tensor_dynamic_infos=_tuple(tensor_dynamic_infos), literals_info=literals_info) + return static_sig, dynamic_sig + + @staticmethod + @nvtx.instrument_nvtx + def replace_sliced_with_static(obj: Any, static_tensors: List[torch.Tensor]) -> Any: + tensor_idx = 0 + + def recursive_replace(o: Any) -> Any: + nonlocal tensor_idx + if isinstance(o, torch.Tensor) and not isinstance(o, torch.nn.Parameter): + # Copy data to the corresponding static tensor slice + static_tensor = static_tensors[tensor_idx] + slices = [slice(None)] * static_tensor.ndim + for i in range(min(o.ndim, static_tensor.ndim)): + slices[i] = slice(0, o.shape[i]) + # Only copy if the data_ptrs are different + if not o.data_ptr() == static_tensor[tuple(slices)].data_ptr(): + static_tensor[tuple(slices)].copy_(o) + tensor_idx += 1 + return static_tensor[tuple(slices)] + + elif isinstance(o, dict): + return {k: recursive_replace(v) for k, v in o.items()} + elif isinstance(o, (list, tuple)): + return type(o)(recursive_replace(item) for item in o) + elif is_dataclass(o): + field_values = {f.name: recursive_replace(getattr(o, f.name)) for f in fields(o)} + return type(o)(**field_values) + elif issubclass(o.__class__, InplaceSubstituteFakeClass): + # Do not create a new instance, but modify attributes in place (to keep original initialization logic) + for k, v in o.__dict__.items(): + if not callable(v): + o.__dict__[k] = recursive_replace(v) + return o + elif o is None or isinstance(o, (int, float, str, bool)): + return o # Keep None and basic types + else: + return o + + return recursive_replace(obj) + + @staticmethod + @nvtx.instrument_nvtx + def replace_sliced_with_static_simple( + sliced_tensors: List[torch.Tensor], static_tensors: List[torch.Tensor] + ) -> List[torch.Tensor]: + for sliced_tensor, static_tensor in zip(sliced_tensors, static_tensors): + if not sliced_tensor.data_ptr() == static_tensor.data_ptr(): + slices = [slice(None)] * static_tensor.ndim + for i in range(sliced_tensor.ndim): + slices[i] = slice(0, sliced_tensor.shape[i]) + static_tensor[tuple(slices)].copy_(sliced_tensor) + + @staticmethod + @nvtx.instrument_nvtx + def replace_static_with_sliced(obj: Any, static_tensors: List[torch.Tensor]) -> Any: + tensor_idx = 0 + + def recursive_replace(o: Any) -> Any: + nonlocal tensor_idx + if (isinstance(o, torch.Tensor) and not isinstance(o, torch.nn.Parameter)) or isinstance(o, FakeTensor): + # Replace with the corresponding sliced tensor + static_tensor = static_tensors[tensor_idx] + shape_to_slice = o.shape + slices = [slice(0, dim_size) for dim_size in shape_to_slice] + result_tensor = static_tensor[tuple(slices)] + tensor_idx += 1 + return result_tensor + + elif isinstance(o, dict): + return {k: recursive_replace(v) for k, v in o.items()} + elif isinstance(o, (list, tuple)): + return type(o)(recursive_replace(item) for item in o) + elif is_dataclass(o): + field_values = {f.name: recursive_replace(getattr(o, f.name)) for f in fields(o)} + return type(o)(**field_values) + elif issubclass(o.__class__, InplaceSubstituteFakeClass): + # Do not create a new instance, but modify attributes in place (to keep original initialization logic) + for k, v in o.__dict__.items(): + if not callable(v): + o.__dict__[k] = recursive_replace(v) + return o + elif o is None or isinstance(o, (int, float, str, bool)): + return o # Keep None and basic types + else: + return o + + return recursive_replace(obj) + + @staticmethod + @nvtx.instrument_nvtx + def try_fx_extract_core( + obj: Any, extract_tensors: bool = True, extract_literals: bool = True, with_names: bool = False + ) -> Tuple[List[torch.Tensor], List[str], List[Any]]: + failed_tuple = None, None, None + tensors = [] + names = [] + literals = [] + + if not isinstance(obj, dict) or "args" not in obj or "kwargs" not in obj: + return failed_tuple + args, kwargs = obj["args"], obj["kwargs"] + if kwargs: + return failed_tuple + if not isinstance(args, (list, tuple)): + return failed_tuple + + for idx, item in enumerate(args): + if extract_tensors and isinstance(item, torch.Tensor) and not isinstance(item, torch.nn.Parameter): + tensors.append(item) + elif extract_literals and isinstance(item, (int, float, str, bool)): + literals.append(item) + + names = [""] * len(tensors) + return tensors, names, literals + + @staticmethod + @nvtx.instrument_nvtx + def recursive_extract_core( + obj: Any, extract_tensors: bool = True, extract_literals: bool = True, with_names: bool = False + ) -> Tuple[List[torch.Tensor], List[str], List[Any]]: + tensors = [] + names = [] + literals = [] + + def recursive_traverse(o: Any, prefix: str = ""): + # 1. Extract tensors (if enabled) + if extract_tensors and isinstance(o, torch.Tensor) and not isinstance(o, torch.nn.Parameter): + tensors.append(o) + names.append(prefix) if with_names else None + elif extract_literals and isinstance(o, (int, float, str, bool)): + literals.append(o) if extract_literals else None + elif isinstance(o, dict): + for k, v in o.items(): + new_prefix = f"{prefix}.{k}" if (with_names and extract_tensors) else prefix + recursive_traverse(v, new_prefix) + elif isinstance(o, (list, tuple)): + for idx, item in enumerate(o): + new_prefix = f"{prefix}[{idx}]" if (with_names and extract_tensors) else prefix + recursive_traverse(item, new_prefix) + elif is_dataclass(o): + for f in fields(o): + new_prefix = f"{prefix}.{f.name}" if (with_names and extract_tensors) else prefix + recursive_traverse(getattr(o, f.name), new_prefix) + elif issubclass(o.__class__, InplaceSubstituteFakeClass): + for k, v in o.__dict__.items(): + if not callable(v): + new_prefix = f"{prefix}.{k}" if (with_names and extract_tensors) else prefix + recursive_traverse(v, new_prefix) + elif o is None: + pass + else: + pass + + recursive_traverse(obj) + return tensors, names if with_names else [""] * len(tensors), literals if extract_literals else None + + @staticmethod + @nvtx.instrument_nvtx + def extract_output_template(obj: Any) -> Any: + def recursive_template(o: Any) -> Any: + if isinstance(o, torch.Tensor) and not isinstance(o, torch.nn.Parameter): + return FakeTensor(shape=list(o.shape), dtype=str(o.dtype), device=str(o.device)) + elif isinstance(o, dict): + return {k: recursive_template(v) for k, v in o.items()} + elif isinstance(o, (list, tuple)): + return type(o)(recursive_template(item) for item in o) + elif is_dataclass(o): + field_values = {f.name: recursive_template(getattr(o, f.name)) for f in fields(o)} + return type(o)(**field_values) + elif issubclass(o.__class__, InplaceSubstituteFakeClass): + # 不重新创建实例，直接修改属性（保持原有初始化逻辑） + for k, v in o.__dict__.items(): + if not callable(v): + o.__dict__[k] = recursive_template(v) + return o + elif o is None or isinstance(o, (int, float, str, bool)): + return o + else: + return o + + return recursive_template(obj) + + +class CudaGraphMgr: + """CUDA Graph Manager for caching and managing CUDA Graphs and static tensors.""" + + def __init__(self): + self.cache: Dict[StaticSignature, StaticTensorEntry] = dict() + self.graph_mem_pool: Optional[torch.cuda.graph_pool_handle] = None + self.check_output_inconsistency = False # Not enabled by default + + @property + def graph_count(self) -> int: + count = 0 + for tensor_entry in self.cache.values(): + if tensor_entry.template_entry_dict is not None: + for template_entry in tensor_entry.template_entry_dict.values(): + for graph_entry in template_entry.graph_entry_dict.values(): + if graph_entry.graph is not None and not graph_entry.inconsistent and not graph_entry.invalid: + count += 1 + return count + + @property + def tensor_entry_count(self) -> int: + count = 0 + for tensor_entry in self.cache.values(): + if tensor_entry.input_tensors is not None and tensor_entry.output_tensors is not None: + count += 1 + return count + + @property + def graph_mem_pool_size(self) -> float: + if not hasattr(self, "graph_mem_pool") or self.graph_mem_pool is None: + return 0.0 + pool_stats = torch.cuda.memory.memory_stats(self.graph_mem_pool) + used_mem = pool_stats.get("allocated_bytes.all.current", 0) + return used_mem / (1024 * 1024) # 转换为MB + + @property + def tensor_mem_size(self) -> float: + total_size = 0 # 字节 + for tensor_entry in self.cache.values(): + if tensor_entry.input_tensors is not None: + for t in tensor_entry.input_tensors: + total_size += t.element_size() * t.nelement() + if tensor_entry.output_tensors is not None: + for t in tensor_entry.output_tensors: + total_size += t.element_size() * t.nelement() + return total_size / (1024 * 1024) # 转换为MB + + @nvtx.instrument_nvtx + def formatted_cache_str(self) -> str: + """Format the cache content as a string for debugging.""" + lines = [] + for static_sig, tensor_entry in self.cache.items(): + lines.append(f"StaticSignature: {static_sig}") + s = " Input Static Tensors: " + for it in tensor_entry.input_tensors: + s += f"[shape={list(it.shape)},dtype={str(it.dtype)}] " + lines.append(s) + s = " Output Static Tensors: " + for ot in tensor_entry.output_tensors: + s += f"[shape={list(ot.shape)},dtype={str(ot.dtype)}] " + lines.append(s) + if tensor_entry.template_entry_dict is not None: + for dynamic_sig, template_entry in tensor_entry.template_entry_dict.items(): + lines.append(f" DynamicSignature: {dynamic_sig}") + lines.append(f" Output Template: {template_entry.output_template}") + for layer_number, graph_entry in template_entry.graph_entry_dict.items(): + status = "Valid" + if graph_entry.inconsistent: + status = "Inconsistent" + elif graph_entry.invalid: + status = "Invalid" + lines.append(f" Layer {layer_number}: Graph Status: {status}") + return "\n".join(lines) + + @nvtx.instrument_nvtx + def try_get_cuda_graph( + self, static_sig: StaticSignature, dynamic_sig: DynamicSignature, layer_number: int + ) -> Optional[torch.cuda.CUDAGraph]: + graph_entry = self.try_get_graph_entry(static_sig, dynamic_sig, layer_number) + if ( + graph_entry is not None + and graph_entry.graph is not None + and not graph_entry.inconsistent + and not graph_entry.invalid + ): + return graph_entry.graph + return None + + @nvtx.instrument_nvtx + def get_static_tensors(self, input_static_sig: StaticSignature) -> Optional[Tuple[List[torch.Tensor], List[torch.Tensor]]]: + if input_static_sig in self.cache: + cached_entry = self.cache[input_static_sig] + return cached_entry.input_tensors, cached_entry.output_tensors + raise ValueError("Cached input/output tensors not found for the given static signature.") + + @nvtx.instrument_nvtx + def warmup_run(self, func: Callable, *args, **kwargs) -> Union[torch.Tensor, List[torch.Tensor]]: + warmup_outputs = None + s = torch.cuda.Stream() + s.wait_stream(torch.cuda.current_stream()) + with torch.cuda.stream(s), torch.no_grad(): + for _ in range(1): + warmup_outputs = func(*args, **kwargs) + torch.cuda.current_stream().wait_stream(s) + return warmup_outputs + + @nvtx.instrument_nvtx + def add_static_entry( + self, + static_sig: StaticSignature, + input_tensors: Optional[List[torch.Tensor]] = None, + output_tensors: Optional[List[torch.Tensor]] = None, + ) -> None: + assert static_sig not in self.cache + self.cache[static_sig] = StaticTensorEntry( + input_tensors=input_tensors, output_tensors=output_tensors, template_entry_dict=dict() + ) + + @nvtx.instrument_nvtx + def add_template_entry( + self, input_static_sig: StaticSignature, input_dynamic_sig: DynamicSignature, output_obj: Any = None + ) -> None: + try: + output_template = ArgsUtils.extract_output_template(output_obj) + self.cache[input_static_sig].template_entry_dict[input_dynamic_sig] = OutputTemplateEntry( + graph_entry_dict=dict(), output_template=output_template + ) + except KeyError: + raise ValueError("StaticSignature not found in cache when adding template entry.") + + @nvtx.instrument_nvtx + def add_graph_entry( + self, + input_static_sig: StaticSignature, + input_dynamic_sig: DynamicSignature, + layer_number: int, + graph: torch.cuda.CUDAGraph, + ) -> None: + try: + self.cache[input_static_sig].template_entry_dict[input_dynamic_sig].graph_entry_dict[layer_number] = GraphEntry( + graph=graph, inconsistent=False, invalid=False + ) + except KeyError: + raise ValueError("StaticSignature or DynamicSignature not found in cache when adding graph entry.") + + @nvtx.instrument_nvtx + def try_get_graph_entry( + self, input_static_sig: StaticSignature, input_dynamic_sig: DynamicSignature, layer_number: int + ) -> Optional[GraphEntry]: + try: + return self.cache[input_static_sig].template_entry_dict[input_dynamic_sig].graph_entry_dict[layer_number] + except KeyError: + pass + return None + + @nvtx.instrument_nvtx + def batch_set_graph_invalid(self, static_sig: StaticSignature) -> None: + if static_sig in self.cache: + static_tensor_entry = self.cache[static_sig] + if static_tensor_entry.template_entry_dict is not None: + for template_entry in static_tensor_entry.template_entry_dict.values(): + for graph_entry in template_entry.graph_entry_dict.values(): + graph_entry.invalid = True + + @nvtx.instrument_nvtx + def set_graph_inconsistent( + self, input_static_sig: StaticSignature, input_dynamic_sig: DynamicSignature, layer_number: int + ) -> None: + if input_static_sig not in self.cache: + self.add_static_entry(input_static_sig, None, None) + if input_dynamic_sig not in self.cache[input_static_sig].template_entry_dict: + self.add_template_entry(input_static_sig, input_dynamic_sig, None) + if layer_number not in self.cache[input_static_sig].template_entry_dict[input_dynamic_sig].graph_entry_dict: + self.cache[input_static_sig].template_entry_dict[input_dynamic_sig].graph_entry_dict[layer_number] = GraphEntry( + graph=None, inconsistent=True, invalid=False + ) + self.cache[input_static_sig].template_entry_dict[input_dynamic_sig].graph_entry_dict[layer_number].inconsistent = True + + @nvtx.instrument_nvtx + def wrapped_graph_capture( + self, + func: Callable, + input_obj: Any, + static_input_tensors: List[torch.Tensor], + static_output_tensors: List[torch.Tensor], + ) -> torch.cuda.CUDAGraph: + init_cudagraph_global_pool() + _set_capture_start() + try: + graph = torch.cuda.CUDAGraph() + _static_input_obj = ArgsUtils.replace_sliced_with_static(input_obj, static_input_tensors) + s = None # future: s = GreenCtxManager(0).create_stream() + with torch.cuda.graph(graph, pool=self.graph_mem_pool, stream=s), torch.no_grad(): + _sliced_output_obj = func(*_static_input_obj["args"], **_static_input_obj["kwargs"]) + _static_output_obj = ArgsUtils.replace_sliced_with_static(_sliced_output_obj, static_output_tensors) + except Exception as e: + torch.cuda.synchronize() # 等待所有异步操作完成 + _set_capture_end() + raise e + _set_capture_end() + return graph + + @nvtx.instrument_nvtx + def wrapped_graph_replay( + self, + graph: torch.cuda.CUDAGraph, + static_input_tensors: List[torch.Tensor], + static_output_tensors: List[torch.Tensor], + input_obj: Any, + output_template: Any, + ) -> Any: + _static_input_obj = ArgsUtils.replace_sliced_with_static(input_obj, static_input_tensors) + graph.replay() + output_obj = ArgsUtils.replace_static_with_sliced(output_template, static_output_tensors) + return output_obj + + @nvtx.instrument_nvtx + def replay_graph( + self, input_static_sig: StaticSignature, input_dynamic_sig: DynamicSignature, input_obj: Any, layer_number: int + ) -> Any: + output_template = self.cache[input_static_sig].template_entry_dict[input_dynamic_sig].output_template + static_input_tensors = self.cache[input_static_sig].input_tensors + static_output_tensors = self.cache[input_static_sig].output_tensors + graph = self.try_get_cuda_graph(input_static_sig, input_dynamic_sig, layer_number=layer_number) + assert graph is not None, "CUDA Graph not found for replay." + output_obj = self.wrapped_graph_replay( + graph=graph, + static_input_tensors=static_input_tensors, + static_output_tensors=static_output_tensors, + input_obj=input_obj, + output_template=output_template, + ) + return output_obj + + @nvtx.instrument_nvtx + def capture_and_cache( + self, + func: Callable, + input_obj: Any, + layer_number: int, + input_static_sig: StaticSignature, + input_dynamic_sig: DynamicSignature, + ) -> Any: + """Capture a new CUDA Graph and cache it.""" + # Access static tensors from cache + static_tensor_entry = self.cache[input_static_sig] + assert static_tensor_entry.input_tensors is not None + assert static_tensor_entry.output_tensors is not None + static_input_tensors = static_tensor_entry.input_tensors + static_output_tensors = static_tensor_entry.output_tensors + + # Capture CUDA Graph + graph = self.wrapped_graph_capture( + func=func, + input_obj=input_obj, + static_input_tensors=static_input_tensors, + static_output_tensors=static_output_tensors, + ) + + # Cache the captured graph + graph_entry = self.try_get_graph_entry(input_static_sig, input_dynamic_sig, layer_number) + if graph_entry: + graph_entry.graph = graph + graph_entry.inconsistent = False + graph_entry.invalid = False + else: + self.add_graph_entry( + input_static_sig=input_static_sig, input_dynamic_sig=input_dynamic_sig, layer_number=layer_number, graph=graph + ) + + @nvtx.instrument_nvtx + def if_need_expand_static_tensors( + self, static_tensors: List[torch.Tensor], new_tensors: List[torch.Tensor], input_static_sig: StaticSignature + ) -> bool: + """Judge whether static tensors need to be expanded based on new tensors.""" + res = False + static_infos = input_static_sig.tensor_static_infos + + if len(static_tensors) != len(new_tensors) or len(static_tensors) != len(static_infos): + raise AssertionError( + f"[CUDA Graph] Tensor count mismatch. {len(static_tensors)=}, {len(new_tensors)=}, {len(static_infos)=}" + ) + + for static_t, new_t, static_info in zip(static_tensors, new_tensors, static_infos): + if static_t.ndim != new_t.ndim: + raise AssertionError(f"[CUDA Graph] Rank mismatch. {static_t.shape=}, {new_t.shape=}") + if static_t.dtype != new_t.dtype: + raise AssertionError(f"[CUDA Graph] Dtype mismatch. {static_t.dtype=}, {new_t.dtype=}") + for i in range(static_t.ndim): + if static_info.shapes[i] != -1 and static_info.shapes[i] != new_t.shape[i]: + raise AssertionError( + f"[CUDA Graph] Static dimension mismatch. {static_t.shape=}, {new_t.shape=}, {static_info.shapes=}, dim={i}" + ) + if static_t.shape[i] < new_t.shape[i]: + res = True + return res + + @nvtx.instrument_nvtx + def get_expanded_static_tensors( + self, static_tensors: List[torch.Tensor], new_tensors: List[torch.Tensor] + ) -> List[torch.Tensor]: + """Get expanded static tensors based on new tensors. Reuses existing tensors when possible.""" + expanded_tensors = [] + for static_t, new_t in zip(static_tensors, new_tensors): + if static_t.ndim != new_t.ndim: + raise AssertionError( + f"[CUDA Graph] Rank mismatch during expansion. Static: {static_t.shape}, New: {new_t.shape}" + ) + new_shape = tuple(max(s, n) for s, n in zip(static_t.shape, new_t.shape)) + + if static_t.shape == new_shape: + expanded_tensors.append(static_t) + elif new_shape == new_t.shape: + expanded_tensors.append(new_t) + else: + expanded_tensor = torch.empty(new_shape, dtype=static_t.dtype, device=static_t.device) + expanded_tensors.append(expanded_tensor) + return expanded_tensors + + @nvtx.instrument_nvtx + def try_replay_graph_inline( + self, func: Callable, args: Tuple, kwargs: Dict, layer_number: int + ) -> Tuple[bool, Optional[Union[torch.Tensor, List[torch.Tensor]]]]: + """Try to replay the CUDA Graph inline for fast execution.""" + try: + func_name = func.__qualname__ + input_obj = {"args": args, "kwargs": kwargs} + + input_tensors, input_tensor_names, literals = ArgsUtils.try_fx_extract_core(input_obj) + if None in (input_tensors, input_tensor_names, literals): + input_tensors, input_tensor_names, literals = ArgsUtils.recursive_extract_core(input_obj) + input_static_sig, input_dynamic_sig = ArgsUtils.generate_both_signatures_from_tensors( + func_name, input_tensors, input_tensor_names, literals + ) + static_tensor_entry = self.cache[input_static_sig] + static_input_tensors = static_tensor_entry.input_tensors + static_output_tensors = static_tensor_entry.output_tensors + + template_entry = static_tensor_entry.template_entry_dict[input_dynamic_sig] + output_template = template_entry.output_template + + graph_entry = template_entry.graph_entry_dict[layer_number] + graph = graph_entry.graph + + assert graph is not None, "CUDA Graph not found for inline replay." + assert graph_entry.inconsistent is False, "CUDA Graph marked as inconsistent for inline replay." + assert graph_entry.invalid is False, "CUDA Graph marked as invalid for inline replay." + + ArgsUtils.replace_sliced_with_static_simple(input_tensors, static_input_tensors) + graph.replay() + output_obj = ArgsUtils.replace_static_with_sliced(output_template, static_output_tensors) + + if self.check_output_inconsistency: + cur_output_tensors, cur_output_tensor_names, cur_output_literals = ArgsUtils.recursive_extract_core(output_obj) + cur_output_static_sig, cur_output_dynamic_sig = ArgsUtils.generate_both_signatures_from_tensors( + func.__qualname__, cur_output_tensors, cur_output_tensor_names, cur_output_literals + ) + output_tensors, output_tensor_names, output_literals = ArgsUtils.recursive_extract_core(output_obj) + cached_output_static_sig, cached_output_dynamic_sig = ArgsUtils.generate_both_signatures_from_tensors( + func.__qualname__, output_tensors, output_tensor_names, output_literals + ) + if cur_output_static_sig != cached_output_static_sig or cur_output_dynamic_sig != cached_output_dynamic_sig: + magi_logger.warning( + f"[CUDA Graph] Warning: Output signature changed during inline replay. {func.__qualname__=}, {layer_number=}" + ) + self.set_graph_inconsistent(input_static_sig, input_dynamic_sig, layer_number) + return False, None + return True, output_obj + except KeyError: + return False, None + except AssertionError: + return False, None + except Exception as e: + magi_logger.info( + f"[CUDA Graph] Exception during inline replay: {e=}, {func.__qualname__=}, {layer_number=}", rank="all" + ) + raise e + + @nvtx.instrument_nvtx + def run(self, func: Callable, *args, layer_number: Optional[int], **kwargs) -> Union[torch.Tensor, List[torch.Tensor]]: + """Run the function with CUDA Graph optimization if possible.""" + + # Try inline replay first + success, output_obj = self.try_replay_graph_inline(func=func, args=args, kwargs=kwargs, layer_number=layer_number) + if success: + # print_rank_0(f"[CUDA Graph] Current cache stats: {self.tensor_entry_count=}, {self.graph_count=}.") + return output_obj + + # Extract input signatures + func_name = func.__qualname__ + input_obj = {"args": args, "kwargs": kwargs} + input_tensors, input_tensor_names, literals = ArgsUtils.recursive_extract_core(input_obj) + input_static_sig, input_dynamic_sig = ArgsUtils.generate_both_signatures_from_tensors( + func_name, input_tensors, input_tensor_names, literals + ) + + # Judge if the graph is marked as inconsistent + graph_entry = self.try_get_graph_entry(input_static_sig, input_dynamic_sig, layer_number) + if graph_entry is not None and graph_entry.inconsistent: + return func(*args, **kwargs) + + # Judge if need to expand static tensors + if_need_expand_static_tensors = False + if_cached_tensor_entry = input_static_sig in self.cache + if if_cached_tensor_entry: + static_input_tensors, static_output_tensors = self.get_static_tensors(input_static_sig) + if_need_expand_static_tensors = self.if_need_expand_static_tensors( + static_input_tensors, input_tensors, input_static_sig + ) + + # Warmup run + warmup_output_obj = self.warmup_run(func, *args, **kwargs) + + # Check input signature consistency after warmup + warmup_input_tensors, warmup_input_tensor_names, warmup_literals = ArgsUtils.recursive_extract_core(input_obj) + warmup_input_static_sig, warmup_input_dynamic_sig = ArgsUtils.generate_both_signatures_from_tensors( + func_name, warmup_input_tensors, warmup_input_tensor_names, warmup_literals + ) + if warmup_input_static_sig != input_static_sig or warmup_input_dynamic_sig != input_dynamic_sig: + magi_logger.warning( + f"[CUDA Graph] Warning: Input signature changed during warmup run. {func_name=}, {layer_number=}" + ) + self.set_graph_inconsistent(input_static_sig, input_dynamic_sig, layer_number) + return warmup_output_obj + + # Update cache entries + if if_cached_tensor_entry: + if if_need_expand_static_tensors: + output_tensors, _, _ = ArgsUtils.recursive_extract_core(warmup_output_obj, extract_literals=False) + # Need to expand static tensors + new_static_input_tensors = self.get_expanded_static_tensors(static_input_tensors, input_tensors) + new_static_output_tensors = self.get_expanded_static_tensors(static_output_tensors, output_tensors) + # Register as new cache entries + self.batch_set_graph_invalid(input_static_sig) + self.cache[input_static_sig].input_tensors = new_static_input_tensors + self.cache[input_static_sig].output_tensors = new_static_output_tensors + + self.add_template_entry(input_static_sig, input_dynamic_sig, warmup_output_obj) + else: + # Simply reuse existing static tensor entry + static_tensor_entry = self.cache[input_static_sig] + if input_dynamic_sig not in static_tensor_entry.template_entry_dict: + self.add_template_entry(input_static_sig, input_dynamic_sig, warmup_output_obj) + + else: + # Create new static tensor entry + output_tensors, _, _ = ArgsUtils.recursive_extract_core(warmup_output_obj, extract_literals=False) + self.add_static_entry(input_static_sig, input_tensors, output_tensors) + self.add_template_entry(input_static_sig, input_dynamic_sig, warmup_output_obj) + + # Capture and cache new CUDA Graph + self.capture_and_cache( + func=func, + input_obj=input_obj, + layer_number=layer_number, + input_static_sig=input_static_sig, + input_dynamic_sig=input_dynamic_sig, + ) + + magi_logger.info( + f"[CUDA Graph] Current cache stats: {self.tensor_entry_count=}, {self.graph_count=}, {self.tensor_mem_size=:.2f} MB, {self.graph_mem_pool_size=:.2f} MB" + ) + return warmup_output_obj + + +_IS_GRAPH_CAPTURING = False + + +def _is_graph_capturing(): + """Query if currently capturing.""" + global _IS_GRAPH_CAPTURING + return _IS_GRAPH_CAPTURING + + +def _set_capture_start(): + """Set graph capture has started.""" + global _IS_GRAPH_CAPTURING + _IS_GRAPH_CAPTURING = True + + +def _set_capture_end(): + """Set graph capture has ended.""" + global _IS_GRAPH_CAPTURING + _IS_GRAPH_CAPTURING = False + + +# Singleton instance of CudaGraphMgr +_CUDA_GRAPH_MGR = CudaGraphMgr() + + +def cuda_graph_mgr() -> CudaGraphMgr: + """ + Get the current CudaGraphMgr instance. + Returns: + CudaGraphMgr: The current CudaGraphMgr instance. + Raises: + AssertionError: If the CudaGraphMgr has not been initialized. + """ + assert _CUDA_GRAPH_MGR is not None, "cuda graph manager is not initialized" + return _CUDA_GRAPH_MGR + + +def cuda_graph_enable_if(condition: Callable): + def decorator(func): + """ + Decorator to enable CUDA graph option for a function. The function will be executed using CUDA Graph if the condition func provided outputs True. + Args: + condition (Callable): A callable that returns a bool indicating whether enable CUDA Graph. + """ + + @wraps(func) + def wrapped_func(*args, **kwargs): + enable_cuda_graph = condition() + if not enable_cuda_graph or _is_graph_capturing(): + return func(*args, **kwargs) + + layer_number = getattr(args[0], "layer_number", None) if args else None + + return cuda_graph_mgr().run(func, *args, layer_number=layer_number, **kwargs) + + return wrapped_func + + return decorator + + +def gen_wrap_func_for_cudagraph(func: Callable, mode_prefix: str, target_prefix=None) -> Callable: + """ + Wrap the given function for CUDA Graph: + 1. Generate a unique __qualname__ for caching + 2. Built-in call to cuda_graph_mgr().run + """ + # Generate a unique identifier to avoid cache conflicts + func_id = id(func) if not hasattr(func, "__name__") else func.__name__ + if mode_prefix == "full": + wrapped_func_name = f"Athena_CUDAGraph_{mode_prefix}_{func_id}" + else: # piecewise + wrapped_func_name = f"Athena_CUDAGraph_{mode_prefix}_{target_prefix}_{func_id}" + + @nvtx.instrument_nvtx + def wrapped_func(*args, **kwargs): + layer_number = kwargs.pop("layer_number", None) + res = cuda_graph_mgr().run(func, *args, layer_number=layer_number, **kwargs) + return res + + func.__qualname__ = wrapped_func_name + magi_logger.info(f"Set original function qualname to {wrapped_func_name} for CUDA Graph caching.") + + # Copy attributes from the original function to the wrapped function + wrapped_func.__dict__.update(func.__dict__) + wrapped_func.__qualname__ = wrapped_func_name + for attr in ["__is_first_graph", "__is_last_graph", "__sym_shape_indices"]: + if hasattr(func, attr): + setattr(wrapped_func, attr, getattr(func, attr)) + + return wrapped_func + + +def init_cudagraph_global_pool(): + """Initialize the global CUDA graph memory pool if not already initialized.""" + from magi_compiler.cuda_graph_mgr import cuda_graph_mgr + + if cuda_graph_mgr().graph_mem_pool is None: + cuda_graph_mgr().graph_mem_pool = torch.cuda.graph_pool_handle() + magi_logger.info("Initialized global CUDA graph pool for Athena.") diff --git a/pkgs/MagiCompiler/magi_compiler/joint_graph_partition.py b/pkgs/MagiCompiler/magi_compiler/joint_graph_partition.py new file mode 100644 index 0000000000000000000000000000000000000000..46faf149dad44060da28aa722e0d0c0d15a0b4bd --- /dev/null +++ b/pkgs/MagiCompiler/magi_compiler/joint_graph_partition.py @@ -0,0 +1,180 @@ +# Copyright (c) 2026 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. + +import os +from typing import Any, Optional, Sequence, Tuple +from unittest.mock import patch + +import torch +import torch.fx as fx +from torch._functorch.compile_utils import get_aten_target +from torch._functorch.partitioners import NodeInfo, OpTypes, get_default_op_list, min_cut_rematerialization_partition +from torch._inductor.custom_graph_pass import CustomPartitionerFn +from torch.utils._ordered_set import OrderedSet + +# from magi_compiler.partitioners import min_cut_rematerialization_partition +from .config import RecomputePolicy, get_compile_config +from .utils import compute_code_hash, magi_logger +from .utils.visualize import joint_graph_vis + +SAVE_TENSOR_NODES: Optional[list[fx.Node]] = None + + +def is_memory_increase_by_node(node: fx.Node) -> bool: + # Only support aten.to now + assert get_aten_target(node) == torch.ops.prims.convert_element_type + input_dtype = node.args[0].meta["tensor_meta"].dtype + output_dtype = node.args[1] + assert output_dtype is not None + return output_dtype.itemsize > input_dtype.itemsize + + +def is_primal_contribute_to_bwd_directly(primal_node: fx.Node, node_info: NodeInfo, op_types: OpTypes) -> bool: + """ + FSDP ensures that weights already reside in memory. If there exists a path from the primal to the bwd, and the path does not contain any matmul, then the primal contributes to the bwd directly. + And we should save this primals. + """ + if node_info.is_required_bw(primal_node): + return True + topology_start = set({primal_node}) + + while len(topology_start) > 0: + cur_node = topology_start.pop() + for user in cur_node.users: + if node_info.is_required_bw(user): + return True + if op_types.is_compute_intensive(user): + continue + topology_start.add(user) + return False + + +def is_compute_intensive_and_has_following_recomputable_ops( + intermidiate_node: fx.Node, node_info: NodeInfo, op_types: OpTypes +) -> Tuple[bool, fx.Node]: + """ + If compute-intensive node(CIN) is not the output of fwd graph(has following memory-intensive ops in the fwd graph), then we should save this CIN node. + NOTE: For CIN+aten.to, we should save aten.to op instead of CIN op to save more memory. + """ + if not op_types.is_compute_intensive(intermidiate_node): + return False, None + + save_node = intermidiate_node + topology_start = set({save_node}) + while len(topology_start) > 0: + cur_node = topology_start.pop() + fwd_user_nodes = [] + for user in cur_node.users: + if node_info.is_required_fw(user): + fwd_user_nodes.append(user) + + if len(fwd_user_nodes) > 1: # multiple users, save current node + return True, save_node + elif len(fwd_user_nodes) == 0: # output, return + return False, None + + # save current node if it's user is recomputable + next_node = fwd_user_nodes[0] + if op_types.is_view(next_node): + if save_node == cur_node: + save_node = next_node + topology_start.add(next_node) + # Special case for aten.to, memory efficient case + elif get_aten_target(next_node) == torch.ops.prims.convert_element_type: + is_memory_increase = is_memory_increase_by_node(next_node) + if not is_memory_increase: + save_node = next_node + topology_start.add(next_node) + elif next_node.op == "output": + return False, None + else: + return True, save_node + assert False, f"Should not reach here: {intermidiate_node=} {save_node=}" + + +# TODO: We find an elegant impl to heuristically save nodes, reconstruct this later +def heuristic_choose_saved_values_set(joint_graph: fx.Graph, node_info: NodeInfo, memory_budget=1) -> list[fx.Node]: + output: OrderedSet[fx.Node] = OrderedSet() + op_types = get_default_op_list() + # Select the inputs that are required by the backward pass + for primal_node in node_info.inputs: + if is_primal_contribute_to_bwd_directly(primal_node, node_info, op_types): + output.add(primal_node) + magi_logger.info("MagiCompiler: saved_output forward-input = %s", output) + # Select the compute-intensive nodes that are required by the forward pass + for intermidiate_node in node_info.required_fw_nodes: + is_save, save_node = is_compute_intensive_and_has_following_recomputable_ops(intermidiate_node, node_info, op_types) + if is_save: + output.add(save_node) + magi_logger.info("MagiCompiler: saved_output compute-intensive = %s", output) + global SAVE_TENSOR_NODES + SAVE_TENSOR_NODES = list(output) + return list(output) + + +def custom_joint_graph_partition_fn( + joint_module: fx.GraphModule, + _joint_inputs, + compiler="inductor", + *, + num_fwd_outputs, + static_lifetime_input_indices: Optional[list[int]] = None, +) -> tuple[fx.GraphModule, fx.GraphModule]: + recompute_config = get_compile_config().recompute_config + if recompute_config.recompute_policy == RecomputePolicy.HANDCRAFT: + magi_logger.info("MagiCompiler using handcraft recompute policy") + # TODO: different memory budget definition from torch + with patch("torch._functorch.config.activation_memory_budget", recompute_config.memory_budget): + fwd_module, bwd_module = min_cut_rematerialization_partition( + joint_module, + _joint_inputs, + compiler, + num_fwd_outputs=num_fwd_outputs, + static_lifetime_input_indices=static_lifetime_input_indices, + ) + elif recompute_config.recompute_policy == RecomputePolicy.HEURISTIC: + magi_logger.info("MagiCompiler using heuristic recompute policy") + with patch("torch._functorch.partitioners.choose_saved_values_set", heuristic_choose_saved_values_set): + fwd_module, bwd_module = min_cut_rematerialization_partition( + joint_module, + _joint_inputs, + compiler, + num_fwd_outputs=num_fwd_outputs, + static_lifetime_input_indices=static_lifetime_input_indices, + ) + elif recompute_config.recompute_policy == RecomputePolicy.AUTOSEARCH: + raise ValueError(f"AutoSearch recompute policy is not supported yet") + else: + raise ValueError(f"Invalid recompute policy: {recompute_config.recompute_policy}") + + joint_graph_vis(joint_module, fwd_module, bwd_module, save_tensor_nodes=SAVE_TENSOR_NODES) + + return fwd_module, bwd_module + + +class CustomJointGraphPartitionFn(CustomPartitionerFn): + def __call__( + self, gm: torch.fx.GraphModule, joint_inputs: Sequence[object], **kwargs: Any + ) -> tuple[torch.fx.GraphModule, torch.fx.GraphModule]: + """ + Implementation of the custom partitioner. + """ + return custom_joint_graph_partition_fn(gm, joint_inputs, **kwargs) + + def uuid(self) -> Optional[Any]: + """ + Return an ID to uniquely identify your custom partitioner implementation. + Return None to skip inductor code caching entirely. + """ + return compute_code_hash({os.path.abspath(__file__)}) diff --git a/pkgs/MagiCompiler/magi_compiler/magi_backend.py b/pkgs/MagiCompiler/magi_compiler/magi_backend.py new file mode 100644 index 0000000000000000000000000000000000000000..70f0a3788138596fae12439f8ed783ddae9c2077 --- /dev/null +++ b/pkgs/MagiCompiler/magi_compiler/magi_backend.py @@ -0,0 +1,607 @@ +# 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. + + +import ast +import dataclasses +import pprint +import time +from collections.abc import Callable +from contextlib import contextmanager +from pathlib import Path +from typing import Any + +import magi_compiler.utils.envs as envs +import torch +import torch.fx as fx +from torch._dispatch.python import enable_python_dispatcher +from torch._dynamo.utils import lazy_format_graph_code +from torch._guards import detect_fake_mode + +from ._cache_data_cls import CacheEntry, CacheHandle +from .compile_artifacts import MagiSerializableFunction +from .config import CompileConfig, CompileMode, CudaGraphMode +from .cuda_graph_mgr import gen_wrap_func_for_cudagraph +from .joint_graph_partition import CustomJointGraphPartitionFn +from .offload.offload_warpper import OffloadWrapper +from .partition_rules import inductor_partition_rule_context, resolve_defined_ops +from .passes import PostGradPassManager +from .passes.inductor_pass import pass_context +from .passes.replace_pass import FullGraphPassManager +from .piecewise_backend import PiecewiseBackend +from .piecewise_compiler import CompilerInterface, EagerAdaptor, InductorStandaloneAdaptor +from .utils import ( + CompileMonitor, + compilation_counter, + compute_code_hash, + compute_hash, + detect_symbolic_tensor_indices, + magi_logger, +) +from .utils.envs import MAGI_CUSTOM_PARTITIONER_FN, MAGI_MODEL_TAG, MAGI_POST_GRAD_PASS +from .utils.visualize import save_fx_graph_visualization + +compilation_start_time: float = 0.0 + + +def _print_with_shape_and_time(runtime_shape: int | None, prefix: str = ""): + elapsed = time.time() - compilation_start_time + if runtime_shape is None: + magi_logger.info("%s for dynamic shape, took %.3f s", prefix, elapsed) + else: + magi_logger.info("%s for shape %s, took %.3f s", prefix, str(runtime_shape), elapsed) + + +@dataclasses.dataclass +class SplitItem: + submod_name: str + graph_id: int + is_splitting_graph: bool + graph: fx.GraphModule + + +def make_compiler(compile_config: CompileConfig) -> CompilerInterface: + if compile_config.backend == "inductor": + # Use standalone_compile with PyTorch 2.8+ + assert hasattr(torch._inductor, "standalone_compile"), "standalone_compile not found in PyTorch Inductor" + magi_logger.info("Using InductorStandaloneAdaptor") + return InductorStandaloneAdaptor() + else: + assert compile_config.backend == "eager", f"Invalid backend for MagiCompiler: {compile_config.backend}" + magi_logger.info("Using EagerAdaptor") + return EagerAdaptor() + + +class CompilerManager: + """ + Manage the compilation process, including graph compilation, compile artifacts caching and loading. + + The cache is a dict mapping `(runtime_shape, graph_index, backend_name)` to `any_data` returned from the compiler. + + When serializing the cache, we save it to a Python file for readability. We don't use json here because json doesn't support int as key. + """ + + def __init__(self, compile_config: CompileConfig): + self.cache: dict[CacheEntry, CacheHandle] = dict() + self.compile_config = compile_config + self.compiler = make_compiler(compile_config) + self.disable_cache = envs.MAGI_DISABLE_COMPILE_CACHE + + @property + def hash(self) -> str: + return self.compiler.hash + + @contextmanager + def compile_context(self, runtime_shape: int | None = None): + """Provide compilation context for the duration of compilation to set + any torch global properties we want to scope to a single Inductor + compilation (e.g. partition rules, pass context).""" + with pass_context(runtime_shape): + if self.compile_config.use_inductor_graph_partition: + inductor_partition_ops = resolve_defined_ops(self.compile_config.splitting_ops) + with inductor_partition_rule_context(inductor_partition_ops): + yield + else: + yield + + def initialize_cache(self, cache_dir: Path, prefix: str = ""): + """ + Initialize the cache directory for the compiler. + + The organization of the cache directory is as follows: + cache_dir=/path/to/torch_compile_cache/rank_i_j/hash_str/prefix/ + inside cache_dir, there will be: + - magi_compile_cache.py + - computation_graph.py + + for multiple prefixes, they can share the same base cache dir of + /path/to/torch_compile_cache/rank_i_j/hash_str/ to store some + common compilation artifacts. + """ + + self.cache_dir: Path = cache_dir + self.cache_file_path: Path = cache_dir / "magi_compile_cache.py" + + if self.disable_cache: + magi_logger.info("MagiCompiler's cache is disabled.") + return + + magi_logger.info("Using cache directory: %s for MagiCompiler", cache_dir) + if self.cache_file_path.exists(): + # load the cache from the file + with self.cache_file_path.open() as f: + # Parse Python literals using ast.literal_eval, which is a safe alternative to eval(). + raw = ast.literal_eval(f.read()) + self.cache = {CacheEntry(*entry): CacheHandle(*handle) for entry, handle in raw.items()} + + self.compiler.initialize_cache(cache_dir=self.cache_dir, prefix=prefix) + + def save_to_file(self): + if self.disable_cache: + return + # serialize to a literal-friendly dict + serializable = {(e.runtime_shape, e.graph_index, e.backend_name): (h.key, h.path) for e, h in self.cache.items()} + printer = pprint.PrettyPrinter(indent=4) + data = printer.pformat(serializable) + with self.cache_file_path.open("w") as f: + f.write(data) + + def load(self, graph: fx.GraphModule, example_inputs: list[Any], cache_entry: CacheEntry) -> Callable | None: + if cache_entry not in self.cache: + return None + cache_handle = self.cache[cache_entry] + _print_with_shape_and_time( + cache_entry.runtime_shape, + f"Directly load the {cache_entry.graph_index}-th graph from {cache_entry.backend_name} via handle {cache_handle}", + ) + return self.compiler.load(graph, example_inputs, cache_entry, cache_handle) + + # TODO(hongyu): Support training mode here + def compile( + self, + graph: fx.GraphModule, + example_inputs: tuple[torch.fx.node.Argument, ...], + compile_config: CompileConfig, + graph_index: int = 0, + num_graphs: int = 1, + runtime_shape: int | None = None, + ) -> Callable: + # Step0: update some global metrics + compilation_counter.num_backend_compilations += 1 + if graph_index == 0: + global compilation_start_time + compilation_start_time = time.time() + + # Step1: Try loading from the cache + cache_entry = CacheEntry(runtime_shape, graph_index, self.compiler.name) + compiled_graph = self.load(graph, example_inputs, cache_entry) + if compiled_graph is not None: + return compiled_graph + + # Step2: Compile the graph + key = f"artifact_shape_{runtime_shape}_subgraph_{graph_index}" + with self.compile_context(runtime_shape): + compiled_graph, cache_handle = self.compiler.compile( + graph, example_inputs, compile_config.inductor_compile_config, runtime_shape, key + ) + assert compiled_graph is not None, "Failed to compile the graph" + + # Step3: Store the artifact in the cache + if not self.disable_cache and cache_handle is not None: + assert cache_entry not in self.cache, "Cache entry already exists" + self.cache[cache_entry] = cache_handle + compilation_counter.num_cache_entries += 1 + _print_with_shape_and_time(runtime_shape, f"Compile the {graph_index}/{num_graphs} graph") + + return compiled_graph + + +# TODO(hongyu): Support training mode here +class PiecewiseCompileInterpreter(torch.fx.Interpreter): + """ + Code adapted from `torch.fx.passes.shape_prop.ShapeProp`. + It runs the given graph with fake inputs, and compile some submodules specified by `compile_submod_names` with compilation configs. + + NOTE: the order in `compile_submod_names` matters, because it will be used to determine the order of the compiled piecewise graphs. + The first graph will handle logging, and the last graph has some special cudagraph output handling. + """ + + def __init__( + self, + module: torch.fx.GraphModule, + compiler_manager: CompilerManager, + compile_submod_names: list[str], + compile_config: CompileConfig, + ): + super().__init__(module) + + self.fake_mode = detect_fake_mode() + self.compiler_manager = compiler_manager + self.compile_submod_names = compile_submod_names + self.compile_config = compile_config + # extra_traceback is attribute of torch.fx.Interpreter, when it is True, it annoyingly dumps the torch.fx.Graph on errors. + self.extra_traceback = False + + def _fix_graph_device_placement(self, module: torch.nn.Module): + for name, child in module.named_children(): + self._fix_graph_device_placement(child) + + if isinstance(module, torch.fx.GraphModule): + needs_recompile = False + target_device = torch.cuda.current_device() + + factory_functions = [ + torch.empty, + torch.zeros, + torch.ones, + torch.full, + torch.rand, + torch.randn, + torch.arange, + torch.tensor, + torch.ops.aten.empty.memory_format, + ] + + for node in module.graph.nodes: + if node.op == 'call_function': + is_factory = node.target in factory_functions or ( + hasattr(node.target, '__name__') and node.target.__name__ in ['empty', 'zeros', 'ones', 'full'] + ) + + if is_factory: + if 'device' in node.kwargs: + current_dev = node.kwargs['device'] + if str(current_dev) == 'cpu' or current_dev == torch.device('cpu'): + node.update_kwarg('device', target_device) + needs_recompile = True + + if needs_recompile: + module.recompile() + + def run(self, *args): + fake_args = [self.fake_mode.from_tensor(t) if isinstance(t, torch.Tensor) else t for t in args] + if self.compile_config.offload_config.model_cpu_offload: + self._fix_graph_device_placement(self.module) + for i, arg in enumerate(fake_args): + if isinstance(arg, torch.Tensor): + fake_args[i] = arg.cuda() + + with self.fake_mode, enable_python_dispatcher(): + return super().run(*fake_args) + + def call_module( + self, target: torch.fx.node.Target, args: tuple[torch.fx.node.Argument, ...], kwargs: dict[str, Any] + ) -> Any: + assert isinstance(target, str) + output = super().call_module(target, args, kwargs) + if target not in self.compile_submod_names: + return output + + index = self.compile_submod_names.index(target) + submod = self.fetch_attr(target) + sym_shape_indices = [i for i, x in enumerate(args) if isinstance(x, torch.SymInt)] + magi_logger.info(f"Compiling {target=}, {sym_shape_indices=}, {args=}") + + compiled_graph_for_dynamic_shape = self.compiler_manager.compile( + submod, args, self.compile_config, graph_index=index, num_graphs=len(self.compile_submod_names), runtime_shape=None + ) + + piecewise_backend = PiecewiseBackend( + submod, + compiled_graph_for_dynamic_shape, + self.compile_config, + index, + len(self.compile_submod_names), + sym_shape_indices, + self.compiler_manager, + ) + + if self.compile_config.use_inductor_graph_partition or self.compile_config.cudagraph_mode != CudaGraphMode.PIECEWISE: + self.module.__dict__[target] = piecewise_backend + else: + wrapped_backend = gen_wrap_func_for_cudagraph( + func=piecewise_backend, mode_prefix=CudaGraphMode.PIECEWISE.name.lower(), target_prefix=target + ) + + self.module.__dict__[target] = wrapped_backend + magi_logger.info( + f"Wrapped piecewise submodule {target} (index {index}) with CUDA Graph " + f"[PIECEWISE mode, first_graph={piecewise_backend.is_first_graph}, last_graph={piecewise_backend.is_last_graph}]" + ) + + return output + + +class MagiBackend: + """ + The compilation backend for `torch.compile` with MagiCompiler. + It is used for compilation mode of `CompileMode.MAGI_COMPILE`, + where we customize the compilation. + + The major work of this backend is to split the graph into + piecewise graphs, and pass them to the piecewise backend. + + This backend also adds the PostGradPassManager to Inductor config, + which handles the post-grad passes. + """ + + compile_config: CompileConfig + _called_once: bool = False + # for the graph we compiled + graph: fx.GraphModule + compiler_manager: CompilerManager + # for cudagraph + sym_tensor_indices: list[int] # indices for tensors that have symbolic shapes + input_buffers: list[torch.Tensor] # buffers for input tensors that have symbolic shapes + + def __init__(self, compile_config: CompileConfig, model_tag: str = ""): + self.model_tag = model_tag or MAGI_MODEL_TAG + self.compile_config = compile_config + self._configure_custom_passes() + self.compiler_manager: CompilerManager = CompilerManager(self.compile_config) + + self.sym_tensor_indices = [] + self.input_buffers = [] + + def _configure_custom_passes(self): + # Custom pass 1: full graph passes between Dynamo and AOTAutograd + self.full_graph_pass_manager = FullGraphPassManager(self.compile_config.pass_config) + + # Custom pass 2: custom partitioner function + custom_partitioner_fn = CustomJointGraphPartitionFn() + if MAGI_CUSTOM_PARTITIONER_FN in self.compile_config.inductor_compile_config: + existing_fn = self.compile_config.inductor_compile_config[MAGI_CUSTOM_PARTITIONER_FN] + assert isinstance(existing_fn, CustomJointGraphPartitionFn) + assert existing_fn.uuid() == custom_partitioner_fn.uuid() + self.compile_config.inductor_compile_config[MAGI_CUSTOM_PARTITIONER_FN] = custom_partitioner_fn + + # Custom pass 3: post-grad passes after AOTAutograd + post_grad_pass_manager = PostGradPassManager() + post_grad_pass_manager.configure(self.compile_config) + + # Run post-grad custom passes with post_grad_custom_post_pass hook + if MAGI_POST_GRAD_PASS in self.compile_config.inductor_compile_config: + existing_pass = self.compile_config.inductor_compile_config[MAGI_POST_GRAD_PASS] + assert isinstance(existing_pass, PostGradPassManager) + assert existing_pass.uuid() == post_grad_pass_manager.uuid() + + self.compile_config.inductor_compile_config[MAGI_POST_GRAD_PASS] = post_grad_pass_manager + + def _init_cache(self) -> str: + hash_key = compute_hash( + [self.compile_config.hash, self.compiler_manager.hash, compute_code_hash(self.compile_config.traced_files)] + ) + self.compile_config.traced_files.clear() + + # Path: .../model_{idx}_{model_tag}_rank_{rank}/{hash}/{model_tag}/ (last segment = class name or user tag) + self.local_cache_dir: Path = self.compile_config.cache_dump_path() / hash_key / self.model_tag + self.local_cache_dir.mkdir(parents=True, exist_ok=True) + + self.compiler_manager.initialize_cache(self.local_cache_dir, self.model_tag) + + def _save_partitioned_graph(self, split_gm: fx.GraphModule): + graph_path = self.local_cache_dir / "computation_graph.py" + if not graph_path.exists(): + # code adapted from + # https://github.com/thuml/depyf/blob/dab831108a752d1facc00acdd6d4243891845c37/depyf/explain/patched_lazy_format_graph_code.py#L30 + # use `print_readable` because it can include submodules + src = "from __future__ import annotations\nimport torch\n" + split_gm.print_readable(print_output=False) + src = src.replace("", "GraphModule") + with open(graph_path, "w") as f: + f.write(src) + magi_logger.info("Computation graph saved to %s", graph_path) + + def _split_graph(self, graph: fx.GraphModule) -> tuple[fx.GraphModule, list[SplitItem]]: + # Step 1: resolve the splitting ops + if self.compile_config.use_inductor_graph_partition: + # Let Inductor decide partitioning; avoid FX-level pre-splitting. + fx_split_ops: list[str] = [] + else: + fx_split_ops = self.compile_config.splitting_ops or [] + resolved_ops: list[torch._ops.OpOverload] = resolve_defined_ops(fx_split_ops) + magi_logger.info(f"Setting up FX-level graph split with ops: {fx_split_ops=}") + magi_logger.info(f"Resolved splitting ops for FX-level graph split: {resolved_ops=}") + + # Step 2: split graph by ops, we split graph based on resolved_ops, which becomes the partitioned single graph. + subgraph_id = 0 + node_to_subgraph_id = {} + split_op_graphs = [] + for node in graph.graph.nodes: + if node.op in ("output", "placeholder"): + continue + # Match node.target against resolved_ops, node.target can be OpOverloadPacket, need to check .default + if node.op == "call_function" and ( + node.target in resolved_ops or (hasattr(node.target, "default") and node.target.default in resolved_ops) + ): + magi_logger.info(f"Splitting graph at {node=} with {node.target=}") + subgraph_id += 1 + node_to_subgraph_id[node] = subgraph_id + split_op_graphs.append(subgraph_id) + subgraph_id += 1 + else: + node_to_subgraph_id[node] = subgraph_id + + # Step 3: split the graph based on node_to_subgraph_id + # pytorch might reorder the nodes and the semantics of the graph will change when we have mutations in the graph, if we don't set keep_original_order=True + split_gm = torch.fx.passes.split_module.split_module( + graph, None, lambda node: node_to_subgraph_id[node], keep_original_order=True + ) + + def _extract_example_values(args) -> list: + example_values = [] + + def _recurse_extract(arg): + if isinstance(arg, (list, tuple)): + for sub_arg in arg: + _recurse_extract(sub_arg) + else: + example_value = arg.meta.get("example_value") + assert example_value is not None, f"Output arg {arg} has no example_value for tensor_meta recovery" + example_values.append(example_value) + + _recurse_extract(args) + return example_values + + def _format_output_values(values: list): + if not values: + return None + return tuple(values) if len(values) > 1 else values[0] + + def _recursive_recover_tensor_meta(gm: fx.GraphModule): + """ + 递归恢复指定 GraphModule 及其所有嵌套 submodule 中所有 node 的 example_value + 支持任意层级的 submodule 嵌套 + """ + for node in gm.graph.nodes: + if node.meta.get("example_value") is not None: + continue + + if node.op == "call_module": + submod: fx.GraphModule = getattr(gm, node.target) + _recursive_recover_tensor_meta(submod) # 递归调用，处理嵌套 submodule + output_node = next(n for n in submod.graph.nodes if n.op == "output") + assert output_node is not None, f"Output node not found in submodule {node.target}" + output_values = _extract_example_values(output_node.args) + node.meta["example_value"] = _format_output_values(output_values) + elif node.op == "call_function": + if "getitem" in str(node.target): + prev_node, getitem_index = node.args + prev_example_value = prev_node.meta.get("example_value") + assert ( + prev_example_value is not None + ), f"Previous node {prev_node} has no example_value for tensor_meta recovery of node {node}" + node.meta["example_value"] = prev_example_value[getitem_index] + elif "grad" in str(node.target) or "device" in str(node.target): # 暂时不做处理 + node.meta["example_value"] = None + elif node.op == "output": + output_values = _extract_example_values(node.args[0]) + node.meta["example_value"] = _format_output_values(output_values) + + else: + raise ValueError(f"Unsupported node op for tensor_meta recovery: {node.op} for node {node}") + + magi_logger.info(f"Recovered example_value for node {node.name}: {node.meta['example_value']=}") + + # Recover tensor_meta for all nodes in split_gm and its submodules + if envs.MAGI_ENABLE_PROFILE: + _recursive_recover_tensor_meta(split_gm) + + # Step 4: fetch all the submodules + piecewise_graphs = [] + names = [name for (name, module) in split_gm.named_modules()] + for name in names: + # Only keep the top-level modules, skip recursive child modules or the root module + if "." in name or name == "": + continue + + module = getattr(split_gm, name) + assert isinstance(module, fx.GraphModule), f"Expected fx.GraphModule, got {type(module)}" + + graph_id = int(name.replace("submod_", "")) + piecewise_graphs.append(SplitItem(name, graph_id, (graph_id in split_op_graphs), module)) + # sort by integer graph_id, rather than string name + piecewise_graphs.sort(key=lambda x: x.graph_id) + + # Step 5: visualize the split graph + # depyf already hooks lazy_format_graph_code and dumps the graph, we do not print the graph here + lazy_format_graph_code("Before split", graph, print_output=True, include_stride=True, include_device=True) + lazy_format_graph_code("After split", split_gm, print_output=True, include_stride=True, include_device=True) + + if envs.MAGI_ENABLE_FX_GRAPH_VIZ: + save_fx_graph_visualization(split_gm.graph, sub_dir="after_split", filename="split_gm_root") + for item in piecewise_graphs: + save_fx_graph_visualization(item.graph.graph, sub_dir="after_split", filename=item.submod_name) + + return split_gm, piecewise_graphs + + def __call__(self, graph: fx.GraphModule, example_inputs) -> MagiSerializableFunction: + assert not self._called_once, "MagiBackend can only be called once cause compilation is a one-time process" + self._called_once = True + magi_logger.info("Dynamo traced files (for compilation cache):\n%s", "\n".join(self.compile_config.traced_files)) + compilation_counter.num_graphs_seen += 1 + CompileMonitor().mark("Dynamo bytecode transform") + + self._init_cache() + + self.full_graph_pass_manager(graph) + + split_gm, piecewise_graphs = self._split_graph(graph) + + submod_names_to_compile = [item.submod_name for item in piecewise_graphs if not item.is_splitting_graph] + compilation_counter.num_piecewise_graphs_seen += len(piecewise_graphs) + compilation_counter.num_piecewise_capturable_graphs_seen += len(submod_names_to_compile) + magi_logger.info(f"Piecewise modules waiting for compilation: {submod_names_to_compile}") + + # propagate the split graph to the piecewise backend, compile submodules with symbolic shapes + try: + PiecewiseCompileInterpreter(split_gm, self.compiler_manager, submod_names_to_compile, self.compile_config).run( + *example_inputs + ) + except Exception as e: + # Magi compile 的集中失败入口：直接打印 ERROR，方便在大模型日志中 grep + magi_logger.error("Magi compile failed while compiling piecewise submodules %s: %s", submod_names_to_compile, e) + raise + self._save_partitioned_graph(split_gm) + + # TODO: Support DBO and NAT here + # split_gm = DBOGraphModule(split_gm, self.compile_config) + if self.compile_config.offload_config.model_cpu_offload: + split_gm = OffloadWrapper(split_gm, self.compile_config) + + # if envs.MAGI_ENABLE_TOKENFLOW: + # from magi_compiler.tokenflow.graph_fork import GraphForkWrapper + + if envs.MAGI_ENABLE_PROFILE: + from magi_compiler.tokenflow.graph_profile import gen_profile_wrap_func + + split_gm = gen_profile_wrap_func(split_gm) + + if self.compile_config.cudagraph_mode == CudaGraphMode.FULL and self.compile_config.cudagraph_copy_inputs: + return self._serialize_func_with_cudagraph(graph, split_gm, example_inputs) + + return MagiSerializableFunction(graph, example_inputs, self.model_tag, split_gm) + + def _serialize_func_with_cudagraph( + self, graph: fx.GraphModule, split_gm: fx.GraphModule, example_inputs: list[Any] + ) -> MagiSerializableFunction: + fake_mode = detect_fake_mode() + fake_args = [fake_mode.from_tensor(t) if isinstance(t, torch.Tensor) else t for t in example_inputs] + + self.sym_tensor_indices = detect_symbolic_tensor_indices(fake_args) + + wrapped_split_gm = gen_wrap_func_for_cudagraph(func=split_gm, mode_prefix=CudaGraphMode.FULL.name.lower()) + + return MagiSerializableFunction(graph, example_inputs, self.model_tag, wrapped_split_gm) + + +def init_backend(compile_config: CompileConfig) -> str | Callable: + """ + Initialize the backend based on CompileConfig. + """ + if compile_config.compile_mode is None or compile_config.compile_mode == CompileMode.NONE: + raise ValueError("No compilation mode is set.") + + from torch._dynamo.backends.registry import list_backends + + torch_backends = list_backends(exclude_tags=tuple()) + magi_logger.info("Supported torch backends: %s", torch_backends) + if compile_config.compile_mode == CompileMode.TORCH_COMPILE: + assert compile_config.backend in torch_backends, f"Invalid backend for torch compilation: {compile_config.backend}" + return compile_config.backend + elif compile_config.compile_mode == CompileMode.MAGI_COMPILE: + assert compile_config.backend in ["eager", "inductor"], f"Invalid backend for MagiCompiler: {compile_config.backend}" + model_tag = getattr(compile_config, "model_tag", None) or MAGI_MODEL_TAG + return MagiBackend(compile_config, model_tag=model_tag) + else: + raise ValueError(f"Invalid compile mode: {compile_config.compile_mode}") diff --git a/pkgs/MagiCompiler/magi_compiler/magi_compiler_base.py b/pkgs/MagiCompiler/magi_compiler/magi_compiler_base.py new file mode 100644 index 0000000000000000000000000000000000000000..e2e77cc5528de57417b36a40583ac058293981d8 --- /dev/null +++ b/pkgs/MagiCompiler/magi_compiler/magi_compiler_base.py @@ -0,0 +1,219 @@ +# 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. + +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +import inspect +import os +import sys +from abc import abstractmethod +from contextlib import contextmanager +from types import CodeType +from typing import Callable, Literal + +import magi_compiler.utils.envs as envs +import torch +from magi_compiler.utils import compute_hash, get_git_version +from magi_compiler.utils.compile_time_monitor import CompileMonitor + +from .config import CompileConfig, CompileMode +from .magi_backend import init_backend +from .utils import compute_code_hash, compute_code_hash_with_content, magi_logger + + +def _verify_source_unchanged(source_info, compile_config: CompileConfig) -> None: + file_contents = {} + for source in source_info.inlined_sources: + module = sys.modules[source.module] + file = inspect.getfile(module) + file_contents[file] = source.content + compile_config.traced_files.add(file) + expected_checksum = compute_code_hash_with_content(file_contents) + actual_checksum = compute_code_hash(set(file_contents.keys())) + if expected_checksum != actual_checksum: + raise RuntimeError("Source code has changed since the last compilation. Recompiling the model.") + + +class MagiCompilerBase: + compile_config: CompileConfig + """ + A wrapper class for torch.compile, with a custom dispatch logic. + Subclasses should: + 1. Implement the forward method + 2. Implement the dispatch logic in the __call__ method + It can use `self.compiled_codes` to access the compiled bytecode, + and `with self.dispatch_to_compiled_code:` to dispatch to + the compiled code. + 3. Implement the `__init__` method to determine how to call + `torch.compile` over the forward method. + """ + + def __init__(self, compile_config: CompileConfig): + backend = init_backend(compile_config) + options = None + if isinstance(backend, str) and backend == "inductor": + options = compile_config.inductor_compile_config + if envs.MAGI_AOT_COMPILE: + options = options or {} + # Drop all the guards in the AOT compile mode as bytecode hook is not used anymore. + options["guard_filter_fn"] = lambda guards: [False for _ in guards] + assert hasattr(torch._dynamo.config, "enable_aot_compile"), "enable_aot_compile config not available" + torch._dynamo.config.enable_aot_compile = True + + self.compiled_callable = torch.compile(self.forward, fullgraph=True, backend=backend, options=options) + self.original_code_object: CodeType = self.__class__.forward.__code__ + self.compiled_code: CodeType | None = None + self.aot_compiled_fn: Callable | None = None + + @property + def aot_compilation_path(self) -> str: + """ + When using torch.compile in AOT mode, we store the cache artifacts + under cache_root_dir/torch_aot_compile/{hash}/rank_i_j. The {hash} + contains all of the factors except for the source files being + traced through, because we don't actually know which source files + to check at this point (before dynamo runs). + On loading we will actually look at the source files being traced + through. If any source file have changed (compared with the + serialized backend artifacts), then we need to generate a new AOT + compile artifact from scratch. + """ + hash_key = compute_hash([self.forward, self.compile_config.model_idx, self.compile_config.hash, get_git_version()]) + cache_dir = os.path.join(self.compile_config.cache_root_dir, "torch_aot_compile", hash_key) + rank = torch.distributed.get_rank() if torch.distributed.is_initialized() else 0 + cache_dir = os.path.join(cache_dir, f"rank_{rank}") + os.makedirs(cache_dir, exist_ok=True) + aot_compilation_path = os.path.join(cache_dir, "model") + return aot_compilation_path + + def try_load_aot_compile_artifacts(self) -> Callable | None: + if self.aot_compiled_fn is not None: + return self.aot_compiled_fn + if not os.path.exists(self.aot_compilation_path): + return None + with open(self.aot_compilation_path, "rb") as f: + CompileMonitor().start( + self.compile_config.compile_mode == CompileMode.MAGI_COMPILE, self.compile_config.debug_dump_path() + ) + loaded_fn = torch.compiler.load_compiled_function(f) + _verify_source_unchanged(loaded_fn.source_info(), self.compile_config) + return loaded_fn + + def aot_compile(self, *args, **kwargs): + """ + Run the model in AOT (Ahead-Of-Time) compile mode. + + All compilation work is completed before execution, suitable for production environment. + This results in longer compilation time but superior runtime performance. + """ + assert hasattr(self.compiled_callable, "aot_compile"), "aot_compile is not supported by the current configuration" + return self.compiled_callable.aot_compile((args, kwargs)) + + def jit_compile(self, *args, **kwargs): + """ + Run the model in JIT (Just-In-Time) compile mode. + + Compilation occurs at runtime, first run may be slower due to compilation overhead. + """ + handle = torch._dynamo.convert_frame.register_bytecode_hook(self.bytecode_hook) + output = self.compiled_callable(*args, **kwargs) + handle.remove() + return output + + @abstractmethod + def forward(self, *args, **kwargs): + ... + + def bytecode_hook(self, old_code: CodeType, new_code: CodeType): + """Hook to save the compiled bytecode for direct execution.""" + if old_code is not self.original_code_object: + return + # Step1: Check if the old bytecode is from the compiled code + # code borrowed from depyf enable_debugging.py + frame = sys._getframe() + while frame and frame.f_back: + frame = frame.f_back + code_name = frame.f_code.co_name + file_name = frame.f_code.co_filename.split(os.path.sep)[-1] + if code_name == "_compile" and file_name == "convert_frame.py": + break + frame = frame.f_locals["frame"] + assert frame.f_code == old_code + + if hasattr(frame.f_locals, "self") and frame.f_locals["self"] is not self: + return + + # Step2: Save the compiled bytecode + self.compiled_code = new_code + + # Step3: Save the decompiled code + path = self.compile_config.debug_dump_path() + decompiled_file = os.path.join(path, "decompiled_code.py") + if os.path.exists(decompiled_file): + return + try: + # usually the decompilation will succeed for most models, as we guarantee a full-graph compilation in Dynamo. + # but there's no 100% guarantee, since decompliation is not a reversible process. + from magi_compiler.magi_depyf import decompile as magi_decompile + + src = magi_decompile(new_code) + with open(decompiled_file, "w") as f: + f.write(src) + magi_logger.info("Dynamo transformed code saved to %s", decompiled_file) + except Exception: + pass + + @contextmanager + def dispatch_to_compiled_fwd(self, mode: Literal["jit", "aot"] = "jit"): + """ + Context manager to dispatch to the compiled code. + Why does this work? Because Dynamo guarantees that the compiled + bytecode has exactly the same arguments, cell variables, and free + variables as the original code. Therefore we can directly switch + the code object in the function and call it. + + See https://dev-discuss.pytorch.org/t/what-is-the-relationship-requirement-among-original-bytecode-transformed-bytecode-and-bytecode-returned-by-hooks-in-dynamo/1693/7 + for more details. + + NOTE: Why compile `forward` but invoke through `old_call`? + + In torch.nn.Module, `__call__` wraps `forward` with critical runtime logic: + - Pre/post forward hooks + - FSDP parameter sharding/gathering and device placement + + Our strategy: use this context manager to temporarily replace `self.forward` + with the compiled version, then invoke `old_call(self, *args, **kwargs)`. + + This way: + 1. `old_call` executes hooks and FSDP mechanics normally + 2. When `old_call` internally calls `self.forward`, it hits our compiled code + 3. Compiled code runs within the proper FSDP/hook context + + Calling `self.forward()` directly would bypass FSDP (seeing sharded/invalid + params) and skip hooks that other components may rely on. + """ + if mode == "jit": + assert self.compiled_code is not None + self.__class__.forward.__code__ = self.compiled_code + yield + self.__class__.forward.__code__ = self.original_code_object + elif mode == "aot": + assert self.aot_compiled_fn is not None + old_forward = self.forward + self.forward = lambda *args, **kwargs: self.aot_compiled_fn(self, *args, **kwargs) + yield + self.forward = old_forward + else: + raise ValueError(f"Invalid mode: {mode}") diff --git a/pkgs/MagiCompiler/magi_compiler/magi_depyf/__init__.py b/pkgs/MagiCompiler/magi_compiler/magi_depyf/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..f708aafb6d6d6dc069d2c576a2c99686650d5c2f --- /dev/null +++ b/pkgs/MagiCompiler/magi_compiler/magi_depyf/__init__.py @@ -0,0 +1,21 @@ +# Copyright (c) 2026 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. + +"""magi_depyf — a modern bytecode decompiler and torch.compile inspector.""" + +from .decompile import DecompilationError, Decompiler, decompile, safe_decompile + +__version__ = "0.1.0" + +__all__ = ["Decompiler", "decompile", "safe_decompile", "DecompilationError", "__version__"] diff --git a/pkgs/MagiCompiler/magi_compiler/magi_depyf/decompile/__init__.py b/pkgs/MagiCompiler/magi_compiler/magi_depyf/decompile/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..0610bd1b36f98238b914e1e4793c8b88b1f2c8ef --- /dev/null +++ b/pkgs/MagiCompiler/magi_compiler/magi_depyf/decompile/__init__.py @@ -0,0 +1,19 @@ +# Copyright (c) 2026 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. + +"""Decompilation: bytecode → Python source, plus recompile/fix/postprocess.""" + +from .decompiler import DecompilationError, Decompiler, decompile, safe_decompile + +__all__ = ["Decompiler", "decompile", "safe_decompile", "DecompilationError"] diff --git a/pkgs/MagiCompiler/magi_compiler/magi_depyf/decompile/bytecode/__init__.py b/pkgs/MagiCompiler/magi_compiler/magi_depyf/decompile/bytecode/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..cdc565928d0873e6ead5a2f92aae37e9c7a7ede8 --- /dev/null +++ b/pkgs/MagiCompiler/magi_compiler/magi_depyf/decompile/bytecode/__init__.py @@ -0,0 +1,22 @@ +# Copyright (c) 2026 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. + +"""Bytecode processing — pure Python, no torch dependency.""" + +from .decompile_context import DecompileContext +from .handler_registry import HandlerRegistry, registry +from .instruction import Instruction +from .source_emitter import SourceEmitter + +__all__ = ["Instruction", "SourceEmitter", "HandlerRegistry", "DecompileContext", "registry"] diff --git a/pkgs/MagiCompiler/magi_compiler/magi_depyf/decompile/bytecode/decompile_context.py b/pkgs/MagiCompiler/magi_compiler/magi_depyf/decompile/bytecode/decompile_context.py new file mode 100644 index 0000000000000000000000000000000000000000..1b63e01d80e221d37c22073f67ce6246666e928d --- /dev/null +++ b/pkgs/MagiCompiler/magi_compiler/magi_depyf/decompile/bytecode/decompile_context.py @@ -0,0 +1,53 @@ +# Copyright (c) 2026 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. + +"""DecompileContext — read-only bag passed to every handler. + +Handlers receive ``(emitter, inst, ctx)`` — they mutate *emitter* +and call *ctx* methods but never touch the ``Decompiler`` directly. +""" + +from __future__ import annotations + +from types import CodeType +from typing import TYPE_CHECKING, Callable, Dict, Tuple + +if TYPE_CHECKING: + from .instruction import Instruction + + +class DecompileContext: + """Read-only context providing handlers with instructions, code object, + and the ``decompile_range`` callback for recursive sub-block decompilation.""" + + def __init__( + self, + code: CodeType, + instructions: Tuple["Instruction", ...], + indentation: int, + decompile_range: Callable, + offset_to_index: Dict[int, int], + ) -> None: + self.code = code + self.instructions = instructions + self.indentation = indentation + self.decompile_range = decompile_range + self._offset_to_index = offset_to_index + + def index_of(self, offset: int) -> int: + """Return the index of the instruction at *offset* (O(1) lookup).""" + try: + return self._offset_to_index[offset] + except KeyError: + raise ValueError(f"No instruction at offset {offset}") from None diff --git a/pkgs/MagiCompiler/magi_compiler/magi_depyf/decompile/bytecode/handler_registry.py b/pkgs/MagiCompiler/magi_compiler/magi_depyf/decompile/bytecode/handler_registry.py new file mode 100644 index 0000000000000000000000000000000000000000..3d2a694fc01f8babec248bac7dedf0297297fdb4 --- /dev/null +++ b/pkgs/MagiCompiler/magi_compiler/magi_depyf/decompile/bytecode/handler_registry.py @@ -0,0 +1,62 @@ +# Copyright (c) 2026 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. + +"""HandlerRegistry — opcode-to-handler dispatch. + +A *handler* is a plain function with signature:: + + (emitter: SourceEmitter, inst: Instruction, ctx: DecompileContext) -> Optional[int] + +Returning ``None`` advances to the next instruction. +Returning an ``int`` jumps to that instruction index. +""" + +from __future__ import annotations + +from typing import TYPE_CHECKING, Callable, List, Optional + +if TYPE_CHECKING: + pass + +HandlerFn = Callable[..., Optional[int]] + + +class HandlerRegistry: + """Maps opcode names -> handler functions.""" + + def __init__(self) -> None: + self._handlers: dict[str, HandlerFn] = {} + + def register(self, *opnames: str) -> Callable[[HandlerFn], HandlerFn]: + """Decorator that registers *fn* for one or more opcode names.""" + + def decorator(fn: HandlerFn) -> HandlerFn: + for name in opnames: + self._handlers[name] = fn + return fn + + return decorator + + def get(self, opname: str) -> Optional[HandlerFn]: + return self._handlers.get(opname) + + def __contains__(self, opname: str) -> bool: + return opname in self._handlers + + def supported_opnames(self) -> List[str]: + return sorted(self._handlers.keys()) + + +# Singleton registry — handlers register against this at import time. +registry = HandlerRegistry() diff --git a/pkgs/MagiCompiler/magi_compiler/magi_depyf/decompile/bytecode/handlers/__init__.py b/pkgs/MagiCompiler/magi_compiler/magi_depyf/decompile/bytecode/handlers/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..213b41aa91ccf3bce70e694d48ff3f147de12c8d --- /dev/null +++ b/pkgs/MagiCompiler/magi_compiler/magi_depyf/decompile/bytecode/handlers/__init__.py @@ -0,0 +1,22 @@ +# Copyright (c) 2026 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. + +"""Import every handler module so they register against the global registry.""" + +from . import arithmetic # noqa: F401 +from . import calls # noqa: F401 +from . import containers # noqa: F401 +from . import control_flow # noqa: F401 +from . import load_store # noqa: F401 +from . import stack_ops # noqa: F401 diff --git a/pkgs/MagiCompiler/magi_compiler/magi_depyf/decompile/bytecode/handlers/arithmetic.py b/pkgs/MagiCompiler/magi_compiler/magi_depyf/decompile/bytecode/handlers/arithmetic.py new file mode 100644 index 0000000000000000000000000000000000000000..ae5c6951f8141fddcaf1a887a6a5af0b7155a0eb --- /dev/null +++ b/pkgs/MagiCompiler/magi_compiler/magi_depyf/decompile/bytecode/handlers/arithmetic.py @@ -0,0 +1,144 @@ +# Copyright (c) 2026 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. + +"""Handlers for unary, binary, inplace, and comparison operations.""" + +from __future__ import annotations + +from ..decompile_context import DecompileContext +from ..handler_registry import registry +from ..instruction import Instruction +from ..source_emitter import SourceEmitter + +_reg = registry.register + +# ── Unary ───────────────────────────────────────────────────────────────── + +_UNARY_SYMBOLS = {"UNARY_NEGATIVE": "-", "UNARY_POSITIVE": "+", "UNARY_INVERT": "~", "UNARY_NOT": "not"} + + +@_reg(*_UNARY_SYMBOLS) +def _unary(em: SourceEmitter, inst: Instruction, ctx: DecompileContext) -> None: + em.push(f"({_UNARY_SYMBOLS[inst.opname]} {em.pop()})") + + +@_reg("GET_LEN") +def _get_len(em: SourceEmitter, inst: Instruction, ctx: DecompileContext) -> None: + em.push(f"len({em.peek()})") + + +# ── Binary ──────────────────────────────────────────────────────────────── + +_BINARY_SYMBOLS = { + "BINARY_MULTIPLY": "*", + "BINARY_ADD": "+", + "BINARY_SUBTRACT": "-", + "BINARY_TRUE_DIVIDE": "/", + "BINARY_FLOOR_DIVIDE": "//", + "BINARY_MODULO": "%", + "BINARY_POWER": "**", + "BINARY_AND": "&", + "BINARY_OR": "|", + "BINARY_XOR": "^", + "BINARY_LSHIFT": "<<", + "BINARY_RSHIFT": ">>", + "BINARY_MATRIX_MULTIPLY": "@", +} + + +@_reg(*_BINARY_SYMBOLS) +def _binary(em: SourceEmitter, inst: Instruction, ctx: DecompileContext) -> None: + rhs = em.pop() + lhs = em.pop() + em.push(f"({lhs} {_BINARY_SYMBOLS[inst.opname]} {rhs})") + + +@_reg("BINARY_SUBSCR") +def _subscr(em: SourceEmitter, inst: Instruction, ctx: DecompileContext) -> None: + rhs = em.pop() + lhs = em.pop() + em.push(f"{lhs}[{rhs}]") + + +@_reg("BINARY_SLICE") +def _slice(em: SourceEmitter, inst: Instruction, ctx: DecompileContext) -> None: + end = em.pop() + start = em.pop() + container = em.pop() + em.push(f"{container}[{start}:{end}]") + + +# ── Inplace ─────────────────────────────────────────────────────────────── + +_INPLACE_SYMBOLS = { + "INPLACE_MULTIPLY": "*", + "INPLACE_ADD": "+", + "INPLACE_SUBTRACT": "-", + "INPLACE_TRUE_DIVIDE": "/", + "INPLACE_FLOOR_DIVIDE": "//", + "INPLACE_MODULO": "%", + "INPLACE_POWER": "**", + "INPLACE_AND": "&", + "INPLACE_OR": "|", + "INPLACE_XOR": "^", + "INPLACE_LSHIFT": "<<", + "INPLACE_RSHIFT": ">>", + "INPLACE_MATRIX_MULTIPLY": "@", +} + + +@_reg(*_INPLACE_SYMBOLS) +def _inplace(em: SourceEmitter, inst: Instruction, ctx: DecompileContext) -> None: + rhs = em.pop() + lhs = em.pop() + em.emit(f"{lhs} {_INPLACE_SYMBOLS[inst.opname]}= {rhs}") + em.push(lhs) + + +@_reg("BINARY_OP") +def _binary_op(em: SourceEmitter, inst: Instruction, ctx: DecompileContext) -> None: + """Python 3.12+ unified BINARY_OP.""" + rhs = em.pop() + lhs = em.pop() + if "=" in inst.argrepr: + em.emit(f"{lhs} {inst.argrepr} {rhs}") + em.push(lhs) + else: + em.push(f"({lhs} {inst.argrepr} {rhs})") + + +# ── Comparison ──────────────────────────────────────────────────────────── + + +@_reg("COMPARE_OP") +def _compare(em: SourceEmitter, inst: Instruction, ctx: DecompileContext) -> None: + rhs = em.pop() + lhs = em.pop() + em.push(f"({lhs} {inst.argval} {rhs})") + + +@_reg("IS_OP") +def _is_op(em: SourceEmitter, inst: Instruction, ctx: DecompileContext) -> None: + rhs = em.pop() + lhs = em.pop() + op = "is" if inst.argval == 0 else "is not" + em.push(f"({lhs} {op} {rhs})") + + +@_reg("CONTAINS_OP") +def _contains(em: SourceEmitter, inst: Instruction, ctx: DecompileContext) -> None: + rhs = em.pop() + lhs = em.pop() + op = "in" if inst.argval == 0 else "not in" + em.push(f"({lhs} {op} {rhs})") diff --git a/pkgs/MagiCompiler/magi_compiler/magi_depyf/decompile/bytecode/handlers/calls.py b/pkgs/MagiCompiler/magi_compiler/magi_depyf/decompile/bytecode/handlers/calls.py new file mode 100644 index 0000000000000000000000000000000000000000..cf14fb4612530ccf1e8bd738b9010ed90384d0a8 --- /dev/null +++ b/pkgs/MagiCompiler/magi_compiler/magi_depyf/decompile/bytecode/handlers/calls.py @@ -0,0 +1,200 @@ +# Copyright (c) 2026 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. + +"""Handlers for function-call and function-creation opcodes.""" + +from __future__ import annotations + +import sys +from typing import Optional + +from ..decompile_context import DecompileContext +from ..handler_registry import registry +from ..instruction import Instruction +from ..source_emitter import SourceEmitter + +_reg = registry.register + + +@_reg("KW_NAMES") +def _kw_names(em: SourceEmitter, inst: Instruction, ctx: DecompileContext) -> None: + # Python 3.11+ instruction that passes keyword argument names to the subsequent CALL. + # inst.arg indexes into co_consts for the key-name tuple, e.g. ('y', 'z'). + # Push repr so it becomes the string "('y', 'z')"; the CALL handler later eval()s it back to a tuple. + names = ctx.code.co_consts[inst.arg] + em.push(repr(names)) + + +@_reg("CALL") +def _call(em: SourceEmitter, inst: Instruction, ctx: DecompileContext) -> None: + """Python 3.11+ unified CALL. + + 3.12 stack layout: [NULL, callable, arg0, ..., argN-1] (KW_NAMES precedes) + 3.11 stack layout: [NULL, callable, arg0, ..., argN-1] (KW_NAMES → PRECALL → CALL) + """ + # Check whether KW_NAMES precedes CALL (indicating keyword arguments exist). + # 3.12: KW_NAMES → CALL; 3.11: KW_NAMES → PRECALL → CALL + preceding = [x for x in ctx.instructions if x.offset < inst.offset] + has_kw = False + if preceding: + if preceding[-1].opname == "KW_NAMES" or ( + len(preceding) > 1 + and preceding[-2].opname == "KW_NAMES" + and preceding[-1].opname == "PRECALL" # 3.11 transitional opcode, removed in 3.12 + ): + has_kw = True + + kw_names: tuple = () + if has_kw: + kw_names = eval(em.pop()) # retrieve the tuple stored by KW_NAMES from the stack + args = [em.pop() for _ in range(inst.argval)][::-1] + pos_args = args[: len(args) - len(kw_names)] + kw_args = args[len(args) - len(kw_names) :] + kwcalls = [f"{n}={v}" for n, v in zip(kw_names, kw_args)] + func = em.pop() + # 3.11+ PUSH_NULL / LOAD_GLOBAL(NULL+name) pushes a NULL sentinel before the call. + # After popping the callable, the top of stack may be NULL (represented as None); clear it. + if em.stack_size and em.peek() is None: + em.pop() + # GET_ITER produces "iter(x)"; if func happens to be "iter(x)" it is actually an argument + # (e.g. in the next(iter(x)) pattern), and the real callable is further down the stack. + if "iter(" in str(func): + pos_args = [func] + func = em.pop() + em.push(f"{func}({', '.join(pos_args + kwcalls)})") + # replace_tos_with_temp: the call result may be referenced multiple times (assignment, passing, + # method call), so store it in a temp to avoid repeated evaluation and side effects. + em.replace_tos_with_temp() + + +@_reg("CALL_FUNCTION", "CALL_METHOD") +def _call_legacy(em: SourceEmitter, inst: Instruction, ctx: DecompileContext) -> None: + """CALL_FUNCTION / CALL_METHOD (Python ≤3.10).""" + args = [em.pop() for _ in range(inst.argval)][::-1] + func = em.pop() + em.push(f"{func}({', '.join(args)})") + em.replace_tos_with_temp() + + +@_reg("CALL_FUNCTION_KW") +def _call_function_kw(em: SourceEmitter, inst: Instruction, ctx: DecompileContext) -> None: + kw_args = eval(em.pop()) + kw_vals = [em.pop() for _ in range(len(kw_args))] + kw_vals.reverse() + kwcalls = [f"{n}={v}" for n, v in zip(kw_args, kw_vals)] + pos_args = [em.pop() for _ in range(inst.argval - len(kw_args))][::-1] + func = em.pop() + em.push(f"{func}({', '.join(pos_args + kwcalls)})") + em.replace_tos_with_temp() + + +@_reg("CALL_FUNCTION_EX") +def _call_function_ex(em: SourceEmitter, inst: Instruction, ctx: DecompileContext) -> None: + # 3.11+ stack: [NULL, func, args (, kwargs)] + # After popping func, clear the NULL sentinel before pushing the result + if inst.argval == 0: + a = em.pop() + f = em.pop() + if em.stack_size and em.peek() is None: + em.pop() + em.push(f"{f}(*{a})") + elif inst.argval == 1: + kw = em.pop() + a = em.pop() + f = em.pop() + if em.stack_size and em.peek() is None: + em.pop() + em.push(f"{f}(*{a}, **{kw})") + em.replace_tos_with_temp() + + +@_reg("CALL_INTRINSIC_1") +def _intrinsic_1(em: SourceEmitter, inst: Instruction, ctx: DecompileContext) -> None: + """Python 3.12 instruction replacing some internal C-level calls. + argrepr identifies the specific operation, e.g. INTRINSIC_PRINT, INTRINSIC_UNARY_POSITIVE. + Most are compiler-internal operations (import *, typealias) rarely triggered by user code.""" + _SKIP = { + "INTRINSIC_1_INVALID", + "INTRINSIC_IMPORT_STAR", + "INTRINSIC_STOPITERATION_ERROR", + "INTRINSIC_ASYNC_GEN_WRAP", + "INTRINSIC_TYPEVAR", + "INTRINSIC_PARAMSPEC", + "INTRINSIC_TYPEVARTUPLE", + "INTRINSIC_SUBSCRIPT_GENERIC", + "INTRINSIC_TYPEALIAS", + } + if inst.argrepr in _SKIP: + return + if inst.argrepr == "INTRINSIC_PRINT": + em.emit(f"print({em.pop()})") + em.push("None") + elif inst.argrepr == "INTRINSIC_UNARY_POSITIVE": + em.set_at(0, f"+{em.peek()}") + elif inst.argrepr == "INTRINSIC_LIST_TO_TUPLE": + em.push(f"tuple({em.pop()})") + + +# ── MAKE_FUNCTION ───────────────────────────────────────────────────────── + + +@_reg("MAKE_FUNCTION") +def _make_function(em: SourceEmitter, inst: Instruction, ctx: DecompileContext) -> Optional[int]: + """Handle bytecode for def inner(...) and lambda. + + Bytecode: LOAD_CONST → MAKE_FUNCTION → STORE_FAST name + The handler recursively decompiles the inner code object and emits the full def statement. + """ + if sys.version_info < (3, 11): + # 3.10: qualified_name string is still on the stack + qual_name = em.pop() + try: + qual_name = eval(qual_name) + except Exception: + pass + func_name = qual_name.split(".")[-1] + if "<" in func_name: # , , etc. — invalid identifiers + em.emit(f'"original function name {func_name} is illegal, use a temp name."') + func_name = em.make_temp() + else: + func_name = em.make_temp() + + code = em.pop() # inner CodeType object pushed by LOAD_CONST + # argval bit flags indicate whether extra function components remain on the stack + if inst.argval & 0x08: + em.pop() # closure tuple (cell references for freevars) + if inst.argval & 0x04: + em.pop() # annotations dict + if inst.argval & 0x02: + em.pop() # keyword-only defaults tuple + if inst.argval & 0x01: + em.pop() # positional defaults tuple + + # If the next instruction is STORE_FAST, use the target variable name as the function name + this_idx = ctx.index_of(inst.offset) + immediately_used = False + if ctx.instructions[this_idx + 1].opname == "STORE_FAST": + func_name = ctx.instructions[this_idx + 1].argval + immediately_used = True + + # Recurse: create a new Decompiler instance for the inner code object + from ...decompiler import Decompiler + + inner = Decompiler(code).decompile(overwrite_fn_name=func_name) + em.emit_raw(inner) + + if immediately_used: + return this_idx + 2 # skip the MAKE_FUNCTION + STORE_FAST pair + em.push(func_name) # not immediately assigned — push onto stack for later use (e.g. as an argument) + return None diff --git a/pkgs/MagiCompiler/magi_compiler/magi_depyf/decompile/bytecode/handlers/containers.py b/pkgs/MagiCompiler/magi_compiler/magi_depyf/decompile/bytecode/handlers/containers.py new file mode 100644 index 0000000000000000000000000000000000000000..f57266074132e2a6402eed66edaa97541b2ceeb9 --- /dev/null +++ b/pkgs/MagiCompiler/magi_compiler/magi_depyf/decompile/bytecode/handlers/containers.py @@ -0,0 +1,200 @@ +# Copyright (c) 2026 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. + +"""Handlers for BUILD_*, UNPACK_*, LIST_EXTEND/APPEND, SET_ADD, MAP_ADD, +FORMAT_VALUE, and BUILD_SLICE / BUILD_STRING.""" + +from __future__ import annotations + +import sys + +from ..decompile_context import DecompileContext +from ..handler_registry import registry +from ..instruction import Instruction +from ..source_emitter import SourceEmitter + +_reg = registry.register + + +# ── BUILD tuple / list / set ────────────────────────────────────────────── + + +def _safe_str(val) -> str: + """Convert a stack value to string, handling None sentinels from PUSH_NULL.""" + return "None" if val is None else str(val) + + +@_reg("BUILD_TUPLE", "BUILD_TUPLE_UNPACK", "BUILD_TUPLE_UNPACK_WITH_CALL") +def _build_tuple(em: SourceEmitter, inst: Instruction, ctx: DecompileContext) -> None: + args = [_safe_str(em.pop()) for _ in range(inst.argval)][::-1] + if "UNPACK" in inst.opname: + args = [f"*{a}" for a in args] + em.push(f"({args[0]},)" if inst.argval == 1 else f"({', '.join(args)})") + + +@_reg("BUILD_LIST", "BUILD_LIST_UNPACK") +def _build_list(em: SourceEmitter, inst: Instruction, ctx: DecompileContext) -> None: + args = [_safe_str(em.pop()) for _ in range(inst.argval)][::-1] + if "UNPACK" in inst.opname: + args = [f"*{a}" for a in args] + em.push(f"[{', '.join(args)}]") + em.replace_tos_with_temp() + + +@_reg("BUILD_SET", "BUILD_SET_UNPACK") +def _build_set(em: SourceEmitter, inst: Instruction, ctx: DecompileContext) -> None: + if inst.argval == 0: + em.push("set()") + else: + args = [em.pop() for _ in range(inst.argval)][::-1] + if "UNPACK" in inst.opname: + args = [f"*{a}" for a in args] + em.push(f"{{{', '.join(args)}}}") + em.replace_tos_with_temp() + + +# ── BUILD map ───────────────────────────────────────────────────────────── + + +@_reg("BUILD_MAP") +def _build_map(em: SourceEmitter, inst: Instruction, ctx: DecompileContext) -> None: + items = [em.pop() for _ in range(inst.argval * 2)][::-1] + keys, vals = items[::2], items[1::2] + em.push(f"{{{', '.join(f'{k}: {v}' for k, v in zip(keys, vals))}}}") + em.replace_tos_with_temp() + + +@_reg("BUILD_MAP_UNPACK", "BUILD_MAP_UNPACK_WITH_CALL") +def _build_map_unpack(em: SourceEmitter, inst: Instruction, ctx: DecompileContext) -> None: + if inst.argval == 0: + em.push("dict()") + else: + args = [em.pop() for _ in range(inst.argval)][::-1] + em.push(f"{{{', '.join(f'**{a}' for a in args)}}}") + em.replace_tos_with_temp() + + +@_reg("BUILD_CONST_KEY_MAP") +def _const_key_map(em: SourceEmitter, inst: Instruction, ctx: DecompileContext) -> None: + keys = eval(em.pop()) + vals = [em.pop() for _ in range(inst.argval)][::-1] + em.push(f"{{{', '.join(f'{k!r}: {v}' for k, v in zip(keys, vals))}}}") + em.replace_tos_with_temp() + + +@_reg("BUILD_STRING") +def _build_string(em: SourceEmitter, inst: Instruction, ctx: DecompileContext) -> None: + args = [em.pop() for _ in range(inst.argval)][::-1] + em.push(" + ".join(args)) + + +@_reg("BUILD_SLICE") +def _build_slice(em: SourceEmitter, inst: Instruction, ctx: DecompileContext) -> None: + tos = em.pop() + tos1 = em.pop() + if inst.argval == 2: + em.push(f"slice({tos1}, {tos})") + elif inst.argval == 3: + tos2 = em.pop() + em.push(f"slice({tos2}, {tos1}, {tos})") + + +@_reg("LIST_TO_TUPLE") +def _list_to_tuple(em: SourceEmitter, inst: Instruction, ctx: DecompileContext) -> None: + em.push(f"tuple({em.pop()})") + + +# ── Mutating container ops ──────────────────────────────────────────────── + + +@_reg("LIST_EXTEND") +def _list_extend(em: SourceEmitter, inst: Instruction, ctx: DecompileContext) -> None: + values = em.pop() + temp = em.replace_tos_with_temp(depth=inst.argval) + em.emit(f"{temp}.extend({values})") + + +@_reg("LIST_APPEND") +def _list_append(em: SourceEmitter, inst: Instruction, ctx: DecompileContext) -> None: + argval = inst.argval if inst.argval != 1 else 2 + container = em.stack[-argval] + value = em.pop() + em.emit(f"{container}.append({value})") + + +@_reg("SET_UPDATE", "DICT_UPDATE", "DICT_MERGE") +def _generic_update(em: SourceEmitter, inst: Instruction, ctx: DecompileContext) -> None: + assert inst.argval == 1, "Only tested for argval==1" + values = em.pop() + temp = em.replace_tos_with_temp() + em.emit(f"{temp}.update({values})") + + +@_reg("SET_ADD") +def _set_add(em: SourceEmitter, inst: Instruction, ctx: DecompileContext) -> None: + argval = inst.argval if inst.argval != 1 else 2 + container = em.stack[-argval] + value = em.pop() + em.emit(f"{container}.add({value})") + + +@_reg("MAP_ADD") +def _map_add(em: SourceEmitter, inst: Instruction, ctx: DecompileContext) -> None: + container = em.stack[-inst.argval - 1] + if sys.version_info >= (3, 8): + value = em.pop() + key = em.pop() + else: + key = em.pop() + value = em.pop() + em.emit(f"{container}.__setitem__({key}, {value})") + + +# ── Unpack ──────────────────────────────────────────────────────────────── + + +@_reg("UNPACK_SEQUENCE") +def _unpack_seq(em: SourceEmitter, inst: Instruction, ctx: DecompileContext) -> None: + varname = em.pop() + tmps = [em.make_temp() for _ in range(inst.argval)] + em.emit("".join(f"{t}, " for t in tmps) + f"= {varname}") + for t in reversed(tmps): + em.push(t) + + +@_reg("UNPACK_EX") +def _unpack_ex(em: SourceEmitter, inst: Instruction, ctx: DecompileContext) -> None: + varname = em.pop() + tmps = [em.make_temp() for _ in range(inst.argval)] + star = em.make_temp() + em.emit(f"{', '.join(tmps)}, *{star} = {varname}") + em.push(star) + for t in reversed(tmps): + em.push(t) + + +# ── Format ──────────────────────────────────────────────────────────────── + + +@_reg("FORMAT_VALUE") +def _format_value(em: SourceEmitter, inst: Instruction, ctx: DecompileContext) -> None: + func, spec = inst.argval + if spec: + form_spec = em.pop() + value = em.pop() + em.push(f"format({value}, {form_spec})") + else: + value = em.pop() + fn = str if func is None else func + em.push(f"{fn.__name__}({value})") diff --git a/pkgs/MagiCompiler/magi_compiler/magi_depyf/decompile/bytecode/handlers/control_flow.py b/pkgs/MagiCompiler/magi_compiler/magi_depyf/decompile/bytecode/handlers/control_flow.py new file mode 100644 index 0000000000000000000000000000000000000000..904951f278d8802c45d6c6bc3af8c6ecfda48835 --- /dev/null +++ b/pkgs/MagiCompiler/magi_compiler/magi_depyf/decompile/bytecode/handlers/control_flow.py @@ -0,0 +1,273 @@ +# Copyright (c) 2026 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. + +"""Handlers for control-flow opcodes: jumps, if/else, for, return, yield, raise.""" + +from __future__ import annotations + +from typing import Optional + +from ..decompile_context import DecompileContext +from ..handler_registry import registry +from ..instruction import Instruction +from ..source_emitter import LoopContext, SourceEmitter + +_reg = registry.register + + +# ── Simple returns / yield / raise ──────────────────────────────────────── + + +@_reg("RETURN_VALUE") +def _return_value(em: SourceEmitter, inst: Instruction, ctx: DecompileContext) -> None: + em.emit(f"return {em.peek()}") + em.pop() + + +@_reg("RETURN_CONST") +def _return_const(em: SourceEmitter, inst: Instruction, ctx: DecompileContext) -> None: + em.emit(f"return {repr(inst.argval)}") + + +@_reg("YIELD_VALUE") +def _yield_value(em: SourceEmitter, inst: Instruction, ctx: DecompileContext) -> None: + import sys + + if sys.version_info >= (3, 12): + raise NotImplementedError("YIELD_VALUE is not supported in Python 3.12+") + em.emit(f"yield {em.peek()}") + + +@_reg("RETURN_GENERATOR") +def _return_generator(em: SourceEmitter, inst: Instruction, ctx: DecompileContext) -> None: + """Python 3.11+ generator function prologue. Each generator has its own stack frame; + RETURN_GENERATOR creates the generator object and returns it to the caller, + subsequent next(gen) resumes from RESUME. Push None as a placeholder during decompilation.""" + em.push(None) + + +@_reg("GEN_START") +def _gen_start(em: SourceEmitter, inst: Instruction, ctx: DecompileContext) -> None: + """Python 3.11 marks generator start (replaced by RESUME in 3.12).""" + assert inst.argval == 0, "Only generator expression is supported" + + +@_reg("RAISE_VARARGS") +def _raise_varargs(em: SourceEmitter, inst: Instruction, ctx: DecompileContext) -> None: + if inst.argval == 0: + em.emit("raise") + elif inst.argval == 1: + em.emit(f"raise {em.pop()}") + elif inst.argval == 2: + tos = em.pop() + tos1 = em.pop() + em.emit(f"raise {tos1} from {tos}") + + +@_reg("BREAK_LOOP") +def _break_loop(em: SourceEmitter, inst: Instruction, ctx: DecompileContext) -> None: + em.emit("break") + + +# ── Unconditional jumps ─────────────────────────────────────────────────── + + +@_reg("JUMP_ABSOLUTE") +@_reg("JUMP_FORWARD") +@_reg("JUMP_BACKWARD") +@_reg("JUMP_BACKWARD_NO_INTERRUPT") +def _abs_jump(em: SourceEmitter, inst: Instruction, ctx: DecompileContext) -> Optional[int]: + """Unconditional jump. Returns len(instructions) to make decompile_range stop immediately.""" + target = inst.jump_target_offset() + idx = ctx.index_of(target) + loop = em.loop + if loop is not None: + if idx >= loop.end_index: + em.emit("break") + return len(ctx.instructions) + if idx == loop.start_index: + em.emit("continue") + return len(ctx.instructions) + return idx + + +# ── Conditional jumps (if / else) ───────────────────────────────────────── + + +@_reg("POP_JUMP_IF_TRUE", "POP_JUMP_IF_FALSE") +@_reg("POP_JUMP_FORWARD_IF_TRUE", "POP_JUMP_FORWARD_IF_FALSE") +@_reg("POP_JUMP_BACKWARD_IF_TRUE", "POP_JUMP_BACKWARD_IF_FALSE") +@_reg("POP_JUMP_FORWARD_IF_NONE", "POP_JUMP_FORWARD_IF_NOT_NONE") +@_reg("POP_JUMP_BACKWARD_IF_NONE", "POP_JUMP_BACKWARD_IF_NOT_NONE") +@_reg("JUMP_IF_TRUE_OR_POP", "JUMP_IF_FALSE_OR_POP") +@_reg("POP_JUMP_IF_NOT_NONE", "POP_JUMP_IF_NONE") +def _jump_if(em: SourceEmitter, inst: Instruction, ctx: DecompileContext) -> Optional[int]: + """Decompile if/else structure. + + Standard if/else bytecode: + POP_JUMP_IF_FALSE else_start ← this_idx + (if-body) + JUMP_FORWARD after_else ← last instruction of if-body + >> else_start: ← jump_idx + (else-body) + >> after_else: ← merge point (end) + """ + + jump_offset = inst.jump_target_offset() + jump_idx = ctx.index_of(jump_offset) + this_idx = ctx.index_of(inst.offset) + + # ── Step 1: condition expression and branch stack state ── + cond = em.peek() + fall_stack = list(em.stack) + jump_stack = list(em.stack) + + if "IF_NOT_NONE" in inst.opname: + cond = f"({cond} is None)" + elif "IF_NONE" in inst.opname: + cond = f"({cond} is not None)" + elif "IF_TRUE" in inst.opname: + cond = f"(not {cond})" + else: + cond = f"{cond}" + + if "POP_JUMP" in inst.opname: + jump_stack.pop() + fall_stack.pop() + elif "OR_POP" in inst.opname: + fall_stack.pop() + + # ── Step 2: merge point candidate upper bounds ── + merge_upper_bounds = [len(ctx.instructions)] + if em.loop is not None: + merge_upper_bounds.append(em.loop.end_index) + + # ── Step 3: find "skip else" JUMPs in the if-body ── + def _is_forward_past_else(i: Instruction) -> bool: + return i.is_jump and i.jump_target_offset() >= jump_offset + + forward_targets = [i.jump_target_offset() for i in ctx.instructions[this_idx:jump_idx] if _is_forward_past_else(i)] + + # ── Step 4: compute merge point by case ── + if not forward_targets: + if jump_idx <= this_idx: + # Case C: backward jump (inside loop), emit if cond: continue + rev_cond = em.peek() + if "IF_NOT_NONE" in inst.opname: + rev_cond = f"({rev_cond} is not None)" + elif "IF_NONE" in inst.opname: + rev_cond = f"({rev_cond} is None)" + elif "IF_TRUE" in inst.opname: + rev_cond = f"{rev_cond}" + elif "IF_FALSE" in inst.opname: + rev_cond = f"(not {rev_cond})" + em.emit(f"if {rev_cond}:") + em.emit(em.indent("continue\n").rstrip("\n")) + return None + # Case B: both branches terminate with RETURN/RAISE + end = jump_idx + else: + # Case A: standard if/else, infer merge point from forward_targets + max_jump = max(forward_targets) + max_idx = ctx.index_of(max_jump) + all_targets = [i.jump_target_offset() for i in ctx.instructions[this_idx:max_idx] if _is_forward_past_else(i)] + max_idx = ctx.index_of(max(all_targets)) + + last = ctx.instructions[max_idx - 1] + if not ("RAISE" in last.opname or "RETURN" in last.opname or "STORE" in last.opname): + old = max_idx + while max_idx < len(ctx.instructions): + op = ctx.instructions[max_idx].opname + if "STORE" in op or "RETURN" in op: + max_idx += 1 + break + if ("JUMP" in op and max_idx > old) or "FOR_ITER" in op: + break + max_idx += 1 + + merge_upper_bounds.append(max_idx) + end = min(merge_upper_bounds) + + # ── Step 5: else-body end position (PR#91 fix) ── + else_end = end + if end == jump_idx and jump_idx < len(ctx.instructions): + last_if = ctx.instructions[jump_idx - 1] + if "RETURN" in last_if.opname or "RAISE" in last_if.opname: + else_end = len(ctx.instructions) + if em.loop is not None: + else_end = min(else_end, em.loop.end_index) + + # ── Step 6: decompile both branches ── + with em.fork(stack=fall_stack) as if_em: + ctx.decompile_range(this_idx + 1, jump_idx, if_em) + if_body = em.indent(if_em.get_source()) + if_end_stack = list(if_em.stack) + em.emit_raw(f"if {cond}:\n{if_body}") + + with em.fork(stack=jump_stack) as else_em: + ctx.decompile_range(jump_idx, else_end, else_em) + else_body = else_em.get_source() + if else_body: + em.emit_raw(f"else:\n{em.indent(else_body)}") + + em.stack[:] = if_end_stack + return else_end + + +# ── FOR_ITER ────────────────────────────────────────────────────────────── + + +@_reg("FOR_ITER") +def _for_iter(em: SourceEmitter, inst: Instruction, ctx: DecompileContext) -> Optional[int]: + """Decompile for loop. + + Bytecode layout (3.12): + FOR_ITER target ← get next value; jump to target when exhausted (END_FOR) + (loop body) + JUMP_BACKWARD for_iter ← normal back-jump (not continue) + >> target: END_FOR + + Loop body range excludes the trailing JUMP_BACKWARD to avoid emitting a spurious continue. + """ + start_idx = ctx.index_of(inst.offset) + end_idx = ctx.index_of(inst.jump_target_offset()) + + temp = em.make_temp() + iterator = em.pop() + em.push(temp) + + # Determine the actual end position of the loop body: + # if the instruction at end_idx is a back-jump to FOR_ITER, extend end_idx so + # the LoopContext boundary is correct (break needs to jump past end_idx) + if end_idx < len(ctx.instructions): + at_end = ctx.instructions[end_idx] + if at_end.is_jump and at_end.jump_target_offset() == inst.offset: + end_idx += 1 + + # Exclude the trailing JUMP_BACKWARD: it is the normal loop back-jump mechanism, not continue. + # Only JUMP_BACKWARDs in the middle of the loop body are continue (handled by _abs_jump). + body_end = end_idx + if body_end > start_idx + 1: + back_jump = ctx.instructions[body_end - 1] + if back_jump.is_jump and back_jump.jump_target_offset() == inst.offset: + body_end -= 1 + + loop = LoopContext(start_index=start_idx, end_index=end_idx) + with em.fork(stack=list(em.stack), loop=loop) as body_em: + ctx.decompile_range(start_idx + 1, body_end, body_em) + + body_src = em.indent(body_em.get_source()) + em.emit_raw(f"for {temp} in {iterator}:\n{body_src}") + em.stack[:] = body_em.stack + return end_idx diff --git a/pkgs/MagiCompiler/magi_compiler/magi_depyf/decompile/bytecode/handlers/load_store.py b/pkgs/MagiCompiler/magi_compiler/magi_depyf/decompile/bytecode/handlers/load_store.py new file mode 100644 index 0000000000000000000000000000000000000000..6438dc858b24e8f693c686eead723f15a0da9e77 --- /dev/null +++ b/pkgs/MagiCompiler/magi_compiler/magi_depyf/decompile/bytecode/handlers/load_store.py @@ -0,0 +1,262 @@ +# Copyright (c) 2026 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. + +"""Handlers for LOAD_*, STORE_*, DELETE_*, IMPORT_*, PUSH_NULL, GET_ITER.""" + +from __future__ import annotations + +from types import CodeType + +from ..decompile_context import DecompileContext +from ..handler_registry import registry +from ..instruction import Instruction +from ..source_emitter import SourceEmitter + +_reg = registry.register + + +# ── NOP / unsupported sentinels ────────────────────────────────────────── + + +@_reg("NOP", "RESUME", "EXTENDED_ARG", "SETUP_LOOP", "POP_BLOCK") +@_reg("PRECALL", "BEGIN_FINALLY", "END_FINALLY", "MAKE_CELL") +@_reg("RERAISE", "END_FOR", "COPY_FREE_VARS") +def _nop(em: SourceEmitter, inst: Instruction, ctx: DecompileContext) -> None: + pass + + +@_reg("GET_YIELD_FROM_ITER") +@_reg("POP_EXCEPT", "WITH_EXCEPT_START", "JUMP_IF_NOT_EXC_MATCH") +@_reg("CHECK_EG_MATCH", "PUSH_EXC_INFO", "PREP_RERAISE_STAR") +@_reg("WITH_CLEANUP_FINISH", "CALL_FINALLY", "POP_FINALLY") +@_reg("WITH_CLEANUP_START", "SETUP_EXCEPT", "CHECK_EXC_MATCH") +@_reg("CLEANUP_THROW") +@_reg("GET_AWAITABLE", "GET_AITER", "GET_ANEXT", "END_ASYNC_FOR") +@_reg("BEFORE_ASYNC_WITH", "SETUP_ASYNC_WITH", "SEND", "ASYNC_GEN_WRAP") +@_reg("CACHE") +@_reg("PRINT_EXPR", "COPY_DICT_WITHOUT_KEYS") +@_reg("IMPORT_STAR") +@_reg("YIELD_FROM", "SETUP_ANNOTATIONS", "LOAD_BUILD_CLASS") +@_reg("MATCH_MAPPING", "MATCH_SEQUENCE", "MATCH_KEYS", "MATCH_CLASS") +@_reg("CALL_INTRINSIC_2") +@_reg("SETUP_FINALLY", "SETUP_WITH", "BEFORE_WITH") +def _unsupported(em: SourceEmitter, inst: Instruction, ctx: DecompileContext) -> None: + from ...decompiler import DecompilationError + + raise DecompilationError(f"Unsupported opcode: {inst.opname}", instruction=inst) + + +# ── LOAD instructions ──────────────────────────────────────────────────── + + +@_reg("LOAD_CONST") +def _load_const(em: SourceEmitter, inst: Instruction, ctx: DecompileContext) -> None: + """Load a constant. Branches: can_repr → direct repr / type → importlib / + torch prefix → import torch / CodeType → push as-is for MAKE_FUNCTION.""" + can_repr = False + try: + can_repr = eval(repr(inst.argval)) == inst.argval + except BaseException: + pass + if can_repr: + em.push(repr(inst.argval)) + elif isinstance(inst.argval, type): + module = inst.argval.__module__ + name = inst.argval.__name__ + em.emit("import importlib") + tmp = em.make_temp() + em.emit(f'{tmp} = importlib.import_module("{module}").{name}') + em.push(tmp) + elif inst.argrepr.startswith("torch."): + em.emit("import torch") + tmp = em.make_temp() + em.emit(f"{tmp} = {inst.argval}") + em.push(tmp) + elif isinstance(inst.argval, CodeType): + em.push(inst.argval) + else: + from ...decompiler import DecompilationError + + raise DecompilationError( + f"LOAD_CONST: cannot represent co_consts[{inst.arg}] = {repr(inst.argval)!r} " + f"(type {type(inst.argval).__name__}) as source code", + instruction=inst, + ) + + +@_reg("LOAD_FAST", "LOAD_FAST_CHECK") +@_reg("LOAD_GLOBAL", "LOAD_DEREF", "LOAD_NAME") +@_reg("LOAD_CLASSDEREF", "LOAD_CLOSURE") +def _generic_load(em: SourceEmitter, inst: Instruction, ctx: DecompileContext) -> None: + """Generic load. 3.11+ LOAD_GLOBAL argrepr "NULL + name" pushes a NULL sentinel first. + Python <3.12 comprehension parameter name ".0" is replaced with "comp_arg_0".""" + if "NULL + " in inst.argrepr: + em.push(None) + if inst.argrepr.startswith("."): + em.push(inst.argval.replace(".", "comp_arg_")) + else: + em.push(inst.argval) + + +# Python 3.12 comprehension variable protection: LOAD_FAST_AND_CLEAR saves old value + STORE_FAST restores. +# During decompilation, temp variables used for loops don't need save/restore; push a sentinel so STORE_FAST skips. +_CLEAR_SENTINEL = object() + + +@_reg("LOAD_FAST_AND_CLEAR") +def _load_fast_and_clear(em: SourceEmitter, inst: Instruction, ctx: DecompileContext) -> None: + em.push(_CLEAR_SENTINEL) + + +@_reg("LOAD_LOCALS") +def _load_locals(em: SourceEmitter, inst: Instruction, ctx: DecompileContext) -> None: + """3.12 class body: locals() returns a new dict snapshot, cached in a temp to avoid repeated calls.""" + em.push("locals()") + em.replace_tos_with_temp() + + +@_reg("LOAD_FROM_DICT_OR_GLOBALS", "LOAD_FROM_DICT_OR_DEREF") +def _load_from_dict(em: SourceEmitter, inst: Instruction, ctx: DecompileContext) -> None: + """3.12 class body: look up in locals dict first, fall back to globals if not found.""" + tos = em.pop() + em.push(f"{tos}[{inst.argval}] if '{inst.argval}' in {tos} else {inst.argval}") + em.replace_tos_with_temp() + + +@_reg("LOAD_ATTR") +def _load_attr(em: SourceEmitter, inst: Instruction, ctx: DecompileContext) -> None: + """Attribute access. isidentifier() checks if the attr name is valid; if not, use getattr().""" + lhs = str(em.pop()) + rhs = inst.argval + em.push(f"{lhs}.{rhs}" if rhs.isidentifier() else f"getattr({lhs}, {rhs!r})") + + +@_reg("LOAD_SUPER_ATTR") +def _load_super_attr(em: SourceEmitter, inst: Instruction, ctx: DecompileContext) -> None: + self_obj = em.pop() + cls_obj = em.pop() + super_obj = em.pop() + em.push(f"{super_obj}({cls_obj}, {self_obj}).{inst.argval}") + em.replace_tos_with_temp() + + +@_reg("LOAD_METHOD") +def _load_method(em: SourceEmitter, inst: Instruction, ctx: DecompileContext) -> None: + em.push(f"{em.pop()}.{inst.argval}") + + +@_reg("LOAD_ASSERTION_ERROR") +def _load_assertion_error(em: SourceEmitter, inst: Instruction, ctx: DecompileContext) -> None: + em.push("AssertionError") + + +@_reg("PUSH_NULL") +def _push_null(em: SourceEmitter, inst: Instruction, ctx: DecompileContext) -> None: + """3.11+ pushes a NULL sentinel before function calls; the CALL handler will clear it.""" + em.push(None) + + +@_reg("GET_ITER") +def _get_iter(em: SourceEmitter, inst: Instruction, ctx: DecompileContext) -> None: + em.push(f"iter({em.pop()})") + + +# ── STORE instructions ─────────────────────────────────────────────────── + + +@_reg("STORE_FAST", "STORE_GLOBAL", "STORE_DEREF", "STORE_NAME") +def _generic_store(em: SourceEmitter, inst: Instruction, ctx: DecompileContext) -> None: + """Generic store. Skips _CLEAR_SENTINEL and self-assignment, protects variable names on the stack that are about to be overwritten.""" + left = inst.argval + right = em.pop() + if right is _CLEAR_SENTINEL: + return + if left != right: + if isinstance(left, str) and left in em.stack: + tmp = em.make_temp() + em.emit(f"{tmp} = {left}") + em.stack[:] = [tmp if x == left else x for x in em.stack] + em.emit(f"{left} = {right}") + + +@_reg("STORE_SUBSCR") +def _store_subscr(em: SourceEmitter, inst: Instruction, ctx: DecompileContext) -> None: + index = em.pop() + obj = em.pop() + value = em.pop() + em.emit(f"{obj}[{index}] = {value}") + + +@_reg("STORE_SLICE") +def _store_slice(em: SourceEmitter, inst: Instruction, ctx: DecompileContext) -> None: + end = em.pop() + start = em.pop() + container = em.pop() + value = em.pop() + em.emit(f"{container}[{start}:{end}] = {value}") + + +@_reg("STORE_ATTR") +def _store_attr(em: SourceEmitter, inst: Instruction, ctx: DecompileContext) -> None: + obj = em.pop() + value = em.pop() + em.emit(f"{obj}.{inst.argval} = {value}") + + +# ── DELETE instructions ────────────────────────────────────────────────── + + +@_reg("DELETE_SUBSCR") +def _delete_subscr(em: SourceEmitter, inst: Instruction, ctx: DecompileContext) -> None: + index = em.pop() + obj = em.pop() + if f"{obj}[{index}]" not in em.stack: + em.emit(f"del {obj}[{index}]") + + +@_reg("DELETE_NAME", "DELETE_GLOBAL", "DELETE_DEREF") +def _generic_delete(em: SourceEmitter, inst: Instruction, ctx: DecompileContext) -> None: + em.emit(f"del {inst.argval}") + + +@_reg("DELETE_FAST") +def _delete_fast(em: SourceEmitter, inst: Instruction, ctx: DecompileContext) -> None: + """Dynamo cleans up temp variables; no explicit del needed after decompilation.""" + pass + + +@_reg("DELETE_ATTR") +def _delete_attr(em: SourceEmitter, inst: Instruction, ctx: DecompileContext) -> None: + em.emit(f"del {em.pop()}.{inst.argval}") + + +# ── IMPORT instructions ────────────────────────────────────────────────── + + +@_reg("IMPORT_NAME") +def _import_name(em: SourceEmitter, inst: Instruction, ctx: DecompileContext) -> None: + """import os.path → binds 'os' (top-level module), accesses submodules via os.path.sep.""" + name = inst.argval.split(".")[0] + fromlist = em.pop() + level = em.pop() + em.emit(f"{name} = __import__({inst.argval!r}, fromlist={fromlist}, level={level})") + em.push(name) + + +@_reg("IMPORT_FROM") +def _import_from(em: SourceEmitter, inst: Instruction, ctx: DecompileContext) -> None: + name = inst.argval + module = em.peek() + em.emit(f"{name} = {module}.{name}") + em.push(name) diff --git a/pkgs/MagiCompiler/magi_compiler/magi_depyf/decompile/bytecode/handlers/stack_ops.py b/pkgs/MagiCompiler/magi_compiler/magi_depyf/decompile/bytecode/handlers/stack_ops.py new file mode 100644 index 0000000000000000000000000000000000000000..f3040f2440e0bdc78cbe4421683f84f5448c1f71 --- /dev/null +++ b/pkgs/MagiCompiler/magi_compiler/magi_depyf/decompile/bytecode/handlers/stack_ops.py @@ -0,0 +1,84 @@ +# Copyright (c) 2026 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. + +"""Handlers for stack-manipulation opcodes: ROT, SWAP, COPY, POP, DUP. + +See bytecode_explained.py §16 for details. +""" + +from __future__ import annotations + +from ..decompile_context import DecompileContext +from ..handler_registry import registry +from ..instruction import Instruction +from ..source_emitter import SourceEmitter + +_reg = registry.register + + +# ── ROT_N family (Python ≤3.10, replaced by SWAP/COPY in 3.11+) ─────────── + + +@_reg("ROT_N") +@_reg("ROT_TWO") +@_reg("ROT_THREE") +@_reg("ROT_FOUR") +def _rot_n(em: SourceEmitter, inst: Instruction, ctx: DecompileContext) -> None: + """Top-n stack rotation: [a, b, c] → [c, a, b] (n=3). + + ROT_TWO: a, b → b, a (swap, used for a, b = b, a) + ROT_THREE: a, b, c → c, a, b (3-element rotation) + ROT_FOUR: a, b, c, d → d, a, b, c (4-element rotation) + ROT_N: generic n-element rotation (argval = n) + """ + n = inst.argval if inst.opname == "ROT_N" else {"ROT_TWO": 2, "ROT_THREE": 3, "ROT_FOUR": 4}[inst.opname] + vals = em.stack[-n:] + em.stack[-n:] = [vals[-1]] + vals[:-1] + + +@_reg("SWAP") +def _swap(em: SourceEmitter, inst: Instruction, ctx: DecompileContext) -> None: + """Python 3.11+: swap stack[-1] and stack[-n].""" + n = inst.argval + em.stack[-1], em.stack[-n] = em.stack[-n], em.stack[-1] + + +@_reg("COPY") +def _copy(em: SourceEmitter, inst: Instruction, ctx: DecompileContext) -> None: + """Python 3.11+: copy stack[-n] to top of stack (COPY 1 = DUP_TOP).""" + n = inst.argval + if n == 0: + return + em.push(em.stack[-1 - (n - 1)]) + + +@_reg("POP_TOP") +def _pop_top(em: SourceEmitter, inst: Instruction, ctx: DecompileContext) -> None: + if em.stack_size > 0: + em.pop() + + +@_reg("DUP_TOP") +def _dup_top(em: SourceEmitter, inst: Instruction, ctx: DecompileContext) -> None: + """Python ≤3.10: duplicate top of stack. Replaced by COPY 1 in 3.11+.""" + em.push(em.peek()) + + +@_reg("DUP_TOP_TWO") +def _dup_top_two(em: SourceEmitter, inst: Instruction, ctx: DecompileContext) -> None: + """Python ≤3.10: duplicate top two stack items. Replaced by two COPYs in 3.11+.""" + tos = em.peek(0) + tos1 = em.peek(1) + em.push(tos1) + em.push(tos) diff --git a/pkgs/MagiCompiler/magi_compiler/magi_depyf/decompile/bytecode/instruction.py b/pkgs/MagiCompiler/magi_compiler/magi_depyf/decompile/bytecode/instruction.py new file mode 100644 index 0000000000000000000000000000000000000000..7315adf0a0bed3e5d7477d0e5859974a5b1495b8 --- /dev/null +++ b/pkgs/MagiCompiler/magi_compiler/magi_depyf/decompile/bytecode/instruction.py @@ -0,0 +1,129 @@ +# Copyright (c) 2026 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. + +"""Enhanced Instruction dataclass with rich querying properties.""" + +from __future__ import annotations + +import dataclasses +import dis +import sys +from typing import Any, Optional + +_ALL_JUMP_OPCODES = frozenset(dis.hasjabs) | frozenset(dis.hasjrel) +_PY311 = sys.version_info >= (3, 11) + +_LOAD_OPCODES = frozenset(n for n in dis.opname if n.startswith("LOAD_") or n in ("PUSH_NULL", "GET_ITER")) +_STORE_OPCODES = frozenset(n for n in dis.opname if n.startswith("STORE_")) +_DELETE_OPCODES = frozenset(n for n in dis.opname if n.startswith("DELETE_")) + + +@dataclasses.dataclass +class Instruction: + """Mutable mirror of ``dis.Instruction`` with convenience queries. + + Unlike the stdlib version this is mutable so cleanup passes can + modify instructions in-place (e.g. NOP-ing unreachable bytecode). + """ + + opcode: int + opname: str + # arg: raw integer argument (the number in the bytecode), may be an index into co_consts/co_varnames + # argval: Python object resolved by the dis module (value of co_consts[arg], or a variable name string) + # argrepr: human-readable string of argval (e.g. "NULL + print", "to 20") + # See bytecode_explained.py §1 for details + arg: Optional[int] + argval: Any + argrepr: str + offset: Optional[int] = None + starts_line: Optional[int] = None + is_jump_target: bool = False + + # -- identity / hashing (by object id, not value) ---------------------- + + def __hash__(self) -> int: + return id(self) + + def __eq__(self, other: object) -> bool: + return self is other + + def __repr__(self) -> str: + return f"Instruction({self.opname}, offset={self.offset}, argval={self.argrepr!r})" + + # -- category queries --------------------------------------------------- + + @property + def is_load(self) -> bool: + return self.opname in _LOAD_OPCODES + + @property + def is_store(self) -> bool: + return self.opname in _STORE_OPCODES + + @property + def is_delete(self) -> bool: + return self.opname in _DELETE_OPCODES + + @property + def is_jump(self) -> bool: + return self.opcode in _ALL_JUMP_OPCODES + + @property + def is_conditional_jump(self) -> bool: + return self.is_jump and ("IF" in self.opname or "FOR_ITER" in self.opname) + + @property + def is_unconditional_jump(self) -> bool: + return self.is_jump and not self.is_conditional_jump + + @property + def is_return(self) -> bool: + return self.opname in ("RETURN_VALUE", "RETURN_CONST") + + @property + def is_nop(self) -> bool: + return self.opname == "NOP" + + # -- jump target -------------------------------------------------------- + + def jump_target_offset(self) -> Optional[int]: + """Return the absolute bytecode offset this instruction jumps to, + or ``None`` if it is not a jump instruction.""" + if not self.is_jump: + return None + if "to " in self.argrepr: + return int(self.argrepr.replace("to ", "").strip()) + if self.opcode in dis.hasjabs: + return self.argval + if self.opcode in dis.hasjrel: + return self.argval if _PY311 else self.offset + self.argval + return None + + # -- mutation helpers (for cleanup passes) ------------------------------ + + def nop_(self) -> None: + """In-place convert this instruction to a NOP.""" + self.opname = "NOP" + self.opcode = dis.opmap["NOP"] + self.arg = 0 + self.argval = 0 + self.argrepr = "" + self.is_jump_target = False + + # -- factory ------------------------------------------------------------ + + @staticmethod + def from_dis(i: dis.Instruction) -> "Instruction": + """Create from a stdlib ``dis.Instruction``.""" + return Instruction(i.opcode, i.opname, i.arg, i.argval, i.argrepr, i.offset, i.starts_line, i.is_jump_target) diff --git a/pkgs/MagiCompiler/magi_compiler/magi_depyf/decompile/bytecode/source_emitter.py b/pkgs/MagiCompiler/magi_compiler/magi_depyf/decompile/bytecode/source_emitter.py new file mode 100644 index 0000000000000000000000000000000000000000..26a2834784bb482b1e6d58f37dc6cf9ad114becf --- /dev/null +++ b/pkgs/MagiCompiler/magi_compiler/magi_depyf/decompile/bytecode/source_emitter.py @@ -0,0 +1,153 @@ +# Copyright (c) 2026 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. + +"""SourceEmitter: manages the evaluation stack and source-code emission. + +This replaces the bare ``DecompilerState`` (just ``source_code: str`` +and ``stack: list``) with a proper class that owns *all* mutable state +touched during decompilation, including the temp-variable counter +(instance-level, not class-level — thread-safe by design). +""" + +from __future__ import annotations + +import contextlib +import dataclasses +from typing import Any, Iterator, List, Optional + + +@dataclasses.dataclass +class LoopContext: + """Current loop boundaries, used for break/continue determination. + + Range semantics similar to range(start, end): + start_index: index of FOR_ITER itself (inclusive — part of the loop) + end_index: index of the first instruction outside the loop (exclusive — not part of the loop) + + break/continue determination (used in _abs_jump): + jump target >= end_index → break (jump out of the loop) + jump target == start_index → continue (jump back to loop head) + """ + + start_index: int # index of FOR_ITER, inclusive (part of the loop) + end_index: int # first instruction outside the loop, exclusive (not part of the loop) + + +class SourceEmitter: + """Stateful accumulator for the decompiler's output. + + Improvements over depyf's ``DecompilerState``: + * ``_temp_counter`` is **instance-level** (no thread-safety issues). + * Stack operations (``push / pop / peek``) are proper methods. + * ``emit()`` appends with a trailing newline automatically. + * ``fork()`` context-manager creates a child emitter for sub-blocks + (if-else branches, loop bodies, etc.) and returns it so the caller + can inspect the generated source and final stack. + """ + + def __init__(self, indent_size: int = 4, temp_prefix: str = "__temp_", *, _parent_counter: Optional[list] = None) -> None: + self._lines: List[str] = [] + self._stack: List[Any] = [] + self._indent_size = indent_size + self._temp_prefix = temp_prefix + # Share counter across forks so names are globally unique within + # one Decompiler invocation, but still instance-scoped. + self._counter: list = _parent_counter if _parent_counter is not None else [0] + self.loop: Optional[LoopContext] = None + + # -- source emission ---------------------------------------------------- + + def emit(self, line: str) -> None: + """Append *line* (with auto newline) to accumulated source.""" + self._lines.append(line + "\n") + + def emit_raw(self, text: str) -> None: + """Append pre-formatted *text* verbatim (e.g. nested function defs).""" + self._lines.append(text) + + def get_source(self) -> str: + return "".join(self._lines) + + # -- stack operations --------------------------------------------------- + + def push(self, value: Any) -> None: + self._stack.append(value) + + def pop(self) -> Any: + return self._stack.pop() + + def peek(self, depth: int = 0) -> Any: + """Return item at ``stack[-(depth+1)]`` without popping.""" + return self._stack[-(depth + 1)] + + def set_at(self, depth: int, value: Any) -> None: + """Set ``stack[-(depth+1)]`` to *value*.""" + self._stack[-(depth + 1)] = value + + @property + def stack(self) -> List[Any]: + """Direct access (for complex multi-item operations).""" + return self._stack + + @property + def stack_size(self) -> int: + return len(self._stack) + + # -- temp variables (instance-scoped counter) --------------------------- + + def make_temp(self) -> str: + """Return a unique temporary variable name.""" + self._counter[0] += 1 + return f"{self._temp_prefix}{self._counter[0]}" + + def replace_tos_with_temp(self, depth: int = 1) -> str: + """Replace ``stack[-depth]`` with a fresh temp, emitting the + assignment ``__temp_N = ``. Returns the temp name.""" + old = self._stack[-depth] + name = self.make_temp() + self.emit(f"{name} = {old}") + self._stack[-depth] = name + return name + + # -- sub-block forking -------------------------------------------------- + + @contextlib.contextmanager + def fork(self, stack: Optional[List[Any]] = None, loop: Optional[LoopContext] = None) -> Iterator["SourceEmitter"]: + """Create a child emitter for a sub-block (if-branch, loop body …). + + The child shares the temp counter but has its own ``_lines`` and + ``_stack``. If *loop* is ``None`` the parent's loop context is + inherited (matching depyf's ``new_state`` semantics). + + Usage:: + + with emitter.fork(stack=my_stack) as child: + decompile_range(start, end, child) + child_source = child.get_source() + child_final_stack = child.stack + """ + child = SourceEmitter(indent_size=self._indent_size, temp_prefix=self._temp_prefix, _parent_counter=self._counter) + child._stack = list(stack) if stack is not None else list(self._stack) + if loop is not None: + child.loop = loop + elif self.loop is not None: + child.loop = self.loop + yield child + + # -- indentation helpers ------------------------------------------------ + + def indent(self, text: str) -> str: + """Add one level of indentation to every line in *text*.""" + prefix = " " * self._indent_size + return "".join(prefix + line + "\n" for line in text.splitlines()) diff --git a/pkgs/MagiCompiler/magi_compiler/magi_depyf/decompile/decompiler.py b/pkgs/MagiCompiler/magi_compiler/magi_depyf/decompile/decompiler.py new file mode 100644 index 0000000000000000000000000000000000000000..013430eef96b1479babb497d349e7464221d232f --- /dev/null +++ b/pkgs/MagiCompiler/magi_compiler/magi_depyf/decompile/decompiler.py @@ -0,0 +1,230 @@ +# Copyright (c) 2026 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. + +"""Decompiler — the orchestrator that ties everything together. + +This module is the only place that coordinates ``SourceEmitter``, +``HandlerRegistry``, and ``DecompileContext``. +Individual handler functions never import from here (except for +``DecompilationError`` and recursive ``Decompiler`` usage in +``MAKE_FUNCTION``). +""" + +from __future__ import annotations + +import dis +import inspect +import os +from types import CodeType +from typing import Callable, List, Optional, Union + +# Force handler registration by importing the package. +import magi_compiler.magi_depyf.decompile.bytecode.handlers # noqa: F401 + +from .bytecode.decompile_context import DecompileContext +from .bytecode.handler_registry import registry +from .bytecode.instruction import Instruction +from .bytecode.source_emitter import SourceEmitter + +# --------------------------------------------------------------------------- +# Errors +# --------------------------------------------------------------------------- + + +class DecompilationError(Exception): + """Raised when decompilation fails. + + Carries optional ``instruction`` context so callers can produce + actionable error messages. + """ + + def __init__(self, message: str = "", *, instruction: Optional[Instruction] = None): + self.message = message + self.instruction = instruction + super().__init__(message) + + def __str__(self) -> str: + loc = "" + if self.instruction is not None: + loc = f" at {self.instruction}" + return f"DecompilationError: {self.message}{loc}" + + +# --------------------------------------------------------------------------- +# Signature builder (lives here — it's a Decompiler concern, not a util) +# --------------------------------------------------------------------------- + + +class SignatureBuilder: + """Build the ``def fn(args):`` header from a ``CodeType``.""" + + @staticmethod + def build(code: CodeType, overwrite_name: Optional[str] = None) -> str: + n = code.co_argcount + code.co_kwonlyargcount + names = [x.replace(".", "comp_arg_") if x.startswith(".") else x for x in code.co_varnames[:n]] + if code.co_flags & inspect.CO_VARARGS: + names.append("*" + code.co_varnames[n]) + n += 1 + if code.co_flags & inspect.CO_VARKEYWORDS: + names.append("**" + code.co_varnames[n]) + n += 1 + fn_name = overwrite_name or code.co_name + return f"def {fn_name}({', '.join(names)}):\n" + + +# --------------------------------------------------------------------------- +# Decompiler +# --------------------------------------------------------------------------- + + +class Decompiler: + """Decompile a ``CodeType`` into Python source code. + + Design differences from depyf's ``Decompiler``: + + * Handlers live in separate modules and receive + ``(emitter, inst, ctx)`` — they never reference this class. + * All mutable state is on ``SourceEmitter`` (instance-scoped counter). + * ``decompile_range`` is delegated *through* ``DecompileContext`` + so handlers can recurse without importing this class (except + ``MAKE_FUNCTION`` which needs a fresh ``Decompiler`` instance). + """ + + _TERMINATORS = frozenset({"RETURN_VALUE", "RETURN_CONST", "RAISE_VARARGS"}) + + def __init__(self, code: Union[CodeType, Callable]) -> None: + if callable(code) and not isinstance(code, CodeType): + code = _get_code_owner(code).__code__ + self.code: CodeType = code + self.instructions = [Instruction.from_dis(i) for i in dis.get_instructions(code)] + self._cleanup() + + # -- bytecode cleanup --------------------------------------------------- + + def _cleanup(self) -> None: + """Propagate line numbers and NOP dead code after unconditional exits.""" + cur: Optional[int] = None + for inst in self.instructions: + if inst.starts_line is not None: + cur = inst.starts_line + inst.starts_line = cur + + in_dead = False + for inst in self.instructions: + if in_dead: + if inst.is_jump_target: + in_dead = False + else: + inst.nop_() + elif inst.opname in self._TERMINATORS: + in_dead = True + + # -- core loop ---------------------------------------------------------- + + def decompile_range(self, start: int, end: int, emitter: SourceEmitter) -> None: + """Execute instruction handlers from *start* to *end* (exclusive).""" + idx = start + try: + while idx < end: + inst = self.instructions[idx] + handler = registry.get(inst.opname) + if handler is None: + raise DecompilationError(f"No handler for opcode {inst.opname}", instruction=inst) + ctx = self._make_context(emitter) + result = handler(emitter, inst, ctx) + idx = result if result is not None else idx + 1 + except DecompilationError: + raise + except Exception as e: + raise DecompilationError(f"Failed at {inst!r} in {self.code.co_name}", instruction=inst) from e + + def _make_context(self, emitter: SourceEmitter) -> DecompileContext: + return DecompileContext( + code=self.code, + instructions=tuple(self.instructions), + indentation=emitter._indent_size, + decompile_range=lambda start, end, em: self.decompile_range(start, end, em), + offset_to_index={inst.offset: idx for idx, inst in enumerate(self.instructions)}, + ) + + # -- public API --------------------------------------------------------- + + def decompile(self, indentation: int = 4, temp_prefix: str = "__temp_", overwrite_fn_name: Optional[str] = None) -> str: + """Return decompiled Python source code.""" + try: + emitter = SourceEmitter(indent_size=indentation, temp_prefix=temp_prefix) + self.decompile_range(0, len(self.instructions), emitter) + body = emitter.get_source() + + if os.environ.get("DEPYF_REMOVE_TEMP", "1") == "1": + from .postprocess import run_all as _postprocess + + body = _postprocess(body, temp_prefix, indentation) + + header = SignatureBuilder.build(self.code, overwrite_fn_name) + + global_names = {i.argval for i in dis.get_instructions(self.code) if i.opname == "STORE_GLOBAL"} + preamble = "" + if global_names: + preamble += "global " + ", ".join(global_names) + "\n" + if self.code.co_freevars: + preamble += "nonlocal " + ", ".join(self.code.co_freevars) + "\n" + + body = preamble + body + return header + emitter.indent(body) + except DecompilationError: + raise + except Exception as e: + raise DecompilationError(f"Failed to decompile {self.code.co_name}") from e + + @staticmethod + def supported_opnames() -> List[str]: + return registry.supported_opnames() + + +# --------------------------------------------------------------------------- +# Module-level convenience +# --------------------------------------------------------------------------- + + +def decompile(code: Union[CodeType, Callable]) -> str: + """One-liner: decompile a code object or callable to source.""" + return Decompiler(code).decompile() + + +def safe_decompile(code: CodeType) -> str: + """Decompile *code* without raising; fall back to depyf then placeholder.""" + try: + return Decompiler(code).decompile() + except Exception: + try: + from depyf import decompile as _depyf_decompile + + return _depyf_decompile(code) + except Exception: + return f"# Failed to decompile {code.co_name}\n" + + +# --------------------------------------------------------------------------- +# Private helpers +# --------------------------------------------------------------------------- + + +def _get_code_owner(fn): + """Walk through wrappers to find the object that owns ``__code__``.""" + if hasattr(fn, "__func__"): + return fn.__func__ + if hasattr(fn, "__wrapped__"): + return _get_code_owner(fn.__wrapped__) + return fn diff --git a/pkgs/MagiCompiler/magi_compiler/magi_depyf/decompile/postprocess/__init__.py b/pkgs/MagiCompiler/magi_compiler/magi_depyf/decompile/postprocess/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..0d0cb557f2bf393ca92f8a7ff2d5c2d948e26589 --- /dev/null +++ b/pkgs/MagiCompiler/magi_compiler/magi_depyf/decompile/postprocess/__init__.py @@ -0,0 +1,35 @@ +# Copyright (c) 2026 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. + +"""Source-level post-processing pipeline for decompiled code. + +Each pass is a function ``(source, ...) -> source`` that performs one +semantics-preserving transformation. ``run_all`` applies them in order. +All passes are best-effort: on any exception they return the input unchanged. +""" + +from .branch_dedup import dedup_branch_tails +from .for_temps import eliminate_for_temps +from .inline_temps import eliminate_inline_temps + + +def run_all(source: str, temp_prefix: str = "__temp_", indent: int = 4) -> str: + """Apply all post-processing passes in sequence.""" + source = eliminate_for_temps(source, temp_prefix, indent) + source = eliminate_inline_temps(source, temp_prefix, indent) + source = dedup_branch_tails(source, indent) + return source + + +__all__ = ["run_all", "eliminate_for_temps", "eliminate_inline_temps", "dedup_branch_tails"] diff --git a/pkgs/MagiCompiler/magi_compiler/magi_depyf/decompile/postprocess/branch_dedup.py b/pkgs/MagiCompiler/magi_compiler/magi_depyf/decompile/postprocess/branch_dedup.py new file mode 100644 index 0000000000000000000000000000000000000000..11864a9ee36adc06b0cf038c3ab51684f0677d7e --- /dev/null +++ b/pkgs/MagiCompiler/magi_compiler/magi_depyf/decompile/postprocess/branch_dedup.py @@ -0,0 +1,99 @@ +# Copyright (c) 2026 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. + +"""Pass 3: if/else branch tail deduplication. + +Move identical trailing statements from if/else branches to after the block. + +Example:: + + if cond: if cond: + x = 1 x = 1 + return x → else: + else: x = 2 + x = 2 return x + return x +""" + +from __future__ import annotations + +import ast +from typing import List, Tuple + +import astor + + +def dedup_branch_tails(source: str, indent: int = 4) -> str: + """Move identical trailing statements from if/else branches to after the block.""" + try: + tree = ast.parse(source) + changed = False + for node in ast.walk(tree): + if hasattr(node, "body") and isinstance(node.body, list): + new_body, c = _dedup_stmts(node.body) + if c: + node.body = new_body + changed = True + if not changed: + return source + ast.fix_missing_locations(tree) + return astor.to_source(tree, indent_with=" " * indent) + except Exception: + return source + + +def _dedup_stmts(stmts: List[ast.stmt]) -> Tuple[List[ast.stmt], bool]: + """Process a statement list, extracting common if/else tails.""" + result: List[ast.stmt] = [] + changed = False + + for stmt in stmts: + for attr in ("body", "orelse", "handlers", "finalbody"): + sub = getattr(stmt, attr, None) + if isinstance(sub, list) and sub: + new_sub, c = _dedup_stmts(sub) + if c: + setattr(stmt, attr, new_sub) + changed = True + + if isinstance(stmt, ast.If) and stmt.orelse: + n = _common_tail_length(stmt.body, stmt.orelse) + if n > 0: + common = stmt.body[-n:] + stmt.body = stmt.body[:-n] or [ast.Pass()] + stmt.orelse = stmt.orelse[:-n] or [] + result.append(stmt) + result.extend(common) + changed = True + continue + + result.append(stmt) + + return result, changed + + +def _common_tail_length(body: List[ast.stmt], orelse: List[ast.stmt]) -> int: + """Count identical trailing statements (by AST dump equality).""" + count = 0 + i, j = len(body) - 1, len(orelse) - 1 + while i >= 0 and j >= 0: + if ast.dump(body[i]) == ast.dump(orelse[j]): + count += 1 + i -= 1 + j -= 1 + else: + break + if count >= len(body) or count >= len(orelse): + count = min(len(body), len(orelse)) - 1 + return max(count, 0) diff --git a/pkgs/MagiCompiler/magi_compiler/magi_depyf/decompile/postprocess/for_temps.py b/pkgs/MagiCompiler/magi_compiler/magi_depyf/decompile/postprocess/for_temps.py new file mode 100644 index 0000000000000000000000000000000000000000..8947a6b57984f63acbe6974c8a39433d19e700e5 --- /dev/null +++ b/pkgs/MagiCompiler/magi_compiler/magi_depyf/decompile/postprocess/for_temps.py @@ -0,0 +1,57 @@ +# Copyright (c) 2026 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. + +"""Pass 1: for-loop temp elimination. + +``for __temp in iter: var = __temp; ...`` → ``for var in iter: ...`` +""" + +from __future__ import annotations + +import ast + +import astor + + +def eliminate_for_temps(source: str, temp_prefix: str = "__temp_", indent: int = 4) -> str: + """Only applies when the first body statement is a plain assignment + from the temp to a real variable.""" + try: + tree = ast.parse(source) + tree = _ForTempEliminator(temp_prefix).visit(tree) + ast.fix_missing_locations(tree) + return astor.to_source(tree, indent_with=" " * indent) + except Exception: + return source + + +class _ForTempEliminator(ast.NodeTransformer): + def __init__(self, prefix: str): + self._prefix = prefix + + def visit_For(self, node: ast.For) -> ast.For: + self.generic_visit(node) + if not ( + isinstance(node.target, ast.Name) + and node.target.id.startswith(self._prefix) + and node.body + and isinstance(node.body[0], ast.Assign) + and len(node.body[0].targets) == 1 + and isinstance(node.body[0].value, ast.Name) + and node.body[0].value.id == node.target.id + ): + return node + node.target = node.body[0].targets[0] + node.body = node.body[1:] or [ast.Pass()] + return node diff --git a/pkgs/MagiCompiler/magi_compiler/magi_depyf/decompile/postprocess/inline_temps.py b/pkgs/MagiCompiler/magi_compiler/magi_depyf/decompile/postprocess/inline_temps.py new file mode 100644 index 0000000000000000000000000000000000000000..bca2680806b4d40c2546aea6af80d3a3e12a6b33 --- /dev/null +++ b/pkgs/MagiCompiler/magi_compiler/magi_depyf/decompile/postprocess/inline_temps.py @@ -0,0 +1,165 @@ +# Copyright (c) 2026 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. + +"""Pass 2: single-use temp inlining. + +``__temp = expr; use(__temp)`` → ``use(expr)`` for single-use temps. +""" + +from __future__ import annotations + +import ast +from collections import defaultdict +from typing import List, Optional + +import astor + + +def eliminate_inline_temps(source: str, temp_prefix: str = "__temp_", indent: int = 4) -> str: + """Inline single-use temporaries into their use site.""" + try: + tree = ast.parse(source) + _set_parents(tree) + + occurrences: dict[str, list] = defaultdict(list) + for node in ast.walk(tree): + if isinstance(node, ast.Name) and node.id.startswith(temp_prefix): + occurrences[node.id].append(node) + + _INDENT_NODES = ( + ast.FunctionDef, + ast.AsyncFunctionDef, + ast.For, + ast.AsyncFor, + ast.While, + ast.If, + ast.Try, + ast.With, + ast.AsyncWith, + ast.ClassDef, + ) + + for name in occurrences: + occ = occurrences[name] + if len(occ) == 2: + n1, n2 = occ + _, p1, p2 = _lowest_common_parent(n1, n2) + ap = p1 if isinstance(getattr(n1, "parent", None), ast.Assign) else p2 + can = not isinstance(ap, _INDENT_NODES) + if can: + can = _safe_to_inline(tree, n1, n2) + occ.append(can) + tree = _RemoveAssign(name, occurrences).visit(tree) + tree = _InlineTemp(name, occurrences).visit(tree) + + return astor.to_source(tree, indent_with=" " * indent) + except Exception: + return source + + +# --------------------------------------------------------------------------- +# AST helpers +# --------------------------------------------------------------------------- + + +def _set_parents(node: ast.AST, parent: Optional[ast.AST] = None) -> None: + for child in ast.iter_child_nodes(node): + child.parent = parent # type: ignore[attr-defined] + _set_parents(child, child) + + +def _get_parents(node: ast.AST) -> List[ast.AST]: + out = [] + while node: + out.append(node) + node = getattr(node, "parent", None) + return out + + +def _lowest_common_parent(n1: ast.AST, n2: ast.AST): + p1 = _get_parents(n1) + p2 = _get_parents(n2) + p1.reverse() + p2.reverse() + last = c1 = c2 = None + for a, b in zip(p1, p2): + if a is b: + last = a + else: + c1, c2 = a, b + break + return last, c1, c2 + + +def _safe_to_inline(tree: ast.AST, def_node: ast.AST, use_node: ast.AST) -> bool: + """Verify the RHS variable is not reassigned between definition and use.""" + assign_parent = getattr(def_node, "parent", None) + if not isinstance(assign_parent, ast.Assign): + return True + rhs = assign_parent.value + if not isinstance(rhs, ast.Name): + return True + + rhs_name = rhs.id + stmts: List[ast.stmt] = [] + for node in ast.walk(tree): + if hasattr(node, "body") and isinstance(node.body, list): + stmts = node.body + break + try: + def_idx = next(i for i, s in enumerate(stmts) if s is assign_parent) + use_stmt = getattr(use_node, "parent", None) + while use_stmt and use_stmt not in stmts: + use_stmt = getattr(use_stmt, "parent", None) + use_idx = next(i for i, s in enumerate(stmts) if s is use_stmt) + except StopIteration: + return True + + for stmt in stmts[def_idx + 1 : use_idx]: + if isinstance(stmt, ast.Assign): + for t in stmt.targets: + if isinstance(t, ast.Name) and t.id == rhs_name: + return False + return True + + +class _RemoveAssign(ast.NodeTransformer): + def __init__(self, name: str, occ: dict): + self._name = name + self._occ = occ + + def visit_Assign(self, node: ast.Assign): + if len(node.targets) == 1 and isinstance(node.targets[0], ast.Name): + n = node.targets[0].id + if n == self._name: + o = self._occ[n] + if len(o) == 1: + return ast.Expr(value=node.value) + if len(o) == 3 and isinstance(o[-1], bool): + o.append(node.value) + if o[-2]: + return None + return node + + +class _InlineTemp(ast.NodeTransformer): + def __init__(self, name: str, occ: dict): + self._name = name + self._occ = occ + + def visit_Name(self, node: ast.Name): + o = self._occ.get(node.id, []) + if node.id == self._name and len(o) == 4 and isinstance(o[-2], bool) and o[-2]: + return o[-1] + return node diff --git a/pkgs/MagiCompiler/magi_compiler/magi_depyf/decompile/recompiler.py b/pkgs/MagiCompiler/magi_compiler/magi_depyf/decompile/recompiler.py new file mode 100644 index 0000000000000000000000000000000000000000..29a4e54fbdc513a2bed6ae9e6648c546f59f8379 --- /dev/null +++ b/pkgs/MagiCompiler/magi_compiler/magi_depyf/decompile/recompiler.py @@ -0,0 +1,53 @@ +# Copyright (c) 2026 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. + +"""CodeRecompiler: round-trip decompile -> compile -> extract target CodeType. + +Pipeline: CodeType -> decompile -> compile -> find target. +""" + +from __future__ import annotations + +from types import CodeType +from typing import List + +from .decompiler import Decompiler + + +class CodeRecompiler: + """Decompile *code*, recompile, and produce a compatible ``CodeType``.""" + + @staticmethod + def recompile( + code_to_decompile: CodeType, reference_code: CodeType, indentation: int = 4, temp_prefix: str = "__temp_" + ) -> CodeType: + """Full round-trip: decompile -> compile -> find target.""" + fn_name = reference_code.co_name + + src = Decompiler(code_to_decompile).decompile( + indentation=indentation, temp_prefix=temp_prefix, overwrite_fn_name=fn_name + ) + + compiled = compile(src, "noname", "exec") + all_codes = CodeRecompiler.collect_code_objects(compiled) + return [c for c in all_codes if c.co_name == fn_name][0] + + @staticmethod + def collect_code_objects(code: CodeType) -> List[CodeType]: + """Recursively collect all ``CodeType`` objects from *code*.""" + result = [code] + for c in code.co_consts: + if isinstance(c, CodeType): + result.extend(CodeRecompiler.collect_code_objects(c)) + return result diff --git a/pkgs/MagiCompiler/magi_compiler/magi_depyf/demo_toy_example.py b/pkgs/MagiCompiler/magi_compiler/magi_depyf/demo_toy_example.py new file mode 100644 index 0000000000000000000000000000000000000000..5f03c5f0d6441b3f9e945112a9c3aae1e0b0e45c --- /dev/null +++ b/pkgs/MagiCompiler/magi_compiler/magi_depyf/demo_toy_example.py @@ -0,0 +1,54 @@ +# Copyright (c) 2026 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. + +"""Demo: magi_depyf.dump_src with the depyf tutorial toy_example. + +Run: PYTHONPATH=. python demo_toy_example.py +""" +import torch +from magi_compiler.magi_depyf.inspect import dump_src + +device = "cuda" if torch.cuda.is_available() else "cpu" +torch.set_default_device(device) + + +@torch.compile +def toy_example(a, b): + x = a / (torch.abs(a) + 1) + if b.sum() < 0: + b = b * -1 + return x * b + + +def main(): + for _ in range(100): + toy_example(torch.randn(10), torch.randn(10)) + + +if __name__ == "__main__": + import os + import shutil + + out = "./magi_dump_src_dir" + if os.path.exists(out): + shutil.rmtree(out) + with dump_src(out): + main() + + print("\n=== Generated files ===") + for root, dirs, files in os.walk(out): + level = root.replace(out, "").count(os.sep) + print(f"{' ' * level}{os.path.basename(root)}/") + for f in files: + print(f"{' ' * (level + 1)}{f}") diff --git a/pkgs/MagiCompiler/magi_compiler/magi_depyf/inspect/__init__.py b/pkgs/MagiCompiler/magi_compiler/magi_depyf/inspect/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..4af536f18a00157eac13eca2a1c9806cbd87c492 --- /dev/null +++ b/pkgs/MagiCompiler/magi_compiler/magi_depyf/inspect/__init__.py @@ -0,0 +1,57 @@ +# Copyright (c) 2026 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. + +"""Inspection layer: capture torch.compile events, introspect artifacts, write structured output.""" + +from typing import Optional + +from .dump_src import dump_src +from .introspect import Introspector +from .model import CompiledFnInfo, EntryInfo, FunctionInfo, GuardInfo, GuardNode, SubgraphInfo +from .result import CaptureResult +from .session import CaptureSession +from .writer import FunctionWriter, write_function + + +def debug_compiled(fn, output_dir: Optional[str] = None) -> FunctionInfo: + """Introspect a compiled function and optionally write debug output. + + Args: + fn: The original (uncompiled) function. + output_dir: If provided, write organized files to this directory. + + Returns: + FunctionInfo with full compilation state. + """ + info = Introspector.build_function_info(fn) + if output_dir is not None: + write_function(info, output_dir) + return info + + +__all__ = [ + "CompiledFnInfo", + "EntryInfo", + "FunctionInfo", + "GuardInfo", + "GuardNode", + "SubgraphInfo", + "Introspector", + "FunctionWriter", + "write_function", + "dump_src", + "debug_compiled", + "CaptureSession", + "CaptureResult", +] diff --git a/pkgs/MagiCompiler/magi_compiler/magi_depyf/inspect/dump_src.py b/pkgs/MagiCompiler/magi_compiler/magi_depyf/inspect/dump_src.py new file mode 100644 index 0000000000000000000000000000000000000000..1a5230f19bb3a6cf59575b4a5bdc36815eae7c5f --- /dev/null +++ b/pkgs/MagiCompiler/magi_compiler/magi_depyf/inspect/dump_src.py @@ -0,0 +1,78 @@ +# Copyright (c) 2026 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. + +"""dump_src: context manager that captures torch.compile artifacts and +writes structured source output to disk. + +Usage:: + + from magi_compiler.magi_depyf.inspect import dump_src + + @torch.compile + def my_fn(x): + return x.sum() + + with dump_src("./output_dir"): + my_fn(torch.randn(10)) + +Internally uses ``CaptureSession`` to intercept compilation events, +then runs ``Introspector`` post-hoc for full introspection. +""" + +from __future__ import annotations + +import contextlib +from pathlib import Path +from typing import Set + +from magi_compiler.utils import magi_logger + +from .introspect import Introspector +from .session import CaptureSession +from .writer import write_function + + +@contextlib.contextmanager +def dump_src(dump_src_dir: str): + """Context manager that captures torch.compile artifacts and writes output. + + Uses CaptureSession for hook management and post-hoc introspection + of CacheEntries after execution completes. + """ + dump_dir = Path(dump_src_dir) + dump_dir.mkdir(parents=True, exist_ok=True) + + with CaptureSession() as session: + yield + + seen: Set[str] = set() + overview_paths: list[Path] = [] + for r in session.results: + name = r.original_code.co_name + if name in seen: + continue + if name.startswith("torch_dynamo_resume_in_"): + continue + seen.add(name) + + try: + info = Introspector.build_function_info(r.original_code, fn_globals=r.fn_globals) + root = write_function(info, dump_dir) + overview_paths.append(root / "overview.md") + except Exception as e: + magi_logger.warning("[magi_depyf] failed to process '%s': %s", name, e) + + for p in overview_paths: + if p.exists(): + magi_logger.info("[magi_depyf] %s", p) diff --git a/pkgs/MagiCompiler/magi_compiler/magi_depyf/inspect/introspect.py b/pkgs/MagiCompiler/magi_compiler/magi_depyf/inspect/introspect.py new file mode 100644 index 0000000000000000000000000000000000000000..a1cee3b2be1294eedea92f2b88439d12a79d2e93 --- /dev/null +++ b/pkgs/MagiCompiler/magi_compiler/magi_depyf/inspect/introspect.py @@ -0,0 +1,524 @@ +# Copyright (c) 2026 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. + +"""Runtime introspection of torch.compile artifacts. + +Walk actual runtime state (CacheEntry chain, guard trees, __compiled_fn +objects) to build the structured model. All torch imports are lazy so +this module can be imported without torch. +""" + +from __future__ import annotations + +import io +from pathlib import Path +from typing import Any, Dict, Optional + +from ..decompile import safe_decompile +from .model import CompiledFnInfo, EntryInfo, FunctionInfo, GuardInfo, GuardNode, SubgraphInfo + + +class Introspector: + """Namespace for runtime introspection helpers (all static methods).""" + + @staticmethod + def get_cache_entries(fn) -> list: + """Return CacheEntry list for *fn* (function or code object).""" + from torch._dynamo.eval_frame import _debug_get_cache_entry_list + + code = fn.__code__ if hasattr(fn, "__code__") else fn + return _debug_get_cache_entry_list(code) + + @staticmethod + def build_guard_tree(node, max_depth: int = 32, _depth: int = 0) -> GuardNode: + """Recursively build a GuardNode from a GuardManager C++ object.""" + type_name = type(node).__name__ + leaf_guards = [] + for lg in node.get_leaf_guards(): + for part in lg.verbose_code_parts(): + leaf_guards.append(part.strip()[:120]) + children = [] + if _depth < max_depth: + for child in node.get_child_managers(): + children.append(Introspector.build_guard_tree(child, max_depth, _depth + 1)) + return GuardNode(type_name=type_name, leaf_guards=leaf_guards, children=children) + + @staticmethod + def extract_guard_info(entry) -> Optional[GuardInfo]: + """Extract structured guard info from a CacheEntry (post-hoc introspection). + + Operates on the persisted CacheEntry and builds a full GuardNode tree. + """ + try: + gm = entry.guard_manager + tree = Introspector.build_guard_tree(gm.root) + closure_vars: Dict[str, str] = {} + if hasattr(gm, "closure_vars") and gm.closure_vars: + for k, v in list(gm.closure_vars.items())[:10]: + closure_vars[k] = repr(v)[:100] + return GuardInfo(tree=tree, closure_vars=closure_vars or None) + except Exception: + return None + + @staticmethod + def extract_compiled_fn_info(name: str, fn_globals: dict) -> Optional[CompiledFnInfo]: + """Inspect a __compiled_fn_xxx from fn.__globals__. + + Handles three backend types: + eager: wrapper -> closure[0]=GraphModule.forward (bound method) + inductor: wrapper -> closure[0]=aot_forward -> ... -> CompiledFxGraph + magi_compile: MagiSerializableFunction -> split_gm -> PiecewiseBackend(s) + """ + obj = fn_globals.get(name) + if obj is None: + return None + + magi_info = Introspector._try_extract_magi_info(name, obj) + if magi_info is not None: + return magi_info + + info = CompiledFnInfo(name=name, backend="eager") + + gm = Introspector._find_graph_module(obj) + if gm is not None: + Introspector._fill_graph_module_info(info, gm) + + cfx = Introspector._find_compiled_fx_graph(obj) + if cfx is not None: + info.backend = "inductor" + Introspector._fill_compiled_fx_graph_info(info, cfx) + + return info + + @staticmethod + def _fill_graph_module_info(info: CompiledFnInfo, gm) -> None: + try: + info.readable_code = gm.print_readable(print_output=False) + except Exception: + pass + try: + info.graph_module_code = str(gm.code) if hasattr(gm, "code") else None + except Exception: + pass + try: + buf = io.StringIO() + gm.graph.print_tabular(file=buf) + info.fx_graph_tabular = buf.getvalue() + except Exception: + pass + + @staticmethod + def _fill_compiled_fx_graph_info(info: CompiledFnInfo, cfx) -> None: + try: + info.source_code = cfx.source_code + except Exception: + pass + try: + info.inductor_post_grad_graph = cfx.inductor_post_grad_graph_str + except Exception: + pass + try: + info.cache_key = cfx.cache_key + except Exception: + pass + try: + info.runnable_graph_str = cfx.runnable_graph_str + except Exception: + pass + + # -- Magi backend introspection ---------------------------------------- + + @staticmethod + def _try_extract_magi_info(name: str, obj) -> Optional[CompiledFnInfo]: + """Detect MagiSerializableFunction and walk its hierarchy. + + MagiSerializableFunction hierarchy: + .graph_module → fx.GraphModule (full graph before splitting) + .optimized_call → split_gm (fx.GraphModule with PiecewiseBackend submodules) + .submod_N → PiecewiseBackend + .graph → fx.GraphModule (the subgraph) + .compiled_graph_for_general_shape → inductor compiled output + + Dynamo wraps the backend result in a DisableContext closure, so the + MagiSerializableFunction may live one level deep in the closure chain. + """ + msf = obj if (hasattr(obj, "graph_module") and hasattr(obj, "optimized_call")) else None + if msf is None and callable(obj) and getattr(obj, "__closure__", None): + for cell in obj.__closure__: + try: + val = cell.cell_contents + except ValueError: + continue + if hasattr(val, "graph_module") and hasattr(val, "optimized_call"): + msf = val + break + if msf is None: + return None + obj = msf + + import torch.fx + + info = CompiledFnInfo(name=name, backend="magi_compile") + + full_gm = getattr(obj, "graph_module", None) + if isinstance(full_gm, torch.fx.GraphModule): + Introspector._fill_graph_module_info(info, full_gm) + + split_gm = getattr(obj, "optimized_call", None) + + # In FULL cudagraph mode, optimized_call is a wrapper function whose + # __dict__ carries the GraphModule's attributes (via __dict__.update). + # Unwrap to find the actual GraphModule for print_readable / named_children. + actual_gm = split_gm if isinstance(split_gm, torch.fx.GraphModule) else None + if actual_gm is None and split_gm is not None: + actual_gm = Introspector._find_graph_module_deep(split_gm) + + info.cudagraph_mode = Introspector._detect_cudagraph_mode(split_gm, actual_gm) + + if actual_gm is not None: + try: + info.split_graph_readable = actual_gm.print_readable(print_output=False) + except Exception: + pass + + # PiecewiseCompileInterpreter replaces submodules via __dict__, + # so named_children() still sees the original GraphModules while + # __dict__ contains the PiecewiseBackend (or cudagraph wrapper). + # In FULL cudagraph mode, those __dict__ entries are copied onto + # the wrapper function, so we look up runtime objects from + # split_gm (the wrapper) rather than actual_gm. + runtime_source = split_gm if split_gm is not None else actual_gm + for sub_name, original_gm in actual_gm.named_children(): + runtime_obj = runtime_source.__dict__.get(sub_name, original_gm) + sg_info = Introspector._extract_subgraph_info(sub_name, runtime_obj, original_gm) + if sg_info is not None: + info.subgraph_infos.append(sg_info) + + info.subgraph_infos.sort(key=lambda s: s.name) + + return info + + @staticmethod + def _extract_subgraph_info(sub_name: str, runtime_obj, original_gm=None) -> Optional[SubgraphInfo]: + """Extract info from one submodule of the split graph. + + Args: + sub_name: The submodule name (e.g. "submod_0"). + runtime_obj: The actual runtime object — PiecewiseBackend, + cudagraph wrapper, or the original GraphModule. + original_gm: The original GraphModule before replacement (from _modules). + """ + import torch.fx + + piecewise = Introspector._unwrap_piecewise_backend(runtime_obj) + + if piecewise is not None: + sg = SubgraphInfo(name=sub_name, is_splitting_graph=False) + inner_gm = getattr(piecewise, "graph", None) + if isinstance(inner_gm, torch.fx.GraphModule): + Introspector._fill_subgraph_gm_info(sg, inner_gm) + + compiled = getattr(piecewise, "compiled_graph_for_general_shape", None) + if compiled is not None: + sg.inductor_code = Introspector._try_extract_inductor_source(compiled) + + if sg.inductor_code is None: + sg.inductor_code = Introspector._read_artifact_source_from_piecewise(piecewise) + + return sg + + gm = original_gm if isinstance(original_gm, torch.fx.GraphModule) else None + if gm is None and isinstance(runtime_obj, torch.fx.GraphModule): + gm = runtime_obj + + if gm is not None: + sg = SubgraphInfo(name=sub_name, is_splitting_graph=True) + Introspector._fill_subgraph_gm_info(sg, gm) + return sg + + return None + + @staticmethod + def _unwrap_piecewise_backend(obj): + """Find a PiecewiseBackend from obj, unwrapping closures/wrappers if needed.""" + if hasattr(obj, "graph") and hasattr(obj, "compiled_graph_for_general_shape"): + return obj + + if callable(obj) and hasattr(obj, "__closure__") and obj.__closure__: + for cell in obj.__closure__: + try: + val = cell.cell_contents + except ValueError: + continue + if hasattr(val, "graph") and hasattr(val, "compiled_graph_for_general_shape"): + return val + return None + + @staticmethod + def _fill_subgraph_gm_info(sg: SubgraphInfo, gm) -> None: + try: + sg.readable_code = gm.print_readable(print_output=False) + except Exception: + pass + try: + sg.graph_module_code = str(gm.code) if hasattr(gm, "code") else None + except Exception: + pass + try: + buf = io.StringIO() + gm.graph.print_tabular(file=buf) + sg.fx_graph_tabular = buf.getvalue() + except Exception: + pass + + @staticmethod + def _try_extract_inductor_source(compiled) -> Optional[str]: + """Try to extract inductor kernel source from a compiled graph object. + + Handles CompiledFxGraph, CompiledArtifact, and closure-wrapped variants. + """ + for attr in ("source_code", "_source_code"): + val = getattr(compiled, attr, None) + if isinstance(val, str) and val: + return val + + cfx = Introspector._find_compiled_fx_graph(compiled) + if cfx is not None: + try: + return cfx.source_code + except Exception: + pass + + if hasattr(compiled, "print_readable"): + try: + return compiled.print_readable(print_output=False) + except Exception: + pass + + return None + + @staticmethod + def _read_artifact_source_from_piecewise(piecewise) -> Optional[str]: + """Read Inductor-generated source from the saved artifact directory. + + PiecewiseBackend stores a compiler_manager whose cache maps + CacheEntry(runtime_shape, graph_index, backend_name) → CacheHandle(key, path). + The artifact at CacheHandle.path is an unpacked directory containing + ``py/*.py`` — the full Inductor output code. + """ + try: + compiler_manager = getattr(piecewise, "compiler_manager", None) + if compiler_manager is None: + return None + cache = getattr(compiler_manager, "cache", None) + if not cache: + return None + index = getattr(piecewise, "piecewise_compile_index", None) + if index is None: + return None + + for cache_entry, cache_handle in cache.items(): + if cache_entry.graph_index == index and cache_entry.runtime_shape is None: + artifact_path = getattr(cache_handle, "path", None) + if artifact_path: + return Introspector._read_py_from_artifact(artifact_path) + return None + except Exception: + return None + + @staticmethod + def _read_py_from_artifact(artifact_path: str) -> Optional[str]: + """Read the Inductor-generated Python wrapper from an artifact directory. + + The unpacked artifact layout varies across PyTorch versions; the + wrapper ``.py`` file has been observed under ``yb/`` and ``py/``. + We try known directories first, then fall back to scanning all + immediate subdirectories. + """ + root = Path(artifact_path) + if not root.is_dir(): + return None + + for candidate_dir in ("yb", "py"): + d = root / candidate_dir + if d.is_dir(): + py_files = sorted(d.glob("*.py")) + if py_files: + try: + return py_files[0].read_text(encoding="utf-8") + except Exception: + pass + + for d in sorted(root.iterdir()): + if d.is_dir(): + py_files = sorted(d.glob("*.py")) + if py_files: + try: + return py_files[0].read_text(encoding="utf-8") + except Exception: + pass + return None + + @staticmethod + def _detect_cudagraph_mode(split_gm, actual_gm) -> str: + """Detect cudagraph wrapping mode from the split graph structure. + + - FULL: split_gm itself is a cudagraph wrapper (not a GraphModule), + with __qualname__ containing 'Athena_CUDAGraph_full'. + - PIECEWISE: split_gm is a GraphModule, but its __dict__ submodules are + cudagraph wrappers with __qualname__ 'Athena_CUDAGraph_piecewise'. + - NONE: no cudagraph wrapping detected. + """ + _CG_PREFIX = "Athena_CUDAGraph_" + + qualname = getattr(split_gm, "__qualname__", "") or "" + if qualname.startswith(f"{_CG_PREFIX}full"): + return "FULL" + + if actual_gm is not None: + for key, val in actual_gm.__dict__.items(): + if not key.startswith("submod_"): + continue + sub_qualname = getattr(val, "__qualname__", "") or "" + if sub_qualname.startswith(f"{_CG_PREFIX}piecewise"): + return "PIECEWISE" + + return "NONE" + + @staticmethod + def _find_graph_module_deep(obj, _depth: int = 0, _max_depth: int = 4) -> Optional[Any]: + """Recursively walk closure chain to find a ``torch.fx.GraphModule``. + + This is needed for FULL cudagraph mode where the split GraphModule is + wrapped by ``gen_wrap_func_for_cudagraph`` (+ ``@instrument_nvtx``), + placing the GraphModule 2-3 levels deep in the closure chain. + """ + import torch.fx + + if isinstance(obj, torch.fx.GraphModule): + return obj + if _depth >= _max_depth: + return None + if not callable(obj) or not getattr(obj, "__closure__", None): + return None + for cell in obj.__closure__: + try: + val = cell.cell_contents + except ValueError: + continue + if isinstance(val, torch.fx.GraphModule): + return val + if callable(val): + found = Introspector._find_graph_module_deep(val, _depth + 1, _max_depth) + if found is not None: + return found + return None + + @staticmethod + def _find_graph_module(obj) -> Optional[Any]: + """Walk closure chain to find a torch.fx.GraphModule.""" + import torch.fx + + if isinstance(obj, torch.fx.GraphModule): + return obj + if hasattr(obj, "__self__") and isinstance(obj.__self__, torch.fx.GraphModule): + return obj.__self__ + if not callable(obj) or not hasattr(obj, "__closure__") or not obj.__closure__: + return None + for cell in obj.__closure__: + try: + val = cell.cell_contents + except ValueError: + continue + if isinstance(val, torch.fx.GraphModule): + return val + if hasattr(val, "__self__") and isinstance(val.__self__, torch.fx.GraphModule): + return val.__self__ + return None + + @staticmethod + def _find_compiled_fx_graph(obj, _depth: int = 0) -> Optional[Any]: + """Walk closure chain (up to 4 levels) to find a CompiledFxGraph.""" + if _depth > 4: + return None + try: + from torch._inductor.codecache import CompiledFxGraph + except ImportError: + return None + if isinstance(obj, CompiledFxGraph): + return obj + if not callable(obj) or not hasattr(obj, "__closure__") or not obj.__closure__: + return None + for cell in obj.__closure__: + try: + val = cell.cell_contents + except ValueError: + continue + if isinstance(val, CompiledFxGraph): + return val + if callable(val): + found = Introspector._find_compiled_fx_graph(val, _depth + 1) + if found is not None: + return found + return None + + @staticmethod + def build_entry_info(entry, index: int, fn_globals: dict) -> EntryInfo: + """Build an EntryInfo from a CacheEntry.""" + tc = entry.code + decompiled = safe_decompile(tc) + + compiled_names = [n for n in tc.co_names if n.startswith("__compiled")] + compiled_fns = [] + for cn in compiled_names: + cf = Introspector.extract_compiled_fn_info(cn, fn_globals) + if cf: + compiled_fns.append(cf) + + resume_names = [n for n in tc.co_names if n.startswith("__resume")] + resume_fns = [] + for rn in resume_names: + rfn = fn_globals.get(rn) + if rfn is not None and hasattr(rfn, "__code__"): + resume_info = Introspector.build_function_info(rfn, fn_globals=fn_globals) + resume_info.name = rn + resume_fns.append(resume_info) + + guard = Introspector.extract_guard_info(entry) + + return EntryInfo( + index=index, + dynamo_code=tc, + decompiled_source=decompiled, + guard=guard, + compiled_fns=compiled_fns, + resume_fns=resume_fns, + ) + + @staticmethod + def build_function_info(fn, fn_globals: Optional[dict] = None) -> FunctionInfo: + """Build full FunctionInfo by walking CacheEntry chain.""" + if fn_globals is None: + fn_globals = fn.__globals__ if hasattr(fn, "__globals__") else {} + + code = fn.__code__ if hasattr(fn, "__code__") else fn + name = code.co_name + original_source = safe_decompile(code) + + entries_raw = Introspector.get_cache_entries(fn) + entries = [] + for i, raw_entry in enumerate(entries_raw): + entries.append(Introspector.build_entry_info(raw_entry, i, fn_globals)) + + return FunctionInfo(name=name, original_code=code, original_source=original_source, entries=entries) diff --git a/pkgs/MagiCompiler/magi_compiler/magi_depyf/inspect/model.py b/pkgs/MagiCompiler/magi_compiler/magi_depyf/inspect/model.py new file mode 100644 index 0000000000000000000000000000000000000000..ad88c75cc81a32808ed6900183e5eedd906f4409 --- /dev/null +++ b/pkgs/MagiCompiler/magi_compiler/magi_depyf/inspect/model.py @@ -0,0 +1,241 @@ +# Copyright (c) 2026 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. + +"""Data model for structured compilation output. + +These dataclasses represent the full compilation state that +torch.compile produces, organized to reflect the actual runtime +structure: CacheEntry linked list, fn/resume recursion, +compiled_fn → backend mapping, and guard trees. +""" + +from __future__ import annotations + +import dataclasses +import dis +import inspect +import io +from types import CodeType +from typing import Dict, List, Optional + + +def format_code_info(code: CodeType) -> str: + """Format key attributes of a CodeType for debugging.""" + lines: List[str] = [] + lines.append(f"co_name: {code.co_name}") + if hasattr(code, "co_qualname"): + lines.append(f"co_qualname: {code.co_qualname}") + lines.append(f"co_filename: {code.co_filename}") + lines.append(f"co_firstlineno: {code.co_firstlineno}") + lines.append(f"co_argcount: {code.co_argcount}") + lines.append(f"co_kwonlyargcount:{code.co_kwonlyargcount}") + lines.append(f"co_varnames: {code.co_varnames}") + lines.append(f"co_freevars: {code.co_freevars}") + lines.append(f"co_cellvars: {code.co_cellvars}") + lines.append(f"co_names: {code.co_names}") + flags = code.co_flags + flag_strs = [name for name, val in _CODE_FLAGS.items() if flags & val] + lines.append(f"co_flags: 0x{flags:04x} ({' | '.join(flag_strs) if flag_strs else 'none'})") + lines.append(f"co_stacksize: {code.co_stacksize}") + lines.append("") + lines.append("co_consts:") + for i, c in enumerate(code.co_consts): + lines.append(f" [{i:3d}] {type(c).__name__:12s} {_safe_repr(c)}") + lines.append("") + lines.append("dis:") + buf = io.StringIO() + dis.dis(code, file=buf) + lines.append(buf.getvalue()) + return "\n".join(lines) + + +_CODE_FLAGS = { + "CO_OPTIMIZED": inspect.CO_OPTIMIZED, + "CO_NEWLOCALS": inspect.CO_NEWLOCALS, + "CO_VARARGS": inspect.CO_VARARGS, + "CO_VARKEYWORDS": inspect.CO_VARKEYWORDS, + "CO_NESTED": inspect.CO_NESTED, + "CO_GENERATOR": inspect.CO_GENERATOR, + "CO_COROUTINE": inspect.CO_COROUTINE, + "CO_ASYNC_GENERATOR": inspect.CO_ASYNC_GENERATOR, +} + + +def _safe_repr(obj, max_len: int = 120) -> str: + try: + r = repr(obj) + except Exception: + r = f"" + if len(r) > max_len: + r = r[: max_len - 3] + "..." + return r + + +@dataclasses.dataclass +class GuardNode: + """One node in the guard tree (mirrors RootGuardManager / GuardManager).""" + + type_name: str + leaf_guards: List[str] + children: List["GuardNode"] = dataclasses.field(default_factory=list) + + def format(self, depth: int = 0, max_depth: int = 32) -> str: + prefix = " " * depth + lines = [f"{prefix}[{self.type_name}] " f"({len(self.leaf_guards)} leaf guards, {len(self.children)} children)"] + for g in self.leaf_guards: + lines.append(f"{prefix} LEAF: {g}") + if depth < max_depth: + for i, child in enumerate(self.children): + lines.append(f"{prefix} child[{i}]:") + lines.append(child.format(depth + 2, max_depth)) + elif self.children: + lines.append(f"{prefix} ... ({len(self.children)} children omitted)") + return "\n".join(lines) + + +@dataclasses.dataclass +class SubgraphInfo: + """One piecewise subgraph in the magi split pipeline.""" + + name: str + is_splitting_graph: bool = False + readable_code: Optional[str] = None + graph_module_code: Optional[str] = None + fx_graph_tabular: Optional[str] = None + inductor_code: Optional[str] = None + + def format(self) -> str: + if self.inductor_code: + return self.inductor_code + if self.readable_code: + return self.readable_code + if self.graph_module_code: + return self.graph_module_code + tag = "splitting_op" if self.is_splitting_graph else "compiled" + return f"# {self.name} ({tag})\n" + + +@dataclasses.dataclass +class CompiledFnInfo: + """What __compiled_fn_xxx actually points to in the backend.""" + + name: str + backend: str # "eager", "inductor", or "magi_compile" + cudagraph_mode: Optional[str] = None # "NONE", "PIECEWISE", "FULL" (magi_compile only) + readable_code: Optional[str] = None + graph_module_code: Optional[str] = None + fx_graph_tabular: Optional[str] = None + source_code: Optional[str] = None + inductor_post_grad_graph: Optional[str] = None + runnable_graph_str: Optional[str] = None + cache_key: Optional[str] = None + split_graph_readable: Optional[str] = None + subgraph_infos: List["SubgraphInfo"] = dataclasses.field(default_factory=list) + + def format(self) -> str: + """Full content for writing to file (compiled output).""" + if self.source_code: + return self.source_code + if self.readable_code: + return self.readable_code + if self.graph_module_code: + return self.graph_module_code + return f"# {self.name} (backend={self.backend})\n" + + def format_summary(self) -> str: + """Short summary for overview / full_code.""" + header = f"{self.name} (backend={self.backend}" + if self.cudagraph_mode: + header += f", cudagraph={self.cudagraph_mode}" + header += ")" + lines = [header] + if self.cache_key: + lines.append(f" cache_key: {self.cache_key}") + if self.graph_module_code: + lines.append(" GraphModule.code:") + for l in self.graph_module_code.strip().splitlines(): + lines.append(f" {l}") + if self.subgraph_infos: + lines.append(f" piecewise subgraphs: {len(self.subgraph_infos)}") + for sg in self.subgraph_infos: + tag = "splitting_op" if sg.is_splitting_graph else "compiled" + lines.append(f" {sg.name} ({tag})") + return "\n".join(lines) + + +@dataclasses.dataclass +class GuardInfo: + """Guard information for a CacheEntry.""" + + tree: Optional[GuardNode] = None + closure_vars: Optional[Dict[str, str]] = None + + def format(self) -> str: + lines = [] + if self.tree: + lines.append(self.tree.format()) + if self.closure_vars: + lines.append(" closure_vars:") + for k, v in list(self.closure_vars.items())[:8]: + lines.append(f" {k} = {v}") + return "\n".join(lines) + + +@dataclasses.dataclass +class EntryInfo: + """One CacheEntry in the linked list.""" + + index: int + dynamo_code: Optional[CodeType] = None + decompiled_source: str = "" + guard: Optional[GuardInfo] = None + compiled_fns: List[CompiledFnInfo] = dataclasses.field(default_factory=list) + resume_fns: List["FunctionInfo"] = dataclasses.field(default_factory=list) + + def format(self, indent: int = 0) -> str: + pfx = " " * indent + lines = [f"{pfx}entry[{self.index}]:"] + if self.decompiled_source: + lines.append(f"{pfx} dynamo_code (decompiled):") + for l in self.decompiled_source.splitlines(): + lines.append(f"{pfx} {l}") + if self.compiled_fns: + lines.append(f"{pfx} compiled functions:") + for cf in self.compiled_fns: + lines.append(cf.format_summary()) + if self.guard: + lines.append(f"{pfx} guards:") + lines.append(self.guard.format()) + if self.resume_fns: + lines.append(f"{pfx} resume functions:") + for rf in self.resume_fns: + lines.append(rf.format(indent + 2)) + return "\n".join(lines) + + +@dataclasses.dataclass +class FunctionInfo: + """A compiled function and its CacheEntry chain.""" + + name: str + original_code: Optional[CodeType] = None + original_source: str = "" + entries: List[EntryInfo] = dataclasses.field(default_factory=list) + + def format(self, indent: int = 0) -> str: + pfx = " " * indent + lines = [f"{pfx}{self.name}: {len(self.entries)} cache entries"] + for entry in self.entries: + lines.append(entry.format(indent + 1)) + return "\n".join(lines) diff --git a/pkgs/MagiCompiler/magi_compiler/magi_depyf/inspect/result.py b/pkgs/MagiCompiler/magi_compiler/magi_depyf/inspect/result.py new file mode 100644 index 0000000000000000000000000000000000000000..fe20ff932e7298efd627f7c9c8b69c2a5a93098a --- /dev/null +++ b/pkgs/MagiCompiler/magi_compiler/magi_depyf/inspect/result.py @@ -0,0 +1,51 @@ +# Copyright (c) 2026 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. + +"""CaptureResult — structured data model for one compilation event.""" + +from __future__ import annotations + +import dataclasses +import time +from types import CodeType +from typing import List, Optional + + +@dataclasses.dataclass +class CaptureResult: + """Data captured from a single ``torch.compile`` bytecode event. + + - original_code: the user's original function code + - dynamo_code: the code after Dynamo transformation (with __compiled_fn / __resume calls) + - decompiled_source: dynamo_code decompiled back to Python source + - fn_globals: the function's global namespace (for post-hoc introspection) + """ + + function_name: str + original_code: CodeType + dynamo_code: CodeType + decompiled_source: str + fn_globals: Optional[dict] = None + guards: List[str] = dataclasses.field(default_factory=list) + graph_source: Optional[str] = None + timestamp: float = dataclasses.field(default_factory=time.time) + + def summary(self) -> str: + n_guards = len(self.guards) + return ( + f"[{self.function_name}] " + f"original={self.original_code.co_name}, " + f"guards={n_guards}, " + f"graph={'yes' if self.graph_source else 'no'}" + ) diff --git a/pkgs/MagiCompiler/magi_compiler/magi_depyf/inspect/session.py b/pkgs/MagiCompiler/magi_compiler/magi_depyf/inspect/session.py new file mode 100644 index 0000000000000000000000000000000000000000..77982e951497934172454fb3201e8e4344a21da5 --- /dev/null +++ b/pkgs/MagiCompiler/magi_compiler/magi_depyf/inspect/session.py @@ -0,0 +1,118 @@ +# Copyright (c) 2026 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. + +"""CaptureSession -- lifecycle-managed torch.compile interception. + +All state is instance-scoped (no global dictionaries). Only the +official ``register_bytecode_hook`` API is used -- no monkey-patching. +""" + +from __future__ import annotations + +import os +import sys +from types import CodeType, FrameType +from typing import List, Optional + +from ..decompile import safe_decompile +from .result import CaptureResult + +_MAX_FRAME_WALK = 64 + + +class CaptureSession: + """Context-manager that intercepts ``torch.compile`` bytecode events. + + Usage:: + + with CaptureSession() as session: + compiled_fn(input_tensor) + for r in session.results: + print(r.summary()) + print(r.decompiled_source) + """ + + def __init__(self) -> None: + self._results: List[CaptureResult] = [] + self._hook_handle = None + + @property + def results(self) -> List[CaptureResult]: + return list(self._results) + + def __enter__(self) -> "CaptureSession": + try: + import torch._dynamo.convert_frame as cf + except ImportError as e: + raise ImportError("CaptureSession requires PyTorch (torch._dynamo). " "Install with: pip install torch") from e + + hook = self._make_hook(self._results) + self._hook_handle = cf.register_bytecode_hook(hook) + return self + + def __exit__(self, exc_type, exc_val, exc_tb) -> None: + if self._hook_handle is not None: + self._hook_handle.remove() + self._hook_handle = None + + def clear(self) -> None: + """Discard all captured results.""" + self._results.clear() + + # ------------------------------------------------------------------ + # Hook internals + # ------------------------------------------------------------------ + + @staticmethod + def _find_compile_frame(start: Optional[FrameType] = None) -> Optional[FrameType]: + """Walk the call stack to find the ``_compile`` frame inside + ``convert_frame.py``. Returns ``None`` if not found within + ``_MAX_FRAME_WALK`` steps. + """ + frame = start or sys._getframe(1) + for _ in range(_MAX_FRAME_WALK): + if frame is None: + return None + name = frame.f_code.co_name + filename = os.path.basename(frame.f_code.co_filename) + if name == "_compile" and filename == "convert_frame.py": + return frame + frame = frame.f_back + return None + + @staticmethod + def _make_hook(results: List[CaptureResult]): + """Return a bytecode-hook callable that appends ``CaptureResult`` + objects to *results*. + """ + + def hook(old_code: CodeType, new_code: CodeType) -> CodeType: + compile_frame = CaptureSession._find_compile_frame() + fn_globals = None + if compile_frame is not None: + inner_frame = compile_frame.f_locals.get("frame") + if inner_frame is not None: + fn_globals = inner_frame.f_globals + + result = CaptureResult( + function_name=old_code.co_name, + original_code=old_code, + dynamo_code=new_code, + decompiled_source=safe_decompile(new_code), + fn_globals=fn_globals, + ) + results.append(result) + return new_code + + return hook diff --git a/pkgs/MagiCompiler/magi_compiler/magi_depyf/inspect/writer.py b/pkgs/MagiCompiler/magi_compiler/magi_depyf/inspect/writer.py new file mode 100644 index 0000000000000000000000000000000000000000..0148d24830352fb7c6c6519b562e315f95a85d0f --- /dev/null +++ b/pkgs/MagiCompiler/magi_compiler/magi_depyf/inspect/writer.py @@ -0,0 +1,280 @@ +# Copyright (c) 2026 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. + +"""File writer: serialize FunctionInfo tree to organized output files. + +Output structure: + output_dir/ + {fn_name}/ + decompiled_code.py (original bytecode decompiled) + overview.md + entry_0/ + decompiled_code.py (dynamo-transformed bytecode decompiled) + guards.txt + compiled_fns/ + __compiled_fn_xxx.py + resume_fns/ + {resume_name}/ + decompiled_code.py + overview.md + entry_0/ + ... + entry_1/ + ... +""" + +from __future__ import annotations + +from pathlib import Path +from typing import List + +from .model import CompiledFnInfo, EntryInfo, FunctionInfo, GuardNode, SubgraphInfo, format_code_info + +_DECOMPILED = "decompiled_code.py" +_BYTECODE_INFO = "bytecode_info.txt" +_OVERVIEW = "overview.md" + +# ── Convenience function (preserves existing public API) ───────────── + + +def write_function(fn_info: FunctionInfo, output_dir: str | Path) -> Path: + """Write *fn_info* to a directory tree under *output_dir*. + Returns the root directory created.""" + return FunctionWriter(fn_info, output_dir).write() + + +# ── Main class ─────────────────────────────────────────────────────── + + +class FunctionWriter: + """Serialize a :class:`FunctionInfo` tree into an organized directory.""" + + def __init__(self, fn_info: FunctionInfo, output_dir: str | Path) -> None: + self.fn_info = fn_info + self.root = (Path(output_dir) / self._sanitize(fn_info.name)).resolve() + + def write(self) -> Path: + """Write all output files. Returns the root directory.""" + self.root.mkdir(parents=True, exist_ok=True) + if self.fn_info.original_source: + self._write_text(self.root / _DECOMPILED, self.fn_info.original_source) + if self.fn_info.original_code is not None: + self._write_text(self.root / _BYTECODE_INFO, format_code_info(self.fn_info.original_code)) + self._write_text(self.root / _OVERVIEW, self._format_overview(self.fn_info, self.root)) + for entry in self.fn_info.entries: + self._write_entry(entry, self.root / f"entry_{entry.index}") + return self.root + + # ── Entry writing ──────────────────────────────────────────────── + + def _write_entry(self, entry: EntryInfo, entry_dir: Path) -> None: + entry_dir.mkdir(parents=True, exist_ok=True) + self._write_text(entry_dir / _DECOMPILED, entry.decompiled_source) + if entry.dynamo_code is not None: + self._write_text(entry_dir / _BYTECODE_INFO, format_code_info(entry.dynamo_code)) + + if entry.guard: + self._write_text(entry_dir / "guards.txt", entry.guard.format()) + + if entry.compiled_fns: + cfns_dir = entry_dir / "compiled_fns" + cfns_dir.mkdir(parents=True, exist_ok=True) + for cf in entry.compiled_fns: + base = self._sanitize(cf.name) + self._write_text(cfns_dir / f"{base}.py", cf.format()) + if cf.inductor_post_grad_graph: + self._write_text(cfns_dir / f"{base}_post_grad.py", cf.inductor_post_grad_graph) + if cf.runnable_graph_str: + self._write_text(cfns_dir / f"{base}_runnable.py", cf.runnable_graph_str) + if cf.split_graph_readable: + self._write_text(cfns_dir / f"{base}_split_graph.py", cf.split_graph_readable) + for sg in cf.subgraph_infos: + self._write_subgraph(cfns_dir / base, sg) + + if entry.resume_fns: + rfns_dir = entry_dir / "resume_fns" + for rf in entry.resume_fns: + sub_writer = FunctionWriter(rf, rfns_dir) + sub_writer.write() + + # ── Subgraph writing (magi piecewise) ─────────────────────────── + + def _write_subgraph(self, parent_dir: Path, sg: SubgraphInfo) -> None: + sg_dir = parent_dir / sg.name + sg_dir.mkdir(parents=True, exist_ok=True) + tag = "splitting_op" if sg.is_splitting_graph else "compiled" + if sg.readable_code: + self._write_text(sg_dir / "graph_module.py", sg.readable_code) + if sg.graph_module_code: + self._write_text(sg_dir / "graph_module_code.py", sg.graph_module_code) + if sg.fx_graph_tabular: + self._write_text(sg_dir / "fx_graph_tabular.txt", sg.fx_graph_tabular) + if sg.inductor_code: + self._write_text(sg_dir / "inductor_output.py", sg.inductor_code) + summary = f"# {sg.name} ({tag})\n" + if sg.readable_code: + summary += f"# graph_module.py: GraphModule.print_readable()\n" + if sg.inductor_code: + summary += f"# inductor_output.py: inductor kernel source\n" + self._write_text(sg_dir / "README.txt", summary) + + # ── Overview (markdown) ─────────────────────────────────────────── + + def _format_overview(self, fn_info: FunctionInfo, root: Path) -> str: + lines: List[str] = [] + lines.append(f"# {fn_info.name}") + lines.append("") + lines.append(f"**Root:** `{self.root}` ") + lines.append(f"**Cache entries:** {len(fn_info.entries)}") + lines.append("") + + if fn_info.original_source: + lines.append(f"[decompiled code (before dynamo)](./{_DECOMPILED})") + if fn_info.original_code is not None: + lines.append(f"[bytecode info](./{_BYTECODE_INFO})") + lines.append("") + + for entry in fn_info.entries: + lines.extend(self._format_entry_md(entry, root)) + + return "\n".join(lines) + "\n" + + def _format_entry_md(self, entry: EntryInfo, root: Path) -> List[str]: + entry_dir = root / f"entry_{entry.index}" + lines: List[str] = [] + lines.append(f"## entry\\[{entry.index}\\]") + lines.append("") + + items = self._build_entry_items(entry, entry_dir, root) + lines.extend(self._render_tree_md(items, depth=0)) + lines.append("") + return lines + + def _build_entry_items(self, entry: EntryInfo, entry_dir: Path, root: Path) -> "List[_TreeItem]": + items: List[_TreeItem] = [] + + items.append(_TreeItem("decompiled code", self._rel(entry_dir / _DECOMPILED, root))) + if entry.dynamo_code is not None: + items.append(_TreeItem("bytecode info", self._rel(entry_dir / _BYTECODE_INFO, root))) + + if entry.guard: + n = len(self._collect_leaf_guards(entry.guard.tree)) if entry.guard.tree else 0 + label = f"guards ({n} leaf)" if n else "guards" + items.append(_TreeItem(label, self._rel(entry_dir / "guards.txt", root))) + + if entry.compiled_fns: + cfns_dir = entry_dir / "compiled_fns" + cf_children = self._build_compiled_fn_items(entry.compiled_fns, cfns_dir, root) + items.append(_TreeItem("compiled_fns/", "", cf_children)) + + if entry.resume_fns: + rfns_dir = entry_dir / "resume_fns" + rf_children: List[_TreeItem] = [] + for rf in entry.resume_fns: + resume_dir = rfns_dir / self._sanitize(rf.name) + fn_children: List[_TreeItem] = [] + if rf.original_source: + fn_children.append(_TreeItem("decompiled code", self._rel(resume_dir / _DECOMPILED, root))) + for re in rf.entries: + sub_entry_dir = resume_dir / f"entry_{re.index}" + entry_children = self._build_entry_items(re, sub_entry_dir, root) + fn_children.append( + _TreeItem(f"entry\\[{re.index}\\]", self._rel(sub_entry_dir / _DECOMPILED, root), entry_children) + ) + rf_children.append(_TreeItem(f"`{rf.name}`/", self._rel(resume_dir / _OVERVIEW, root), fn_children)) + items.append(_TreeItem("resume_fns/", "", rf_children)) + + return items + + def _build_compiled_fn_items(self, compiled_fns: List[CompiledFnInfo], cfns_dir: Path, root: Path) -> "List[_TreeItem]": + items: List[_TreeItem] = [] + for cf in compiled_fns: + base = self._sanitize(cf.name) + label = f"`{cf.name}` ({cf.backend}" + if cf.cudagraph_mode and cf.cudagraph_mode != "NONE": + label += f", cudagraph={cf.cudagraph_mode}" + label += ")" + items.append(_TreeItem(label, self._rel(cfns_dir / f"{base}.py", root))) + if cf.inductor_post_grad_graph: + items.append( + _TreeItem( + f"`{cf.name}` (inductor_post_grad_graph_str)", self._rel(cfns_dir / f"{base}_post_grad.py", root) + ) + ) + if cf.runnable_graph_str: + items.append(_TreeItem(f"`{cf.name}` (runnable_graph_str)", self._rel(cfns_dir / f"{base}_runnable.py", root))) + if cf.split_graph_readable: + items.append(_TreeItem(f"`{cf.name}` (split_graph)", self._rel(cfns_dir / f"{base}_split_graph.py", root))) + if cf.subgraph_infos: + sg_children: List[_TreeItem] = [] + for sg in cf.subgraph_infos: + sg_dir = cfns_dir / base / sg.name + tag = "splitting_op" if sg.is_splitting_graph else "compiled" + sg_sub: List[_TreeItem] = [] + if sg.readable_code: + sg_sub.append(_TreeItem("graph_module", self._rel(sg_dir / "graph_module.py", root))) + if sg.inductor_code: + sg_sub.append(_TreeItem("inductor_output", self._rel(sg_dir / "inductor_output.py", root))) + if sg.fx_graph_tabular: + sg_sub.append(_TreeItem("fx_graph_tabular", self._rel(sg_dir / "fx_graph_tabular.txt", root))) + sg_children.append(_TreeItem(f"`{sg.name}` ({tag})", self._rel(sg_dir / "README.txt", root), sg_sub)) + items.append(_TreeItem("piecewise_subgraphs/", "", sg_children)) + return items + + @staticmethod + def _render_tree_md(items: "List[_TreeItem]", depth: int = 0) -> List[str]: + lines: List[str] = [] + indent = " " * depth + for item in items: + if item.rel_path: + lines.append(f"{indent}- [{item.label}](./{item.rel_path})") + else: + lines.append(f"{indent}- {item.label}") + if item.children: + lines.extend(FunctionWriter._render_tree_md(item.children, depth + 1)) + return lines + + # ── Helpers ────────────────────────────────────────────────────── + + @staticmethod + def _collect_leaf_guards(node: GuardNode) -> List[str]: + guards = list(node.leaf_guards) + for child in node.children: + guards.extend(FunctionWriter._collect_leaf_guards(child)) + return guards + + @staticmethod + def _rel(path: Path, base: Path) -> str: + try: + return str(path.relative_to(base)) + except ValueError: + return str(path) + + @staticmethod + def _write_text(path: Path, content: str) -> None: + path.parent.mkdir(parents=True, exist_ok=True) + path.write_text(content, encoding="utf-8") + + @staticmethod + def _sanitize(name: str) -> str: + return name.replace("<", "_").replace(">", "_").replace(" ", "_").replace(".", "_") + + +class _TreeItem: + __slots__ = ("label", "rel_path", "children") + + def __init__(self, label: str, rel_path: str, children: "List[_TreeItem] | None" = None): + self.label = label + self.rel_path = rel_path + self.children = children or [] diff --git a/pkgs/MagiCompiler/magi_compiler/offload/offload_warpper.py b/pkgs/MagiCompiler/magi_compiler/offload/offload_warpper.py new file mode 100644 index 0000000000000000000000000000000000000000..8f29e3754386f4c825650d9d57b7ef27cfd8457f --- /dev/null +++ b/pkgs/MagiCompiler/magi_compiler/offload/offload_warpper.py @@ -0,0 +1,212 @@ +# Copyright (c) 2026 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. + +import collections +import operator +from typing import Any, Dict + +import torch +from magi_compiler.config import CompileConfig +from magi_compiler.offload.profiler import OffloadProfiler +from magi_compiler.offload.scheduler import OffloadRuntimeContext, SchedulerFactory +from magi_compiler.utils.nvtx import add_nvtx_event +from torch.fx import GraphModule, Node +from torch.fx.node import map_arg + + +class OffloadExecutor: + def __init__(self, graph_module: GraphModule, compile_config: CompileConfig): + self.graph_module = graph_module + self.compile_config = compile_config + + self.compute_stream = torch.cuda.current_stream() + self.h2d_stream = torch.cuda.Stream() + + self.warmup = True + self.second_call = False + self.buffers: Dict[str, torch.Tensor] = {} + self.persistent_weights: Dict[str, torch.Tensor] = {} + self.submod_0_weight_handoff: Dict[Node, torch.Tensor] = {} + + self._analyze_graph() + self.profiler = OffloadProfiler() + + common_args = { + "submod_nodes": self.submod_nodes, + "submod_weights_map": self.submod_weights_map, + "name_node_map": self.name_node_map, + "weight_sizes": self.submod_weight_sizes, + } + + self.scheduler = SchedulerFactory.create(self.compile_config, common_args) + + def _analyze_graph(self): + self.submod_nodes = [n for n in self.graph_module.graph.nodes if n.op == "call_module"] + + self.placeholder_nodes = [] + self.arg_index_weight = {} + self.user_counts = collections.defaultdict(int) + self.name_node_map = {} + + placeholder_idx = 0 + for node in self.graph_module.graph.nodes: + for input_node in node.all_input_nodes: + self.user_counts[input_node] += 1 + + if node.op == "placeholder": + is_w = isinstance(node.meta.get("example_value"), torch.nn.Parameter) + self.arg_index_weight[placeholder_idx] = is_w + self.placeholder_nodes.append(node) + self.name_node_map[node.name] = node + placeholder_idx += 1 + + self.submod_weights_map = {} + self.submod_weight_sizes = {} + + for node in self.submod_nodes: + weight_names = [] + size = 0 + for arg in node.args: + if isinstance(arg, Node) and self._is_weight_node(arg): + if arg.name in self.name_node_map: + weight_names.append(arg.name) + val = arg.meta.get("example_value") + if val is not None: + size += val.numel() * val.element_size() + + self.submod_weights_map[node.name] = weight_names + self.submod_weight_sizes[node.name] = size + + def _is_weight_node(self, node: Node) -> bool: + return node.op == "placeholder" and isinstance(node.meta.get("example_value"), torch.nn.Parameter) + + def _prepare_inputs(self, args) -> Dict[Node, Any]: + env = {} + args = list(args) + submod_0 = self.submod_nodes[0] + + for i, node in enumerate(self.placeholder_nodes): + arg_val = args[i] + is_weight = self.arg_index_weight[i] + + # case 1: input tensor + if not is_weight: + if isinstance(arg_val, torch.Tensor): + arg_val = arg_val.to("cuda", non_blocking=False) + env[node] = arg_val + continue + + # case 2: kept weight + if self.scheduler.is_kept(node.name): + if node.name not in self.persistent_weights: + t = arg_val.to("cuda", non_blocking=False) if arg_val.device.type == "cpu" else arg_val + self.persistent_weights[node.name] = t + env[node] = self.persistent_weights[node.name] + continue + + # case 3: submod 0 weight + if submod_0 in node.users: + if self.warmup and arg_val.device.type == "cpu": + self.buffers[node.name] = arg_val + arg_val = arg_val.to("cuda", non_blocking=False) + elif not self.warmup: + if node in self.submod_0_weight_handoff: + arg_val = self.submod_0_weight_handoff[node] + del self.submod_0_weight_handoff[node] + + env[node] = arg_val + + return env + + def _finalize_warmup(self): + profile_results = self.profiler.summarize() + self.scheduler.schedule_kept_weights(profile_results) + self.warmup = False + + def __call__(self, *args): + env = self._prepare_inputs(args) + current_user_counts = self.user_counts.copy() + runtime_ctx = OffloadRuntimeContext( + env=env, + h2d_stream=self.h2d_stream, + compute_stream=self.compute_stream, + buffers=self.buffers, + submod_0_handoff=self.submod_0_weight_handoff, + need_profile=self.second_call or self.warmup, + ) + need_profile = self.second_call + + for node in self.graph_module.graph.nodes: + if node.op == "placeholder": + continue + + elif node.op == "call_module": + self.scheduler.prefetch(node.name, runtime_ctx) + + if need_profile: + if torch.distributed.is_initialized(): + torch.distributed.barrier() + self.profiler.start_compute_profile(node.name, self.compute_stream) + + with add_nvtx_event(node.name): + with torch.cuda.stream(self.compute_stream): + s_args = map_arg(node.args, lambda n: env[n]) + s_kwargs = map_arg(node.kwargs, lambda n: env[n]) + env[node] = getattr(self.graph_module, node.target)(*s_args, **s_kwargs) + del s_args, s_kwargs + + if need_profile: + if torch.distributed.is_initialized(): + torch.distributed.barrier() + self.profiler.end_compute_profile(node.name, self.compute_stream) + + elif node.op == "call_function": + # ... (Standard execution logic same as before) + if node.target == operator.getitem: + parent_node, idx = node.args + env[node] = env[parent_node][idx] + else: + with torch.cuda.stream(self.compute_stream): + f_args = map_arg(node.args, lambda n: env[n]) + f_kwargs = map_arg(node.kwargs, lambda n: env[n]) + env[node] = node.target(*f_args, **f_kwargs) + + elif node.op == "output": + if self.second_call: + self._finalize_warmup() + self.second_call = False + if self.warmup: + self.second_call = True + self.warmup = False + + return map_arg(node.args[0], lambda n: env[n]) + + # Memory Management + for input_node in node.all_input_nodes: + current_user_counts[input_node] -= 1 + if current_user_counts[input_node] == 0: + if input_node in env: + tensor_obj = env[input_node] + if isinstance(tensor_obj, torch.Tensor) and tensor_obj.is_cuda: + tensor_obj.record_stream(self.compute_stream) + del env[input_node] + return None + + +class OffloadWrapper: + def __init__(self, graph_module: torch.fx.GraphModule, compile_config: CompileConfig): + self.executor = OffloadExecutor(graph_module, compile_config) + + def __call__(self, *args): + return self.executor(*args) diff --git a/pkgs/MagiCompiler/magi_compiler/offload/profiler.py b/pkgs/MagiCompiler/magi_compiler/offload/profiler.py new file mode 100644 index 0000000000000000000000000000000000000000..4a014788afe12979d4e30b5c062995cee906c5c0 --- /dev/null +++ b/pkgs/MagiCompiler/magi_compiler/offload/profiler.py @@ -0,0 +1,91 @@ +# Copyright (c) 2026 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. + +from typing import Dict + +import torch + + +class OffloadProfiler: + def __init__(self): + self.compute_events: Dict[str, Dict[str, torch.cuda.Event]] = {} + self.timings: Dict[str, Dict[str, float]] = {} + + def start_compute_profile(self, name: str, stream: torch.cuda.Stream): + if name not in self.compute_events: + self.compute_events[name] = {} + start_event = torch.cuda.Event(enable_timing=True) + start_event.record(stream) + self.compute_events[name]["start"] = start_event + + def end_compute_profile(self, name: str, stream: torch.cuda.Stream): + end_event = torch.cuda.Event(enable_timing=True) + end_event.record(stream) + self.compute_events[name]["end"] = end_event + + def get_h2d_bandwidth(self, size_mb=1024, iters=3, warmup=3, dtype=torch.float32, device=torch.device("cuda")): + torch.cuda.synchronize() + + num_elements = size_mb * 1024 * 1024 // torch.tensor([], dtype=dtype).element_size() + + cpu_tensor = torch.empty(num_elements, dtype=dtype, pin_memory=True) + gpu_tensor = torch.empty(num_elements, dtype=dtype, device=device) + + # warmup + for _ in range(warmup): + gpu_tensor.copy_(cpu_tensor, non_blocking=True) + torch.cuda.synchronize() + + start_event = torch.cuda.Event(enable_timing=True) + end_event = torch.cuda.Event(enable_timing=True) + + start_event.record() + + for _ in range(iters): + gpu_tensor.copy_(cpu_tensor, non_blocking=True) + + end_event.record() + torch.cuda.synchronize() + + elapsed_ms = start_event.elapsed_time(end_event) + elapsed_s = elapsed_ms / 1000.0 + + total_bytes = size_mb * 1024 * 1024 * iters + bandwidth = total_bytes / elapsed_s / 1e9 # GB/s + + return bandwidth + + def broadcast_obj(self, obj, src=0): + obj_list = [obj] + torch.distributed.broadcast_object_list(obj_list, src=src) + return obj_list[0] + + def summarize(self) -> Dict[str, Dict[str, float]]: + torch.cuda.synchronize() + results = {} + for name, evs in self.compute_events.items(): + if name not in results: + results[name] = {} + if "start" in evs and "end" in evs: + results[name]["compute"] = evs["start"].elapsed_time(evs["end"]) + + h2d_bandwidth = self.get_h2d_bandwidth() + results["h2d_bandwidth"] = h2d_bandwidth + + if torch.distributed.is_initialized(): + h2d_bandwidth = self.broadcast_obj(h2d_bandwidth) + results = self.broadcast_obj(results) + + self.timings = results + return results diff --git a/pkgs/MagiCompiler/magi_compiler/offload/scheduler.py b/pkgs/MagiCompiler/magi_compiler/offload/scheduler.py new file mode 100644 index 0000000000000000000000000000000000000000..99d91b1012f5c5098355e13e9b27cc34fa718571 --- /dev/null +++ b/pkgs/MagiCompiler/magi_compiler/offload/scheduler.py @@ -0,0 +1,466 @@ +# Copyright (c) 2026 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. + +import abc +import collections +from dataclasses import dataclass +from typing import Any, Dict, List, Set, Type + +import torch +from magi_compiler.config import CompileConfig +from magi_compiler.utils import magi_logger +from torch.fx import Node + + +@dataclass +class OffloadRuntimeContext: + """ + Wrap the dynamic runtime context needed during Offload + """ + + env: Dict[Node, Any] # Tensor environment of the current computation graph + h2d_stream: torch.cuda.Stream # CUDA stream for data transfer + compute_stream: torch.cuda.Stream # CUDA stream for computation + buffers: Dict[str, torch.Tensor] # GPU copies stored during Warmup + submod_0_handoff: Dict[Node, torch.Tensor] # Container for weights used by Submod 0 for the next iteration + need_profile: bool = False + + +SchedulerType = Type["OffloadScheduler"] + + +class SchedulerFactory: + _REGISTRY: Dict[str, SchedulerType] = {} + + @classmethod + def register(cls, name: str): + def decorator(scheduler_cls: SchedulerType): + cls._REGISTRY[name] = scheduler_cls + return scheduler_cls + + return decorator + + @classmethod + def create(cls, compile_config: CompileConfig, common_args: Dict[str, Any]) -> "OffloadScheduler": + offload_policy_name = compile_config.offload_config.offload_policy.value + + scheduler_cls = cls._REGISTRY.get(offload_policy_name) + if not scheduler_cls: + raise ValueError(f"Unknown offload policy: {offload_policy_name}. " f"Available: {list(cls._REGISTRY.keys())}") + + return scheduler_cls(compile_config=compile_config, **common_args) + + +class OffloadScheduler(abc.ABC): + def __init__( + self, + compile_config: CompileConfig, + submod_nodes: List[Node], + submod_weights_map: Dict[str, List[str]], + name_node_map: Dict[str, Node], + weight_sizes: Dict[str, int], + ): + """ + :param config: Compile configuration + :param submod_nodes: List of submodule nodes in execution order + :param submod_weights_map: {submod_name: [weight_names...]} + :param name_node_map: {weight_name: weight_node} for looking up Node objects + """ + self.compile_config = compile_config + self.submod_nodes = submod_nodes + self.submod_num = len(submod_nodes) + self.submod_weights_map = submod_weights_map + self.name_node_map = name_node_map + self.weight_sizes = weight_sizes + + self.kept_weights: Set[str] = set() + + @abc.abstractmethod + def schedule_kept_weights(self, profile_data: Dict[str, float]): + """ + Static decision: determine which weights to keep on GPU based on profile data + """ + pass + + @abc.abstractmethod + def prefetch(self, current_node_name: str, ctx: OffloadRuntimeContext): + """ + + :param env: Tensor environment of the current execution environment + :param h2d_stream: CUDA stream for data transfer + :param buffers: GPU copies stored during Warmup + :param submod_0_next_iter_weights: Container for weights used by Submod 0 for the next iteration + """ + pass + + def is_kept(self, weight_name: str) -> bool: + return weight_name in self.kept_weights + + def get_keep_weight_size(self) -> int: + return sum( + [ + self.name_node_map[w_name].meta.get("example_value").numel() + * self.name_node_map[w_name].meta.get("example_value").element_size() + for w_name in self.kept_weights + ] + ) + + +@SchedulerFactory.register("BASE") +class BaseScheduler(OffloadScheduler): + """ + Basic strategy implementation: + 1. schedule_kept_weights: Default to not keep any weights on GPU (gpu_resident_weight_ratio=0.0) + 2. prefetch: Inherit the original "Odd/Even + Next-Next" logic + """ + + def __init__( + self, + compile_config: CompileConfig, + submod_nodes: List[Node], + submod_weights_map: Dict[str, List[str]], + name_node_map: Dict[str, Node], + weight_sizes: Dict[str, int], + ): + super().__init__(compile_config, submod_nodes, submod_weights_map, name_node_map, weight_sizes) + self.submod_cpu_names = set() + + def schedule_kept_weights(self, profile_data: Dict[str, float]): + self.kept_weights.clear() + # offload all submodules to CPU + self.submod_cpu_names = set(self.submod_nodes[i].name for i in range(len(self.submod_nodes))) + + def prefetch(self, current_node_name: str, ctx: OffloadRuntimeContext): + env = ctx.env + h2d_stream = ctx.h2d_stream + buffers = ctx.buffers + submod_0_next_iter_weights = ctx.submod_0_handoff + compute_stream = ctx.compute_stream + + try: + idx = int(current_node_name.split('_')[-1]) + except ValueError: + return + + max_lookahead = 2 + target_node = None + is_next_iter = False + + for offset in range(1, max_lookahead + 1): + candidate_idx = (idx + offset) % self.submod_num + candidate_node = self.submod_nodes[candidate_idx] + + if self.weight_sizes.get(candidate_node.name, 0) > 0: + target_node = candidate_node + + if (idx + offset) >= self.submod_num: + is_next_iter = True + else: + is_next_iter = False + + break + + if target_node is None: + return + + last_target = getattr(self, "last_prefetched_target", None) + + if idx == 0: + last_target = None + + if target_node.name == last_target: + return + + self.last_prefetched_target = target_node.name + + compute_stream.wait_stream(h2d_stream) + + weight_names = self.submod_weights_map.get(target_node.name, []) + + with torch.cuda.stream(h2d_stream): + for w_name in weight_names: + if w_name in self.kept_weights: + continue + + w_node = self.name_node_map.get(w_name) + if not w_node: + continue + + if is_next_iter: + if w_name in buffers: + submod_0_next_iter_weights[w_node] = buffers[w_name].to("cuda", non_blocking=True) + + elif w_node in env and hasattr(env[w_node], 'device') and env[w_node].device.type == "cpu": + env[w_node] = env[w_node].to("cuda", non_blocking=True) + + +@SchedulerFactory.register("COST_EFFECTIVE") +class CostEffectiveScheduler(BaseScheduler): + """ + Cost-effective strategy implementation: + Determine which weights to keep on GPU based on the ratio of time to size + """ + + def schedule_kept_weights(self, profile_data: Dict[str, float]): + weight_sizes = self.weight_sizes + gpu_resident_weight_ratio = self.compile_config.offload_config.gpu_resident_weight_ratio + if gpu_resident_weight_ratio <= 0.0: + return + + candidates = [] + total_weight_size = 0 + + for name, timing in profile_data.items(): + if name == "h2d_bandwidth": + continue + duration = timing.get("compute", 0.0) + size = weight_sizes.get(name, 0) + ratio = duration / size if size > 0 else float('inf') + candidates.append((name, ratio, size)) + total_weight_size += size + + candidates.sort(key=lambda x: x[1]) + + current_kept_size = 0 + limit = total_weight_size * gpu_resident_weight_ratio + + self.kept_weights.clear() + + for name, ratio, s in candidates: + if name == "submod_0": + continue + if current_kept_size + s <= limit: + current_kept_size += s + for w_name in self.submod_weights_map.get(name, []): + self.kept_weights.add(w_name) + self.submod_cpu_names.add(name) + else: + break + + magi_logger.info(f"schedule_kept_weights size {current_kept_size} keep_submod_names: {self.submod_cpu_names}") + + +@SchedulerFactory.register("HEURISTIC") +class HeuristicScheduler(BaseScheduler): + def __init__( + self, + compile_config: CompileConfig, + submod_nodes: List[Node], + submod_weights_map: Dict[str, List[str]], + name_node_map: Dict[str, Node], + weight_sizes: Dict[str, int], + ): + super().__init__(compile_config, submod_nodes, submod_weights_map, name_node_map, weight_sizes) + + # safety margin (ms), default 0.1ms, to prevent pipeline bubbles due to profile fluctuations + self.safety_margin = getattr(compile_config, "prefetch_margin_ms", 0.1) + + # prefetch schedule: { trigger_node_name : target_submod_index } + # when execute the trigger_node_name, prefetch the target_submod_index + self.prefetch_schedule = collections.defaultdict() + self.submod_load_events: Dict[str, torch.cuda.Event] = collections.defaultdict(torch.cuda.Event) + for i in range(len(self.submod_nodes)): + self.submod_load_events[self.submod_nodes[i].name] = torch.cuda.Event() + + def schedule_kept_weights(self, profile_data: Dict[str, Any]): + """ + Core logic: build the prefetch schedule based on the profile data. + profile_data format: + { + "submod_name": { + "compute": 10.5, + "h2d": 5.2 + }, + "h2d_bandwidth": 5.2 + } + """ + self.kept_weights.clear() + self.prefetch_schedule.clear() + weight_sizes = self.weight_sizes + + len(self.submod_nodes) + compute_times = [] + h2d_times = [] + + # GiB/s + EST_BANDWIDTH = profile_data.get("h2d_bandwidth") * self.compile_config.offload_config.bandwidth_safety_factor + + for i, node in enumerate(self.submod_nodes): + data = profile_data.get(node.name, {}) + c_time = data.get("compute", 0.0) + + w_size = weight_sizes.get(node.name, 0) + # w_size is in bytes, EST_BANDWIDTH is in GB/s (1e9 bytes/s) + # h_time should be in milliseconds + h_time = (w_size / (EST_BANDWIDTH * 1e6)) if EST_BANDWIDTH > 0 else 0 + + compute_times.append(c_time) + h2d_times.append(h_time) + + keep_submod_names = collections.defaultdict(list) + + schedule_nodes = collections.defaultdict(dict) + + submod_cpu_name = set() + cpu_weight_size = 0 + + for i in range(len(self.submod_nodes) - 1, -1, -1): + if weight_sizes.get(self.submod_nodes[i].name, 0.0) > 0.0: + offload_idx = i + break + # traverse submod_nodes in reverse order to ensure the latest start of transmission + i = offload_idx - 1 + schedule_node_compute_time = 0 + while i >= 0: + end_idx = i + schedule_node_compute_time = 0 + + while schedule_node_compute_time < self.safety_margin + h2d_times[offload_idx] and i >= 0: + schedule_node_compute_time += compute_times[i] + i -= 1 + + if i < 0 and schedule_node_compute_time < self.safety_margin + h2d_times[offload_idx]: + break + + submod_cpu_name.add(self.submod_nodes[offload_idx].name) + cpu_weight_size += weight_sizes.get(self.submod_nodes[offload_idx].name, 0) + + schedule_nodes[offload_idx].update( + { + "compute_start_idx": i + 1, + "compute_end_idx": end_idx, + "h2d_idx": offload_idx, + "compute_time": schedule_node_compute_time, + "h2d_time": h2d_times[offload_idx], + "ratio": schedule_node_compute_time / (self.safety_margin + h2d_times[offload_idx]), + } + ) + + offload_idx = i + 1 + + while weight_sizes.get(self.submod_nodes[offload_idx].name, 0.0) == 0.0 and offload_idx <= end_idx: + offload_idx += 1 + + gpu_resident_weight_ratio = self.compile_config.offload_config.gpu_resident_weight_ratio + + total_weight_size = sum(weight_sizes.values()) + limit = total_weight_size * gpu_resident_weight_ratio + keep_gpu_size = total_weight_size - cpu_weight_size + + while limit < keep_gpu_size: + schedule_nodes_list = [(i, s_node["ratio"]) for i, s_node in schedule_nodes.items()] + schedule_nodes_list.sort(key=lambda x: x[1], reverse=True) + changed = False + for s_node_idx, _ in schedule_nodes_list: + s_node = schedule_nodes[s_node_idx] + compute_start_idx = s_node["compute_start_idx"] + + for i in range(s_node["compute_start_idx"] + 1, s_node["compute_end_idx"] + 1): + if self.submod_nodes[i].name in submod_cpu_name or weight_sizes.get(self.submod_nodes[i].name, 0.0) == 0.0: + continue + + offload_idx = i + schedule_node_compute_time = sum(compute_times[compute_start_idx:i]) + schedule_nodes[offload_idx].update( + { + "compute_start_idx": compute_start_idx, + "compute_end_idx": max(i - 1, compute_start_idx), + "h2d_idx": offload_idx, + "compute_time": schedule_node_compute_time, + "h2d_time": h2d_times[offload_idx], + "ratio": schedule_node_compute_time / (self.safety_margin + h2d_times[offload_idx]), + } + ) + + submod_cpu_name.add(self.submod_nodes[i].name) + keep_gpu_size -= weight_sizes.get(self.submod_nodes[i].name, 0) + + schedule_nodes[s_node_idx].update( + { + "compute_start_idx": i, + "compute_end_idx": s_node["compute_end_idx"], + "h2d_idx": s_node_idx, + "compute_time": s_node["compute_time"] - schedule_node_compute_time, + "h2d_time": h2d_times[s_node_idx], + "ratio": (s_node["compute_time"] - schedule_node_compute_time) + / (self.safety_margin + h2d_times[s_node_idx]), + } + ) + + changed = True + break + + if changed: + break + + if not changed: + break + + self.submod_cpu_names = submod_cpu_name + + for i, s_node in schedule_nodes.items(): + self.prefetch_schedule[self.submod_nodes[s_node["compute_start_idx"]].name] = s_node["h2d_idx"] + + # add weights of submod that are not on CPU + keep_submod_names = set() + for i in range(len(self.submod_nodes)): + if self.submod_nodes[i].name not in submod_cpu_name: + for w_name in self.submod_weights_map.get(self.submod_nodes[i].name, []): + self.kept_weights.add(w_name) + keep_submod_names.add(self.submod_nodes[i].name) + + def prefetch(self, current_node_name: str, ctx: OffloadRuntimeContext): + """ + 1. check if the current node is on CPU, if yes, wait for the submod load event + 2. if not on CPU, skip + 3. prefetch the weights of the submod that the current node is pointing to + """ + if ctx.need_profile: + return super().prefetch(current_node_name, ctx) + + env = ctx.env + h2d_stream = ctx.h2d_stream + buffers = ctx.buffers + submod_0_next_iter_weights = ctx.submod_0_handoff + compute_stream = ctx.compute_stream + + if current_node_name in self.submod_cpu_names: + compute_stream.wait_event(self.submod_load_events[current_node_name]) + + h2d_idx = self.prefetch_schedule.get(current_node_name, None) + if h2d_idx is None: + return + + compute_complete_event = torch.cuda.Event() + compute_complete_event.record(compute_stream) + h2d_stream.wait_event(compute_complete_event) + + h2d_node = self.submod_nodes[h2d_idx] + is_last_step = h2d_idx == len(self.submod_nodes) - 1 + + weight_names = self.submod_weights_map.get(h2d_node.name, []) + + with torch.cuda.stream(h2d_stream): + for w_name in weight_names: + if w_name in self.kept_weights: + continue + + w_node = self.name_node_map.get(w_name) + + if w_node in env and hasattr(env[w_node], 'device') and env[w_node].device.type == "cpu": + env[w_node] = env[w_node].to("cuda", non_blocking=True) + elif is_last_step and w_name in buffers: + submod_0_next_iter_weights[w_node] = buffers[w_name].to("cuda", non_blocking=True) + + self.submod_load_events[h2d_node.name].record(h2d_stream) diff --git a/pkgs/MagiCompiler/magi_compiler/partition_rules.py b/pkgs/MagiCompiler/magi_compiler/partition_rules.py new file mode 100644 index 0000000000000000000000000000000000000000..59c9ca769cfa91355a6e46adf4c4c0a47f5a3586 --- /dev/null +++ b/pkgs/MagiCompiler/magi_compiler/partition_rules.py @@ -0,0 +1,93 @@ +# 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. + + +import contextlib + +import torch +from torch._library.utils import lookup_op + +from .utils import magi_logger + + +def resolve_defined_ops(op_names: list[str]) -> list["torch._ops.OpOverload"]: + """Resolve operator names to OpOverload objects. + + Skips operators that fail to resolve (e.g., operators not registered or + model-specific operators not present in the current model). + + Note: Users should inspect the operator graph before lowering and ensure + the specified operators are present in the final graph. Built-in PyTorch + operators (aten::*, torch::*) may be decomposed, fused, or transformed + during Inductor's compilation passes, so use them with caution. + + Args: + op_names: List of operator names in PyTorch format + (e.g., "vllm::unified_attention") + + Returns: + List of successfully resolved operator overloads + """ + resolved = [] + for op_name in op_names: + try: + resolved.append(lookup_op(op_name)) + except Exception: + # Skip operators that don't exist (e.g., model-specific ops) + magi_logger.info(f"Failed to resolve operator for graph partition: {op_name}") + continue + + return resolved + + +@contextlib.contextmanager +def inductor_partition_rule_context(overloads: list["torch._ops.OpOverload"]): + """Context manager to temporarily register Inductor partition rules. + + Registers custom partition rules for specified operators, forcing the + Inductor scheduler to partition the graph at these operators. The rules + are automatically restored to their previous state on exit. + + Note: Callers should use resolve_defined_ops() to convert operator names + to OpOverload objects before calling this function. + + Args: + overloads: List of resolved operator overload objects. + """ + if not overloads: + magi_logger.info("No partition ops provided; skipping rule registration.") + yield + return + + # NOTE: May cause error in this version of PyTorch + from torch._inductor.scheduler import _custom_should_partition_fns, register_should_partition_rule # type: ignore + + def _always_partition(*_args, **_kwargs): + return True + + # Save current state before registering + saved_rules = _custom_should_partition_fns.copy() + + for overload in overloads: + register_should_partition_rule(overload, _always_partition) + + magi_logger.info("Registered inductor partition rules for %d operators", len(overloads)) + + try: + yield + finally: + # Clear and restore previous state + _custom_should_partition_fns.clear() + _custom_should_partition_fns.update(saved_rules) + magi_logger.info("Restored previous partition rules state.") diff --git a/pkgs/MagiCompiler/magi_compiler/passes/__init__.py b/pkgs/MagiCompiler/magi_compiler/passes/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..092933634030cce28690b6b7842eac3306871c2f --- /dev/null +++ b/pkgs/MagiCompiler/magi_compiler/passes/__init__.py @@ -0,0 +1,18 @@ +# 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. + +from .inductor_pass import InductorPass +from .pass_manager import PostGradPassManager + +__all__ = ["PostGradPassManager", "InductorPass"] diff --git a/pkgs/MagiCompiler/magi_compiler/passes/fix_functionalization.py b/pkgs/MagiCompiler/magi_compiler/passes/fix_functionalization.py new file mode 100644 index 0000000000000000000000000000000000000000..8703f97db5ca7842a93c39ae67a1caa221b8e81c --- /dev/null +++ b/pkgs/MagiCompiler/magi_compiler/passes/fix_functionalization.py @@ -0,0 +1,210 @@ +# 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. + +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +import operator +from collections.abc import Iterable + +import torch +from torch._higher_order_ops.auto_functionalize import auto_functionalized + +from ..utils import is_func, magi_logger +from .magi_inductor_pass import MagiInductorPass + + +class FixFunctionalizationPass(MagiInductorPass): + """ + This pass defunctionalizes certain nodes to avoid redundant tensor copies. + After this pass, DCE (dead-code elimination) should never be run, + as de-functionalized nodes may appear as dead code. + + To add new nodes to defunctionalize, add to the if-elif chain in __call__. + """ + + @MagiInductorPass.time_and_log + def __call__(self, graph: torch.fx.Graph): + self.nodes_to_remove: list[torch.fx.Node] = [] + count = 0 + for node in graph.nodes: + if not is_func(node, auto_functionalized): + continue # Avoid deep if-elif nesting + + kwargs = node.kwargs + at_target = node.args[0] + + if at_target == torch.ops._C.rotary_embedding.default: + query = kwargs["query"] + key = kwargs["key"] + getitem_nodes = self.getitem_users(node) + + if ( + is_func(query, operator.getitem) + and is_func(key, operator.getitem) + and query.args[0] == key.args[0] + and is_func(query.args[0], torch.ops.aten.split_with_sizes.default) + and all( + is_func(user, torch.ops.aten.slice_scatter.default) + for getitem_node in getitem_nodes.values() + for user in getitem_node.users + ) + ): + # Pattern where query and key are slices of an mm_node. + # While functionalized, results at [1] and [2] are scattered + # back into mm_node. So after de-functionalization, we can + # just use mm_node directly. + + mm_node = query.args[0].args[0] + for user in getitem_nodes.values(): + for user_of_getitem in user.users: + if is_func(user_of_getitem, torch.ops.aten.slice_scatter.default): + user_of_getitem.replace_all_uses_with(mm_node) + self._remove(user_of_getitem) + self._remove(user) + + self.insert_defunctionalized(graph, node) + self._remove(node) + + else: + # Directly replace the auto_functionalize(rotary_embedding) + # with the inplace rotary_embedding. In theory, we shouldn't + # do this blindly, but in practice in vLLM it's ok. The best + # solution is to use auto_functionalization_v2 and then use + # inductor's builtin defunctionalization (reinplacing) pass. + mutated_args = {1: "query", 2: "key"} + self.defunctionalize(graph, node, mutated_args) + + # rms_norm replacements avoid the most copies for LLaMa. + elif at_target == torch.ops._C.fused_add_rms_norm.default: + mutated_args = {1: "input", 2: "residual"} + self.defunctionalize(graph, node, mutated_args) + elif at_target == torch.ops._C.fused_add_rms_norm_static_fp8_quant.default: # noqa: E501 + mutated_args = {1: "result", 2: "residual"} + self.defunctionalize(graph, node, mutated_args) + elif at_target == torch.ops._C.rms_norm_dynamic_per_token_quant.default: # noqa: E501 + mutated_args = {1: "result", 2: "scale", 3: "residual"} + self.defunctionalize(graph, node, mutated_args) + elif at_target in [torch.ops._C.rms_norm.default, torch.ops._C.rms_norm_static_fp8_quant.default]: + mutated_args = {1: "result"} + self.defunctionalize(graph, node, mutated_args) + # For some reason we need to specify the args for both + # silu_and_mul and silu_and_mul_quant. The kwargs + # pathway gets the wrong answer. + elif at_target == torch.ops._C.silu_and_mul.default: + mutated_args = {1: "result"} + self.defunctionalize(graph, node, mutated_args, args=("result", "input")) + elif at_target == torch.ops._C.silu_and_mul_quant.default: + mutated_args = {1: "result"} + self.defunctionalize(graph, node, mutated_args, args=("result", "input", "scale")) + elif ( + hasattr(torch.ops._C, "silu_and_mul_nvfp4_quant") + and at_target == torch.ops._C.silu_and_mul_nvfp4_quant.default + ): + mutated_args = {1: "result", 2: "result_block_scale"} + self.defunctionalize( + graph, node, mutated_args, args=("result", "result_block_scale", "input", "input_global_scale") + ) + else: + continue # skip the count + + count += 1 + + self.dump_graph(graph, "before_cleanup") + + # Remove the nodes all at once + count_removed = len(self.nodes_to_remove) + for node in self.nodes_to_remove: + graph.erase_node(node) + + magi_logger.info("De-functionalized %s nodes, removed %s nodes", count, count_removed) + self.nodes_to_remove.clear() + + def _remove(self, node_or_nodes: torch.fx.Node | Iterable[torch.fx.Node]): + """ + Stage a node (or nodes) for removal at the end of the pass. + """ + if isinstance(node_or_nodes, torch.fx.Node): + self.nodes_to_remove.append(node_or_nodes) + else: + self.nodes_to_remove.extend(node_or_nodes) + + def defunctionalize( + self, + graph: torch.fx.Graph, + node: torch.fx.Node, + mutated_args: dict[int, torch.fx.Node | str], + args: tuple[torch.fx.Node | str, ...] | None = None, + ): + """ + De-functionalize a node by replacing it with a call to the original. + It also replaces the getitem users with the mutated arguments. + See replace_users_with_mutated_args and insert_defunctionalized. + """ + self.replace_users_with_mutated_args(node, mutated_args) + self.insert_defunctionalized(graph, node, args=args) + self._remove(node) + + def replace_users_with_mutated_args(self, node: torch.fx.Node, mutated_args: dict[int, torch.fx.Node | str]): + """ + Replace all getitem users of the auto-functionalized node with the + mutated arguments. + :param node: The auto-functionalized node + :param mutated_args: The mutated arguments, indexed by getitem index. + If the value of an arg is a string, `node.kwargs[arg]` is used. + """ + for idx, user in self.getitem_users(node).items(): + arg = mutated_args[idx] + arg = node.kwargs[arg] if isinstance(arg, str) else arg + user.replace_all_uses_with(arg) + self._remove(user) + + def getitem_users(self, node: torch.fx.Node) -> dict[int, torch.fx.Node]: + """ + Returns the operator.getitem users of the auto-functionalized node, + indexed by the index they are getting. + """ + users = {} + for user in node.users: + if is_func(user, operator.getitem): + idx = user.args[1] + # NOTE: Corner case: Maybe multiple users for the same index? + users[idx] = user + return users + + def insert_defunctionalized( + self, graph: torch.fx.Graph, node: torch.fx.Node, args: tuple[torch.fx.Node | str, ...] | None = None + ): + """ + Insert a new defunctionalized node into the graph before node. + If one of the kwargs is 'out', provide args directly, + as node.kwargs cannot be used. + See https://github.com/pytorch/pytorch/blob/a00faf440888ffb724bad413f329a49e2b6388e7/torch/_inductor/lowering.py#L351 + + :param graph: Graph to insert the defunctionalized node into + :param node: The auto-functionalized node to defunctionalize + :param args: If we cannot use kwargs, specify args directly. + If an arg is a string, `node.kwargs[arg]` is used. + """ # noqa: E501 + assert is_func(node, auto_functionalized), f"node must be auto-functionalized, is {node} instead" + + # Create a new call to the original function + with graph.inserting_before(node): + function = node.args[0] + if args is None: + graph.call_function(function, kwargs=node.kwargs) + else: + # Args passed as strings refer to items in node.kwargs + args = tuple(node.kwargs[arg] if isinstance(arg, str) else arg for arg in args) + graph.call_function(function, args=args) diff --git a/pkgs/MagiCompiler/magi_compiler/passes/inductor_pass.py b/pkgs/MagiCompiler/magi_compiler/passes/inductor_pass.py new file mode 100644 index 0000000000000000000000000000000000000000..6b341f0a26dfb2cdb8c5665d4b61a1455eaa14f5 --- /dev/null +++ b/pkgs/MagiCompiler/magi_compiler/passes/inductor_pass.py @@ -0,0 +1,99 @@ +# 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. + +import hashlib +import inspect +import json +import types +from contextlib import contextmanager +from typing import Any + +from torch._inductor.custom_graph_pass import CustomGraphPass + +from ..utils import compute_hash + +_pass_context = None + + +class PassContext: + def __init__(self, runtime_shape: int | None): + self.runtime_shape = runtime_shape + + +def get_pass_context() -> PassContext: + """Get the current pass context.""" + assert _pass_context is not None + return _pass_context + + +@contextmanager +def pass_context(runtime_shape: int | None): + """ + A context manager that stores the current pass context, usually it is a list of sizes to specialize. + """ + global _pass_context + prev_context = _pass_context + _pass_context = PassContext(runtime_shape) + try: + yield + finally: + _pass_context = prev_context + + +class InductorPass(CustomGraphPass): + """ + A custom graph pass that uses a hash of its source as the UUID. + This is defined as a convenience and should work in most cases. + """ + + def uuid(self) -> Any: + """ + Provide a unique identifier for the pass, used in Inductor code cache. + This should depend on the pass implementation, so that changes to the + pass result in recompilation. + By default, the object source is hashed. + """ + return InductorPass.hash_source(self) + + @staticmethod + def hash_source(*srcs: str | Any): + """ + Utility method to hash the sources of functions or objects. + :param srcs: strings or objects to add to the hash. + Objects and functions have their source inspected. + :return: + """ + hasher = hashlib.sha256() + for src in srcs: + if isinstance(src, str): + src_str = src + elif isinstance(src, (types.FunctionType, type)): + src_str = inspect.getsource(src) + else: + # object instance + src_str = inspect.getsource(src.__class__) + hasher.update(src_str.encode("utf-8")) + return hasher.hexdigest() + + @staticmethod + def hash_dict(dict_: dict[Any, Any]): + """ + Utility method to hash a dictionary, can alternatively be used for uuid. + :return: A sha256 hash of the json rep of the dictionary. + """ + encoded = json.dumps(dict_, sort_keys=True).encode("utf-8") + return compute_hash(encoded) + + def is_applicable(self, shape: int | None): + return True diff --git a/pkgs/MagiCompiler/magi_compiler/passes/magi_inductor_pass.py b/pkgs/MagiCompiler/magi_compiler/passes/magi_inductor_pass.py new file mode 100644 index 0000000000000000000000000000000000000000..635177b9560c6c6493c556c7eb27ac5195ba3f0d --- /dev/null +++ b/pkgs/MagiCompiler/magi_compiler/passes/magi_inductor_pass.py @@ -0,0 +1,74 @@ +# 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. + +import functools +import time +from dataclasses import dataclass +from typing import ClassVar + +import torch +from magi_compiler.config import CompileConfig +from magi_compiler.utils import magi_logger +from torch._dynamo.utils import lazy_format_graph_code + +from .inductor_pass import InductorPass + + +@dataclass +class InductorCompilationConfig: + splitting_ops: list[str] | None = None + use_inductor_graph_partition: bool = False + + +class MagiInductorPass(InductorPass): + """ + An inductor pass with access to MagiCompiler PassConfig. It provides timing, logging, and dumping utilities. + """ + + # Keep track of pass index for debug dump ordering. + dump_prefix: ClassVar[int | None] = None + + def __init__(self, config: CompileConfig): + # Get only the necessary CompilationConfig for the inductor pass, since + # full `CompilationConfig` contains pointer to model which is unsafe. + self.compilation_config = InductorCompilationConfig( + splitting_ops=config.splitting_ops, use_inductor_graph_partition=config.use_inductor_graph_partition + ) + self.pass_config = config.pass_config + self.pass_name = self.__class__.__name__ + + @staticmethod + def time_and_log(call_fn): + @functools.wraps(call_fn) + def wrapped(self: MagiInductorPass, graph: torch.fx.Graph): + self.begin() + self.dump_graph(graph, "before") + call_fn(self, graph) + self.dump_graph(graph, "after") + self.end_and_log() + + return wrapped + + def dump_graph(self, graph: torch.fx.Graph, stage: str): + i = MagiInductorPass.dump_prefix + i_str = "" if i is None else f".{i}" + lazy_format_graph_code(f"post_grad{i_str}.{self.pass_name}.{stage}", graph.owning_module) + + def begin(self): + self._start_time = time.perf_counter_ns() + + def end_and_log(self): + self._end_time = time.perf_counter_ns() + duration_ms = float(self._end_time - self._start_time) / 1.0e6 + magi_logger.info("%s completed in %.1f ms", self.pass_name, duration_ms) diff --git a/pkgs/MagiCompiler/magi_compiler/passes/pass_manager.py b/pkgs/MagiCompiler/magi_compiler/passes/pass_manager.py new file mode 100644 index 0000000000000000000000000000000000000000..89dd06f49d2ef3402982b0776f7b2dd2f43251ae --- /dev/null +++ b/pkgs/MagiCompiler/magi_compiler/passes/pass_manager.py @@ -0,0 +1,150 @@ +# 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. + +import functools + +from torch import fx as fx + +from ..config import CompileConfig +from ..utils import magi_logger, set_env_var +from ..utils.envs import MAGI_PATTERN_MATCH_DEBUG +from .fix_functionalization import FixFunctionalizationPass +from .inductor_pass import CustomGraphPass, InductorPass, get_pass_context +from .magi_inductor_pass import MagiInductorPass +from .post_cleanup import PostCleanupPass + + +def with_pattern_match_debug(fn): + """ + Function decorator that turns on inductor pattern match debug + for the duration of the call. + Used to avoid logging builtin Inductor pattern matching. + """ + + @functools.wraps(fn) + def wrapper(*args, **kwargs): + if (debug_val := MAGI_PATTERN_MATCH_DEBUG) is not None: + # optionally check rank here + with set_env_var("TORCHINDUCTOR_PATTERN_MATCH_DEBUG", debug_val): + return fn(*args, **kwargs) + return fn(*args, **kwargs) + + return wrapper + + +class PostGradPassManager(CustomGraphPass): + """ + The pass manager for post-grad passes. + It handles configuration, adding custom passes, and running passes. + It supports uuid for the Inductor code cache. + + The order of the post-grad post-passes is: + 1. passes (constructor parameter) + 2. default passes (NoopEliminationPass, FusionPass) + 3. config["post_grad_custom_post_pass"] (if it exists) + 4. fix_functionalization + This way, all passes operate on a functionalized graph. + """ + + def __init__(self): + self.passes: list[InductorPass] = [] + + @with_pattern_match_debug + def __call__(self, graph: fx.Graph): + magi_logger.info("Run PostGradPassManager") + MagiInductorPass.dump_prefix = 0 # reset dump index + + shape = get_pass_context().runtime_shape + for pass_ in self.passes: + # NOTE: while better than if? + if pass_.is_applicable(shape): + pass_(graph) + MagiInductorPass.dump_prefix += 1 + else: + magi_logger.info("Skipping %s with shape %s", pass_, shape) + + # post-cleanup goes before fix_functionalization because it requires a functional graph + self.post_cleanup(graph) + MagiInductorPass.dump_prefix += 1 + + # always run fix_functionalization last + self.fix_functionalization(graph) + MagiInductorPass.dump_prefix = None # Cleanup index + + def configure(self, config: CompileConfig): + self.pass_config = config.pass_config + + # TODO: Support custom passes later, add UlyssesOverlapPass here. + # if self.pass_config.enable_noop: + # self.passes += [NoOpEliminationPass(config)] + + # if self.pass_config.enable_sequence_parallelism: + # self.passes += [SequenceParallelismPass(config)] + # if self.pass_config.enable_async_tp: + # self.passes += [AsyncTPPass(config)] + + # if self.pass_config.enable_fi_allreduce_fusion: + # self.passes += [AllReduceFusionPass(config)] + + # if self.pass_config.enable_fusion: + # self.passes += [RMSNormQuantFusionPass(config)] + # self.passes += [ActivationQuantFusionPass(config)] + + # if self.pass_config.enable_attn_fusion: + # self.passes += [AttnFusionPass(config)] + + # needs a functional graph + self.post_cleanup = PostCleanupPass(config) + self.fix_functionalization = FixFunctionalizationPass(config) + + # [HACK: Bug with Inductor graph partition and torch.compile cache] + # In PyTorch 2.9, torch.compile has a bug where the graph + # partition is not taken into account during caching. + # Because CompileMode.MAGI_COMPILE is the only mode that uses + # Inductor graph partition, and MAGI_COMPILE implies there + # is a PostGradPassManager, we put the list of operators to graph + # partition into the PostGradPassManager's uuid (which + # then gets incorporated into Inductor's FX graph cache key). + # Remove this hack whenever torch.compile fixes it. + + # This is the list of operators that vLLM asks Inductor to split. + self.inductor_splitting_ops = [] + if config.use_inductor_graph_partition and config.splitting_ops is not None: + # Sort them so we're not dependent on the ordering. + self.inductor_splitting_ops = sorted(config.splitting_ops) + + def add(self, pass_: InductorPass): + assert isinstance(pass_, InductorPass) + self.passes.append(pass_) + + def uuid(self): + """ + The PostGradPassManager is set as a custom pass in the Inductor and + affects compilation caching. Its uuid depends on the UUIDs of all + dependent passes and the pass config. See InductorPass for more info. + """ + state = { + "pass_config": self.pass_config.uuid(), # TODO: Find a way to hash the pass config. + "passes": [], + "inductor_splitting_ops": [], + } + for pass_ in self.passes: + state["passes"].append(pass_.uuid()) + state["passes"].append(self.post_cleanup.uuid()) + state["passes"].append(self.fix_functionalization.uuid()) + + # See [HACK: Bug with Inductor graph partition and torch.compile cache] + state["inductor_splitting_ops"].extend(self.inductor_splitting_ops) + + return InductorPass.hash_dict(state) diff --git a/pkgs/MagiCompiler/magi_compiler/passes/post_cleanup.py b/pkgs/MagiCompiler/magi_compiler/passes/post_cleanup.py new file mode 100644 index 0000000000000000000000000000000000000000..edc12f9781fc609373e6f18543e1d939a645bb83 --- /dev/null +++ b/pkgs/MagiCompiler/magi_compiler/passes/post_cleanup.py @@ -0,0 +1,32 @@ +# 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. + +from torch import fx +from torch._inductor.pattern_matcher import stable_topological_sort + +from .magi_inductor_pass import MagiInductorPass + + +class PostCleanupPass(MagiInductorPass): + """ + This pass performs cleanup after custom passes. + It topologically sorts the graph and removes unused nodes. + This is needed because the pattern matcher does not guarantee producing + a topologically sorted graph, and there may be unused nodes left around. + """ + + @MagiInductorPass.time_and_log + def __call__(self, graph: fx.Graph) -> None: + stable_topological_sort(graph) + graph.eliminate_dead_code() diff --git a/pkgs/MagiCompiler/magi_compiler/passes/replace_pass.py b/pkgs/MagiCompiler/magi_compiler/passes/replace_pass.py new file mode 100644 index 0000000000000000000000000000000000000000..1eada95b3abb45f18235aa0d68efae23df22f696 --- /dev/null +++ b/pkgs/MagiCompiler/magi_compiler/passes/replace_pass.py @@ -0,0 +1,34 @@ +# 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. + +import torch + + +def replace_fa_with_sage(gm: torch.fx.GraphModule) -> torch.fx.GraphModule: + for node in gm.graph.nodes: + if node.op == "call_function" and node.target == torch.ops.athena.flash_attn_func: + node.target = torch.ops.athena.sage_attn_func + + +class FullGraphPassManager: + """ + A manager to apply various graph passes on the full graph before splitting. + """ + + def __init__(self, pass_config): + self.pass_config = pass_config + + def __call__(self, gm: torch.fx.GraphModule): + if self.pass_config.enable_sage_attn: + replace_fa_with_sage(gm) diff --git a/pkgs/MagiCompiler/magi_compiler/piecewise_backend.py b/pkgs/MagiCompiler/magi_compiler/piecewise_backend.py new file mode 100644 index 0000000000000000000000000000000000000000..f73ef7c5d783b182b39863737c2197f5d923615b --- /dev/null +++ b/pkgs/MagiCompiler/magi_compiler/piecewise_backend.py @@ -0,0 +1,115 @@ +# 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. + +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +import dataclasses +from collections.abc import Callable +from typing import TYPE_CHECKING, Any + +import torch.fx as fx + +from .config import CompileConfig +from .utils.compile_time_monitor import CompileMonitor + +if TYPE_CHECKING: + from .magi_backend import CompilerManager + + +@dataclasses.dataclass +class ConcreteSizeEntry: + runtime_shape: int + compiled: bool = False + runnable: Callable = None # type: ignore + + +class PiecewiseBackend: + def __init__( + self, + graph: fx.GraphModule, + compiled_graph_for_general_shape: Callable, + compile_config: CompileConfig, + piecewise_compile_index: int, + piecewise_submodule_number: int, + sym_shape_indices: list[int], + compiler_manager: "CompilerManager", + ): + """ + The backend for piecewise compilation. It mainly handles the compilation of static shapes and dispatching based on runtime shape. + + We will compile `self.graph` once for the general shape, and then compile for different shapes specified in `compile_config.compile_sizes`. + """ + self.graph = graph + self.compiled_graph_for_general_shape = compiled_graph_for_general_shape + self.compile_config = compile_config + self.piecewise_compile_index = piecewise_compile_index + self.piecewise_submodule_number = piecewise_submodule_number + self.compiler_manager = compiler_manager + self.sym_shape_indices = sym_shape_indices + + self.is_first_graph = piecewise_compile_index == 0 + self.is_last_graph = piecewise_compile_index == piecewise_submodule_number - 1 + self.is_first_run = True + + # to_be_compiled_sizes tracks the remaining sizes to compile, + # and updates during the compilation process, so we need to copy it + self.to_be_compiled_sizes: set[int] = set(self.compile_config.compile_sizes) + + # the entries for different shapes that we need to compile + self.concrete_size_entries: dict[int, ConcreteSizeEntry] = {} + for shape in self.to_be_compiled_sizes: + self.concrete_size_entries[shape] = ConcreteSizeEntry( + runtime_shape=shape, runnable=self.compiled_graph_for_general_shape + ) + + def check_for_ending_compilation(self): + if self.is_last_graph and not self.to_be_compiled_sizes: + # no specific sizes to compile, save the cache for the next run + self.compiler_manager.save_to_file() + CompileMonitor().end() + + def __call__(self, *args) -> Any: + if self.is_first_run: + self.is_first_run = False + self.check_for_ending_compilation() + return self.compiled_graph_for_general_shape(*args) + + assert len(self.sym_shape_indices) != 0, "No symbolic shape indices found" + runtime_shape = args[self.sym_shape_indices[0]] + if runtime_shape not in self.concrete_size_entries: + # we don't need to do anything for this shape + return self.compiled_graph_for_general_shape(*args) + + entry = self.concrete_size_entries[runtime_shape] + + if not entry.compiled: + entry.compiled = True + self.to_be_compiled_sizes.remove(runtime_shape) + # args are real arguments + entry.runnable = self.compiler_manager.compile( + self.graph, + args, + self.compile_config.inductor_compile_config, + self.compile_config, + graph_index=self.piecewise_compile_index, + num_graphs=self.piecewise_submodule_number, + runtime_shape=runtime_shape, + ) + + # finished compilations for all required shapes + if self.is_last_graph and not self.to_be_compiled_sizes: + self.check_for_ending_compilation() + + return entry.runnable(*args) diff --git a/pkgs/MagiCompiler/magi_compiler/piecewise_compiler.py b/pkgs/MagiCompiler/magi_compiler/piecewise_compiler.py new file mode 100644 index 0000000000000000000000000000000000000000..c5d727a6cce166f072c50957de3941c7dff50359 --- /dev/null +++ b/pkgs/MagiCompiler/magi_compiler/piecewise_compiler.py @@ -0,0 +1,216 @@ +# 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. + +from abc import abstractmethod +from collections.abc import Callable +from pathlib import Path +from typing import Any + +import magi_compiler.utils.envs as envs +import torch +import torch._inductor.compile_fx +import torch.fx as fx + +from ._cache_data_cls import CacheEntry, CacheHandle +from .utils import compilation_counter, compute_hash, magi_logger + + +class CompilerInterface: + """ + The interface for a compiler that can be used by MagiCompiler. + """ + + # The name of the compiler, e.g. inductor. This is a class-level attribute. + name: str + + @abstractmethod + def initialize_cache(self, cache_dir: Path, prefix: str = ""): + """ + when the MagiCompiler process uses `cache_dir` as the cache directory, + the compiler should initialize itself with the cache directory, + e.g. by re-directing its own cache directory to a sub-directory. + + prefix can be used in combination with cache_dir to figure out the base + cache directory, e.g. there're multiple parts of model being compiled, + but we want to share the same cache directory for all of them. + + e.g. + cache_dir = "/path/to/dir/backbone", prefix = "backbone" + cache_dir = "/path/to/dir/eagle_head", prefix = "eagle_head" + """ + pass + + @property + @abstractmethod + def hash(self) -> str: + """ + Gather all the relevant information from the config, to compute a hash so that we can cache the compiled model. + + This function should only consider the information that is specific to the compiler. + """ + return "" + + @abstractmethod + def compile( + self, + graph: fx.GraphModule, + example_inputs: list[Any], + inductor_compile_config: dict[str, Any], + runtime_shape: int | None = None, + key: str | None = None, + ) -> tuple[Callable | None, Any | None]: + """ + Compile the graph with the given example inputs and compiler config, + with a runtime shape. If the `runtime_shape` is None, it means + the `example_inputs` have a dynamic shape. Otherwise, the + `runtime_shape` specifies the shape of the inputs. + Right now we only support one variable shape for all inputs, which is the sequence length + (number of tokens) during inference. + + Dynamo will make sure `graph(*example_inputs)` is valid. + + The function should return a compiled callable function, as well as + a handle that can be used to directly load the compiled function. + + The handle should be a plain Python object, preferably a string or a + file path for readability. + + If the compiler doesn't support caching, it should return None for the + handle. If the compiler fails to compile the graph, it should return + None for the compiled function as well. + + `key` is required for StandaloneInductorAdapter, it specifies where to + save the compiled artifact. The compiled artifact gets saved to + `cache_dir/key`. + """ + return None, None + + @abstractmethod + def load( + self, graph: fx.GraphModule, example_inputs: list[Any], cache_entry: CacheEntry, cache_handle: CacheHandle + ) -> Callable: + """ + Load the compiled function from the handle. Raises an error if the handle is invalid. + + The handle is the second return value of the `compile` function. + """ + raise NotImplementedError("caching is not supported") + + +class InductorStandaloneAdaptor(CompilerInterface): + """ + The adaptor for the Inductor compiler, which requires PyTorch 2.8+. + + Mainly reuses the standalone_compile function from PyTorch Inductor. + """ + + name = "inductor_standalone" + + @property + def hash(self) -> str: + # summarize system state and pytorch state + from torch._inductor.codecache import CacheBase, torch_key + + factors: list[Any] = [CacheBase.get_system(), torch_key()] + return compute_hash(factors) + + def initialize_cache(self, cache_dir: Path, prefix: str = ""): + self.cache_dir: Path = cache_dir + + def compile( + self, + graph: fx.GraphModule, + example_inputs: list[Any], + inductor_compile_config: dict[str, Any], + runtime_shape: int | None = None, + key: str | None = None, + ) -> tuple[Callable | None, CacheHandle | None]: + # Step1: Update compile settings + compilation_counter.num_inductor_compiles += 1 + current_config = {} + if inductor_compile_config is not None: + current_config.update(inductor_compile_config) + if isinstance(runtime_shape, int): + # for a specific sequence length, tuning triton kernel parameters can be beneficial + current_config.update( + { + "max_autotune": envs.MAGI_ENABLE_INDUCTOR_MAX_AUTOTUNE, + "coordinate_descent_tuning": envs.MAGI_ENABLE_INDUCTOR_COORDINATE_DESCENT_TUNING, + } + ) + dynamic_shapes = "from_example_inputs" + else: + dynamic_shapes = "from_tracing_context" + + # Step2: Compile the graph + from torch._inductor import standalone_compile + + compiled_graph = standalone_compile( + graph, example_inputs, dynamic_shapes=dynamic_shapes, options={"config_patches": current_config} + ) + + # Step3: Save the compiled artifact + # When standalone_compile is invoked from within a torch.compile backend, + # AOTAutograd's cache key computation may be silently bypassed, leaving + # aot_autograd artifacts empty. In that case save() will raise an + # AssertionError, so we fall back to running without a cache handle. + # TODO(hongyu): Support caching the compiled artifact. + assert key is not None + if hasattr(self, "cache_dir") and self.cache_dir is not None: + try: + path: Path = self.cache_dir / key + compiled_graph.save(path=path.as_posix(), format="unpacked") + compilation_counter.num_compiled_artifacts_saved += 1 + return compiled_graph, CacheHandle(key, path.as_posix()) + except (AssertionError, RuntimeError) as e: + magi_logger.warning("Failed to save compiled artifact for key '%s', skipping cache: %s", key, e) + + return compiled_graph, None + + def load( + self, graph: fx.GraphModule, example_inputs: list[Any], cache_entry: CacheEntry, cache_handle: CacheHandle + ) -> Callable: + assert isinstance(cache_handle.key, str) and cache_handle.key is not None + assert isinstance(cache_handle.path, str) and cache_handle.path is not None + inductor_compiled_graph = torch._inductor.CompiledArtifact.load(path=cache_handle.path, format="unpacked") + + from torch._inductor.compile_fx import graph_returns_tuple + + is_return_tuple = graph_returns_tuple(graph) + + def compiled_graph_wrapper(*args): + graph_output = inductor_compiled_graph(*args) + if is_return_tuple: + return graph_output + else: + return graph_output[0] + + return compiled_graph_wrapper + + +class EagerAdaptor(CompilerInterface): + name = "eager" + + def compile( + self, + graph: fx.GraphModule, + example_inputs: list[Any], + inductor_compile_config: dict[str, Any], + runtime_shape: int | None = None, + key: str | None = None, + ) -> tuple[Callable | None, CacheHandle | None]: + compilation_counter.num_eager_compiles += 1 + # we don't need to compile the graph, just return the graph itself. + # It does not support caching, return None for the handle. + return graph, None diff --git a/pkgs/MagiCompiler/magi_compiler/tokenflow/__init__.py b/pkgs/MagiCompiler/magi_compiler/tokenflow/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..3eaa44adb1bd11bb4c8d48c6f00d7a08f292f395 --- /dev/null +++ b/pkgs/MagiCompiler/magi_compiler/tokenflow/__init__.py @@ -0,0 +1,13 @@ +# Copyright (c) 2026 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. diff --git a/pkgs/MagiCompiler/magi_compiler/tokenflow/graph_executor.py b/pkgs/MagiCompiler/magi_compiler/tokenflow/graph_executor.py new file mode 100644 index 0000000000000000000000000000000000000000..43ef429187c5342a099ff165341d90ffc1e1036e --- /dev/null +++ b/pkgs/MagiCompiler/magi_compiler/tokenflow/graph_executor.py @@ -0,0 +1,359 @@ +# Copyright (c) 2026 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. + +from enum import Enum +from typing import Any, Dict, List + +import torch +from magi_compiler.tokenflow.green_ctx import GreenCtxManager +from torch import fx + + +class FX_NODE_OP(Enum): + PLACEHOLDER = "placeholder" + GET_ATTR = "get_attr" + CALL_FUNCTION = "call_function" + CALL_METHOD = "call_method" + CALL_MODULE = "call_module" + OUTPUT = "output" + + +class LaneType(Enum): + COMPUTE = "compute" + OTHERS = "others" + + +class NodeOnlyExecutor: + def __init__(self, graph_module: fx.GraphModule, device: torch.device): + self.graph_module = graph_module + self.device = device + self.name_to_node = {node.name: node for node in graph_module.graph.nodes} + + @classmethod + def replace_nodes_in_args(cls, args, value_map): + if isinstance(args, torch.fx.Node): + return value_map[args.name] + elif isinstance(args, (list, tuple)): + return type(args)(cls.replace_nodes_in_args(a, value_map) for a in args) + elif isinstance(args, dict): + return {k: cls.replace_nodes_in_args(v, value_map) for k, v in args.items()} + else: + return args + + def execute(self, node: fx.Node, value_map: Dict[str, Any], stream: torch.cuda.Stream = None) -> None: + node_name = node.name + + if node.op == FX_NODE_OP.PLACEHOLDER.value: + if node_name in value_map: + return value_map[node_name] + raise RuntimeError("PLACEHOLDER节点不应被执行。") + + args = self.replace_nodes_in_args(node.args, value_map) if node.args and node.args != () else () + kwargs = self.replace_nodes_in_args(node.kwargs, value_map) if node.kwargs and node.kwargs != {} else {} + + with torch.cuda.stream(stream): + # with nullcontext(): + if node.op == FX_NODE_OP.GET_ATTR.value: + attr_val = self.graph_module + for attr in node.target.split("."): + attr_val = getattr(attr_val, attr) + result = attr_val + + elif node.op == FX_NODE_OP.CALL_FUNCTION.value: + result = node.target(*args, **kwargs) + + elif node.op == FX_NODE_OP.CALL_METHOD.value: + obj = args[0] + method = getattr(obj, node.target) + result = method(*args[1:], **kwargs) + + elif node.op == FX_NODE_OP.CALL_MODULE.value: + submod = self.graph_module + for mod_name in node.target.split("."): + submod = getattr(submod, mod_name) + result = submod(*args, **kwargs) + + elif node.op == FX_NODE_OP.OUTPUT.value: + result = args[0] if len(args) == 1 else args + + else: + raise NotImplementedError(f"不支持的op类型: {node.op}") + + assert result is not None, f"节点 {node_name} 执行未返回结果。" + value_map[node_name] = result + + +class GraphRawExecutor: + def __init__(self, graph_module: fx.GraphModule, device: torch.device = None): + # 基础属性初始化 + self.graph = graph_module.graph + self.module = graph_module + self.device = device + + self.topological_nodes = [node for node in self.graph.nodes] + self.name_to_node = {node.name: node for node in self.topological_nodes} + + self.node_executor = NodeOnlyExecutor(self.module, self.device) + + self.value_map: Dict[str, Any] = {} + self.stream_map: Dict[str, torch.cuda.Stream] = {} + + for node in self.topological_nodes: + if node.op == FX_NODE_OP.PLACEHOLDER.value: + continue + self.stream_map[node.name] = torch.cuda.default_stream(device=self.device) + + # cuda_graph_mgr().run(func, *args, layer_number=layer_number, **kwargs) + # @cuda_graph_enable_if(condition=lambda: True) + def execute(self, *inputs) -> Any: + for idx, node in enumerate(self.topological_nodes): + if node.op == FX_NODE_OP.PLACEHOLDER.value: + self.value_map[node.name] = inputs[idx] + continue + + self.node_executor.execute(node=node, value_map=self.value_map, stream=self.stream_map[node.name]) + + if node.op == FX_NODE_OP.OUTPUT.value: + output_result = self.value_map[node.name] + break + + return output_result + + def synchronize(self): + for stream in self.stream_map.values(): + if stream is not None: + stream.synchronize() + + def cleanup(self): + pass + + +class GraphNormalExecutor: + def __init__(self, graph_module: torch.fx.GraphModule, device: torch.device = None): + self.graph = graph_module.graph + self.module = graph_module + self.device = device or torch.device("cuda" if torch.cuda.is_available() else "cpu") + self.name_to_node = {node.name: node for node in self.graph.nodes} + self.topological_names = [node.name for node in self.graph.nodes] + self.value_map = {node.name: None for node in self.graph.nodes} + self.dependencies = {node.name: [] for node in self.graph.nodes} + self.rev_dependencies = {node.name: [] for node in self.graph.nodes} + + self._resolve_node_dependencies() + + self.stream_map: dict[str, torch.cuda.Stream] = {} + self.event_map: dict[str, torch.cuda.Event] = {} + for node_name in self.topological_names: + if self.name_to_node[node_name].op == FX_NODE_OP.PLACEHOLDER.value: + continue + self.stream_map[node_name] = torch.cuda.Stream(device=self.device) + self.event_map[node_name] = torch.cuda.Event(enable_timing=False, blocking=False) + + self.node_executor = NodeOnlyExecutor(self.module, self.device) + + def _resolve_node_dependencies(self): + for node in self.graph.nodes: + dep_node_names = [] + + def extract_deps(arg): + if isinstance(arg, torch.fx.Node): + dep_node_names.append(arg.name) + elif isinstance(arg, (tuple, list)): + for a in arg: + extract_deps(a) + elif isinstance(arg, dict): + for v in arg.values(): + extract_deps(v) + + extract_deps(node.args) + extract_deps(node.kwargs) + dep_node_names = list(set(dep_node_names)) + self.dependencies[node.name] = dep_node_names + + for dep_name in dep_node_names: + self.rev_dependencies[dep_name].append(node.name) + + def wait_for_dependencies(self, node_name: str): + for dep_name in self.dependencies[node_name]: + if self.name_to_node[dep_name].op == FX_NODE_OP.PLACEHOLDER.value: + continue + self.stream_map[node_name].wait_event(self.event_map[dep_name]) + + def _replace_nodes_in_args(self, args): + if isinstance(args, torch.fx.Node): + return self.value_map[args.name] + elif isinstance(args, (tuple, list)): + return type(args)(self._replace_nodes_in_args(a) for a in args) + elif isinstance(args, dict): + return {k: self._replace_nodes_in_args(v) for k, v in args.items()} + else: + return args + + def execute(self, *inputs) -> Any: + self.value_map = {} + + placeholder_names = [node.name for node in self.graph.nodes if node.op == FX_NODE_OP.PLACEHOLDER.value] + assert len(placeholder_names) == len(inputs), f"输入数量不匹配：图需要 {len(placeholder_names)} 个输入，但提供了 {len(inputs)} 个。" + + for i, node_name in enumerate(placeholder_names): + self.value_map[node_name] = inputs[i] + + for idx, node_name in enumerate(self.topological_names): + if self.name_to_node[node_name].op == FX_NODE_OP.PLACEHOLDER.value: + continue + + node = self.name_to_node[node_name] + stream = self.stream_map[node_name] + + self.wait_for_dependencies(node_name) + + self.node_executor.execute(node=node, value_map=self.value_map, stream=stream) + self.event_map[node_name].record(stream) + + if node.op == FX_NODE_OP.OUTPUT.value: + return self.value_map[node_name] + + raise RuntimeError("图中未找到OUTPUT节点，执行未完成。") + + def synchronize(self): + for stream in self.stream_map.values(): + stream.synchronize() + + def cleanup(self): + pass + + +class GraphStageConfig: + def __init__(self, name: str, sm_dict: Dict[str, int], lane_node_dict: Dict[str, List[str]]): + self.name = name + self.lane_sm_dict = sm_dict + self.lane_node_dict = lane_node_dict + + +class GraphOptimizer: + @staticmethod + def generate_stages_per_op(graph: torch.fx.Graph) -> List[GraphStageConfig]: + stages = [] + for idx, node in enumerate(graph.nodes): + if node.op == FX_NODE_OP.PLACEHOLDER.value: + continue + stage_name = f"stage_{idx}_{node.name}" + stages.append( + GraphStageConfig( + name=stage_name, + sm_dict={LaneType.COMPUTE.value: GreenCtxManager(0).max_sm}, + lane_node_dict={LaneType.COMPUTE.value: [node.name]}, + ) + ) + return stages + + @staticmethod + def generate_stages_all_in_one(graph: torch.fx.Graph) -> List[GraphStageConfig]: + all_node_names = [node.name for node in graph.nodes if node.op != FX_NODE_OP.PLACEHOLDER.value] + return [ + GraphStageConfig( + name="stage_all_in_one", + sm_dict={LaneType.COMPUTE.value: 132}, + lane_node_dict={LaneType.COMPUTE.value: all_node_names}, + ) + ] + + +class GraphStageExecutor: + def __init__(self, graph_module: torch.fx.GraphModule, stage_configs: List[GraphStageConfig], device: torch.device = None): + self.graph_module = graph_module + self.graph = graph_module.graph + self.stage_configs = stage_configs + self.device = device or torch.device("cuda" if torch.cuda.is_available() else "cpu") + self.name_to_node = {node.name: node for node in self.graph.nodes} + self.value_map = {node.name: None for node in self.graph.nodes} + + self.stage_green_manager: Dict[str, GreenCtxManager] = {} + self.stage_lane_stream: Dict[str, Dict[str, torch.cuda.Stream]] = {} + self.stage_lane_event: Dict[str, Dict[str, torch.cuda.Event]] = {} + + for idx, stage_config in enumerate(stage_configs): + stage_name = stage_config.name + self.stage_green_manager[stage_name] = GreenCtxManager(device_index=self.device.index) + self.stage_lane_stream[stage_name] = {} + self.stage_lane_event[stage_name] = {} + for lane_type in stage_config.lane_node_dict.keys(): + cur_green_manager = self.stage_green_manager[stage_name] + cur_sm_count = stage_config.lane_sm_dict.get(lane_type, 0) + assert cur_sm_count >= 0, f"阶段 {stage_name} 的泳道 {lane_type} 的 SM 数量不能为负数" + self.stage_lane_stream[stage_name][lane_type] = ( + cur_green_manager.create_stream(sm_count=cur_sm_count) if cur_sm_count else None + ) + self.stage_lane_event[stage_name][lane_type] = torch.cuda.Event(blocking=False) + + self.input_nodes = [node for node in self.graph.nodes if node.op == FX_NODE_OP.PLACEHOLDER.value] + self.output_node = next(node for node in self.graph.nodes if node.op == FX_NODE_OP.OUTPUT.value) + + self.node_executor = NodeOnlyExecutor(self.graph_module, self.device) + + def _replace_nodes_in_args(self, args): + if isinstance(args, torch.fx.Node): + return self.value_map[args.name] + elif isinstance(args, (list, tuple)): + return type(args)(self._replace_nodes_in_args(a) for a in args) + elif isinstance(args, dict): + return {k: self._replace_nodes_in_args(v) for k, v in args.items()} + else: + return args + + def wait_for_stage_dependencies(self, stage_name: str): + stage_idx = next(i for i, sc in enumerate(self.stage_configs) if sc.name == stage_name) + if stage_idx == 0: + return + + curr_stage = self.stage_configs[stage_idx] + prev_stage = self.stage_configs[stage_idx - 1] + + for curr_lane in curr_stage.lane_node_dict.keys(): + curr_stream = self.stage_lane_stream[curr_stage.name][curr_lane] + for prev_lane in prev_stage.lane_node_dict.keys(): + prev_event = self.stage_lane_event[prev_stage.name][prev_lane] + curr_stream.wait_event(prev_event) + + def execute(self, *inputs) -> Any: + assert len(inputs) == len(self.input_nodes), "输入数量与PLACEHOLDER节点数量不匹配。" + for idx, input_node in enumerate(self.input_nodes): + self.value_map[input_node.name] = inputs[idx] + + for idx, stage in enumerate(self.stage_configs): + stage_name = stage.name + + self.wait_for_stage_dependencies(stage_name) + for lane_type, node_names in stage.lane_node_dict.items(): + stream = self.stage_lane_stream[stage_name][lane_type] + for node_name in node_names: + node = self.name_to_node[node_name] + self.node_executor.execute(node=node, value_map=self.value_map, stream=stream) + event = self.stage_lane_event[stage_name][lane_type] + event.record(stream) + + self.synchronize() + return self.value_map[self.output_node.name] + + def synchronize(self): + for stage_config in self.stage_configs: + stage_name = stage_config.name + for lane_type in stage_config.lane_node_dict.keys(): + stream = self.stage_lane_stream[stage_name][lane_type] + if stream is not None: + stream.synchronize() + + def cleanup(self): + for stage_manager in self.stage_green_manager.values(): + stage_manager.cleanup() diff --git a/pkgs/MagiCompiler/magi_compiler/tokenflow/graph_profile.py b/pkgs/MagiCompiler/magi_compiler/tokenflow/graph_profile.py new file mode 100644 index 0000000000000000000000000000000000000000..51b286f20dfbb54846ca35fdf35d5b5e395a2d2b --- /dev/null +++ b/pkgs/MagiCompiler/magi_compiler/tokenflow/graph_profile.py @@ -0,0 +1,427 @@ +# Copyright (c) 2026 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. + +import json +import os +import re +import weakref +from typing import Any, Dict, Tuple + +import matplotlib.pyplot as plt +import numpy as np +import seaborn as sns +import torch +from magi_compiler.config import get_compile_config +from magi_compiler.tokenflow.graph_executor import FX_NODE_OP # 导入你的枚举类 +from magi_compiler.tokenflow.green_ctx import GreenCtxManager, GreenStreamPool +from magi_compiler.tokenflow.sampler import exponential_aligned_sampler +from magi_compiler.utils import magi_logger +from torch._subclasses.fake_tensor import FakeTensor +from torch.fx import GraphModule, Node + +MIN_LATENCY_MS = 0.005 # 最小延迟下限（5微秒） + + +class GraphProfileWrapper: + def __init__(self, graph_module: GraphModule): + # 核心属性初始化 + self.graph_module = graph_module + self.noncompute_node_names = set() + self.profile_results: Dict[int, Dict[str, Dict[int, float]]] = {} + + # 初始化SM测试列表 + max_sm = GreenCtxManager.get_max_sm() + min_sm, align = GreenCtxManager.get_min_and_align_sm() + high_sm_counts = exponential_aligned_sampler(min_sm, max_sm, num_samples=5, align=align) + low_sm_counts = [max_sm - sm for sm in high_sm_counts] + self.sm_samples = sorted(list(set(high_sm_counts).union(set(low_sm_counts)).union(set([max_sm])))) + self.sm_samples.remove(0) if 0 in self.sm_samples else None + + # 初始化SeqLen测试列表 + self.seqlen_samples = exponential_aligned_sampler(2, 4096, num_samples=10, align=1) + self.seqlen_samples.reverse() # 从大到小测试 + + # 结果保存路径与测试配置 + rank_idx = torch.distributed.get_rank() if torch.distributed.is_initialized() else 0 + self.save_base_path = os.path.join(get_compile_config().cache_root_dir, "graph_profile", f"rank_{rank_idx}") + self.profiled = False + + # 初始化流池 + self.green_stream_pool = GreenStreamPool() + for sm in self.sm_samples: + self.green_stream_pool.get_stream(sm) + + self.node_param_ref_dict = {} + + def _resolve_symint_expression(self, sym_expr: Any, seq_len: int) -> int: + # 普通整数直接返回 + if isinstance(sym_expr, int): + return sym_expr + + # 解析SymInt表达式 + if isinstance(sym_expr, torch.SymInt): + expr_str = str(sym_expr) + symbols = re.findall(r's\d+', expr_str) + + if not symbols: + return int(sym_expr) + + eval_env = {sym: seq_len for sym in symbols} + res_str = expr_str + for sym, val in eval_env.items(): + res_str = res_str.replace(sym, str(val)) + + # print(f"解析 SymInt: {expr_str} -> {res_str}") + return int(eval(res_str)) + + raise TypeError(f"不支持的表达式类型: {type(sym_expr)}") + + def _generate_real_tensor( + self, shape: Tuple[int, ...], stride: Tuple[int, ...], dtype: torch.dtype, device: torch.device, seq_len: int + ) -> torch.Tensor: + # 解析shape中的SymInt表达式 + resolved_shape = [] + for dim in shape: + resolved_dim = self._resolve_symint_expression(dim, seq_len) + resolved_shape.append(resolved_dim) + resolved_shape = tuple(resolved_shape) + + # 解析stride中的SymInt表达式 + resolved_stride = [] + for s in stride: + resolved_s = self._resolve_symint_expression(s, seq_len) + resolved_stride.append(resolved_s) + resolved_stride = tuple(resolved_stride) + + # 验证shape与stride长度匹配 + assert len(resolved_stride) == len( + resolved_shape + ), f"stride长度({len(resolved_stride)})与shape长度({len(resolved_shape)})不匹配" + + # 创建指定布局的空Tensor并填充随机数据 + tensor = torch.empty_strided( + size=resolved_shape, stride=resolved_stride, dtype=dtype, device=device, requires_grad=False + ) + tensor.normal_(mean=0.0, std=1.0) + + # 验证布局（调试用，保留断言） + assert tensor.shape == resolved_shape, f"生成Tensor形状({tensor.shape})与预期({resolved_shape})不匹配" + assert tensor.stride() == resolved_stride, f"生成Tensor stride({tensor.stride()})与预期({resolved_stride})不匹配" + + return tensor + + def _prepare_node_inputs(self, node: Node, seq_len: int) -> Tuple[Tuple, Dict]: + # 处理位置参数与关键字参数 + args = [self._process_arg(arg, seq_len) for arg in node.args] + kwargs = {key: self._process_arg(value, seq_len) for key, value in node.kwargs.items()} + return tuple(args), kwargs + + def _process_arg(self, arg: Any, seq_len: int) -> Any: + """递归处理参数，替换节点引用为真实随机Tensor""" + # 处理fx.Node类型参数 + if isinstance(arg, Node): + input_node = arg + example_value = arg.meta.get("example_value") + assert example_value is not None, f"无法获取节点 {input_node.name} 的元信息" + + # 处理标量与SymInt类型 + if isinstance(example_value, (int, float, str, bool)): + return example_value + if isinstance(example_value, torch.SymInt): + return self._resolve_symint_expression(example_value, seq_len) + + if isinstance(example_value, torch.nn.Parameter): + assert input_node.name in self.node_param_ref_dict, f"无法找到节点 {input_node.name} 对应的参数" + res = self.node_param_ref_dict[input_node.name]() + assert res is not None, f"节点 {input_node.name} 对应的参数已被释放" + return res + + # 处理Tensor/FakeTensor类型 + if isinstance(example_value, (torch.Tensor, FakeTensor)): + shape = example_value.shape + dtype = example_value.dtype + device = example_value.device + stride = example_value.stride() + + # 生成真实Tensor并返回 + assert shape is not None, f"无法获取节点 {input_node.name} 的输出形状" + return self._generate_real_tensor(shape, stride, dtype, device, seq_len) + + # 嵌套结构递归处理 + elif isinstance(arg, (list, tuple)): + return type(arg)(self._process_arg(item, seq_len) for item in arg) + + # 其他类型直接返回 + else: + return arg + + def run_batch_profile(self): + """批量执行节点性能分析""" + # 打印批量分析核心配置 + magi_logger.info(f"===== 开始批量性能分析 =====", rank=0) + magi_logger.info( + f"测试SeqLen: {self.seqlen_samples} | SM粒度: {self.sm_samples} | 节点总数: {len(list(self.graph_module.graph.nodes))}", + rank=0, + ) + + # 初始化seq_len结果存储 + for seq_len in self.seqlen_samples: + if seq_len not in self.profile_results: + self.profile_results[seq_len] = {} + + # 遍历所有节点进行测试 + for node_idx, node in enumerate(self.graph_module.graph.nodes): + node_name = node.name + node_op = node.op + node_perf = None + + skip_record = False + if node_op in ["placeholder", "output"]: + skip_record = True + elif node_op == "call_function" and any(k in str(node.target) for k in ["getitem", "list"]): + skip_record = True + if skip_record: + for seq_len in self.seqlen_samples: + self.profile_results[seq_len][node_name] = {sm: MIN_LATENCY_MS for sm in self.sm_samples} # 5ns + magi_logger.info(f"📌 跳过非计算节点: {node_name} ({node_op})", rank=0) + self.noncompute_node_names.add(node_name) + continue + + # 匹配重复节点名称（b\d+s\d+_xxx） + elif re.match(r"b\d+s\d+_.*", node_name) and node_perf is None: + base_name = re.sub(r"^b\d+s\d+_", "", node_name) + for seq_len in self.seqlen_samples: + if base_name in self.profile_results[seq_len]: + node_perf = self.profile_results[seq_len][base_name] + self.profile_results[seq_len][node_name] = node_perf + skip_record = True + if skip_record: + magi_logger.info(f"📌 重用节点性能: {node_name} ({node_op}) -> 基准节点: {base_name}", rank=0) + continue + + # 遍历所有SeqLen进行节点测试 + for seq_len in self.seqlen_samples: + magi_logger.info(f"--- Profile 节点 {node_name} (SeqLen={seq_len}) ---", rank=0) + try: + # 准备输入并执行性能分析 + real_args, real_kwargs = self._prepare_node_inputs(node, seq_len) + node_perf = self._profile_node_with_data(node=node, args=real_args, kwargs=real_kwargs) + + # 释放内存并保存结果 + del real_args, real_kwargs + assert isinstance(node_perf, dict), f"节点性能结果格式错误" + self.profile_results[seq_len][node_name] = node_perf + + torch.cuda.empty_cache() + torch.cuda.synchronize() + + except Exception as e: + magi_logger.info(f"⚠️ 节点 {node_name} (SeqLen={seq_len}) 测试失败: {str(e)=}", rank=0) + import traceback + + traceback.print_exc() + self.profile_results[seq_len][node_name] = {sm: -1.0 for sm in self.sm_samples} + torch.cuda.empty_cache() + raise e + + self._incremental_save_results() + + # 打印汇总结果 + magi_logger.info(f"===== 批量性能分析完成 =====", rank=0) + + # 打印各SeqLen节点性能汇总 + for seq_len in self.seqlen_samples: + magi_logger.info(f"===== SeqLen={seq_len} 性能汇总（单位：ms） =====", rank=0) + header = f"{'节点名称':<20}" + "".join([f"SM={sm:<10}" for sm in self.sm_samples]) + magi_logger.info(header, rank=0) + magi_logger.info("-" * len(header), rank=0) + + for node in self.graph_module.graph.nodes: + if node.name not in self.profile_results[seq_len]: + continue + if node.name in self.noncompute_node_names: + continue + perf_dict = self.profile_results[seq_len].get(node.name, {}) + assert isinstance(perf_dict, dict), f"节点性能结果格式错误" + row = f"{node.name:<20}" + "".join([f"{perf_dict.get(sm, -1.0):<10.4f}" for sm in self.sm_samples]) + magi_logger.info(row, rank=0) + + self.profiled = True + + def _profile_node_with_data( + self, node: Node, args: Tuple, kwargs: Dict, warmup_steps: int = 30, test_steps: int = 30 + ) -> Dict[int, float]: + """单节点性能测试（指定输入数据）""" + node_perf = {} + node_name = node.name + node_op = node.op + + # 遍历所有SM数量进行测试 + for sm in self.sm_samples: + try: + stream = self.green_stream_pool.get_stream(sm) + + # 节点执行函数定义 + def execute_node_once(): + with torch.cuda.stream(stream): + if node_op == FX_NODE_OP.CALL_FUNCTION.value: + return node.target(*args, **kwargs) + elif node_op == FX_NODE_OP.CALL_METHOD.value: + obj = args[0] + return getattr(obj, node.target)(*args[1:], **kwargs) + elif node_op == FX_NODE_OP.CALL_MODULE.value: + submod = self.graph_module + for mod_name in node.target.split("."): + submod = getattr(submod, mod_name) + return submod(*args, **kwargs) + else: + raise NotImplementedError(f"不支持的节点类型: {node_op}") + + # 预热与计时 + torch.cuda.synchronize() + for _ in range(warmup_steps): + execute_node_once() + torch.cuda.synchronize() + + start_event = torch.cuda.Event(enable_timing=True) + end_event = torch.cuda.Event(enable_timing=True) + + start_event.record(stream) + for _ in range(test_steps): + execute_node_once() + end_event.record(stream) + torch.cuda.synchronize() + + # 计算平均延迟 + total_time = start_event.elapsed_time(end_event) + node_perf[sm] = total_time / test_steps + + except Exception as e: + magi_logger.info(f"⚠️ 节点 {node_name} (SM={sm}) 测试失败: {str(e)=}", rank=0) + node_perf[sm] = -1.0 + torch.cuda.empty_cache() + torch.cuda.synchronize() + + return node_perf + + def _incremental_save_results(self): + """增量保存性能结果并生成热力图""" + # 创建保存目录 + os.makedirs(self.save_base_path, exist_ok=True) + + # 增量保存JSON结果 + json_path = os.path.join(self.save_base_path, "profile_results.json") + existing_results = {} + if os.path.exists(json_path): + with open(json_path, "r", encoding="utf-8") as f: + existing_results = json.load(f) + existing_results = {int(k): v for k, v in existing_results.items()} + + # 格式化新数据并合并 + new_data = { + int(seq_len): { + node_name: {int(sm): float(latency) for sm, latency in sm_data.items()} + for node_name, sm_data in node_data.items() + if sm_data + } + for seq_len, node_data in self.profile_results.items() + } + existing_results.update(new_data) + + # 写入JSON文件 + with open(json_path, "w", encoding="utf-8") as f: + json.dump(existing_results, f, indent=4, ensure_ascii=False) + magi_logger.info(f"📁 性能结果已保存至: {json_path}", rank=0) + + # 生成热力图（有有效数据时） + if len(existing_results) > 0 and all(len(v) > 0 for v in existing_results.values()): + self._generate_heatmaps(existing_results) + + def _generate_heatmaps(self, results: Dict[int, Dict[str, Dict[int, float]]]): + """生成节点性能热力图""" + # 配置绘图参数 + plt.rcParams.update( + { + 'font.sans-serif': ['DejaVu Sans', 'Arial', 'Helvetica'], + 'axes.unicode_minus': False, + 'font.family': 'sans-serif', + 'figure.dpi': 300, + 'savefig.dpi': 300, + } + ) + + magi_logger.info(f"📊 开始生成热力图（有效SeqLen: {list(results.keys())}）", rank=0) + + # 提取所有唯一数据 + all_seq_lens = sorted([int(k) for k in results.keys()]) + all_sms = sorted( + {int(sm) for seq_data in results.values() for node_data in seq_data.values() for sm in node_data.keys()} + ) + all_nodes = sorted({node for seq_data in results.values() for node in seq_data.keys()}) + + # 为每个节点生成热力图 + for node_name in all_nodes: + if node_name in self.noncompute_node_names: + continue + # 构建数据矩阵 + data_matrix = np.array( + [[results.get(seq_len, {}).get(node_name, {}).get(sm, -1.0) for sm in all_sms] for seq_len in all_seq_lens] + ) + + # 绘制并保存热力图 + plt.figure(figsize=(12, 8)) + sns.heatmap( + data_matrix, + annot=True, + fmt=".4f", + cmap="RdYlBu_r", + xticklabels=all_sms, + yticklabels=all_seq_lens, + cbar_kws={"label": "Latency (ms)"}, + mask=(data_matrix < 0), + annot_kws={"size": 8}, + ) + + plt.title(f"Node {node_name} Latency Heatmap", fontsize=14, pad=20) + plt.xlabel("SM Count", fontsize=12) + plt.ylabel("Sequence Length", fontsize=12) + plt.tight_layout() + + img_path = os.path.join(self.save_base_path, f"{node_name}_latency_heatmap.png") + plt.savefig(img_path, dpi=300, bbox_inches="tight", facecolor='white', edgecolor='none') + plt.close() + + magi_logger.info(f"📊 热力图生成完成", rank=0) + + def record_node_param_refs(self, real_args: Tuple): + """记录节点参数引用，供后续输入数据生成使用""" + for node in self.graph_module.graph.nodes: + if node.op == "placeholder": + input_idx = list(self.graph_module.graph.nodes).index(node) + if input_idx < len(real_args): + arg_value = real_args[input_idx] + if isinstance(arg_value, torch.nn.Parameter): + self.node_param_ref_dict[node.name] = weakref.ref(arg_value) + + def __call__(self, *args, **kwargs): + """类可调用入口，未分析则先执行批量分析""" + if not self.profiled: + self.record_node_param_refs(args) + self.run_batch_profile() + return self.graph_module(*args, **kwargs) + + +def gen_profile_wrap_func(graph_module: GraphModule) -> GraphProfileWrapper: + """生成性能分析封装实例""" + return GraphProfileWrapper(graph_module) diff --git a/pkgs/MagiCompiler/magi_compiler/tokenflow/green_ctx.py b/pkgs/MagiCompiler/magi_compiler/tokenflow/green_ctx.py new file mode 100644 index 0000000000000000000000000000000000000000..88d5cc948f41cc54ede561f8349e7299adff9eb3 --- /dev/null +++ b/pkgs/MagiCompiler/magi_compiler/tokenflow/green_ctx.py @@ -0,0 +1,234 @@ +# Copyright (c) 2026 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. + +from typing import List, Tuple + +import cuda.bindings.driver as driver +import cuda.bindings.runtime as runtime +import torch +from cuda.bindings.driver import CUdevResource +from magi_compiler.utils import magi_logger + + +class GreenStreamPool: + def __init__(self, device_index: int = 0): + self.device_index = device_index + self.sm_stream_dict = {} + self.sm_green_ctx_manager_dict = {} + self.max_sm = GreenCtxManager.get_max_sm(device_index) + self.min_sm, self.align_sm = GreenCtxManager.get_min_and_align_sm(device_index) + + def get_stream(self, sm: int) -> torch.cuda.Stream: + target_sm = sm + if sm in self.sm_stream_dict: + return self.sm_stream_dict[sm] + + assert sm > 0, "请求的 SM 数量必须为正整数" + assert sm <= self.max_sm, f"请求的 SM 数量超过硬件限制: {sm} > {self.max_sm}" + sm_remain = self.max_sm - sm + if sm % self.align_sm != 0 and sm_remain % self.align_sm != 0: + raise RuntimeError(f"请求的 SM 数量不符合对齐要求: {sm} (最小对齐单位: {self.align_sm}), " f"剩余 SM 数量: {sm_remain} 也不符合对齐要求") + align_sm, remain_sm = (sm, sm_remain) if sm % self.align_sm == 0 or sm == self.max_sm else (sm_remain, sm) + del sm, sm_remain + + g_ctx_manager = GreenCtxManager(self.device_index) + + if align_sm > 0: + self.sm_green_ctx_manager_dict[align_sm] = g_ctx_manager + align_stream = g_ctx_manager.create_stream(align_sm) + self.sm_stream_dict[align_sm] = align_stream + + if remain_sm > 0: + self.sm_green_ctx_manager_dict[remain_sm] = g_ctx_manager + remain_stream = g_ctx_manager.create_stream(remain_sm) + self.sm_stream_dict[remain_sm] = remain_stream + + return self.sm_stream_dict[target_sm] + + def cleanup(self): + for sm_count, g_ctx_manager in self.sm_green_ctx_manager_dict.items(): + g_ctx_manager.cleanup() + self.sm_stream_dict = {} + self.sm_green_ctx_manager_dict = {} + + +class GreenCtxManager: + def __init__(self, device_index: int = 0): + self.device_index = device_index + self.device = torch.device(f'cuda:{device_index}') + + self.resources: List[CUdevResource] = [] + self.green_contexts: List = [] + self.green_streams: List[torch.cuda.Stream] = [] + + runtime.cudaInitDevice(self.device_index, 0, 0) + self.cu_dev = self.check_errors(driver.cuDeviceGet(self.device_index)) + + self.max_sm = self.get_max_sm(self.device_index) + self.min_sm, self.align_sm = self.get_min_and_align_sm(self.device_index) + self.used_sm = 0 + + self.remaining_resource = self.check_errors( + driver.cuDeviceGetDevResource(self.cu_dev, driver.CUdevResourceType.CU_DEV_RESOURCE_TYPE_SM) + ) + + @classmethod + def get_max_sm(cls, device_index: int = 0) -> int: + runtime.cudaInitDevice(device_index, 0, 0) + cu_dev = cls.check_errors(driver.cuDeviceGet(device_index)) + max_sm = cls.check_errors( + driver.cuDeviceGetAttribute(driver.CUdevice_attribute.CU_DEVICE_ATTRIBUTE_MULTIPROCESSOR_COUNT, cu_dev) + ) + return max_sm + + @classmethod + def get_min_and_align_sm(cls, device_index: int = 0) -> Tuple[int, int]: + runtime.cudaInitDevice(device_index, 0, 0) + cu_dev = cls.check_errors(driver.cuDeviceGet(device_index)) + major = cls.check_errors( + driver.cuDeviceGetAttribute(driver.CUdevice_attribute.CU_DEVICE_ATTRIBUTE_COMPUTE_CAPABILITY_MAJOR, cu_dev) + ) + if major >= 9: + return (8, 8) # Hopper 架构 + elif major >= 8: + return (4, 2) # Ampere 架构 + elif major == 7: + return (2, 2) # Volta/Turing 架构 + elif major == 6: + return (1, 1) # Pascal 架构 + else: + raise RuntimeError(f"不支持的计算能力: {major}") + + @classmethod + def check_errors(cls, result): + if result[0].value: + err, name = driver.cuGetErrorName(result[0]) + msg = name if err == driver.CUresult.CUDA_SUCCESS else f"Code {result[0].value}" + raise RuntimeError(f"CUDA Error: {msg}") + return result[1] if len(result) == 2 else result[1:] + + def align_sm_counts(self, sm_counts: List[int]) -> List[int]: + adjusted = [] + for idx, c in enumerate(sm_counts): + assert type(c) == int and c > 0, f"请求的 SM 数量必须为正整数, 实际值: {c}" + if idx == len(sm_counts) - 1: + # 如果这次准备分配完所有剩余SM，那么是可以允许一个不对齐的count的出现的，但得在最后（128+4=132） + if sum(adjusted) + self.used_sm + c == self.max_sm: + adjusted.append(c) + break + val = ((max(c, self.min_sm) + self.align_sm - 1) // self.align_sm) * self.align_sm + adjusted.append(val) + + total_requested = sum(adjusted) + if total_requested > self.max_sm: + raise RuntimeError( + f"请求 SM 总数溢出! 硬件总数: {self.max_sm}, " + f"原始请求: {sum(sm_counts)}, 对齐后请求: {total_requested}\n" + f"对齐后的分配列表: {adjusted}" + ) + if adjusted != sm_counts: + magi_logger.debug( + "DEBUG: 硬件总 SM: %s, 最终分配方案: %s, 剩余: %s", self.max_sm, adjusted, self.max_sm - total_requested, rank=0 + ) + return adjusted + + def _create_stream_from_resource(self, resource: CUdevResource) -> torch.cuda.Stream: + desc = self.check_errors(driver.cuDevResourceGenerateDesc([resource], 1)) + green_ctx = self.check_errors( + driver.cuGreenCtxCreate(desc, self.cu_dev, driver.CUgreenCtxCreate_flags.CU_GREEN_CTX_DEFAULT_STREAM) + ) + cu_stream = self.check_errors( + driver.cuGreenCtxStreamCreate(green_ctx, driver.CUstream_flags.CU_STREAM_NON_BLOCKING, 0) + ) + return torch.cuda.get_stream_from_external(cu_stream) + + def _refresh_remaining_resource(self, remainder): + # 如果 == 0 那就直接返回None + if remainder is None or remainder.sm.smCount == 0: + self.remaining_resource = None + return + desc = self.check_errors(driver.cuDevResourceGenerateDesc([remainder], 1)) + tmp_g_ctx = self.check_errors( + driver.cuGreenCtxCreate(desc, self.cu_dev, driver.CUgreenCtxCreate_flags.CU_GREEN_CTX_DEFAULT_STREAM) + ) + remaining_resource = self.check_errors( + driver.cuGreenCtxGetDevResource(tmp_g_ctx, driver.CUdevResourceType.CU_DEV_RESOURCE_TYPE_SM) + ) + + driver.cuGreenCtxDestroy(tmp_g_ctx) + self.remaining_resource = remaining_resource + + def batch_create_streams(self, sm_counts: List[int]) -> List[torch.cuda.Stream]: + assert type(sm_counts) == list and len(sm_counts) > 0, "sm_counts 必须为非空列表" + assert all(isinstance(c, int) and c > 0 for c in sm_counts), "sm_counts 中的每个元素必须为正整数" + + if self.remaining_resource is None: + raise RuntimeError("没有剩余的 SM 资源可供分配！") + + # 对齐SM数量（最后一项已适配剩余资源） + sm_counts = self.align_sm_counts(sm_counts) + streams = [] + + for idx, sm_count in enumerate(sm_counts): + # 最后一项且剩余资源刚好等于请求值：直接使用剩余资源（跳过切分） + if idx == len(sm_counts) - 1 and self.remaining_resource and self.remaining_resource.sm.smCount == sm_count: + target_resource = self.remaining_resource + self.remaining_resource = None + remainder = None + else: + result = self.check_errors(driver.cuDevSmResourceSplitByCount(1, self.remaining_resource, 0, sm_count)) + splits, cnt, remainder = result + assert cnt == 1, f"每次切分只能得到 1 个子集, 实际得到: {cnt}" + assert len(splits) == 1, f"每次切分只能得到 1 个子集, 实际得到: {len(splits)}" + assert splits[0].sm.smCount == sm_count, f"切分得到的子集大小不匹配请求大小: {splits[0].sm.smCount} != {sm_count}" + target_resource = splits[0] + + desc = self.check_errors(driver.cuDevResourceGenerateDesc([target_resource], 1)) + green_ctx = self.check_errors( + driver.cuGreenCtxCreate(desc, self.cu_dev, driver.CUgreenCtxCreate_flags.CU_GREEN_CTX_DEFAULT_STREAM) + ) + cu_stream = self.check_errors( + driver.cuGreenCtxStreamCreate(green_ctx, driver.CUstream_flags.CU_STREAM_NON_BLOCKING, 0) + ) + stream = torch.cuda.get_stream_from_external(cu_stream) + streams.append(stream) + self.green_streams.append(stream) + self.green_contexts.append(green_ctx) + self.resources.append(target_resource) + self._refresh_remaining_resource(remainder) + self.used_sm += sm_count + + return streams + + def create_stream(self, sm_count: int = None) -> torch.cuda.Stream: + if sm_count is None: + sm_count = self.max_sm + assert isinstance(sm_count, int) and sm_count > 0, "sm_count 必须为正整数" + return self.batch_create_streams([sm_count])[0] + + def cleanup(self): + cnt = len(self.green_streams) + for i in range(cnt): + stream = self.green_streams[i] + # stream.synchronize() + g_ctx = self.green_contexts[i] + driver.cuStreamDestroy(stream.cuda_stream) + driver.cuGreenCtxDestroy(g_ctx) + self.green_streams = [] + self.green_contexts = [] + self.resources = [] + self.remaining_resource = self.check_errors( + driver.cuDeviceGetDevResource(self.cu_dev, driver.CUdevResourceType.CU_DEV_RESOURCE_TYPE_SM) + ) + self.used_sm = 0 diff --git a/pkgs/MagiCompiler/magi_compiler/tokenflow/sampler.py b/pkgs/MagiCompiler/magi_compiler/tokenflow/sampler.py new file mode 100644 index 0000000000000000000000000000000000000000..184692791fcc03eac2317207f5fdadc496469369 --- /dev/null +++ b/pkgs/MagiCompiler/magi_compiler/tokenflow/sampler.py @@ -0,0 +1,49 @@ +# Copyright (c) 2026 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. + +from typing import List + +import numpy as np +from magi_compiler.utils import magi_logger + + +def exponential_aligned_sampler(min_val: int, max_val: int, num_samples: int, align: int = 8) -> List[int]: + if min_val >= max_val: + raise ValueError(f"最小值({min_val})必须小于最大值({max_val})") + if num_samples < 2: + raise ValueError(f"采样个数({num_samples})需≥2（至少包含min/max）") + if align <= 0: + raise ValueError(f"对齐倍数({align})必须为正整数") + if align > (max_val - min_val): + raise ValueError(f"对齐倍数({align})过大，超过范围跨度({max_val - min_val})") + if num_samples > ((max_val - min_val) // align + 1): + raise ValueError(f"采样个数({num_samples})过大，无法在范围内生成足够对齐值") + + aligned_min = ((min_val + align - 1) // align) * align + aligned_max = (max_val // align) * align + + if aligned_min == aligned_max: + raise ValueError(f"对齐后min/max均为{aligned_min}，请调整align或输入范围") + + raw_samples = np.logspace(np.log(aligned_min), np.log(aligned_max), num=num_samples, base=np.e) + aligned_samples = (np.round(raw_samples / align) * align).astype(int) + + final_samples = sorted(list(dict.fromkeys(aligned_samples.tolist()))) + + if len(final_samples) < num_samples: + final_samples = np.linspace(aligned_min, aligned_max, num=num_samples) + final_samples = (np.round(final_samples / align) * align).astype(int).tolist() + + magi_logger.info("生成对齐采样点: %s", final_samples, rank=0) + return final_samples diff --git a/pkgs/MagiCompiler/magi_compiler/tokenflow/utils.py b/pkgs/MagiCompiler/magi_compiler/tokenflow/utils.py new file mode 100644 index 0000000000000000000000000000000000000000..9adf96bc22206cec0d8efa4eb3b37456b96875a2 --- /dev/null +++ b/pkgs/MagiCompiler/magi_compiler/tokenflow/utils.py @@ -0,0 +1,220 @@ +# Copyright (c) 2026 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. + +from typing import Callable + +import torch +from magi_compiler.api import magi_compile +from magi_compiler.utils import magi_logger, nvtx + + +# 先补全依赖的类定义（确保代码可独立运行） +class ModelConfig: + def __init__( + self, + hidden_size, + num_layers, + num_heads_q, + num_heads_kv, + head_dim, + intermediate_size, + activation_type, + params_dtype=torch.float32, + eps=1e-06, + ): + self.hidden_size = hidden_size + self.num_layers = num_layers + self.num_heads_q = num_heads_q + self.num_heads_kv = num_heads_kv + self.head_dim = head_dim + self.intermediate_size = intermediate_size + self.activation_type = activation_type + self.params_dtype = params_dtype + self.eps = eps + + def __repr__(self): + return ( + f"ModelConfig(hidden_size={self.hidden_size}, num_layers={self.num_layers}, " + f"num_heads_q={self.num_heads_q}, num_heads_kv={self.num_heads_kv}, " + f"head_dim={self.head_dim}, intermediate_size={self.intermediate_size}, " + f"activation_type='{self.activation_type}', params_dtype={self.params_dtype}, eps={self.eps})" + ) + + +@magi_compile(dynamic_arg_dims={'x': [0]}) +class CompiledTransformerModel(torch.nn.Module): + def __init__(self, config: ModelConfig): + super().__init__() + self.mod = TransformerModel(config) + + def forward(self, x): + return self.mod(x) + + +class TransformerModel(torch.nn.Module): + def __init__(self, config: ModelConfig): + super().__init__() + self.config = config + self.layers = torch.nn.ModuleList([TransformerLayer(config) for _ in range(config.num_layers)]) + self.final_norm = torch.nn.LayerNorm(config.hidden_size, eps=config.eps, bias=False) + + def forward(self, x): + for layer in self.layers: + x = layer(x) + x = self.final_norm(x) + return x + + +@nvtx.instrument_nvtx +class TransformerLayer(torch.nn.Module): + def __init__(self, config: ModelConfig): + super().__init__() + self.attn_norm = torch.nn.LayerNorm(config.hidden_size, eps=config.eps, bias=False) + self.attention = GroupedQueryAttention(config) + self.mlp_norm = torch.nn.LayerNorm(config.hidden_size, eps=config.eps, bias=False) + self.mlp = MLPLayer(config) + + def forward(self, x): + x = x + self.attention(self.attn_norm(x)) + x = x + self.mlp(self.mlp_norm(x)) + return x + + +class GroupedQueryAttention(torch.nn.Module): + def __init__(self, config: ModelConfig): + super().__init__() + self.n_heads_q = config.num_heads_q # 32 + self.n_heads_kv = config.num_heads_kv # 8 + self.head_dim = config.head_dim # 128 + self.n_rep = self.n_heads_q // self.n_heads_kv # 32//8=4 + self.hidden_size = config.hidden_size # 4096 + + self.q_size = self.n_heads_q * self.head_dim # 32*128=4096 + self.kv_size = self.n_heads_kv * self.head_dim # 8*128=1024 + self.qkv_proj = torch.nn.Linear(config.hidden_size, self.q_size + 2 * self.kv_size, bias=False) + self.o_proj = torch.nn.Linear(self.q_size, config.hidden_size, bias=False) + + def forward(self, x): + qkv = self.qkv_proj(x) + q, k, v = torch.split(qkv, [self.q_size, self.kv_size, self.kv_size], dim=-1) + + q = q.view(1, -1, self.n_heads_q, self.head_dim) + k = k.view(1, -1, self.n_heads_kv, self.head_dim) + v = v.view(1, -1, self.n_heads_kv, self.head_dim) + + if self.n_rep > 1: + k = k.repeat_interleave(self.n_rep, dim=2) + v = v.repeat_interleave(self.n_rep, dim=2) + + q = q.transpose(1, 2) + k = k.transpose(1, 2) + v = v.transpose(1, 2) + + out: torch.Tensor = my_attention(q, k, v) + # out = q + + out = out.transpose(1, 2) + out = out.squeeze(0) + out = out.view(-1, self.q_size) + + out = self.o_proj(out) + + return out + + return x # 临时屏蔽注意力计算，专注测试 MLP 部分的性能 + + +class MLPModel(torch.nn.Module): + def __init__(self, config: ModelConfig): + super().__init__() + self.config = config + self.layers = torch.nn.ModuleList([MLPLayer(config) for _ in range(config.num_layers)]) + + def forward(self, x): + for layer in self.layers: + x = layer(x) + return x + + +@nvtx.instrument_nvtx +class MLPLayer(torch.nn.Module): + def __init__(self, config: ModelConfig): + super().__init__() + self.pre_norm = torch.nn.LayerNorm(config.hidden_size, eps=config.eps, bias=False) + self.fc1 = torch.nn.Linear(config.hidden_size, config.intermediate_size, bias=False) + self.activation = torch.nn.GELU() + self.fc2 = torch.nn.Linear(config.intermediate_size, config.hidden_size, bias=False) + + def forward(self, x): + x = self.pre_norm(x) + # x = self.fc1(x) + x = self.activation(x) + # x = self.fc2(x) + return x + + +@magi_compile(dynamic_arg_dims={'x': [0]}) +class CompiledMiniMLP(torch.nn.Module): + def __init__(self, config: ModelConfig): + super().__init__() + self.mod = MLPModel(config) + + def forward(self, x): + return self.mod(x) + + +def benchmark_func(func: Callable, warmup_steps: int = 10, run_steps: int = 10, desc: str = "测试") -> float: + torch.cuda.synchronize() + start_event = torch.cuda.Event(enable_timing=True) + end_event = torch.cuda.Event(enable_timing=True) + + @nvtx.instrument_nvtx + def warmup(): + for _ in range(warmup_steps): + _ = func() + torch.cuda.synchronize() # 确保预热的 CUDA 操作全部完成 + + warmup() + + total_elapsed_ms = None + + @nvtx.instrument_nvtx + def run(): + nonlocal total_elapsed_ms + total_elapsed_ms = 0.0 # 总耗时（毫秒） + start_event.record() + for _ in range(run_steps): + func() # 要求func内部所有CUDA操作都已提交并完成！ + end_event.record() + end_event.synchronize() # 确保结束事件已完成 + total_elapsed_ms += start_event.elapsed_time(end_event) + + run() + + avg_time = total_elapsed_ms / run_steps / 1000.0 + total_time = total_elapsed_ms / 1000.0 + magi_logger.info("[%s] 完成！平均耗时: %.6f 秒/次 | 总耗时: %.6f 秒 (CUDA Event 精准计时)", desc, avg_time, total_time, rank=0) + + torch.cuda.synchronize() + return avg_time + + +@torch.library.custom_op("athena::my_attention", mutates_args=()) +def my_attention(q: torch.Tensor, k: torch.Tensor, v: torch.Tensor) -> torch.Tensor: + return torch.nn.functional.scaled_dot_product_attention(q, k, v) + + +@my_attention.register_fake +def _(q: torch.Tensor, k: torch.Tensor, v: torch.Tensor) -> torch.Tensor: + return torch.empty_like(q) diff --git a/pkgs/MagiCompiler/magi_compiler/utils/__init__.py b/pkgs/MagiCompiler/magi_compiler/utils/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..de6d959c7a95a381bec24c20dc94675b67ca0e78 --- /dev/null +++ b/pkgs/MagiCompiler/magi_compiler/utils/__init__.py @@ -0,0 +1,41 @@ +# 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. + + +import magi_compiler.utils.nvtx as nvtx + +from ._utils import * +from .compile_counter import compilation_counter +from .compile_time_monitor import CompileMonitor +from .envs import set_env_var +from .hash import compute_code_hash, compute_code_hash_with_content, compute_hash +from .logger import logger, magi_logger +from .ordered_set import OrderedSet +from .singleton_meta import SingletonMeta +from .version import get_git_version + +__all__ = [ + "nvtx", + "compilation_counter", + "CompileMonitor", + "set_env_var", + "compute_code_hash", + "compute_code_hash_with_content", + "compute_hash", + "logger", + "magi_logger", + "OrderedSet", + "SingletonMeta", + "get_git_version", +] diff --git a/pkgs/MagiCompiler/magi_compiler/utils/_utils.py b/pkgs/MagiCompiler/magi_compiler/utils/_utils.py new file mode 100644 index 0000000000000000000000000000000000000000..948b13c743f3feb30ed1a948d2164fcc766a044b --- /dev/null +++ b/pkgs/MagiCompiler/magi_compiler/utils/_utils.py @@ -0,0 +1,36 @@ +# 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. + +from typing import Any + +import torch +from magi_compiler.utils.logger import magi_logger +from torch import fx +from torch.fx.experimental.symbolic_shapes import is_symbolic + + +def is_func(node: fx.Node, target) -> bool: + return node.op == "call_function" and node.target == target + + +def detect_symbolic_tensor_indices(fake_args: list[Any]) -> list[int]: + """Detect indices of input tensors that have symbolic shapes.""" + sym_tensor_indices = [ + i + for i, x in enumerate(fake_args) + if isinstance(x, torch._subclasses.fake_tensor.FakeTensor) and any(is_symbolic(d) for d in x.size()) + ] + if sym_tensor_indices: + magi_logger.info(f"Detected {len(sym_tensor_indices)} symbolic input tensors (dynamic seqlen) for CUDA Graph.") + return sym_tensor_indices diff --git a/pkgs/MagiCompiler/magi_compiler/utils/compile_counter.py b/pkgs/MagiCompiler/magi_compiler/utils/compile_counter.py new file mode 100644 index 0000000000000000000000000000000000000000..d189df258e8c065e5b4c7c9113455cf126d14f48 --- /dev/null +++ b/pkgs/MagiCompiler/magi_compiler/utils/compile_counter.py @@ -0,0 +1,74 @@ +# 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. + +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +import copy +import dataclasses +from contextlib import contextmanager + + +@dataclasses.dataclass +class CompilationCounter: + # How many models are decorated by @magi_compile decorator + num_models_seen: int = 0 + # The number of graphs seen by Dynamo, usually the same as num_models_seen + # NOTE: __init__ updates num_models_seen but __call__ updates num_graphs_seen, + # we also use num_models_seen for cache key + num_graphs_seen: int = 0 + # Total number of subgraphs, which includes the splitting ops + num_piecewise_graphs_seen: int = 0 + # Total number of subgraphs that are captured, which does not include the splitting ops + num_piecewise_capturable_graphs_seen: int = 0 + # Total number of subgraphs that are compiled by the backend, which does not include the splitting ops + num_backend_compilations: int = 0 + # The number of cached graphs + num_cache_entries: int = 0 + # The number of InductorStandaloneAdaptor.compile calls + num_inductor_compiles: int = 0 + # The number of standalone_compile compiled artifacts saved, should be 0 if MAGI_DISABLE_COMPILE_CACHE is true + num_compiled_artifacts_saved: int = 0 + # The number of EagerAdaptor.compile calls + num_eager_compiles: int = 0 + # # Number of gpu_model_runner attempts to trigger CUDAGraphs capture + # num_gpu_runner_capture_triggers: int = 0 + # # Number of CUDAGraphs captured + # num_cudagraph_captured: int = 0 + + def accuracy_check(self): + # check the consistency of the counters + assert self.num_models_seen >= self.num_graphs_seen + assert self.num_piecewise_graphs_seen >= self.num_piecewise_capturable_graphs_seen + assert self.num_piecewise_capturable_graphs_seen == self.num_backend_compilations + assert self.num_cache_entries == (self.num_inductor_compiles + self.num_eager_compiles) + assert self.num_inductor_compiles == 0 or self.num_eager_compiles == 0 + assert self.num_compiled_artifacts_saved == 0 or self.num_compiled_artifacts_saved == self.num_inductor_compiles + + def clone(self) -> "CompilationCounter": + return copy.deepcopy(self) + + @contextmanager + def expect(self, **kwargs): + old = self.clone() + yield + for k, v in kwargs.items(): + assert getattr(self, k) - getattr(old, k) == v, ( + f"{k} not as expected, before it is {getattr(old, k)}" + f", after it is {getattr(self, k)}, " + f"expected diff is {v}" + ) + + +compilation_counter = CompilationCounter() diff --git a/pkgs/MagiCompiler/magi_compiler/utils/compile_time_monitor.py b/pkgs/MagiCompiler/magi_compiler/utils/compile_time_monitor.py new file mode 100644 index 0000000000000000000000000000000000000000..09968fb140f6fbdaf26ad3a8aa5c58707bd77713 --- /dev/null +++ b/pkgs/MagiCompiler/magi_compiler/utils/compile_time_monitor.py @@ -0,0 +1,83 @@ +# 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. + +import time +from pathlib import Path +from typing import Optional + +from .logger import logger +from .singleton_meta import SingletonMeta + + +class CompileMonitor(metaclass=SingletonMeta): + """ + Compile time monitor (singleton pattern). + + This class tracks the compilation time and manages debug output for + torch.compile operations. It uses the SingletonMeta metaclass to ensure + only one instance exists throughout the application lifecycle. + """ + + def __init__(self): + """Initialize the compile monitor with default values.""" + # Timestamp when compilation monitoring starts + self._start_time: float = 0.0 + # Context manager for magi_depyf debug output (if enabled) + self._context_manager = None + + def start(self, enable_debug: bool = True, debug_dump_path: Optional[Path] = None): + """ + Start monitoring compilation time. + + Records the start time and optionally sets up magi_depyf debug output. + + Args: + enable_debug: Whether to save the magi_depyf debug output + debug_dump_path: Path to dump the debug output + """ + self._start_time = time.time() + + if enable_debug and debug_dump_path: + from magi_compiler.magi_depyf.inspect import dump_src + + debug_dump_path.mkdir(parents=True, exist_ok=True) + logger.debug(f"Dumping magi_depyf output to {debug_dump_path}") + self._context_manager = dump_src(debug_dump_path.as_posix()) + self._context_manager.__enter__() + + def mark(self, prefix: str = "") -> float: + """ + Mark the current time point and return elapsed time since start. + + Args: + prefix: Optional prefix string for the log message + + Returns: + Elapsed time in seconds since monitoring started + """ + time_collapsed = time.time() - self._start_time + logger.debug(f"{prefix} collapsed time: {time_collapsed:.2f} s") + return time_collapsed + + def end(self): + """ + End compilation time monitoring. + + Logs the total compilation time and cleans up the magi_depyf context manager + if it was initialized. + """ + logger.debug(f"torch.compile takes {time.time() - self._start_time:.2f} s in total") + if self._context_manager is not None: + self._context_manager.__exit__(None, None, None) + self._context_manager = None diff --git a/pkgs/MagiCompiler/magi_compiler/utils/envs.py b/pkgs/MagiCompiler/magi_compiler/utils/envs.py new file mode 100644 index 0000000000000000000000000000000000000000..48016c106c00fa21a06e5070ba6a977588d10365 --- /dev/null +++ b/pkgs/MagiCompiler/magi_compiler/utils/envs.py @@ -0,0 +1,81 @@ +# 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. + +import contextlib +import os +from typing import Iterator + + +def _env_is_true(env_name: str) -> bool: + return str(os.environ.get(env_name, "0")).lower() in {"1", "true", "yes", "y", "on", "enabled"} + + +@contextlib.contextmanager +def set_env_var(key: str, value: str) -> Iterator[None]: + """Temporarily set an environment variable.""" + old = os.environ.get(key) + os.environ[key] = value + try: + yield + finally: + if old is None: + os.environ.pop(key, None) + else: + os.environ[key] = old + + +# Enable AOT (Ahead-Of-Time) compilation mode for torch.compile. +# When truthy ("1"/"true"/"yes"/"y"/"on"/"enabled"), MagiCompiler: +# - Persists compiled artifacts to {cache_root_dir}/torch_aot_compile/{hash}/rank_*/model. +# - On startup, loads from cache to skip Dynamo tracing/compilation and reduce cold-start latency. +# - Drops shape guards and relies on source consistency checks; any source change triggers recompilation. +# When unset or falsy ("0"/"false"/"no"/"n"/"off"/"disabled"), falls back to the normal JIT path. +MAGI_AOT_COMPILE: bool = _env_is_true("MAGI_AOT_COMPILE") + +# Disable the compile cache for MagiCompiler. +# MagiCompiler will re-compile if we disabled the compile cache. +MAGI_DISABLE_COMPILE_CACHE: bool = _env_is_true("MAGI_DISABLE_COMPILE_CACHE") + +# Enable max autotune, which actually controls `max_autotune` in the Inductor config. +MAGI_ENABLE_INDUCTOR_MAX_AUTOTUNE: bool = _env_is_true("MAGI_ENABLE_INDUCTOR_MAX_AUTOTUNE") + +# Enable coordinate descent tuning, which actually controls `coordinate_descent_tuning` in the Inductor config. +MAGI_ENABLE_INDUCTOR_COORDINATE_DESCENT_TUNING: bool = _env_is_true("MAGI_ENABLE_INDUCTOR_COORDINATE_DESCENT_TUNING") + +# Enable FX graph visualization, which actually controls `enable_fx_graph_viz` in the Inductor config. +MAGI_ENABLE_FX_GRAPH_VIZ: bool = _env_is_true("MAGI_ENABLE_FX_GRAPH_VIZ") + +# FX graph visualization node description mode: simple or detailed. +MAGI_FX_GRAPH_VIZ_NODE_DESC: str = os.getenv("MAGI_FX_GRAPH_VIZ_NODE_DESC", "simple") + +# (Experimental): Enable profiling fx node performance under different sequence lengths and sm allocations. +MAGI_ENABLE_PROFILE: bool = _env_is_true("MAGI_ENABLE_PROFILE") + +# Equal to TORCHINDUCTOR_PATTERN_MATCH_DEBUG environment. +MAGI_PATTERN_MATCH_DEBUG: str | None = os.getenv("MAGI_PATTERN_MATCH_DEBUG") + +# Tag for the part of model being compiled, e.g. backbone/eagle_head, maybe useful in the future +MAGI_MODEL_TAG: str = os.getenv("MAGI_MODEL_TAG", "backbone") + +# temporary path to store shared memory binaries +MAGI_SHARED_BIN_PATH = "/dev/shm" + +# Logging level for MagiCompiler (DEBUG / INFO / WARNING / ERROR). Read once at import time. +MAGI_LOGGING_LEVEL: str = os.getenv("MAGI_LOGGING_LEVEL", "WARNING").upper() + +# Implicit key in torch.compile options, pass name in inductor options +MAGI_POST_GRAD_PASS: str = "post_grad_custom_post_pass" + +# Implicit key in torch.compile options, custom partitioner function name +MAGI_CUSTOM_PARTITIONER_FN: str = "custom_partitioner_fn" diff --git a/pkgs/MagiCompiler/magi_compiler/utils/hash.py b/pkgs/MagiCompiler/magi_compiler/utils/hash.py new file mode 100644 index 0000000000000000000000000000000000000000..035c047445235a4d7a4872cb3997001cfe4e65b6 --- /dev/null +++ b/pkgs/MagiCompiler/magi_compiler/utils/hash.py @@ -0,0 +1,64 @@ +# 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. + +import hashlib +import os +from functools import reduce +from typing import Any, List, Union + +HASH_LENGTH = 10 + + +def _fn_hash_key(fn) -> str: + sha256_hash = hashlib.sha256() + sha256_hash.update(fn.__qualname__.encode()) + sha256_hash.update(str(fn.__code__.co_firstlineno).encode()) + return sha256_hash.hexdigest()[:HASH_LENGTH] + + +def compute_hash(obj: Union[Any, List[Any]]) -> str: + if isinstance(obj, list): + return reduce(lambda x, y: compute_hash(x + y), [compute_hash(_item) for _item in obj], "") + + elif isinstance(obj, dict): + return reduce(lambda x, y: compute_hash(x + y), [compute_hash(_k) + compute_hash(_v) for _k, _v in obj.items()], "") + + elif callable(obj): + return _fn_hash_key(obj) + + return hashlib.md5(str(obj).encode(), usedforsecurity=False).hexdigest()[:HASH_LENGTH] + + +def compute_code_hash_with_content(file_contents: dict[str, str]) -> str: + items = list(sorted(file_contents.items(), key=lambda x: x[0])) + hash_content = [] + for filepath, content in items: + hash_content.append(filepath) + if filepath == "": + # This means the function was dynamically generated, with e.g. exec(). We can't actually check these. + continue + hash_content.append(content) + return compute_hash("\n".join(hash_content)) + + +def compute_code_hash(files: set[str]) -> str: + file_contents = {} + for filepath in files: + # Skip files that don't exist (e.g., , , etc.) + if not os.path.isfile(filepath): + file_contents[filepath] = "" + else: + with open(filepath) as f: + file_contents[filepath] = f.read() + return compute_code_hash_with_content(file_contents) diff --git a/pkgs/MagiCompiler/magi_compiler/utils/logger.py b/pkgs/MagiCompiler/magi_compiler/utils/logger.py new file mode 100644 index 0000000000000000000000000000000000000000..976f36f5c3108de465d183e686c9433e81dce5a8 --- /dev/null +++ b/pkgs/MagiCompiler/magi_compiler/utils/logger.py @@ -0,0 +1,111 @@ +# 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. + +import logging +import os +import sys + +import torch.distributed as dist +from magi_compiler.utils.envs import MAGI_LOGGING_LEVEL + +_FMT = "[%(asctime)s - %(levelname)s] [Rank %(rank)s] %(message)s" +_DATEFMT = "%Y-%m-%d %H:%M:%S" + + +def _get_rank() -> int: + if dist.is_available() and dist.is_initialized(): + try: + return dist.get_rank() + except Exception: + pass + return int(os.getenv("RANK", 0)) + + +def _get_world_size() -> int: + if dist.is_available() and dist.is_initialized(): + try: + return dist.get_world_size() + except Exception: + pass + return 1 + + +def _should_log(rank: int | str) -> bool: + """rank: int (only that rank), 'all' (every rank).""" + if rank == "all": + return True + current = _get_rank() + if isinstance(rank, int): + return current == rank + return False + + +class _RankFormatter(logging.Formatter): + """Inject ``rank`` into every log record so ``%(rank)s`` works in the format string.""" + + def format(self, record: logging.LogRecord) -> str: + record.rank = _get_rank() # type: ignore[attr-defined] + return super().format(record) + + +def _build_logger() -> logging.Logger: + """Create and configure the ``magi_compiler`` logger exactly once.""" + lg = logging.getLogger("magi_compiler") + lg.propagate = False + lg.handlers.clear() + handler = logging.StreamHandler(sys.stderr) + handler.setFormatter(_RankFormatter(fmt=_FMT, datefmt=_DATEFMT)) + lg.addHandler(handler) + lg.setLevel(getattr(logging, MAGI_LOGGING_LEVEL, logging.WARNING)) + return lg + + +_std_logger = _build_logger() + + +class MagiLogger: + """ + Logger for MagiCompiler, backed by the standard-library ``logging`` module + with a dedicated named logger ``"magi_compiler"``. + + Fully isolated from loguru / other logging instances in the same process. + Level is set once at import time via ``MAGI_LOGGING_LEVEL`` (see ``envs.py``). + API: ``.info / .debug / .warning / .error(msg, *args, rank=0)`` + ``rank`` defaults to 0 (only rank-0 logs). Pass an explicit value to override: + int = only that rank, "all" = every rank. + """ + + def info(self, message: str, *args, rank: int | str = 0, **kwargs) -> None: + if not _should_log(rank): + return + _std_logger.info(message, *args, **kwargs) + + def debug(self, message: str, *args, rank: int | str = 0, **kwargs) -> None: + if not _should_log(rank): + return + _std_logger.debug(message, *args, **kwargs) + + def warning(self, message: str, *args, rank: int | str = 0, **kwargs) -> None: + if not _should_log(rank): + return + _std_logger.warning(message, *args, **kwargs) + + def error(self, message: str, *args, rank: int | str = 0, **kwargs) -> None: + if not _should_log(rank): + return + _std_logger.error(message, *args, **kwargs) + + +magi_logger = MagiLogger() +logger = magi_logger # alias diff --git a/pkgs/MagiCompiler/magi_compiler/utils/nvtx.py b/pkgs/MagiCompiler/magi_compiler/utils/nvtx.py new file mode 100644 index 0000000000000000000000000000000000000000..7c36d82684cefb53adea39972012e077a7b8a4b2 --- /dev/null +++ b/pkgs/MagiCompiler/magi_compiler/utils/nvtx.py @@ -0,0 +1,105 @@ +# 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. + +from functools import wraps +from typing import Any, Callable, TypeVar, cast + +import torch + +# fixed the mypy type check missing bug +# when a func is wrapped +# issue: https://stackoverflow.com/questions/65621789/mypy-untyped-decorator-makes-function-my-method-untyped +F = TypeVar("F", bound=Callable[..., Any]) + +# global var for torch.autograd.profiler.emit_nvtx +_EMIT_NVTX_CTX: None | torch.autograd.profiler.emit_nvtx = None + + +@torch.library.custom_op("magi_compile::nvtx_range_push", mutates_args=()) +def nvtx_range_push(event_name: str) -> None: + """torch.ops.magi_compile.nvtx_range_push""" + torch.cuda.nvtx.range_push(event_name) + + +@nvtx_range_push.register_fake +def _(event_name: str) -> None: + pass + + +@torch.library.custom_op("magi_compile::nvtx_range_pop", mutates_args=()) +def nvtx_range_pop() -> None: + """torch.ops.magi_compile.nvtx_range_pop""" + torch.cuda.nvtx.range_pop() + + +@nvtx_range_pop.register_fake +def _() -> None: + pass + + +# NOTE: since torch.compile does not support @contextlib.contextmanager, +# we use the class-based context manager +class add_nvtx_event: + """ + Context manager to add an NVTX event around a code block. + + Args: + event_name (str): The name of the event to be recorded. + """ + + def __init__(self, event_name: str): + self.enter_name = event_name + + def __enter__(self): + if torch.compiler.is_compiling(): + # NOTE: torch.compile supports neither retrieving the attributes from "self" + # nor modifying a variable not in the current scope + # so we have no choice but assign a constant event name when compiling + nvtx_range_push("torch compile region") + else: + torch.cuda.nvtx.range_push(self.enter_name) + return self + + def __exit__(self, *excinfo): + if torch.compiler.is_compiling(): + nvtx_range_pop() + else: + torch.cuda.nvtx.range_pop() + + +def instrument_nvtx(func: F) -> F: + """ + Decorator that records an NVTX range for the duration of the function call. + + Args: + func (Callable): The function to be decorated. + + Returns: + Callable: The wrapped function that is now being profiled. + """ + + @wraps(func) + def wrapped_fn(*args, **kwargs): + if torch.compiler.is_compiling(): + # NOTE: we can not access func.__qualname__ when compiling + # thus use func.__name__ instead + func_name = func.__name__ + else: + func_name = func.__qualname__ + + with add_nvtx_event(func_name): + ret_val = func(*args, **kwargs) + return ret_val + + return cast(F, wrapped_fn) diff --git a/pkgs/MagiCompiler/magi_compiler/utils/ordered_set.py b/pkgs/MagiCompiler/magi_compiler/utils/ordered_set.py new file mode 100644 index 0000000000000000000000000000000000000000..5672ea224ba862f3c0134a14ab6ee1dfdab9b5ad --- /dev/null +++ b/pkgs/MagiCompiler/magi_compiler/utils/ordered_set.py @@ -0,0 +1,135 @@ +# 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. + +from __future__ import annotations + +from collections import OrderedDict +from collections.abc import Iterable, Iterator, MutableSet +from typing import Generic, TypeVar + +T = TypeVar("T") + + +class OrderedSet(MutableSet[T], Generic[T]): + __slots__ = ("_map",) + + def __init__(self, iterable: Iterable[T] | None = None): + self._map: OrderedDict[T, None] = OrderedDict() + if iterable: + self.update(iterable) + + def __contains__(self, x: T) -> bool: + return x in self._map + + def __len__(self) -> int: + return len(self._map) + + def __iter__(self) -> Iterator[T]: + return iter(self._map.keys()) + + def add(self, value: T) -> None: + self._map[value] = None + + def discard(self, value: T) -> None: + self._map.pop(value, None) + + def pop(self, last: bool = True) -> T: + if not self._map: + raise KeyError("pop from an empty OrderedSet") + key = next(reversed(self._map)) if last else next(iter(self._map)) + self._map.pop(key, None) + return key + + def clear(self) -> None: + self._map.clear() + + def to_list(self) -> list[T]: + return list(self._map.keys()) + + def copy(self) -> "OrderedSet[T]": + return OrderedSet(self) + + def __repr__(self) -> str: + cls = self.__class__.__name__ + if not self: + return f"{cls}()" + return f"{cls}([{', '.join(repr(x) for x in self)}])" + + def __eq__(self, other: object) -> bool: + if isinstance(other, OrderedSet): + return list(self) == list(other) + if isinstance(other, set): + return set(self) == other + return NotImplemented + + def union(self, *others: Iterable[T]) -> "OrderedSet[T]": + result = OrderedSet(self) + for other in others: + for x in other: + if x not in result: + result.add(x) + return result + + __or__ = union + + def intersection(self, *others: Iterable[T]) -> "OrderedSet[T]": + if not others: + return self.copy() + common = set(self) + for other in others: + common &= set(other) + return OrderedSet(x for x in self if x in common) + + __and__ = intersection + + def difference(self, *others: Iterable[T]) -> "OrderedSet[T]": + remove = set().union(*(set(o) for o in others)) + return OrderedSet(x for x in self if x not in remove) + + __sub__ = difference + + def symmetric_difference(self, other: Iterable[T]) -> "OrderedSet[T]": + other_set = set(other) + left = [x for x in self if x not in other_set] + right = [x for x in other_set if x not in self] + return OrderedSet([*left, *right]) + + __xor__ = symmetric_difference + + # Pydantic v2 compatible: parse/serialize as list + @classmethod + def __get_pydantic_core_schema__(cls, _source_type, _handler): + from pydantic_core import core_schema as cs + + def validate_from_any(v): + if isinstance(v, OrderedSet): + return v + if v is None: + return cls() + try: + return cls(v) + except TypeError: + raise TypeError("OrderedSet must be built from an iterable") + + return cs.no_info_after_validator_function( + validate_from_any, + cs.list_schema(cs.any_schema()), + serialization=cs.plain_serializer_function_ser_schema(lambda v: list(v)), + ) + + @classmethod + def __get_pydantic_json_schema__(cls, core_schema, handler): + json_schema = handler(core_schema) + json_schema.update({"type": "array"}) + return json_schema diff --git a/pkgs/MagiCompiler/magi_compiler/utils/singleton_meta.py b/pkgs/MagiCompiler/magi_compiler/utils/singleton_meta.py new file mode 100644 index 0000000000000000000000000000000000000000..c683fff8c5b4591f9dc602b9f5bbcb37a01a53ae --- /dev/null +++ b/pkgs/MagiCompiler/magi_compiler/utils/singleton_meta.py @@ -0,0 +1,65 @@ +# 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. + + +import threading + + +class SingletonMeta(type): + """ + Thread-safe singleton metaclass implementation. + + This metaclass ensures that only one instance of a class is created, + even in a multi-threaded environment. It uses a class-level lock to + synchronize instance creation across threads. + + Usage: + class MyClass(metaclass=SingletonMeta): + pass + + # Both instances will be the same object + instance1 = MyClass() + instance2 = MyClass() + assert instance1 is instance2 # True + """ + + # Dictionary to store singleton instances for each class + _instances = {} + # Class-level lock to ensure thread-safe instance creation + _lock = threading.Lock() + + def __call__(cls, *args, **kwargs): + """ + Override __call__ to implement singleton pattern. + + Uses double-checked locking pattern to minimize lock contention: + 1. First check if instance exists (no lock needed) + 2. If not, acquire lock and check again + 3. Create instance only if still doesn't exist + + Args: + *args: Positional arguments passed to the class constructor + **kwargs: Keyword arguments passed to the class constructor + + Returns: + The singleton instance of the class + """ + # Fast path: check if instance already exists (no lock needed) + if cls not in cls._instances: + # Slow path: acquire lock and check again (double-checked locking) + with cls._lock: + # Check again inside the lock to prevent race condition + if cls not in cls._instances: + cls._instances[cls] = super().__call__(*args, **kwargs) + return cls._instances[cls] diff --git a/pkgs/MagiCompiler/magi_compiler/utils/version.py b/pkgs/MagiCompiler/magi_compiler/utils/version.py new file mode 100644 index 0000000000000000000000000000000000000000..89bf19ba9d6246b1e61578f00db0c3b5eb7c7405 --- /dev/null +++ b/pkgs/MagiCompiler/magi_compiler/utils/version.py @@ -0,0 +1,68 @@ +# 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. + +from __future__ import annotations + +import subprocess +from pathlib import Path + + +def get_git_version(short: bool = True, length: int = 7) -> str: + """ + 获取当前仓库的 Git 提交 ID。默认返回短 SHA。 + 当无法通过 git 命令获取时，回退读取 .git/HEAD/packed-refs。 + """ + # 以本文件为锚点推断仓库根目录：.../magi_compiler/magi_compiler/utils/version.py → .../magi_compiler + repo_root = Path(__file__).resolve().parents[2] + git_dir = repo_root / ".git" + + # 优先通过 git 命令获取 + try: + sha = subprocess.check_output( + ["git", "-C", str(repo_root), "rev-parse", "HEAD"], text=True, stderr=subprocess.DEVNULL + ).strip() + return sha[:length] if short else sha + except Exception: + pass + + # 回退：解析 .git/HEAD + try: + head = (git_dir / "HEAD").read_text().strip() + if head.startswith("ref:"): + ref_path = head.split(" ", 1)[1].strip() + ref_file = git_dir / ref_path + if ref_file.exists(): + sha = ref_file.read_text().strip() + else: + # 可能被打包到 packed-refs + packed = git_dir / "packed-refs" + sha = "" + if packed.exists(): + with packed.open() as f: + for line in f: + line = line.strip() + if not line or line.startswith("#") or line.startswith("^"): + continue + parts = line.split(" ") + if len(parts) == 2 and parts[1] == ref_path: + sha = parts[0] + break + if not sha: + return "unknown" + else: + # detached HEAD，HEAD 文件直接保存 SHA + sha = head + return sha[:length] if short else sha + except Exception: + return "unknown" diff --git a/pkgs/MagiCompiler/magi_compiler/utils/visualize/__init__.py b/pkgs/MagiCompiler/magi_compiler/utils/visualize/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..c8d31d39f1817cdf885f3df0778eaeaab75f0288 --- /dev/null +++ b/pkgs/MagiCompiler/magi_compiler/utils/visualize/__init__.py @@ -0,0 +1,18 @@ +# 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. + +from .joint_graph_visualizer import joint_graph_vis +from .visualizer import save_fx_graph_visualization + +__all__ = ["joint_graph_vis", "save_fx_graph_visualization"] diff --git a/pkgs/MagiCompiler/magi_compiler/utils/visualize/joint_graph_visualizer.py b/pkgs/MagiCompiler/magi_compiler/utils/visualize/joint_graph_visualizer.py new file mode 100644 index 0000000000000000000000000000000000000000..c2126c87ee50986c2273502335de0e6c528fb63c --- /dev/null +++ b/pkgs/MagiCompiler/magi_compiler/utils/visualize/joint_graph_visualizer.py @@ -0,0 +1,315 @@ +# 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. + +import os +import textwrap +from typing import Dict, List, Optional, Set, Tuple + +import graphviz +import torch +import torch.fx as fx +from magi_compiler.config import get_compile_config +from magi_compiler.utils import envs, magi_logger + + +class NodeCategory: + FWD = "fwd" + BWD = "bwd" + SAVE_TENSOR = "save_tensor" + INPUT = "input" + FWD_OUTPUT = "fwd_output" + BWD_OUTPUT = "bwd_output" + TANGENT = "tangent" + + +NODE_CATEGORY_STYLES = { + NodeCategory.INPUT: {"shape": "box", "style": "filled,bold", "fillcolor": "#E8F4FD", "color": "#2196F3"}, + NodeCategory.FWD: {"shape": "ellipse", "style": "filled", "fillcolor": "#E8F5E9", "color": "#4CAF50"}, + NodeCategory.BWD: {"shape": "ellipse", "style": "filled", "fillcolor": "#FFF3E0", "color": "#FF9800"}, + NodeCategory.SAVE_TENSOR: {"shape": "hexagon", "style": "filled,bold", "fillcolor": "#FFEB3B", "color": "#D32F2F"}, + NodeCategory.TANGENT: {"shape": "box", "style": "filled,bold", "fillcolor": "#B3E5FC", "color": "#0288D1"}, + NodeCategory.FWD_OUTPUT: {"shape": "box", "style": "filled,bold", "fillcolor": "#C8E6C9", "color": "#388E3C"}, + NodeCategory.BWD_OUTPUT: {"shape": "box", "style": "filled,bold", "fillcolor": "#FFCCBC", "color": "#E64A19"}, +} + + +def get_graph_node_names(graph: fx.Graph) -> Set[str]: + """Extract all node names from a graph.""" + return {node.name for node in graph.nodes} + + +def is_tangent_node(node: fx.Node) -> bool: + """Check if a node is a tangent node (gradient input for backward).""" + return node.name.startswith("tangent") or "tangent" in node.name.lower() + + +def categorize_joint_nodes( + joint_graph: fx.Graph, fwd_graph: fx.Graph, bwd_graph: fx.Graph, save_tensor_nodes: Optional[List[fx.Node]] = None +) -> Tuple[Dict[str, str], Set[str]]: + """ + Categorize nodes in joint graph with priority: save_tensor > fwd > bwd. + + Returns: + - node_categories: dict mapping node name to category + - input_save_tensors: set of input node names that are also save tensors + """ + fwd_names = get_graph_node_names(fwd_graph) + bwd_names = get_graph_node_names(bwd_graph) + save_tensor_names = {node.name for node in save_tensor_nodes} if save_tensor_nodes else set() + + node_categories = {} + input_save_tensors = set() + + for node in joint_graph.nodes: + if node.op == "placeholder": + if is_tangent_node(node): + node_categories[node.name] = NodeCategory.TANGENT + else: + node_categories[node.name] = NodeCategory.INPUT + # Track if this input is also a save tensor + if node.name in save_tensor_names: + input_save_tensors.add(node.name) + elif node.op == "output": + node_categories[node.name] = NodeCategory.BWD_OUTPUT + elif node.name in save_tensor_names: + node_categories[node.name] = NodeCategory.SAVE_TENSOR + elif node.name in fwd_names: + node_categories[node.name] = NodeCategory.FWD + elif node.name in bwd_names: + node_categories[node.name] = NodeCategory.BWD + else: + node_categories[node.name] = NodeCategory.FWD + + return node_categories, input_save_tensors + + +def extract_joint_graph_structure( + graph: fx.Graph, node_categories: Dict[str, str], input_save_tensors: Optional[Set[str]] = None +) -> Tuple[List[Dict], List[Dict]]: + """Extract nodes and edges from joint graph with category-based styling.""" + + def wrap_str(text, width=40): + return textwrap.fill(text, width=width, break_long_words=True, replace_whitespace=False) + + nodes, edges = [], [] + input_save_tensors = input_save_tensors or set() + + for node in graph.nodes: + name_str = wrap_str(str(node.name)) + + if callable(node.target): + target_str = getattr(node.target, "__name__", str(node.target)) + elif hasattr(node.target, "_op"): + target_str = str(node.target._op) + else: + target_str = str(node.target) + target_str = wrap_str(target_str) + + category = node_categories.get(node.name, NodeCategory.FWD) + style_info = NODE_CATEGORY_STYLES.get(category, NODE_CATEGORY_STYLES[NodeCategory.FWD]) + + # Add annotation if input node is also a save tensor + if node.name in input_save_tensors: + node_label = f"{name_str}\n[{target_str}]\n(SaveTensor)" + else: + node_label = f"{name_str}\n[{target_str}]" + + nodes.append({"id": node.name, "style": style_info, "node_label": node_label, "category": category}) + + def traverse_args(args_kwargs): + if isinstance(args_kwargs, (tuple, list)): + for arg in args_kwargs: + traverse_args(arg) + elif isinstance(args_kwargs, dict): + for val in args_kwargs.values(): + traverse_args(val) + elif isinstance(args_kwargs, fx.Node): + d = {"source": args_kwargs.name, "target": node.name} + if d not in edges: + edges.append(d) + + traverse_args(node.args) + traverse_args(node.kwargs) + + return nodes, edges + + +def create_joint_graph_dot(nodes: List[Dict], edges: List[Dict]) -> graphviz.Digraph: + """ + Create a graphviz Digraph for joint graph visualization. + + Layout (using rankdir=BT, bottom-to-top): + - Top: BWD output (gradients) + - Middle-Left: FWD cluster (inputs, fwd ops, save_tensors) + - Middle-Right: BWD cluster (bwd ops) + - Bottom: Tangent inputs + FWD output + """ + dot = graphviz.Digraph( + name="joint_graph", + format="pdf", + graph_attr={ + "rankdir": "BT", + "nodesep": "0.4", + "ranksep": "0.6", + "overlap": "false", + "splines": "true", + "newrank": "true", + "label": "Joint Graph Visualization\\n" + "Blue: Input | Cyan: Tangent | Green: FWD | Yellow: SaveTensor | Orange: BWD | Top: BWD Output", + "labelloc": "t", + "fontsize": "12", + }, + node_attr={"fontname": "Helvetica", "fontsize": "9", "fixedsize": "false", "margin": "0.12"}, + ) + + input_nodes = [] + fwd_nodes = [] + save_tensor_nodes = [] + tangent_nodes = [] + bwd_nodes = [] + bwd_output_nodes = [] + + for node in nodes: + cat = node["category"] + if cat == NodeCategory.INPUT: + input_nodes.append(node) + elif cat == NodeCategory.FWD: + fwd_nodes.append(node) + elif cat == NodeCategory.SAVE_TENSOR: + save_tensor_nodes.append(node) + elif cat == NodeCategory.TANGENT: + tangent_nodes.append(node) + elif cat == NodeCategory.BWD: + bwd_nodes.append(node) + elif cat == NodeCategory.BWD_OUTPUT: + bwd_output_nodes.append(node) + else: + fwd_nodes.append(node) + + with dot.subgraph(name="cluster_fwd") as fwd_cluster: + fwd_cluster.attr(label="Forward Pass", style="rounded,dashed", color="#4CAF50", bgcolor="#F1F8E9", penwidth="2") + for node in input_nodes: + fwd_cluster.node(node["id"], node["node_label"], **node["style"]) + for node in fwd_nodes: + fwd_cluster.node(node["id"], node["node_label"], **node["style"]) + for node in save_tensor_nodes: + fwd_cluster.node(node["id"], node["node_label"], **node["style"]) + + with dot.subgraph(name="cluster_bwd") as bwd_cluster: + bwd_cluster.attr(label="Backward Pass", style="rounded,dashed", color="#FF9800", bgcolor="#FFF8E1", penwidth="2") + for node in bwd_nodes: + bwd_cluster.node(node["id"], node["node_label"], **node["style"]) + + for node in tangent_nodes: + dot.node(node["id"], node["node_label"], **node["style"]) + + for node in bwd_output_nodes: + dot.node(node["id"], node["node_label"], **node["style"]) + + with dot.subgraph() as s: + s.attr(rank="min") + for node in tangent_nodes: + s.node(node["id"]) + for node in save_tensor_nodes: + s.node(node["id"]) + + with dot.subgraph() as s: + s.attr(rank="max") + for node in bwd_output_nodes: + s.node(node["id"]) + + if input_nodes and fwd_nodes: + input_nodes[0]["id"] + first_fwd = fwd_nodes[0]["id"] if fwd_nodes else None + if first_fwd: + pass + + if bwd_nodes and tangent_nodes: + first_bwd = bwd_nodes[0]["id"] + first_tangent = tangent_nodes[0]["id"] + dot.edge(first_tangent, first_bwd, style="invis", constraint="true") + + for edge in edges: + dot.edge(edge["source"], edge["target"]) + + return dot + + +def get_joint_graph_path(sub_dir: str = "", filename: str = "") -> str: + """Get the path for saving joint graph visualization.""" + cache_root_dir = get_compile_config().cache_root_dir + rank = torch.distributed.get_rank() if torch.distributed.is_initialized() else 0 + joint_graph_dir = os.path.join(cache_root_dir, "joint_graph_views", f"rank_{rank}") + os.makedirs(joint_graph_dir, exist_ok=True) + if sub_dir: + joint_graph_sub_dir = os.path.join(joint_graph_dir, sub_dir) + os.makedirs(joint_graph_sub_dir, exist_ok=True) + if filename: + return os.path.join(joint_graph_sub_dir, filename) + return joint_graph_sub_dir + if filename: + return os.path.join(joint_graph_dir, filename) + return joint_graph_dir + + +def joint_graph_vis( + joint_module: fx.GraphModule, + fwd_module: fx.GraphModule, + bwd_module: fx.GraphModule, + save_tensor_nodes: Optional[List[fx.Node]] = None, + file_path: str = None, +): + """ + Visualize joint graph with coloring priority: save_tensor > fwd > bwd. + + Layout (bottom-to-top flow): + - Top: BWD output (gradient outputs) + - Middle-Left: FWD cluster (inputs + fwd ops + save_tensors) + - Middle-Right: BWD cluster (bwd ops) + - Bottom: Tangent inputs + FWD outputs (save_tensors) + + Node colors: + - Blue box: Input nodes (fwd placeholders) + - Cyan box: Tangent nodes (bwd gradient inputs) + - Green ellipse: FWD nodes + - Yellow hexagon: SaveTensor nodes (in FWD cluster) + - Orange ellipse: BWD nodes + - Orange box: BWD output (gradient outputs, at top) + + Args: + joint_module: The joint graph module containing both fwd and bwd + fwd_module: The forward graph module + bwd_module: The backward graph module + save_tensor_nodes: List of nodes that are saved tensors for backward + file_path: Optional path to save the visualization. If None, uses default path. + """ + if not envs.MAGI_ENABLE_FX_GRAPH_VIZ: + magi_logger.info("Joint graph visualization is disabled. Set MAGI_ENABLE_FX_GRAPH_VIZ=true to enable it.") + return + + joint_graph = joint_module.graph if isinstance(joint_module, fx.GraphModule) else joint_module + fwd_graph = fwd_module.graph if isinstance(fwd_module, fx.GraphModule) else fwd_module + bwd_graph = bwd_module.graph if isinstance(bwd_module, fx.GraphModule) else bwd_module + + node_categories, input_save_tensors = categorize_joint_nodes(joint_graph, fwd_graph, bwd_graph, save_tensor_nodes) + + nodes, edges = extract_joint_graph_structure(joint_graph, node_categories, input_save_tensors) + + dot = create_joint_graph_dot(nodes, edges) + + if file_path is None: + file_path = get_joint_graph_path(filename="joint_graph") + + dot.render(filename=file_path, view=False, cleanup=True) + magi_logger.info("Joint graph visualization saved to: %s.pdf", file_path) diff --git a/pkgs/MagiCompiler/magi_compiler/utils/visualize/visualizer.py b/pkgs/MagiCompiler/magi_compiler/utils/visualize/visualizer.py new file mode 100644 index 0000000000000000000000000000000000000000..4df743fde83c5999f8fef3ac28c5facec4948f1d --- /dev/null +++ b/pkgs/MagiCompiler/magi_compiler/utils/visualize/visualizer.py @@ -0,0 +1,234 @@ +# 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. + +import os +from typing import Any, Dict, List, Tuple + +import graphviz +import torch +import torch.fx +from magi_compiler.config import get_compile_config +from magi_compiler.utils import envs, magi_logger + + +class FX_NODE_OP: + PLACEHOLDER = "placeholder" + OUTPUT = "output" + CALL_MODULE = "call_module" + CALL_FUNCTION = "call_function" + CALL_METHOD = "call_method" + GET_ATTR = "get_attr" + DEFAULT = "default" + + +FIXED_NODE_STYLES = { + FX_NODE_OP.PLACEHOLDER: {"shape": "rectangle", "style": "filled,bold", "fillcolor": "#E8F4FD", "color": "#2196F3"}, + FX_NODE_OP.OUTPUT: {"shape": "rectangle", "style": "filled,bold", "fillcolor": "#FCE4EC", "color": "#E91E63"}, + FX_NODE_OP.CALL_MODULE: {"shape": "ellipse", "style": "filled,bold", "fillcolor": "#FFF8E1", "color": "#FFC107"}, + FX_NODE_OP.CALL_FUNCTION: {"shape": "ellipse", "style": "filled", "fillcolor": "#E8F5E9", "color": "#4CAF50"}, + FX_NODE_OP.CALL_METHOD: {"shape": "ellipse", "style": "filled", "fillcolor": "#F3E5F5", "color": "#9C27B0"}, + FX_NODE_OP.GET_ATTR: {"shape": "ellipse", "style": "filled", "fillcolor": "#FFCCBC", "color": "#FF5722"}, + FX_NODE_OP.DEFAULT: {"shape": "ellipse", "style": "filled", "fillcolor": "#F5F5F5", "color": "#666666"}, + "call_function.linear": {"shape": "ellipse", "style": "filled,bold", "fillcolor": "#E8F5E9", "color": "#44FF00"}, +} + + +def build_node_to_code_map(graph: torch.fx.Graph) -> Dict[torch.fx.Node, str]: + node_to_code = {} + + python_code = graph.python_code(root_module="self", verbose=False) + code_lines = python_code.src.strip().split("\n") + lineno_map = python_code._lineno_map + + node_index_map = {idx: node for idx, node in enumerate(graph.nodes)} + + for line_num, node_idx in lineno_map.items(): + if node_idx is None or node_idx not in node_index_map: + continue + node = node_index_map[node_idx] + if 0 <= line_num < len(code_lines): + code_line = code_lines[line_num].strip() + if code_line and not code_line.startswith(("wrap(", "#", "pass")): + node_to_code[node] = code_line + + for node in graph.nodes: + if node not in node_to_code: + assert node.op == FX_NODE_OP.PLACEHOLDER, f"Unexpected missing code for {node.op=}, {node.target=}" + node_to_code[node] = "" + return node_to_code + + +def extract_fx_graph_structure(graph: torch.fx.Graph, simple_desc: bool = False) -> Tuple[List[Dict], List[Dict]]: + import textwrap + + def wrap_str_to_multi_lines(text, width=30): + return textwrap.fill(text, width=width, break_long_words=True, replace_whitespace=False) + + nodes, edges = [], [] + node_to_code = build_node_to_code_map(graph) + + for node in graph.nodes: + name_str = str(node.name) + name_str = wrap_str_to_multi_lines(name_str) + + if node.op == FX_NODE_OP.GET_ATTR: + # torch._inductor.exc.InductorError: DataDependentOutputException: aten._local_scalar_dense.default + meta_str = f"get_attr: {node.target} (skip fake tensor)" + else: + tensor_meta = node.meta.get("tensor_meta") or node.meta.get("val") or node.meta.get("example_value") + meta_str = tensor_meta_to_str(tensor_meta) + meta_str = wrap_str_to_multi_lines(meta_str) + + target_str = target_to_str(node.target) + if hasattr(node, "original_target"): + original_target_str = target_to_str(node.original_target) + target_str += f"\nOriginal: {original_target_str}" + target_str = wrap_str_to_multi_lines(target_str) + + node_code = node_to_code.get(node, "empty") + node_code = wrap_str_to_multi_lines(node_code) + + style_info = FIXED_NODE_STYLES.get(node.op, FIXED_NODE_STYLES[FX_NODE_OP.DEFAULT]) + if node.op == FX_NODE_OP.CALL_FUNCTION and f"call_function.{node.target.__name__}" in FIXED_NODE_STYLES: + style_info = FIXED_NODE_STYLES[f"call_function.{node.target.__name__}"] + + node_label = f"Op: {node.op}\nTarget: {target_str}\nName: {name_str}\nMeta: {meta_str}\nCode: {node_code}" + if simple_desc: + node_label = f"Op: {node.op}\nTarget: {target_str}\nName: {name_str}" + + nodes.append({"id": node.name, "style": style_info, "node_label": node_label}) + + def traverse_args(args_kwargs): + if isinstance(args_kwargs, (tuple, list)): + for arg in args_kwargs: + traverse_args(arg) + elif isinstance(args_kwargs, dict): + for val in args_kwargs.values(): + traverse_args(val) + elif isinstance(args_kwargs, torch.fx.Node): + d = {"source": args_kwargs.name, "target": node.name} + edges.append(d) if d not in edges else None + + traverse_args(node.args) + traverse_args(node.kwargs) + + return nodes, edges + + +def target_to_str(target: Any) -> str: + res = str(target) + if isinstance(target, str): + res = target + elif hasattr(target, "_op"): + res = str(target._op) + elif callable(target): + res = getattr(target, "__name__") + return res + + +def tensor_meta_to_str(tensor_meta: Any) -> str: + if type(tensor_meta) in [int, float, str, bool]: + return str(tensor_meta) + elif isinstance(tensor_meta, (list, tuple)): + return f"[{', '.join([tensor_meta_to_str(t) for t in tensor_meta])}]" + elif isinstance(tensor_meta, torch.Tensor): + d = {} + d["shape"] = tensor_meta.shape if hasattr(tensor_meta, "shape") else "N/A" + d["size"] = tensor_meta.size() if hasattr(tensor_meta, "size") else "N/A" + d["ndim"] = tensor_meta.ndim if hasattr(tensor_meta, "ndim") else "N/A" + d["numel"] = tensor_meta.numel() if hasattr(tensor_meta, "numel") else "N/A" + d["stride"] = tensor_meta.stride if hasattr(tensor_meta, "stride") else "N/A" + d["stride"] = tensor_meta.stride() if hasattr(tensor_meta, "stride") and callable(tensor_meta.stride) else "N/A" + d["is_contiguous"] = tensor_meta.is_contiguous() if hasattr(tensor_meta, "is_contiguous") else "N/A" + d["dtype"] = str(tensor_meta.dtype) if hasattr(tensor_meta, "dtype") else "N/A" + d["device"] = str(tensor_meta.device) if hasattr(tensor_meta, "device") else "N/A" + return ", ".join([f"{k}: {v}" for k, v in d.items()]) + else: + return str(tensor_meta) + + +def create_fx_graph_dot(nodes: list[Dict], edges: list[Dict]) -> graphviz.Digraph: + dot = graphviz.Digraph( + name="fx_graph", + format="pdf", + graph_attr={ + "rankdir": "TD", # Top to Down 布局 + "nodesep": "0.1", # 节点之间的间距 + "ranksep": "0.1", # 不同层级（Rank）之间的间距 + "overlap": "false", # 防止节点重叠 + "splines": "spline", # 使用平滑曲线连线，比 ortho 性能好且易读 + }, + node_attr={ + "fontname": "Helvetica", + "fontsize": "10", + "shape": "rect", + "style": "rounded,filled", + "fixedsize": "false", # 允许节点随文字长度自动撑大 + "margin": "0.2", # 增加文字边距 + }, + ) + + for node in nodes: + dot.node(node["id"], node["node_label"], **node["style"]) + + for edge in edges: + dot.edge(edge["source"], edge["target"]) + + return dot + + +def get_fx_graph_path(sub_dir: str = "", filename: str = "") -> str: + cache_root_dir = get_compile_config().cache_root_dir + rank = torch.distributed.get_rank() if torch.distributed.is_initialized() else 0 + fx_graph_dir = os.path.join(cache_root_dir, "fx_graph_views", f"rank_{rank}") + os.makedirs(fx_graph_dir, exist_ok=True) + if sub_dir: + fx_graph_sub_dir = os.path.join(fx_graph_dir, sub_dir) + os.makedirs(fx_graph_sub_dir, exist_ok=True) + if filename: + return os.path.join(fx_graph_sub_dir, filename) + return fx_graph_sub_dir + if filename: + return os.path.join(fx_graph_dir, filename) + return fx_graph_dir + + +def save_fx_graph_visualization(graph: torch.fx.Graph, sub_dir: str = "", filename: str = "fx_graph"): + """ + Save FX graph visualization as PDF. + + Args: + graph: The FX graph or GraphModule to visualize + sub_dir: Optional subdirectory under the fx_graph_views folder + filename: Filename for the output PDF (without extension) + """ + if not envs.MAGI_ENABLE_FX_GRAPH_VIZ: + magi_logger.info("FX graph visualization is disabled. Set MAGI_ENABLE_FX_GRAPH_VIZ=true to enable it.") + return + + if isinstance(graph, torch.fx.GraphModule): + graph = graph.graph + + assert envs.MAGI_FX_GRAPH_VIZ_NODE_DESC in { + "simple", + "detailed", + }, f"Invalid MAGI_FX_GRAPH_VIZ_NODE_DESC: {envs.MAGI_FX_GRAPH_VIZ_NODE_DESC}" + simple_desc = envs.MAGI_FX_GRAPH_VIZ_NODE_DESC == "simple" + + file_path = get_fx_graph_path(sub_dir=sub_dir, filename=filename) + + nodes, edges = extract_fx_graph_structure(graph, simple_desc=simple_desc) + dot = create_fx_graph_dot(nodes, edges) + dot.render(filename=file_path, view=False, cleanup=True) + magi_logger.info("FX graph visualization saved to: %s.pdf", file_path) diff --git a/pkgs/MagiCompiler/pyproject.toml b/pkgs/MagiCompiler/pyproject.toml new file mode 100644 index 0000000000000000000000000000000000000000..800278ffea456015b553da1b1b733d14fb9c16df --- /dev/null +++ b/pkgs/MagiCompiler/pyproject.toml @@ -0,0 +1,58 @@ +[project] +name = "magi_compiler" +dynamic = ["version"] +description = "A custom compiler that delivers plug-and-play optimizations to accelerate any inference framework. MagiCompiler systematically exposes the capabilities of PyTorch Dynamo, AOTAutograd, Inductor, and Triton, while prioritizing correctness, stability, observability, and maintainability." +authors = [ + {name = "Hongyu Jia", email = "hongyujia@sand.ai"}, + {name = "Zhiyao Cen", email = "2523403608@qq.com"}, + {name = "Taoran Wang", email = "wangtaoran0504nb@outlook.com"}, +] +license = "Apache-2.0" +readme = "README.md" +requires-python = ">=3.10" +classifiers = [ + "Programming Language :: Python :: 3", + "Operating System :: OS Independent", +] + +[project.urls] +"Homepage" = "https://github.com/SandAI-org/MagiCompiler/" +"Bug Tracker" = "https://github.com/SandAI-org/MagiCompiler/issues" + +[tool.isort] +profile = "black" + +[build-system] +requires = [ + "setuptools>=61.0", + "wheel", + "torch", + "ninja", + "packaging", + "cuda-python", + "seaborn" +] +build-backend = "setuptools.build_meta" + +[tool.versioningit] +vcs = "git" +default-version = "0.0.1" + +[tool.versioningit.write] +file = "magi_compiler/_version.py" +encoding = "utf-8" +template = "__version__ = '{version}'" + +[tool.setuptools] +include-package-data = false + +[tool.setuptools.packages.find] +include = ["magi_compiler*"] +exclude = [ + "build*", + "tests*", + "dist*", + "docs*", + "tools*", + "assets*", +] diff --git a/pkgs/MagiCompiler/requirements.txt b/pkgs/MagiCompiler/requirements.txt new file mode 100644 index 0000000000000000000000000000000000000000..5f3a2975941af2355fbb7333cd67f57521a42836 --- /dev/null +++ b/pkgs/MagiCompiler/requirements.txt @@ -0,0 +1,4 @@ +cuda-python +depyf +graphviz +seaborn diff --git a/pkgs/MagiCompiler/setup.py b/pkgs/MagiCompiler/setup.py new file mode 100644 index 0000000000000000000000000000000000000000..a7ddd03c762c2466505ed653c5743569c6c4cd54 --- /dev/null +++ b/pkgs/MagiCompiler/setup.py @@ -0,0 +1,32 @@ +# 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. + +import os + +from setuptools import find_packages, setup + +this_dir = os.path.dirname(os.path.abspath(__file__)) + +PACKAGE_NAME = "magi_compiler" + +with open(os.path.join(this_dir, "README.md"), "r", encoding="utf-8") as fh: + long_description = fh.read() + +setup( + name=PACKAGE_NAME, + packages=find_packages(exclude=["build", "tests", "dist", "docs", "tools", "assets"]), + long_description=long_description, + long_description_content_type="text/markdown", + # ext_modules=[], + # cmdclass={}, +) diff --git a/pkgs/MagiCompiler/tests/__init__.py b/pkgs/MagiCompiler/tests/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..3dbb800db9261283c616f3f0e37218b12a865a34 --- /dev/null +++ b/pkgs/MagiCompiler/tests/__init__.py @@ -0,0 +1,13 @@ +# 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. diff --git a/pkgs/MagiCompiler/tests/conftest.py b/pkgs/MagiCompiler/tests/conftest.py new file mode 100644 index 0000000000000000000000000000000000000000..0cc1a643637a4ec74ddc65d44f282b71645475f0 --- /dev/null +++ b/pkgs/MagiCompiler/tests/conftest.py @@ -0,0 +1,41 @@ +# 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. + +import shutil + +import pytest +import torch +from magi_compiler.config import get_compile_config + +from .model_definition import MLPConfig + + +@pytest.fixture(scope="function") +def device(): + """Device fixture""" + return torch.device("cuda" if torch.cuda.is_available() else "cpu") + + +@pytest.fixture(scope="function") +def mlp_config(): + """MLP configuration fixture""" + return MLPConfig(hidden_size=512, intermediate_size=2048, params_dtype=torch.bfloat16) + + +@pytest.fixture(scope="function", autouse=True) +def cleanup_cache(): + """Auto cleanup cache fixture, executed before and after each test""" + shutil.rmtree(get_compile_config().cache_root_dir, ignore_errors=True) + yield + shutil.rmtree(get_compile_config().cache_root_dir, ignore_errors=True) diff --git a/pkgs/MagiCompiler/tests/magi_depyf/__init__.py b/pkgs/MagiCompiler/tests/magi_depyf/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..3eaa44adb1bd11bb4c8d48c6f00d7a08f292f395 --- /dev/null +++ b/pkgs/MagiCompiler/tests/magi_depyf/__init__.py @@ -0,0 +1,13 @@ +# Copyright (c) 2026 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. diff --git a/pkgs/MagiCompiler/tests/magi_depyf/decompile/__init__.py b/pkgs/MagiCompiler/tests/magi_depyf/decompile/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..3eaa44adb1bd11bb4c8d48c6f00d7a08f292f395 --- /dev/null +++ b/pkgs/MagiCompiler/tests/magi_depyf/decompile/__init__.py @@ -0,0 +1,13 @@ +# Copyright (c) 2026 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. diff --git a/pkgs/MagiCompiler/tests/magi_depyf/decompile/decompile_roundtrip/__init__.py b/pkgs/MagiCompiler/tests/magi_depyf/decompile/decompile_roundtrip/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..3eaa44adb1bd11bb4c8d48c6f00d7a08f292f395 --- /dev/null +++ b/pkgs/MagiCompiler/tests/magi_depyf/decompile/decompile_roundtrip/__init__.py @@ -0,0 +1,13 @@ +# Copyright (c) 2026 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. diff --git a/pkgs/MagiCompiler/tests/magi_depyf/decompile/decompile_roundtrip/helpers.py b/pkgs/MagiCompiler/tests/magi_depyf/decompile/decompile_roundtrip/helpers.py new file mode 100644 index 0000000000000000000000000000000000000000..d9507d58d3f5f4f776d53ad94f3499a58f922dc9 --- /dev/null +++ b/pkgs/MagiCompiler/tests/magi_depyf/decompile/decompile_roundtrip/helpers.py @@ -0,0 +1,95 @@ +# Copyright (c) 2026 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. + +"""Shared helpers for decompile-roundtrip tests.""" + +import dis +from contextlib import contextmanager +from dataclasses import dataclass + +from magi_compiler.magi_depyf import decompile +from magi_compiler.magi_depyf.decompile.recompiler import CodeRecompiler + + +def _assert_uses_op(func, *specs): + """Assert *func*'s bytecode contains at least one instruction matching any spec. + + Each *spec* is either: + - ``str`` -- match ``opname`` exactly + - ``(opname, substr)`` -- match ``opname`` AND ``argrepr`` contains *substr* + """ + for inst in dis.get_instructions(func.__code__): + for spec in specs: + if isinstance(spec, str): + if inst.opname == spec: + return + else: + opname, substr = spec + if inst.opname == opname and substr in str(inst.argrepr): + return + found = sorted({f"{i.opname}({i.argrepr})" if i.argrepr else i.opname for i in dis.get_instructions(func.__code__)}) + spec_str = ", ".join(repr(s) for s in specs) + raise AssertionError(f"Bytecode of {func.__name__} does not contain any of [{spec_str}].\n" f"Found: {found}") + + +def roundtrip_code(func): + """Decompile *func*, compile the result, return the new CodeType.""" + old = func.__code__ + src = decompile(old) + compiled = compile(src, filename=old.co_filename, mode="exec") + codes = CodeRecompiler.collect_code_objects(compiled) + return [c for c in codes if c.co_name == old.co_name][0] + + +@contextmanager +def replaced_code(func, *expected_ops): + """Context-manager: run *func* with decompiled+recompiled code. + + If *expected_ops* are given, assert the bytecode contains at least one + matching instruction **before** the roundtrip (catches constant-folding). + """ + if expected_ops: + _assert_uses_op(func, *expected_ops) + old = func.__code__ + func.__code__ = roundtrip_code(func) + try: + yield + finally: + func.__code__ = old + + +@dataclass +class Point: + x: int + y: int + + def __matmul__(self, other): + return self.x * other.x + self.y * other.y + + def __imatmul__(self, other): + self.x = self.x * other.x + self.y = self.y * other.y + return self + + def __setitem__(self, key, value): + if key == 0: + self.x = value + elif key == 1: + self.y = value + else: + raise IndexError("Point only has two dimensions") + + +point = Point(1, 2) +data_map = {1: 2} diff --git a/pkgs/MagiCompiler/tests/magi_depyf/decompile/decompile_roundtrip/test_comparison.py b/pkgs/MagiCompiler/tests/magi_depyf/decompile/decompile_roundtrip/test_comparison.py new file mode 100644 index 0000000000000000000000000000000000000000..42619dd66c7397797249e0739de094eb0ba35fb5 --- /dev/null +++ b/pkgs/MagiCompiler/tests/magi_depyf/decompile/decompile_roundtrip/test_comparison.py @@ -0,0 +1,44 @@ +# Copyright (c) 2026 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. + +"""Roundtrip tests for comparison / identity / membership operators.""" + +from tests.magi_depyf.decompile.decompile_roundtrip.helpers import replaced_code + + +def test_COMPARE_OP(): + def f(): + return (3 == 3) + (1 < 2) + (2 > 1) + (2 >= 2) + (1 <= 2) + (1 != 2) + + ans = f() + with replaced_code(f, "COMPARE_OP"): + assert f() == ans + + +def test_IS_OP(): + def f(): + return (int is int), (int is not float) + + ans = f() + with replaced_code(f, "IS_OP"): + assert f() == ans + + +def test_CONTAINS_OP(): + def f(): + return (1 in [1, 2, 3]), (5 not in (6, 7, 4)) + + ans = f() + with replaced_code(f, "CONTAINS_OP"): + assert f() == ans diff --git a/pkgs/MagiCompiler/tests/magi_depyf/decompile/decompile_roundtrip/test_containers.py b/pkgs/MagiCompiler/tests/magi_depyf/decompile/decompile_roundtrip/test_containers.py new file mode 100644 index 0000000000000000000000000000000000000000..d261aac5088b91a1ade14a04cfa22f6fe56200bc --- /dev/null +++ b/pkgs/MagiCompiler/tests/magi_depyf/decompile/decompile_roundtrip/test_containers.py @@ -0,0 +1,196 @@ +# Copyright (c) 2026 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. + +"""Roundtrip tests for containers (tuple/list/set/dict), unpack, and comprehensions.""" + +from tests.magi_depyf.decompile.decompile_roundtrip.helpers import replaced_code + +# ====================================================================== +# Containers: tuple / list / set / dict +# ====================================================================== + + +def test_BUILD_TUPLE(): + def f(): + a, b = 1, 2 + return (a, b), (a,) + + ans = f() + with replaced_code(f, "BUILD_TUPLE"): + assert f() == ans + + +def test_BUILD_LIST(): + def f(): + a, b = 1, 2 + return [a, b], [a] + + ans = f() + with replaced_code(f, "BUILD_LIST"): + assert f() == ans + + +def test_BUILD_SET(): + def f(): + a, b = 1, 2 + return {a, b}, {a} + + ans = f() + with replaced_code(f, "BUILD_SET"): + assert f() == ans + + +def test_BUILD_MAP(): + def f(): + a, b = 1, 2 + return {a: 1, 2: 3}, {b: a} + + ans = f() + with replaced_code(f, "BUILD_MAP"): + assert f() == ans + + +def test_BUILD_CONST_KEY_MAP(): + def f(): + return {5: 1, 2: 3} + + ans = f() + with replaced_code(f, "BUILD_CONST_KEY_MAP"): + assert f() == ans + + +def test_BUILD_CONST_KEY_MAP_string_keys(): + def f(): + return {"hello": 1, "world": 2} + + ans = f() + with replaced_code(f, "BUILD_CONST_KEY_MAP"): + assert f() == ans + + +def test_LIST_EXTEND(): + def f(): + return [1, 2, 3] + + ans = f() + with replaced_code(f, "LIST_EXTEND"): + assert f() == ans + + +def test_SET_UPDATE(): + def f(): + return {1, 2, 3} + + ans = f() + with replaced_code(f, "SET_UPDATE"): + assert f() == ans + + +def test_DICT_UPDATE(): + def f(): + a = {1: 2} + b = {'a': 4} + return {**a, **b} + + ans = f() + with replaced_code(f, "DICT_UPDATE"): + assert f() == ans + + +def test_DICT_MERGE(): + def f(): + a = {1: 2} + b = {'a': 4} + a.update(**b) + return a + + ans = f() + with replaced_code(f, "DICT_MERGE"): + assert f() == ans + + +# ====================================================================== +# Unpack +# ====================================================================== + + +def test_UNPACK_SEQUENCE(): + def f(): + a, b = (1, 2) + return a + + ans = f() + with replaced_code(f, "UNPACK_SEQUENCE"): + assert f() == ans + + +def test_UNPACK_SEQUENCE_one(): + def f(): + (a,) = (1,) + return a + + ans = f() + with replaced_code(f, "UNPACK_SEQUENCE"): + assert f() == ans + + +def test_UNPACK_EX(): + def f(): + a, *b = (1, 2, 3) + return b + + ans = f() + with replaced_code(f, "UNPACK_EX"): + assert f() == ans + + +# ====================================================================== +# Comprehensions +# ====================================================================== + + +def test_LIST_COMP(): + def f(a): + return [i**2 for i in range(a)] + + ans = [f(i) for i in range(10)] + with replaced_code(f, "LIST_APPEND"): + assert [f(i) for i in range(10)] == ans + + +def test_SET_COMP(): + def f(a): + return {i**2 for i in range(a)} + + ans = [f(i) for i in range(10)] + with replaced_code(f, "SET_ADD"): + assert [f(i) for i in range(10)] == ans + + +def test_MAP_COMP(): + def f(a): + return {i: i**2 for i in range(a)} + + ans = [f(i) for i in range(10)] + with replaced_code(f, "MAP_ADD"): + assert [f(i) for i in range(10)] == ans + + +def test_NESTED_COMP(): + def f(a): + return [{x: {_ for _ in range(x)} for x in range(i)} for i in range(a)] + + ans = [f(i) for i in range(5)] + with replaced_code(f, "LIST_APPEND"): + assert [f(i) for i in range(5)] == ans diff --git a/pkgs/MagiCompiler/tests/magi_depyf/decompile/decompile_roundtrip/test_control_flow.py b/pkgs/MagiCompiler/tests/magi_depyf/decompile/decompile_roundtrip/test_control_flow.py new file mode 100644 index 0000000000000000000000000000000000000000..d5b7cfb9926d2a52f02690c7a5c6eda59ee563bd --- /dev/null +++ b/pkgs/MagiCompiler/tests/magi_depyf/decompile/decompile_roundtrip/test_control_flow.py @@ -0,0 +1,235 @@ +# Copyright (c) 2026 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. + +"""Roundtrip tests for control flow: if/elif/else, for loops.""" + +from magi_compiler.magi_depyf import decompile +from tests.magi_depyf.decompile.decompile_roundtrip.helpers import replaced_code + +# ====================================================================== +# if / elif / else +# ====================================================================== + + +def test_IF(): + def f(a): + if a == 0: + return 0 + elif a == 1: + return 1 + else: + return 2 + + ans = [f(i) for i in range(10)] + with replaced_code(f, "COMPARE_OP"): + assert [f(i) for i in range(10)] == ans + + +def test_compound_IF_and(): + def f(a, b): + c = 1 + if a > 0 and b > 1: + c += 1 + else: + c += 2 + c += 3 + return c + + ans = [f(a, b) for a in range(-3, 3) for b in range(-3, 3)] + with replaced_code(f, "POP_JUMP_IF_FALSE"): + assert [f(a, b) for a in range(-3, 3) for b in range(-3, 3)] == ans + + +def test_compound_IF_or(): + def f(a, b): + c = 1 + if a > 0 or b > 1: + c += 1 + else: + c += 2 + c += 3 + return c + + ans = [f(a, b) for a in range(-3, 3) for b in range(-3, 3)] + with replaced_code(f, "POP_JUMP_IF_FALSE"): + assert [f(a, b) for a in range(-3, 3) for b in range(-3, 3)] == ans + + +def test_IF_NONE(): + def f(a): + if a is None: + return 0 + elif a is not None: + return 1 + + ans = [f(i) for i in range(10)] + with replaced_code(f, "POP_JUMP_IF_NONE"): + assert [f(i) for i in range(10)] == ans + + +def test_ternary(): + def f(output_hidden_states): + () if output_hidden_states else None + + ans = [f(i) for i in range(2)] + with replaced_code(f, "POP_JUMP_IF_FALSE"): + assert [f(i) for i in range(2)] == ans + + +def test_shortcircuit(): + def f(a, b): + if a > 0 and b > 0: + return a + b + elif a > 1 or b > 2: + return a - b + else: + return 2 + + scope = {} + exec(decompile(f), scope) + for a in [-1, 0, 1, 2]: + for b in [-1, 0, 1, 2]: + assert f(a, b) == scope['f'](a, b) + + +def test_IF_return_in_both_branches(): + def f(x, y): + if x: + return 42 + else: + return [x, y] + + ans = [f(x, y) for x in [True, False] for y in [1, 2]] + with replaced_code(f, "POP_JUMP_IF_FALSE"): + assert [f(x, y) for x in [True, False] for y in [1, 2]] == ans + + +def test_IF_nested(): + def f(a): + if a > 10: + return "big" + elif a > 5: + return "medium" + elif a > 0: + return "small" + else: + return "negative" + + ans = [f(i) for i in [-1, 0, 3, 7, 15]] + with replaced_code(f, "COMPARE_OP"): + assert [f(i) for i in [-1, 0, 3, 7, 15]] == ans + + +def test_IF_assign_then_use(): + def f(flag): + if flag: + x = 10 + else: + x = 20 + return x + 1 + + ans = [f(b) for b in [True, False]] + with replaced_code(f, "POP_JUMP_IF_FALSE"): + assert [f(b) for b in [True, False]] == ans + + +# ====================================================================== +# For loop +# ====================================================================== + + +def test_simple_for(): + def f(a): + for i in range(5): + a += i + return a + + ans = [f(i) for i in range(10)] + with replaced_code(f, "FOR_ITER"): + assert [f(i) for i in range(10)] == ans + + +def test_for_with_break(): + def f(items): + for x in items: + if x < 0: + break + return x + + with replaced_code(f, "FOR_ITER"): + assert f([1, 2, -1, 3]) == -1 + + +def test_for_with_continue(): + def f(n): + total = 0 + for i in range(n): + if i % 2 == 0: + continue + total += i + return total + + ans = f(10) + with replaced_code(f, "FOR_ITER"): + assert f(10) == ans + + +def test_for_nested(): + def f(n): + total = 0 + for i in range(n): + for j in range(i): + total += j + return total + + ans = f(5) + with replaced_code(f, "FOR_ITER"): + assert f(5) == ans + + +def test_for_with_enumerate(): + def f(): + total = 0 + for i, v in enumerate([10, 20, 30]): + total += i * v + return total + + ans = f() + with replaced_code(f, "FOR_ITER"): + assert f() == ans + + +def test_for_dict_items(): + def f(): + d = {"a": 1, "b": 2, "c": 3} + total = 0 + for k, v in d.items(): + total += v + return total + + ans = f() + with replaced_code(f, "FOR_ITER"): + assert f() == ans + + +def test_for_loop_accumulate(): + def f(): + total = 0 + for x in [10, 20, 30]: + total += x + return total + + ans = f() + with replaced_code(f, "FOR_ITER"): + assert f() == ans diff --git a/pkgs/MagiCompiler/tests/magi_depyf/decompile/decompile_roundtrip/test_function_call.py b/pkgs/MagiCompiler/tests/magi_depyf/decompile/decompile_roundtrip/test_function_call.py new file mode 100644 index 0000000000000000000000000000000000000000..5522f44341a4ba3cf9325a4ace7e29f675980301 --- /dev/null +++ b/pkgs/MagiCompiler/tests/magi_depyf/decompile/decompile_roundtrip/test_function_call.py @@ -0,0 +1,83 @@ +# Copyright (c) 2026 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. + +"""Roundtrip tests for function calls (pure — no closures).""" + +from tests.magi_depyf.decompile.decompile_roundtrip.helpers import replaced_code + + +def test_CALL_FUNCTION_NORMAL(): + def f(func): + a = [1, 2, 3] + b = {'a': 4} + return func(1, a, b), func(a=a, b=b) + + def helper(a, b, c=1): + return (a, b, c) + + ans = f(helper) + with replaced_code(f, "CALL"): + assert f(helper) == ans + + +def test_function_signature(): + def func(a, b, c=1, *, d=2): + return (a, b, c, d) + + a, b = [1, 2, 3], {'a': 4} + ans = func(1, a, b, d=5) + with replaced_code(func, "BUILD_TUPLE"): + assert func(1, a, b, d=5) == ans + + +def test_CALL_FUNCTION_EX(): + def f(func): + a = [1, 2, 3] + b = {'a': 4} + return func(*a), func(**b), func(*a, **b) + + def helper(*args, **kwargs): + return (args, kwargs) + + ans = f(helper) + with replaced_code(f, "CALL_FUNCTION_EX"): + assert f(helper) == ans + + +def test_var_args(): + def func(*args, **kwargs): + return (args, kwargs) + + a, b = [1, 2, 3], {'a': 4} + ans = func(1, a, b, d=5) + with replaced_code(func, "BUILD_TUPLE"): + assert func(1, a, b, d=5) == ans + + +def test_complex_signature(): + def func(a, b, *args, **kwargs): + return (a, b, args, kwargs) + + a, b = [1, 2, 3], {'a': 4} + ans = func(1, a, b, d=5) + with replaced_code(func, "BUILD_TUPLE"): + assert func(1, a, b, d=5) == ans + + +def test_call_with_kwargs(): + def f(): + return dict(a=1, b=2) + + with replaced_code(f, "KW_NAMES"): + assert f() == {"a": 1, "b": 2} diff --git a/pkgs/MagiCompiler/tests/magi_depyf/decompile/decompile_roundtrip/test_instruction_api.py b/pkgs/MagiCompiler/tests/magi_depyf/decompile/decompile_roundtrip/test_instruction_api.py new file mode 100644 index 0000000000000000000000000000000000000000..02db08c2164daf7d82a4bfe0049a80b719b5c33c --- /dev/null +++ b/pkgs/MagiCompiler/tests/magi_depyf/decompile/decompile_roundtrip/test_instruction_api.py @@ -0,0 +1,65 @@ +# Copyright (c) 2026 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. + +"""Instruction dataclass API tests: from_dis, category queries, nop_.""" + +import dis + +from magi_compiler.magi_depyf.decompile.bytecode.instruction import Instruction + + +def _make_inst(opname, **kw): + opcode = dis.opmap.get(opname, 0) + return Instruction(opcode=opcode, opname=opname, arg=0, argval=0, argrepr="", **kw) + + +class TestInstructionCategory: + def test_is_load(self): + assert _make_inst("LOAD_FAST").is_load + assert _make_inst("LOAD_GLOBAL").is_load + assert not _make_inst("STORE_FAST").is_load + + def test_is_store(self): + assert _make_inst("STORE_FAST").is_store + assert not _make_inst("LOAD_FAST").is_store + + def test_is_return(self): + assert _make_inst("RETURN_VALUE").is_return + assert _make_inst("RETURN_CONST").is_return + assert not _make_inst("LOAD_FAST").is_return + + def test_is_nop(self): + assert _make_inst("NOP").is_nop + assert not _make_inst("LOAD_FAST").is_nop + + +class TestInstructionNop: + def test_nop_mutates_in_place(self): + inst = _make_inst("LOAD_FAST", offset=10) + inst.nop_() + assert inst.opname == "NOP" + assert inst.is_nop + + +class TestInstructionFromDis: + def test_roundtrip(self): + def f(x): + return x + 1 + + stdlib_insts = list(dis.get_instructions(f.__code__)) + mine = [Instruction.from_dis(i) for i in stdlib_insts] + assert len(mine) == len(stdlib_insts) + for orig, converted in zip(stdlib_insts, mine): + assert converted.opname == orig.opname + assert converted.offset == orig.offset diff --git a/pkgs/MagiCompiler/tests/magi_depyf/decompile/decompile_roundtrip/test_misc.py b/pkgs/MagiCompiler/tests/magi_depyf/decompile/decompile_roundtrip/test_misc.py new file mode 100644 index 0000000000000000000000000000000000000000..79a3527e2f4dcb8981d4476ff3d83fcf0bf546ac --- /dev/null +++ b/pkgs/MagiCompiler/tests/magi_depyf/decompile/decompile_roundtrip/test_misc.py @@ -0,0 +1,389 @@ +# Copyright (c) 2026 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. + +"""Roundtrip tests for miscellaneous opcodes: +LOAD_ATTR, IMPORT, MAKE_FUNCTION, SWAP, FORMAT, SLICE, +constants, STORE_GLOBAL, COPY, RETURN_CONST, etc. +""" + +from magi_compiler.magi_depyf import decompile +from tests.magi_depyf.decompile.decompile_roundtrip.helpers import Point, replaced_code + +# ====================================================================== +# LOAD_ATTR / IMPORT +# ====================================================================== + + +def test_LOAD_ATTR(): + def f(): + p = Point(1, 2) + return p.x + + ans = f() + with replaced_code(f, "LOAD_ATTR"): + assert f() == ans + + +def test_IMPORT_NAME(): + def f(): + import functools + from math import sqrt + + return functools.partial(sqrt, 0.3)() + + ans = f() + with replaced_code(f, "IMPORT_NAME"): + assert f() == ans + + +# ====================================================================== +# MAKE_FUNCTION / nested functions +# ====================================================================== + + +def test_MAKE_FUNCTION(): + def f(a): + def g(b=3): + return a + b + + return g(2) + + ans = [f(i) for i in range(10)] + with replaced_code(f, "MAKE_FUNCTION"): + assert [f(i) for i in range(10)] == ans + + +def test_nested_function(): + def f(): + def inner(x): + return x * 2 + + return inner(21) + + with replaced_code(f, "MAKE_FUNCTION"): + assert f() == 42 + + +# ====================================================================== +# Swap / Rotate +# ====================================================================== + + +def test_ROT_TWO(): + def f(): + a, b = 1, 2 + a, b = b, a + return a, b + + ans = f() + with replaced_code(f, "UNPACK_SEQUENCE"): + assert f() == ans + + +def test_ROT_MULTI(): + def f(): + a, b, c, d = 1, 2, 3, 4 + a, b, c, d = d, c, b, a + return a + + ans = f() + with replaced_code(f, "UNPACK_SEQUENCE"): + assert f() == ans + + +def test_SWAP_triple(): + def f(): + a, b, c = 1, 2, 3 + c, b, a = a, b, c + return a, b, c + + ans = f() + with replaced_code(f, "UNPACK_SEQUENCE"): + assert f() == ans + + +# ====================================================================== +# Format string / BUILD_STRING +# ====================================================================== + + +def test_FORMAT_VALUE(): + def f(): + a, b, c = 1, 2, 3 + return f"{a} {b!r} {b!s} {b!a} {c:.2f}" + + ans = f() + with replaced_code(f, "FORMAT_VALUE"): + assert f() == ans + + +# ====================================================================== +# BUILD_SLICE / BINARY_SLICE / STORE_SLICE +# ====================================================================== + + +def test_BUILD_SLICE(): + def f(): + a = [1, 2, 3, 4, 5] + return a[:] + a[1:] + a[:3] + a[1:3] + a[::-1] + + ans = f() + with replaced_code(f, "BUILD_SLICE"): + assert f() == ans + + +def test_BINARY_SLICE(): + def f(): + a = [10, 20, 30, 40, 50] + return a[1:3], a[:2], a[3:] + + ans = f() + with replaced_code(f, "BINARY_SLICE"): + assert f() == ans + + +def test_STORE_SLICE(): + def f(): + a = [1, 2, 3, 4, 5] + a[1:3] = [20, 30] + return a + + ans = f() + with replaced_code(f, "STORE_SLICE"): + assert f() == ans + + +# ====================================================================== +# Constants (various types) +# ====================================================================== + + +def test_constants(): + def f(): + return (1, 2.5, "hello", True, None, b"bytes", (1, 2)) + + with replaced_code(f, "RETURN_CONST"): + assert f() == (1, 2.5, "hello", True, None, b"bytes", (1, 2)) + + +# ====================================================================== +# GET_LEN +# ====================================================================== + + +def test_GET_LEN(): + def f(): + return len((1, 2, 3)) + + ans = f() + with replaced_code(f, ("LOAD_GLOBAL", "len")): + assert f() == ans + + +# ====================================================================== +# STORE_GLOBAL +# ====================================================================== + + +def test_STORE_GLOBAL(): + def f(): + global len + len = 1 + return len + + with replaced_code(f, "STORE_GLOBAL"): + global len + original_len = len + f() + assert len == 1 + len = original_len + + +# ====================================================================== +# Class method with __class__ +# ====================================================================== + + +class A: + def f(self): + return __class__ + + +def test_class_method(): + """__class__ is an implicit freevar; verify decompilation produces valid text.""" + src = decompile(A.f.__code__) + assert "def f(self):" in src + assert "__class__" in src + + +# ====================================================================== +# COPY / SWAP patterns (Python 3.11+) +# ====================================================================== + + +def test_COPY_simple(): + def f(): + a = [1, 2, 3] + a[0] = a[1] = 99 + return a + + ans = f() + with replaced_code(f, "COPY"): + assert f() == ans + + +# ====================================================================== +# LIST_APPEND (in comprehensions) +# ====================================================================== + + +def test_LIST_APPEND_in_comp(): + def f(n): + return [x * 2 for x in range(n) if x % 2 == 0] + + ans = f(10) + with replaced_code(f, "LIST_APPEND"): + assert f(10) == ans + + +# ====================================================================== +# Mixed complex expressions +# ====================================================================== + + +def test_complex_expression(): + def f(a, b): + return (a + b) * (a - b) // (b + 1) + + ans = [f(a, b) for a in range(1, 5) for b in range(1, 5)] + with replaced_code(f, "BINARY_OP"): + assert [f(a, b) for a in range(1, 5) for b in range(1, 5)] == ans + + +def test_chained_method_calls(): + def f(): + return " hello world ".strip().upper().replace("O", "0") + + ans = f() + with replaced_code(f, "LOAD_ATTR"): + assert f() == ans + + +# ====================================================================== +# RETURN_CONST (Python 3.12+) +# ====================================================================== + + +def test_RETURN_CONST_none(): + def f(): + pass + + ans = f() + with replaced_code(f, "RETURN_CONST"): + assert f() == ans + + +def test_RETURN_CONST_string(): + def f(x): + if x > 0: + return "positive" + else: + return "negative" + + ans = [f(i) for i in [-1, 0, 1]] + with replaced_code(f, "RETURN_CONST"): + assert [f(i) for i in [-1, 0, 1]] == ans + + +def test_RETURN_CONST_number(): + def f(x): + if x: + return 42 + return 0 + + ans = [f(b) for b in [True, False]] + with replaced_code(f, "RETURN_CONST"): + assert [f(b) for b in [True, False]] == ans + + +# ====================================================================== +# LOAD_ASSERTION_ERROR +# ====================================================================== + + +def test_LOAD_ASSERTION_ERROR(): + scope = {} + exec(compile("def f(x):\n assert x > 0, 'must be positive'\n return x", "", "exec"), scope) + f = scope["f"] + src = decompile(f) + exec(src, scope) + g = scope["f"] + assert g(5) == 5 + + +# ====================================================================== +# Import patterns +# ====================================================================== + + +def test_import_dotted(): + def f(): + import os.path + + return os.path.sep + + ans = f() + with replaced_code(f, "IMPORT_NAME"): + assert f() == ans + + +def test_import_from(): + def f(): + from math import sqrt + + return sqrt(4) + + ans = f() + with replaced_code(f, "IMPORT_FROM"): + assert f() == ans + + +# ====================================================================== +# PUSH_NULL + LOAD_GLOBAL pattern (3.11+) +# ====================================================================== + + +def test_global_function_call(): + def f(): + return len([1, 2, 3]) + + ans = f() + with replaced_code(f, ("LOAD_GLOBAL", "len")): + assert f() == ans + + +# ====================================================================== +# LOAD_ATTR normal +# ====================================================================== + + +def test_LOAD_ATTR_normal(): + def f(): + import math + + return math.pi + + ans = f() + with replaced_code(f, "LOAD_ATTR"): + assert f() == ans diff --git a/pkgs/MagiCompiler/tests/magi_depyf/decompile/decompile_roundtrip/test_operators.py b/pkgs/MagiCompiler/tests/magi_depyf/decompile/decompile_roundtrip/test_operators.py new file mode 100644 index 0000000000000000000000000000000000000000..3933de680b8da0ff60f5e6fc71d381ad0e9769d9 --- /dev/null +++ b/pkgs/MagiCompiler/tests/magi_depyf/decompile/decompile_roundtrip/test_operators.py @@ -0,0 +1,353 @@ +# Copyright (c) 2026 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. + +"""Roundtrip tests for unary, binary, and in-place operators.""" + +from tests.magi_depyf.decompile.decompile_roundtrip.helpers import Point, point, replaced_code + +# ====================================================================== +# 1. Unary operators +# ====================================================================== + + +def test_UNARY_POSITIVE(): + def f(): + x = 1 + return +x + + ans = f() + with replaced_code(f, "CALL_INTRINSIC_1"): + assert f() == ans + + +def test_UNARY_NEGATIVE(): + def f(): + x = 1 + return -x + + ans = f() + with replaced_code(f, "UNARY_NEGATIVE"): + assert f() == ans + + +def test_UNARY_NOT(): + def f(): + x = 1 + return not x + + ans = f() + with replaced_code(f, "UNARY_NOT"): + assert f() == ans + + +def test_UNARY_INVERT(): + def f(): + x = 1 + return ~x + + ans = f() + with replaced_code(f, "UNARY_INVERT"): + assert f() == ans + + +# ====================================================================== +# 2. Binary operators +# ====================================================================== + + +def test_BINARY_POWER(): + def f(): + a, b = 2, 3 + return (a**b) ** a + + ans = f() + with replaced_code(f, ("BINARY_OP", "**")): + assert f() == ans + + +def test_BINARY_MULTIPLY(): + def f(): + a, b = 2, 3 + return a * b * a + + ans = f() + with replaced_code(f, ("BINARY_OP", "*")): + assert f() == ans + + +def test_BINARY_MATRIX_MULTIPLY(): + def f(): + return point @ point + + ans = f() + with replaced_code(f, ("BINARY_OP", "@")): + assert f() == ans + + +def test_BINARY_FLOOR_DIVIDE(): + def f(): + a, b = 7, 3 + return (a // b) // a + + ans = f() + with replaced_code(f, ("BINARY_OP", "//")): + assert f() == ans + + +def test_BINARY_TRUE_DIVIDE(): + def f(): + a, b = 7, 3 + return (a / b) / a + + ans = f() + with replaced_code(f, ("BINARY_OP", "/")): + assert f() == ans + + +def test_BINARY_MODULO(): + def f(): + a, b = 10, 3 + return (a % b) % a + + ans = f() + with replaced_code(f, ("BINARY_OP", "%")): + assert f() == ans + + +def test_BINARY_ADD(): + def f(): + a, b = 2, 3 + return (a + b) + a + + ans = f() + with replaced_code(f, ("BINARY_OP", "+")): + assert f() == ans + + +def test_BINARY_SUBTRACT(): + def f(): + a, b = 5, 3 + return (a - b) - a + + ans = f() + with replaced_code(f, ("BINARY_OP", "-")): + assert f() == ans + + +def test_BINARY_SUBSCR(): + def f(): + a = (10, 20, 30) + return a[1] + + ans = f() + with replaced_code(f, "BINARY_SUBSCR"): + assert f() == ans + + +def test_BINARY_LSHIFT(): + def f(): + a, b = 2, 3 + return (a << b) << 1 + + ans = f() + with replaced_code(f, ("BINARY_OP", "<<")): + assert f() == ans + + +def test_BINARY_RSHIFT(): + def f(): + a, b = 16, 2 + return (a >> b) >> 1 + + ans = f() + with replaced_code(f, ("BINARY_OP", ">>")): + assert f() == ans + + +def test_BINARY_AND(): + def f(): + a, b = 0b1100, 0b1010 + return (a & b) & 0b1111 + + ans = f() + with replaced_code(f, ("BINARY_OP", "&")): + assert f() == ans + + +def test_BINARY_XOR(): + def f(): + a, b = 0b1100, 0b1010 + return (a ^ b) ^ 0b0001 + + ans = f() + with replaced_code(f, ("BINARY_OP", "^")): + assert f() == ans + + +def test_BINARY_OR(): + def f(): + a, b = 0b1100, 0b1010 + return (a | b) | 0b0001 + + ans = f() + with replaced_code(f, ("BINARY_OP", "|")): + assert f() == ans + + +# ====================================================================== +# 3. In-place operators +# ====================================================================== + + +def test_INPLACE_POWER(): + def f(): + a = 2 + a **= 3 + return a + + ans = f() + with replaced_code(f, ("BINARY_OP", "**=")): + assert f() == ans + + +def test_INPLACE_MULTIPLY(): + def f(): + a = 2 + a *= 3 + return a + + ans = f() + with replaced_code(f, ("BINARY_OP", "*=")): + assert f() == ans + + +def test_INPLACE_MATRIX_MULTIPLY(): + def f(): + p = Point(1, 2) + p @= p + return p + + ans = f() + with replaced_code(f, ("BINARY_OP", "@=")): + assert f() == ans + + +def test_INPLACE_FLOOR_DIVIDE(): + def f(): + a = 7 + a //= 3 + return a + + ans = f() + with replaced_code(f, ("BINARY_OP", "//=")): + assert f() == ans + + +def test_INPLACE_TRUE_DIVIDE(): + def f(): + a = 7 + a /= 2 + return a + + ans = f() + with replaced_code(f, ("BINARY_OP", "/=")): + assert f() == ans + + +def test_INPLACE_MODULO(): + def f(): + a = 10 + a %= 3 + return a + + ans = f() + with replaced_code(f, ("BINARY_OP", "%=")): + assert f() == ans + + +def test_INPLACE_ADD(): + def f(): + a = 2 + a += 3 + return a + + ans = f() + with replaced_code(f, ("BINARY_OP", "+=")): + assert f() == ans + + +def test_INPLACE_SUBTRACT(): + def f(): + a = 5 + a -= 3 + return a + + ans = f() + with replaced_code(f, ("BINARY_OP", "-=")): + assert f() == ans + + +def test_INPLACE_LSHIFT(): + def f(): + a = 1 + a <<= 3 + return a + + ans = f() + with replaced_code(f, ("BINARY_OP", "<<=")): + assert f() == ans + + +def test_INPLACE_RSHIFT(): + def f(): + a = 16 + a >>= 2 + return a + + ans = f() + with replaced_code(f, ("BINARY_OP", ">>=")): + assert f() == ans + + +def test_INPLACE_AND(): + def f(): + a = 0b1111 + a &= 0b1010 + return a + + ans = f() + with replaced_code(f, ("BINARY_OP", "&=")): + assert f() == ans + + +def test_INPLACE_XOR(): + def f(): + a = 0b1100 + a ^= 0b1010 + return a + + ans = f() + with replaced_code(f, ("BINARY_OP", "^=")): + assert f() == ans + + +def test_INPLACE_OR(): + def f(): + a = 0b1100 + a |= 0b0011 + return a + + ans = f() + with replaced_code(f, ("BINARY_OP", "|=")): + assert f() == ans diff --git a/pkgs/MagiCompiler/tests/magi_depyf/decompile/decompile_roundtrip/test_store_delete.py b/pkgs/MagiCompiler/tests/magi_depyf/decompile/decompile_roundtrip/test_store_delete.py new file mode 100644 index 0000000000000000000000000000000000000000..a5b56db33b7464094c09d71156454a9a63d94713 --- /dev/null +++ b/pkgs/MagiCompiler/tests/magi_depyf/decompile/decompile_roundtrip/test_store_delete.py @@ -0,0 +1,76 @@ +# Copyright (c) 2026 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. + +"""Roundtrip tests for STORE/DELETE operations (subscript, attr, name).""" + +from copy import deepcopy + +from tests.magi_depyf.decompile.decompile_roundtrip.helpers import Point, data_map, replaced_code + + +def test_STORE_SUBSCR(): + def f(): + p = Point(1, 2) + p[0] = 99 + return p + + ans = f() + with replaced_code(f, "STORE_SUBSCR"): + assert f() == ans + + +def test_DELETE_SUBSCR(): + def f(): + a = deepcopy(data_map) + del a[1] + return a + + ans = f() + with replaced_code(f, "DELETE_SUBSCR"): + assert f() == ans + + +def test_STORE_ATTR(): + def f(): + p = Point(1, 2) + p.x = 10 + return p + + ans = f() + with replaced_code(f, "STORE_ATTR"): + assert f() == ans + + +def test_DELETE_ATTR(): + def f(): + p = Point(1, 2) + del p.x + p.x = 99 + return p + + ans = f() + with replaced_code(f, "DELETE_ATTR"): + assert f() == ans + + +def test_DELETE_NAME(): + def f(): + a = 1 + del a + a = 2 + return a + + ans = f() + with replaced_code(f, "DELETE_FAST"): + assert f() == ans diff --git a/pkgs/MagiCompiler/tests/magi_depyf/decompile/dynamo_roundtrip/__init__.py b/pkgs/MagiCompiler/tests/magi_depyf/decompile/dynamo_roundtrip/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..3eaa44adb1bd11bb4c8d48c6f00d7a08f292f395 --- /dev/null +++ b/pkgs/MagiCompiler/tests/magi_depyf/decompile/dynamo_roundtrip/__init__.py @@ -0,0 +1,13 @@ +# Copyright (c) 2026 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. diff --git a/pkgs/MagiCompiler/tests/magi_depyf/decompile/dynamo_roundtrip/conftest.py b/pkgs/MagiCompiler/tests/magi_depyf/decompile/dynamo_roundtrip/conftest.py new file mode 100644 index 0000000000000000000000000000000000000000..f474ad113083c8977542e7dcc34bdc24224509ca --- /dev/null +++ b/pkgs/MagiCompiler/tests/magi_depyf/decompile/dynamo_roundtrip/conftest.py @@ -0,0 +1,30 @@ +# Copyright (c) 2026 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. + +import pytest +import torch + + +@pytest.fixture(autouse=True) +def _no_grad(): + """Disable gradient computation for all dynamo-roundtrip tests. + + Using a fixture (rather than ``with torch.no_grad()`` inside the compiled + function) avoids ``BEFORE_WITH`` / ``POP_EXCEPT`` bytecodes in + Dynamo-generated resume functions, which the decompiler cannot handle. + """ + prev = torch.is_grad_enabled() + torch.set_grad_enabled(False) + yield + torch.set_grad_enabled(prev) diff --git a/pkgs/MagiCompiler/tests/magi_depyf/decompile/dynamo_roundtrip/helpers.py b/pkgs/MagiCompiler/tests/magi_depyf/decompile/dynamo_roundtrip/helpers.py new file mode 100644 index 0000000000000000000000000000000000000000..7192a135a9e069f9a1b9c8b56d88ecdf606245c5 --- /dev/null +++ b/pkgs/MagiCompiler/tests/magi_depyf/decompile/dynamo_roundtrip/helpers.py @@ -0,0 +1,157 @@ +# Copyright (c) 2026 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. + +"""Shared helpers and model definitions for dynamo-roundtrip tests. + +Test methodology: + 1. torch.compile(fn) and run to get expected output + 2. Extract Dynamo cache entry (transformed bytecode) + 3. Decompile -> recompile -> replace fn.__code__ + 4. torch.compile again and verify output matches expected +""" + +from typing import Any + +import torch +from magi_compiler.magi_depyf.decompile.recompiler import CodeRecompiler +from magi_compiler.magi_depyf.inspect.introspect import Introspector + +get_cache_entries = Introspector.get_cache_entries + + +# --------------------------------------------------------------------------- +# Library availability flags +# --------------------------------------------------------------------------- + +_has_timm = False +try: + pass + + _has_timm = True +except ImportError: + pass + +_has_transformers = False +try: + pass + + _has_transformers = True +except ImportError: + pass + +_has_diffusers = False +try: + pass + + _has_diffusers = True +except ImportError: + pass + + +# --------------------------------------------------------------------------- +# Helpers +# --------------------------------------------------------------------------- + + +def _reset(): + torch._dynamo.reset() + + +def _assert_close(actual: Any, expected: Any, atol: float = 1e-5): + """Recursively compare outputs with tolerance.""" + if isinstance(expected, torch.Tensor): + assert isinstance(actual, torch.Tensor), f"Expected tensor, got {type(actual)}" + assert actual.shape == expected.shape, f"Shape mismatch: {actual.shape} vs {expected.shape}" + diff = (actual.float().cpu() - expected.float().cpu()).abs().max().item() + assert diff < atol, f"Tensor mismatch: max diff = {diff}" + elif isinstance(expected, (tuple, list)): + assert type(actual) is type(expected) + assert len(actual) == len(expected), f"Length mismatch: {len(actual)} vs {len(expected)}" + for a, e in zip(actual, expected): + _assert_close(a, e, atol=atol) + elif isinstance(expected, dict): + for k in expected: + if k in actual: + _assert_close(actual[k], expected[k], atol=atol) + elif expected is None: + pass + elif hasattr(expected, "__dict__"): + for k, v in vars(expected).items(): + if isinstance(v, torch.Tensor) and hasattr(actual, k): + _assert_close(getattr(actual, k), v, atol=atol) + else: + assert actual == expected, f"Value mismatch: {actual} vs {expected}" + + +GLOBAL_MODULE = None +GLOBAL_INPUT_KWARGS: dict = {} +GLOBAL_OUTPUT_FN = None + + +def roundtrip_and_verify(fn_or_module, inputs, input_kwargs=None, output_fn=None, backend="eager", atol=1e-5, **compile_kw): + """Compile → decompile → recompile → replace code → re-compile and verify. + + Accepts either a plain function or an ``nn.Module``. When given a module, + it is automatically wrapped via the ``GLOBAL_MODULE`` pattern so that the + wrapper function has no free variables (avoids Dynamo closure issues). + + *input_kwargs*, if given, are forwarded as keyword arguments to the module + call (e.g. ``encoder_hidden_states``). + + *output_fn*, if given, is applied to the module output before returning + (e.g. ``lambda out: out.last_hidden_state``). This is stored as a global + to avoid introducing closure variables. + """ + global GLOBAL_MODULE, GLOBAL_INPUT_KWARGS, GLOBAL_OUTPUT_FN + + if isinstance(fn_or_module, torch.nn.Module): + fn_or_module.eval() + GLOBAL_MODULE = fn_or_module + GLOBAL_INPUT_KWARGS = input_kwargs or {} + GLOBAL_OUTPUT_FN = output_fn + + if GLOBAL_OUTPUT_FN is not None: + + def fn(*args): + return GLOBAL_OUTPUT_FN(GLOBAL_MODULE(*args, **GLOBAL_INPUT_KWARGS)) + + else: + + def fn(*args): + return GLOBAL_MODULE(*args, **GLOBAL_INPUT_KWARGS) + + else: + fn = fn_or_module + + _reset() + torch.manual_seed(42) + compiled = torch.compile(fn, backend=backend, **compile_kw) + expected = compiled(*inputs) + + entries = get_cache_entries(fn) + assert len(entries) >= 1, f"No cache entries for {fn.__code__.co_name}" + + tc = entries[0].code + recompiled = CodeRecompiler.recompile(code_to_decompile=tc, reference_code=tc) + + old_code = fn.__code__ + fn.__code__ = recompiled + try: + _reset() + torch.manual_seed(42) + compiled2 = torch.compile(fn, backend=backend, **compile_kw) + actual = compiled2(*inputs) + _assert_close(actual, expected, atol=atol) + finally: + fn.__code__ = old_code diff --git a/pkgs/MagiCompiler/tests/magi_depyf/decompile/dynamo_roundtrip/test_diffusers.py b/pkgs/MagiCompiler/tests/magi_depyf/decompile/dynamo_roundtrip/test_diffusers.py new file mode 100644 index 0000000000000000000000000000000000000000..f43689b778817b3a8f5fac9ef0369acab12e5e1f --- /dev/null +++ b/pkgs/MagiCompiler/tests/magi_depyf/decompile/dynamo_roundtrip/test_diffusers.py @@ -0,0 +1,53 @@ +# Copyright (c) 2026 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. + +"""Dynamo roundtrip: diffusers attention and transformer blocks.""" + +import pytest +import torch +from tests.magi_depyf.decompile.dynamo_roundtrip.helpers import _has_diffusers, roundtrip_and_verify + + +@pytest.mark.skipif(not _has_diffusers, reason="diffusers not installed") +class TestDiffusersComponents: + """diffusers attention and transformer blocks — directly relevant for + Stable Diffusion / DiT debugging.""" + + def test_diffusers_self_attention(self): + from diffusers.models.attention_processor import Attention + + attn = Attention(query_dim=32, heads=4, dim_head=8) + roundtrip_and_verify(attn, (torch.randn(2, 8, 32),), atol=1e-4) + + def test_diffusers_cross_attention(self): + from diffusers.models.attention_processor import Attention + + attn = Attention(query_dim=32, cross_attention_dim=48, heads=4, dim_head=8) + roundtrip_and_verify( + attn, (torch.randn(2, 8, 32),), input_kwargs={"encoder_hidden_states": torch.randn(2, 6, 48)}, atol=1e-4 + ) + + def test_diffusers_basic_transformer_block(self): + from diffusers.models.attention import BasicTransformerBlock + + block = BasicTransformerBlock(dim=32, num_attention_heads=4, attention_head_dim=8) + roundtrip_and_verify(block, (torch.randn(2, 8, 32),), atol=1e-4) + + def test_diffusers_cross_attn_transformer_block(self): + from diffusers.models.attention import BasicTransformerBlock + + block = BasicTransformerBlock(dim=32, num_attention_heads=4, attention_head_dim=8, cross_attention_dim=48) + roundtrip_and_verify( + block, (torch.randn(2, 8, 32),), input_kwargs={"encoder_hidden_states": torch.randn(2, 6, 48)}, atol=1e-4 + ) diff --git a/pkgs/MagiCompiler/tests/magi_depyf/decompile/dynamo_roundtrip/test_diffusion_blocks.py b/pkgs/MagiCompiler/tests/magi_depyf/decompile/dynamo_roundtrip/test_diffusion_blocks.py new file mode 100644 index 0000000000000000000000000000000000000000..9339bb060255a6f75910b0df4e7ea9e2c335773f --- /dev/null +++ b/pkgs/MagiCompiler/tests/magi_depyf/decompile/dynamo_roundtrip/test_diffusion_blocks.py @@ -0,0 +1,142 @@ +# Copyright (c) 2026 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. + +"""Dynamo roundtrip: diffusion-relevant building blocks (pure PyTorch).""" + +import math + +import torch +import torch.nn as nn +import torch.nn.functional as F +from tests.magi_depyf.decompile.dynamo_roundtrip.helpers import roundtrip_and_verify + + +class _GEGLU(nn.Module): + def __init__(self, dim, out_dim): + super().__init__() + self.proj = nn.Linear(dim, out_dim * 2) + + def forward(self, x): + x, gate = self.proj(x).chunk(2, dim=-1) + return x * F.gelu(gate) + + +class _RMSNorm(nn.Module): + def __init__(self, dim, eps=1e-6): + super().__init__() + self.weight = nn.Parameter(torch.ones(dim)) + self.eps = eps + + def forward(self, x): + norm = torch.rsqrt(x.pow(2).mean(-1, keepdim=True) + self.eps) + return x * norm * self.weight + + +class _SinusoidalEmbedding(nn.Module): + def __init__(self, dim): + super().__init__() + self.dim = dim + + def forward(self, t): + half = self.dim // 2 + freqs = torch.exp(-math.log(10000) * torch.arange(half, device=t.device).float() / half) + args = t[:, None].float() * freqs[None] + return torch.cat([torch.cos(args), torch.sin(args)], dim=-1) + + +class _CrossAttention(nn.Module): + def __init__(self, dim, context_dim, heads=4): + super().__init__() + self.heads = heads + self.head_dim = dim // heads + self.to_q = nn.Linear(dim, dim) + self.to_k = nn.Linear(context_dim, dim) + self.to_v = nn.Linear(context_dim, dim) + self.to_out = nn.Linear(dim, dim) + + def forward(self, x, context): + b, n, _ = x.shape + h = self.heads + q = self.to_q(x).view(b, n, h, self.head_dim).transpose(1, 2) + k = self.to_k(context).view(b, -1, h, self.head_dim).transpose(1, 2) + v = self.to_v(context).view(b, -1, h, self.head_dim).transpose(1, 2) + attn = F.scaled_dot_product_attention(q, k, v) + out = attn.transpose(1, 2).reshape(b, n, -1) + return self.to_out(out) + + +class _AdaLayerNorm(nn.Module): + def __init__(self, dim): + super().__init__() + self.norm = nn.LayerNorm(dim, elementwise_affine=False) + self.scale_shift = nn.Linear(dim, dim * 2) + + def forward(self, x, cond): + scale, shift = self.scale_shift(cond).chunk(2, dim=-1) + return self.norm(x) * (1 + scale) + shift + + +class _SimpleDiTBlock(nn.Module): + def __init__(self, dim, heads=4): + super().__init__() + self.norm1 = _AdaLayerNorm(dim) + self.attn = nn.MultiheadAttention(dim, heads, batch_first=True) + self.norm2 = _AdaLayerNorm(dim) + self.ff = nn.Sequential(nn.Linear(dim, dim * 4), nn.GELU(), nn.Linear(dim * 4, dim)) + + def forward(self, x, cond): + h = self.norm1(x, cond) + h, _ = self.attn(h, h, h) + x = x + h + h = self.norm2(x, cond) + x = x + self.ff(h) + return x + + +class _TimestepMLP(nn.Module): + def __init__(self, time_dim, out_dim): + super().__init__() + self.embed = _SinusoidalEmbedding(time_dim) + self.mlp = nn.Sequential(nn.Linear(time_dim, out_dim), nn.SiLU(), nn.Linear(out_dim, out_dim)) + + def forward(self, t): + return self.mlp(self.embed(t)) + + +class TestDiffusionBlocks: + """Diffusion-relevant building blocks — pure PyTorch, always available.""" + + def test_geglu(self): + roundtrip_and_verify(_GEGLU(32, 64), (torch.randn(2, 8, 32),)) + + def test_rms_norm(self): + roundtrip_and_verify(_RMSNorm(32), (torch.randn(2, 8, 32),)) + + def test_sinusoidal_embedding(self): + roundtrip_and_verify(_SinusoidalEmbedding(32), (torch.arange(4),)) + + def test_cross_attention(self): + model = _CrossAttention(32, context_dim=48, heads=4) + roundtrip_and_verify(model, (torch.randn(2, 8, 32), torch.randn(2, 6, 48))) + + def test_ada_layer_norm(self): + model = _AdaLayerNorm(32) + roundtrip_and_verify(model, (torch.randn(2, 8, 32), torch.randn(2, 8, 32))) + + def test_dit_block(self): + model = _SimpleDiTBlock(32, heads=4) + roundtrip_and_verify(model, (torch.randn(2, 8, 32), torch.randn(2, 8, 32))) + + def test_timestep_mlp(self): + roundtrip_and_verify(_TimestepMLP(32, 64), (torch.arange(4),)) diff --git a/pkgs/MagiCompiler/tests/magi_depyf/decompile/dynamo_roundtrip/test_forward_replacement.py b/pkgs/MagiCompiler/tests/magi_depyf/decompile/dynamo_roundtrip/test_forward_replacement.py new file mode 100644 index 0000000000000000000000000000000000000000..5c5d4dd8c05f37d4c48b27e7d258ea63357e6dee --- /dev/null +++ b/pkgs/MagiCompiler/tests/magi_depyf/decompile/dynamo_roundtrip/test_forward_replacement.py @@ -0,0 +1,113 @@ +# Copyright (c) 2026 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. + +"""Dynamo roundtrip: direct Module.forward.__code__ replacement. + +Mirrors the real magi_compiler pattern: torch.compile(module) then +replace Module.forward.__code__ with decompiled+recompiled code. + +This is the only test file that replaces ``klass.forward.__code__`` +(the production code path in ``magi_compiler_base.py``). All other +dynamo roundtrip tests go through a wrapper function instead. +""" + +import torch +import torch.nn as nn +import torch.nn.functional as F +from magi_compiler.magi_depyf.decompile.recompiler import CodeRecompiler +from tests.magi_depyf.decompile.dynamo_roundtrip.helpers import _assert_close, _reset, get_cache_entries + + +def _roundtrip_forward(module, inputs, backend="eager", atol=1e-5, **compile_kw): + """Compile module directly, decompile forward, replace klass.forward.__code__.""" + _reset() + module.eval() + torch.manual_seed(42) + compiled = torch.compile(module, backend=backend, **compile_kw) + expected = compiled(*inputs) + + entries = get_cache_entries(module.forward) + assert len(entries) >= 1, f"No cache entries for {module.__class__.__name__}.forward" + + tc = entries[0].code + recompiled = CodeRecompiler.recompile(code_to_decompile=tc, reference_code=tc) + + klass = module.__class__ + old_code = klass.forward.__code__ + klass.forward.__code__ = recompiled + try: + _reset() + torch.manual_seed(42) + compiled2 = torch.compile(module, backend=backend, **compile_kw) + actual = compiled2(*inputs) + _assert_close(actual, expected, atol=atol) + finally: + klass.forward.__code__ = old_code + + +class _AdaLayerNorm(nn.Module): + def __init__(self, dim): + super().__init__() + self.norm = nn.LayerNorm(dim, elementwise_affine=False) + self.scale_shift = nn.Linear(dim, dim * 2) + + def forward(self, x, cond): + scale, shift = self.scale_shift(cond).chunk(2, dim=-1) + return self.norm(x) * (1 + scale) + shift + + +class _SimpleDiTBlock(nn.Module): + def __init__(self, dim, heads=4): + super().__init__() + self.norm1 = _AdaLayerNorm(dim) + self.attn = nn.MultiheadAttention(dim, heads, batch_first=True) + self.norm2 = _AdaLayerNorm(dim) + self.ff = nn.Sequential(nn.Linear(dim, dim * 4), nn.GELU(), nn.Linear(dim * 4, dim)) + + def forward(self, x, cond): + h = self.norm1(x, cond) + h, _ = self.attn(h, h, h) + x = x + h + h = self.norm2(x, cond) + x = x + self.ff(h) + return x + + +class TestForwardCodeReplacement: + """Dynamo traces module.forward directly, producing bytecode with ``self`` + in ``co_varnames``, and we replace the class-level ``forward.__code__`` + — exactly like ``magi_compiler_base.py`` does. + """ + + def test_custom_forward(self): + class ResBlock(nn.Module): + def __init__(self): + super().__init__() + self.linear1 = nn.Linear(32, 32) + self.linear2 = nn.Linear(32, 32) + self.norm = nn.LayerNorm(32) + + def forward(self, x): + residual = x + x = F.relu(self.linear1(x)) + x = self.linear2(x) + return self.norm(x + residual) + + _roundtrip_forward(ResBlock(), (torch.randn(2, 8, 32),)) + + def test_multi_input_forward(self): + _roundtrip_forward(_AdaLayerNorm(32), (torch.randn(2, 8, 32), torch.randn(2, 8, 32))) + + def test_dit_block_forward(self): + _roundtrip_forward(_SimpleDiTBlock(32, heads=4), (torch.randn(2, 8, 32), torch.randn(2, 8, 32))) diff --git a/pkgs/MagiCompiler/tests/magi_depyf/decompile/dynamo_roundtrip/test_graph_break.py b/pkgs/MagiCompiler/tests/magi_depyf/decompile/dynamo_roundtrip/test_graph_break.py new file mode 100644 index 0000000000000000000000000000000000000000..329ebeb1766cb888dc07d81150c1317344d396fd --- /dev/null +++ b/pkgs/MagiCompiler/tests/magi_depyf/decompile/dynamo_roundtrip/test_graph_break.py @@ -0,0 +1,182 @@ +# Copyright (c) 2026 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. + +"""Dynamo roundtrip: graph break scenarios and resume function roundtripping. + +When Dynamo encounters an unsupported operation (e.g. ``print()``), it +performs a *graph break*: it splits the function into compiled graph segments +connected by resume functions (``__resume_at_XX``). The transformed bytecode +and the resume functions are all Dynamo-generated code objects that must +roundtrip correctly. + +Gradient computation is disabled globally by the ``_no_grad`` autouse fixture +in ``conftest.py``, rather than via ``with torch.no_grad()`` inside the +compiled function. ``with`` blocks produce ``BEFORE_WITH`` / ``POP_EXCEPT`` +exception-handling bytecode that Dynamo copies verbatim into resume function +code objects; the decompiler cannot process those opcodes. +""" + +import torch +import torch.nn as nn +from magi_compiler.magi_depyf import decompile +from magi_compiler.magi_depyf.decompile.recompiler import CodeRecompiler +from tests.magi_depyf.decompile.dynamo_roundtrip import helpers +from tests.magi_depyf.decompile.dynamo_roundtrip.helpers import _assert_close, _reset, get_cache_entries, roundtrip_and_verify + +# --------------------------------------------------------------------------- +# Helper: recursively roundtrip all code objects (main + resume functions) +# --------------------------------------------------------------------------- + + +def _collect_resume_replacements(transformed_code, fn_globals, replacements, visited): + """Walk *transformed_code*'s ``co_names`` to find ``__resume_at_*`` + functions, decompile→recompile their **original** code, and collect + ``(fn, old_code, new_code)`` tuples for later replacement. + + Recursion: each resume function's *transformed* code (from its cache + entry) may reference further resume functions, so we continue the walk. + """ + for name in transformed_code.co_names: + if not name.startswith("__resume"): + continue + resume_fn = fn_globals.get(name) + if resume_fn is None or not hasattr(resume_fn, "__code__"): + continue + if id(resume_fn) in visited: + continue + visited.add(id(resume_fn)) + + orig_code = resume_fn.__code__ + src = decompile(orig_code) + assert src, f"Empty decompilation for resume fn {orig_code.co_name}" + recompiled = CodeRecompiler.recompile(code_to_decompile=orig_code, reference_code=orig_code) + assert recompiled is not None, f"Recompile failed for {orig_code.co_name}" + replacements.append((resume_fn, orig_code, recompiled)) + + resume_entries = get_cache_entries(resume_fn) + if resume_entries: + _collect_resume_replacements(resume_entries[0].code, fn_globals, replacements, visited) + + +def _roundtrip_all_entries(fn, inputs, backend="eager", atol=1e-5, **compile_kw): + """Compile → decompile → recompile → replace code **recursively** + (top-level fn + all resume functions) → re-compile and verify output. + + Top-level fn: replaced with decompile(TRANSFORMED code from cache entry). + Resume functions: replaced with decompile(ORIGINAL code) — the + Dynamo-generated bytecode that doesn't reference one-shot compiled fns. + """ + _reset() + torch.manual_seed(42) + compiled = torch.compile(fn, backend=backend, **compile_kw) + expected = compiled(*inputs) + + entries = get_cache_entries(fn) + assert len(entries) >= 1, f"No cache entries for {fn.__code__.co_name}" + + tc = entries[0].code + top_recompiled = CodeRecompiler.recompile(code_to_decompile=tc, reference_code=tc) + + resume_replacements = [] + _collect_resume_replacements(tc, fn.__globals__, resume_replacements, visited=set()) + + old_code = fn.__code__ + fn.__code__ = top_recompiled + for resume_fn, orig, new in resume_replacements: + resume_fn.__code__ = new + try: + _reset() + torch.manual_seed(42) + compiled2 = torch.compile(fn, backend=backend, **compile_kw) + actual = compiled2(*inputs) + _assert_close(actual, expected, atol=atol) + finally: + fn.__code__ = old_code + for resume_fn, orig, _ in resume_replacements: + resume_fn.__code__ = orig + + +class TestGraphBreak: + """Functions that cause Dynamo graph breaks, producing resume functions.""" + + def test_print_graph_break(self): + """print() causes a graph break, producing __resume_at functions.""" + layer = nn.Linear(32, 16) + layer.eval() + helpers.GLOBAL_MODULE = layer + + def fn(x): + y = helpers.GLOBAL_MODULE(x) + print("shape:", y.shape) + return y * 2 + + _roundtrip_all_entries(fn, (torch.randn(2, 32),)) + + def test_multi_graph_break(self): + """Multiple graph breaks in one function — same module called twice.""" + layer = nn.Linear(32, 32) + layer.eval() + helpers.GLOBAL_MODULE = layer + + def fn(x): + y = helpers.GLOBAL_MODULE(x) + print("after first:", y.shape) + z = helpers.GLOBAL_MODULE(y) + print("after second:", z.shape) + return z + + _roundtrip_all_entries(fn, (torch.randn(2, 32),)) + + def test_explicit_graph_break(self): + """torch._dynamo.graph_break() explicit break.""" + layer = nn.Linear(16, 16) + layer.eval() + helpers.GLOBAL_MODULE = layer + + def fn(x): + y = helpers.GLOBAL_MODULE(x) + torch._dynamo.graph_break() + return y + 1 + + _roundtrip_all_entries(fn, (torch.randn(2, 16),)) + + def test_conditional_specialization(self): + """Data-independent branch — Dynamo specializes on the bool, no graph break.""" + layer = nn.Linear(16, 16) + layer.eval() + helpers.GLOBAL_MODULE = layer + + def fn(x, flag): + y = helpers.GLOBAL_MODULE(x) + if flag: + return y + 1 + return y - 1 + + roundtrip_and_verify(fn, (torch.randn(2, 16), True)) + + def test_resume_function_recursive_roundtrip(self): + """Verify ALL resume functions in a graph-break chain can be + decompiled, recompiled, and the whole tree re-executed correctly.""" + layer = nn.Linear(16, 16) + layer.eval() + helpers.GLOBAL_MODULE = layer + + def fn(x): + y = helpers.GLOBAL_MODULE(x) + print("break1") + z = y * 2 + print("break2") + return z + 1 + + _roundtrip_all_entries(fn, (torch.randn(2, 16),)) diff --git a/pkgs/MagiCompiler/tests/magi_depyf/decompile/dynamo_roundtrip/test_pytorch_modules.py b/pkgs/MagiCompiler/tests/magi_depyf/decompile/dynamo_roundtrip/test_pytorch_modules.py new file mode 100644 index 0000000000000000000000000000000000000000..bb5b6b5e193de8d914cf25747c31124bcf70174d --- /dev/null +++ b/pkgs/MagiCompiler/tests/magi_depyf/decompile/dynamo_roundtrip/test_pytorch_modules.py @@ -0,0 +1,106 @@ +# Copyright (c) 2026 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. + +"""Dynamo roundtrip: core PyTorch nn.Module tests.""" + +import pytest +import torch +import torch.nn as nn +import torch.nn.functional as F +from tests.magi_depyf.decompile.dynamo_roundtrip.helpers import roundtrip_and_verify + + +class TestPyTorchModules: + """Core PyTorch nn.Module tests — always available.""" + + def test_mlp(self): + """2-layer feedforward network.""" + model = nn.Sequential(nn.Linear(32, 64), nn.ReLU(), nn.Linear(64, 16)) + roundtrip_and_verify(model, (torch.randn(2, 32),)) + + def test_conv_bn_relu(self): + """Conv2d + BatchNorm2d + ReLU stack.""" + model = nn.Sequential( + nn.Conv2d(3, 16, 3, padding=1), nn.BatchNorm2d(16), nn.ReLU(), nn.Conv2d(16, 8, 3, padding=1), nn.BatchNorm2d(8) + ) + roundtrip_and_verify(model, (torch.randn(1, 3, 16, 16),)) + + def test_multihead_attention(self): + """nn.MultiheadAttention self-attention.""" + attn = nn.MultiheadAttention(embed_dim=32, num_heads=4, batch_first=True) + x = torch.randn(2, 8, 32) + roundtrip_and_verify(attn, (x, x, x)) + + def test_transformer_encoder_layer(self): + """nn.TransformerEncoderLayer (self-attn + FFN).""" + layer = nn.TransformerEncoderLayer(d_model=32, nhead=4, dim_feedforward=64, batch_first=True) + roundtrip_and_verify(layer, (torch.randn(2, 8, 32),)) + + @pytest.mark.skip(reason="Dynamo cannot trace LSTM (Graph Count: 0)") + def test_lstm(self): + """nn.LSTM single layer.""" + lstm = nn.LSTM(input_size=16, hidden_size=32, num_layers=1, batch_first=True) + lstm.eval() + + def fn(x): + output, (hn, cn) = lstm(x) + return output + + roundtrip_and_verify(fn, (torch.randn(2, 8, 16),)) + + @pytest.mark.skip(reason="Dynamo cannot trace GRU (Graph Count: 0)") + def test_gru(self): + """nn.GRU single layer.""" + gru = nn.GRU(input_size=16, hidden_size=32, batch_first=True) + gru.eval() + + def fn(x): + output, _ = gru(x) + return output + + roundtrip_and_verify(fn, (torch.randn(2, 8, 16),)) + + def test_embedding_linear(self): + """Embedding -> Linear (language model head pattern).""" + model = nn.Sequential(nn.Embedding(100, 32), nn.Linear(32, 100)) + roundtrip_and_verify(model, (torch.randint(0, 100, (2, 8)),)) + + def test_layernorm_gelu_linear(self): + """LayerNorm -> Linear -> GELU -> Linear.""" + model = nn.Sequential(nn.LayerNorm(32), nn.Linear(32, 64), nn.GELU(), nn.Linear(64, 32)) + roundtrip_and_verify(model, (torch.randn(2, 8, 32),)) + + def test_residual_conv_block(self): + """Residual connection with Conv2d.""" + + class ResBlock(nn.Module): + def __init__(self, ch): + super().__init__() + self.conv1 = nn.Conv2d(ch, ch, 3, padding=1) + self.bn1 = nn.BatchNorm2d(ch) + self.conv2 = nn.Conv2d(ch, ch, 3, padding=1) + self.bn2 = nn.BatchNorm2d(ch) + + def forward(self, x): + residual = x + out = F.relu(self.bn1(self.conv1(x))) + out = self.bn2(self.conv2(out)) + return F.relu(out + residual) + + roundtrip_and_verify(ResBlock(8), (torch.randn(1, 8, 16, 16),)) + + def test_grouped_conv(self): + """Depthwise separable convolution (common in MobileNet/EfficientNet).""" + model = nn.Sequential(nn.Conv2d(16, 16, 3, padding=1, groups=16), nn.Conv2d(16, 32, 1), nn.ReLU()) + roundtrip_and_verify(model, (torch.randn(1, 16, 8, 8),)) diff --git a/pkgs/MagiCompiler/tests/magi_depyf/decompile/dynamo_roundtrip/test_timm.py b/pkgs/MagiCompiler/tests/magi_depyf/decompile/dynamo_roundtrip/test_timm.py new file mode 100644 index 0000000000000000000000000000000000000000..33226aa35a24525764403a96b7aa9f94a9616664 --- /dev/null +++ b/pkgs/MagiCompiler/tests/magi_depyf/decompile/dynamo_roundtrip/test_timm.py @@ -0,0 +1,66 @@ +# Copyright (c) 2026 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. + +"""Dynamo roundtrip: timm models (diverse architectures).""" + +import pytest +import torch +from tests.magi_depyf.decompile.dynamo_roundtrip.helpers import _has_timm, roundtrip_and_verify + + +@pytest.mark.skipif(not _has_timm, reason="timm not installed") +class TestTimmModels: + """Small timm models — diverse architectures.""" + + def test_resnet18(self): + import timm + + model = timm.create_model("resnet18", pretrained=False, num_classes=10) + roundtrip_and_verify(model, (torch.randn(1, 3, 32, 32),), atol=1e-4) + + def test_mobilenetv3_small(self): + import timm + + model = timm.create_model("mobilenetv3_small_050", pretrained=False, num_classes=10) + roundtrip_and_verify(model, (torch.randn(1, 3, 64, 64),), atol=1e-4) + + def test_efficientnet_b0(self): + import timm + + model = timm.create_model("efficientnet_b0", pretrained=False, num_classes=10) + roundtrip_and_verify(model, (torch.randn(1, 3, 64, 64),), atol=1e-4) + + def test_vit_tiny(self): + import timm + + model = timm.create_model("vit_tiny_patch16_224", pretrained=False, img_size=64, num_classes=10) + roundtrip_and_verify(model, (torch.randn(1, 3, 64, 64),), atol=1e-4) + + def test_convnext_tiny(self): + import timm + + model = timm.create_model("convnext_tiny", pretrained=False, num_classes=10) + roundtrip_and_verify(model, (torch.randn(1, 3, 32, 32),), atol=1e-4) + + def test_swin_tiny(self): + import timm + + model = timm.create_model("swin_tiny_patch4_window7_224", pretrained=False, img_size=56, num_classes=10) + roundtrip_and_verify(model, (torch.randn(1, 3, 56, 56),), atol=1e-4) + + def test_deit_tiny(self): + import timm + + model = timm.create_model("deit_tiny_patch16_224", pretrained=False, img_size=64, num_classes=10) + roundtrip_and_verify(model, (torch.randn(1, 3, 64, 64),), atol=1e-4) diff --git a/pkgs/MagiCompiler/tests/magi_depyf/decompile/dynamo_roundtrip/test_transformers.py b/pkgs/MagiCompiler/tests/magi_depyf/decompile/dynamo_roundtrip/test_transformers.py new file mode 100644 index 0000000000000000000000000000000000000000..13e4fd3da8177852a58f6f60eec314db032311ff --- /dev/null +++ b/pkgs/MagiCompiler/tests/magi_depyf/decompile/dynamo_roundtrip/test_transformers.py @@ -0,0 +1,53 @@ +# Copyright (c) 2026 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. + +"""Dynamo roundtrip: HuggingFace transformers models.""" + +import pytest +import torch +from tests.magi_depyf.decompile.dynamo_roundtrip.helpers import _has_transformers, roundtrip_and_verify + + +def _get_last_hidden_state(out): + return out.last_hidden_state + + +@pytest.mark.skipif(not _has_transformers, reason="transformers not installed") +class TestTransformersModels: + """Tiny HuggingFace transformer models.""" + + def test_bert_tiny(self): + from transformers import BertConfig, BertModel + + config = BertConfig( + hidden_size=32, + num_hidden_layers=1, + num_attention_heads=4, + intermediate_size=64, + vocab_size=100, + max_position_embeddings=16, + ) + roundtrip_and_verify(BertModel(config), (torch.randint(0, 100, (1, 8)),), atol=1e-4) + + def test_gpt2_tiny(self): + from transformers import GPT2Config, GPT2Model + + config = GPT2Config(n_embd=32, n_layer=1, n_head=4, vocab_size=100, n_positions=16) + roundtrip_and_verify(GPT2Model(config), (torch.randint(0, 100, (1, 8)),), output_fn=_get_last_hidden_state, atol=1e-4) + + def test_t5_encoder_tiny(self): + from transformers import T5Config, T5EncoderModel + + config = T5Config(d_model=32, d_ff=64, num_heads=4, num_layers=1, vocab_size=100) + roundtrip_and_verify(T5EncoderModel(config), (torch.randint(0, 100, (1, 8)),), atol=1e-4) diff --git a/pkgs/MagiCompiler/tests/magi_depyf/decompile/test_postprocess.py b/pkgs/MagiCompiler/tests/magi_depyf/decompile/test_postprocess.py new file mode 100644 index 0000000000000000000000000000000000000000..e10a055759a9ee91b0dd99e39beb9e52c0b66638 --- /dev/null +++ b/pkgs/MagiCompiler/tests/magi_depyf/decompile/test_postprocess.py @@ -0,0 +1,214 @@ +# Copyright (c) 2026 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 decompile.postprocess pipeline passes.""" + +import textwrap + +from magi_compiler.magi_depyf.decompile.postprocess import ( + dedup_branch_tails, + eliminate_for_temps, + eliminate_inline_temps, + run_all, +) + +# --------------------------------------------------------------------------- +# Pass 1: for-loop temp elimination +# --------------------------------------------------------------------------- + + +class TestForTempElimination: + def test_simple_for(self): + src = textwrap.dedent( + """\ + for __temp_0 in items: + x = __temp_0 + print(x) + """ + ) + result = eliminate_for_temps(src) + assert "for x in items" in result + assert "__temp_0" not in result + + def test_nested_for(self): + src = textwrap.dedent( + """\ + for __temp_0 in outer: + y = __temp_0 + for __temp_1 in inner: + z = __temp_1 + print(z) + """ + ) + result = eliminate_for_temps(src) + assert "for y in outer" in result + assert "for z in inner" in result + assert "__temp" not in result + + def test_no_temp_for(self): + src = textwrap.dedent( + """\ + for x in items: + print(x) + """ + ) + result = eliminate_for_temps(src) + assert "for x in items" in result + + def test_multi_assign_not_eliminated(self): + """If the first statement isn't a simple assignment from the temp, keep it.""" + src = textwrap.dedent( + """\ + for __temp_0 in items: + print(__temp_0) + """ + ) + result = eliminate_for_temps(src) + assert "__temp_0" in result + + +# --------------------------------------------------------------------------- +# Pass 2: inline temp elimination +# --------------------------------------------------------------------------- + + +class TestInlineTempElimination: + def test_simple_inline(self): + src = textwrap.dedent( + """\ + __temp_0 = a + b + result = __temp_0 * 2 + """ + ) + result = eliminate_inline_temps(src) + assert "__temp_0" not in result + assert "result = (a + b) * 2" in result or "result" in result + + def test_no_inline_when_rhs_modified(self): + """Don't inline if the RHS variable is modified between def and use.""" + src = textwrap.dedent( + """\ + __temp_0 = b + b = a + x = __temp_0 + """ + ) + result = eliminate_inline_temps(src) + assert "__temp_0" in result + + +# --------------------------------------------------------------------------- +# Pass 3: branch tail deduplication +# --------------------------------------------------------------------------- + + +class TestBranchTailDedup: + def test_simple_dedup(self): + src = textwrap.dedent( + """\ + if cond: + x = 1 + return x + else: + x = 2 + return x + """ + ) + result = dedup_branch_tails(src) + lines = [l.strip() for l in result.strip().splitlines()] + assert lines.count("return x") == 1 + assert lines[-1] == "return x" + + def test_no_dedup_when_different(self): + src = textwrap.dedent( + """\ + if cond: + return 1 + else: + return 2 + """ + ) + result = dedup_branch_tails(src) + assert "return 1" in result + assert "return 2" in result + + def test_multi_statement_dedup(self): + src = textwrap.dedent( + """\ + if cond: + x = 1 + y = f(x) + return y + else: + x = 2 + y = f(x) + return y + """ + ) + result = dedup_branch_tails(src) + lines = [l.strip() for l in result.strip().splitlines()] + assert lines.count("y = f(x)") == 1 + assert lines.count("return y") == 1 + + def test_elif_chain_dedup(self): + src = textwrap.dedent( + """\ + if a: + x = 1 + return x + elif b: + x = 2 + return x + else: + x = 3 + return x + """ + ) + result = dedup_branch_tails(src) + lines = [l.strip() for l in result.strip().splitlines()] + assert lines.count("return x") == 1 + + def test_keep_one_statement_per_branch(self): + """Both branches have only one (identical) statement — keep it in branches.""" + src = textwrap.dedent( + """\ + if cond: + return 1 + else: + return 1 + """ + ) + result = dedup_branch_tails(src) + assert "return 1" in result + + +# --------------------------------------------------------------------------- +# Full pipeline +# --------------------------------------------------------------------------- + + +class TestRunAll: + def test_for_temp_then_inline(self): + """Pipeline: for-temp elimination runs before inline elimination.""" + src = textwrap.dedent( + """\ + __temp_0 = compute() + for __temp_1 in __temp_0: + x = __temp_1 + print(x) + """ + ) + result = run_all(src) + assert "for x in compute()" in result + assert "__temp" not in result diff --git a/pkgs/MagiCompiler/tests/magi_depyf/decompile/test_source_emitter.py b/pkgs/MagiCompiler/tests/magi_depyf/decompile/test_source_emitter.py new file mode 100644 index 0000000000000000000000000000000000000000..0e3f27da4150b18a612f685a546e5730c4f50676 --- /dev/null +++ b/pkgs/MagiCompiler/tests/magi_depyf/decompile/test_source_emitter.py @@ -0,0 +1,108 @@ +# Copyright (c) 2026 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. + +"""Unit tests for SourceEmitter (the core improvement over depyf).""" + +from magi_compiler.magi_depyf.decompile.bytecode.source_emitter import LoopContext, SourceEmitter + + +class TestStackOperations: + def test_push_pop(self): + em = SourceEmitter() + em.push("a") + em.push("b") + assert em.pop() == "b" + assert em.pop() == "a" + + def test_peek(self): + em = SourceEmitter() + em.push("x") + em.push("y") + assert em.peek(0) == "y" + assert em.peek(1) == "x" + assert em.stack_size == 2 + + def test_set_at(self): + em = SourceEmitter() + em.push("a") + em.push("b") + em.set_at(0, "B") + assert em.peek() == "B" + em.set_at(1, "A") + assert em.stack == ["A", "B"] + + +class TestTempCounter: + def test_instance_scoped(self): + """Verify temp counter is per-instance, not class-level.""" + em1 = SourceEmitter() + em2 = SourceEmitter() + t1 = em1.make_temp() + t2 = em2.make_temp() + assert t1 == "__temp_1" + assert t2 == "__temp_1" + + def test_replace_tos_with_temp(self): + em = SourceEmitter() + em.push("[1, 2, 3]") + name = em.replace_tos_with_temp() + assert name.startswith("__temp_") + assert em.peek() == name + assert f"{name} = [1, 2, 3]\n" in em.get_source() + + +class TestEmission: + def test_emit_appends_newline(self): + em = SourceEmitter() + em.emit("x = 1") + assert em.get_source() == "x = 1\n" + + def test_emit_raw(self): + em = SourceEmitter() + em.emit_raw("def f():\n pass\n") + assert em.get_source() == "def f():\n pass\n" + + def test_indent(self): + em = SourceEmitter(indent_size=4) + assert em.indent("a = 1\nb = 2\n") == " a = 1\n b = 2\n" + + +class TestFork: + def test_fork_shares_counter(self): + em = SourceEmitter() + em.make_temp() # __temp_1 + with em.fork() as child: + t = child.make_temp() + assert t == "__temp_2" + + def test_fork_independent_source(self): + em = SourceEmitter() + em.emit("parent line") + with em.fork(stack=["a"]) as child: + child.emit("child line") + assert "child line" in child.get_source() + assert "child line" not in em.get_source() + + def test_fork_inherits_loop(self): + em = SourceEmitter() + em.loop = LoopContext(start_index=0, end_index=10) + with em.fork() as child: + assert child.loop is not None + assert child.loop.start_index == 0 + + def test_fork_explicit_loop(self): + em = SourceEmitter() + loop = LoopContext(start_index=5, end_index=15) + with em.fork(loop=loop) as child: + assert child.loop.start_index == 5 diff --git a/pkgs/MagiCompiler/tests/magi_depyf/inspect/__init__.py b/pkgs/MagiCompiler/tests/magi_depyf/inspect/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..3eaa44adb1bd11bb4c8d48c6f00d7a08f292f395 --- /dev/null +++ b/pkgs/MagiCompiler/tests/magi_depyf/inspect/__init__.py @@ -0,0 +1,13 @@ +# Copyright (c) 2026 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. diff --git a/pkgs/MagiCompiler/tests/magi_depyf/inspect/test_example.py b/pkgs/MagiCompiler/tests/magi_depyf/inspect/test_example.py new file mode 100644 index 0000000000000000000000000000000000000000..4efa9911c012efd5f05d019277203870011aa6ff --- /dev/null +++ b/pkgs/MagiCompiler/tests/magi_depyf/inspect/test_example.py @@ -0,0 +1,217 @@ +# Copyright (c) 2026 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. + +"""End-to-end inspect pipeline test: simulated torch.compile example. + +Verifies the full pipeline by compiling a function (with and without +graph breaks) and checking: + 1. Hook is called and CaptureResult is produced + 2. fn and __resume functions exist and are decompiled + 3. Guards are obtained from CacheEntry + 4. Backend compiled_fn info is extracted +""" + +from __future__ import annotations + +import shutil +import tempfile + +import pytest + +torch = pytest.importorskip("torch") +import torch.nn as nn +from magi_compiler.magi_depyf.inspect import CaptureSession, write_function +from magi_compiler.magi_depyf.inspect.introspect import Introspector + + +def _reset(): + torch._dynamo.reset() + + +class TestSimpleFunction: + """No graph break — single CacheEntry, single compiled_fn.""" + + def setup_method(self): + _reset() + + def fn(x, y): + return x + y + + self.fn = fn + compiled = torch.compile(fn, backend="eager") + with CaptureSession() as session: + compiled(torch.randn(4), torch.randn(4)) + self.session = session + self.info = Introspector.build_function_info(fn) + + def test_hook_called(self): + assert len(self.session.results) >= 1 + r = self.session.results[0] + assert r.function_name == "fn" + assert r.fn_globals is not None + + def test_fn_decompiled(self): + assert len(self.info.entries) >= 1 + entry = self.info.entries[0] + assert entry.decompiled_source + assert "def" in entry.decompiled_source + assert "__compiled_fn" in entry.decompiled_source + + def test_no_resume(self): + entry = self.info.entries[0] + assert len(entry.resume_fns) == 0 + + def test_guards_obtained(self): + entry = self.info.entries[0] + assert entry.guard is not None + assert entry.guard.tree is not None + assert entry.guard.tree.type_name == "RootGuardManager" + assert len(entry.guard.tree.leaf_guards) > 0 + + def test_compiled_fn_obtained(self): + entry = self.info.entries[0] + assert len(entry.compiled_fns) >= 1 + cf = entry.compiled_fns[0] + assert cf.name.startswith("__compiled_fn") + assert cf.backend in ("eager", "inductor") + assert cf.readable_code is not None or cf.graph_module_code is not None + + +class TestGraphBreakFunction: + """print() causes graph break — produces resume functions.""" + + def setup_method(self): + _reset() + + def fn(x, y): + z = x + y + print("[test] z =", z.shape) + return z * 2 + + self.fn = fn + compiled = torch.compile(fn, backend="eager") + with CaptureSession() as session: + compiled(torch.randn(4), torch.randn(4)) + self.session = session + self.info = Introspector.build_function_info(fn) + + def test_hook_called_multiple_times(self): + assert len(self.session.results) >= 2, f"Graph break should produce >=2 hook events, got {len(self.session.results)}" + + def test_resume_functions_exist(self): + entry = self.info.entries[0] + assert len(entry.resume_fns) >= 1, "Graph break should produce resume functions" + + def test_resume_decompiled(self): + entry = self.info.entries[0] + for rf in entry.resume_fns: + assert rf.name.startswith("__resume") + assert len(rf.entries) >= 1 + re = rf.entries[0] + assert re.decompiled_source + assert "def" in re.decompiled_source + + def test_resume_has_compiled_fn(self): + entry = self.info.entries[0] + for rf in entry.resume_fns: + for re in rf.entries: + assert len(re.compiled_fns) >= 1, f"Resume entry for {rf.name} should have compiled_fn" + + def test_resume_has_guards(self): + entry = self.info.entries[0] + for rf in entry.resume_fns: + for re in rf.entries: + assert re.guard is not None, f"Resume entry for {rf.name} should have guard info" + + +class TestModuleWithNoGrad: + """nn.Module wrapped with torch.no_grad() — common pattern.""" + + def setup_method(self): + _reset() + + layer = nn.Linear(16, 8) + layer.eval() + + def fn(x): + with torch.no_grad(): + return layer(x) + + fn.__name__ = "linear_forward" + self.fn = fn + compiled = torch.compile(fn, backend="eager") + with CaptureSession() as session: + compiled(torch.randn(2, 16)) + self.session = session + self.info = Introspector.build_function_info(fn) + + def test_hook_called(self): + assert len(self.session.results) >= 1 + + def test_fn_decompiled(self): + entry = self.info.entries[0] + assert "def" in entry.decompiled_source + + def test_compiled_fn_has_graph_module(self): + entry = self.info.entries[0] + for cf in entry.compiled_fns: + assert cf.readable_code or cf.graph_module_code, f"compiled_fn {cf.name} should expose GraphModule code" + + +class TestWriteOutput: + """Verify the full write pipeline produces correct file structure.""" + + def test_write_simple(self): + _reset() + tmpdir = tempfile.mkdtemp(prefix="magi_test_") + try: + + def fn(x): + return x.sum() + + compiled = torch.compile(fn, backend="eager") + compiled(torch.randn(5)) + info = Introspector.build_function_info(fn) + root = write_function(info, tmpdir) + + assert (root / "overview.md").exists() + assert (root / "entry_0" / "decompiled_code.py").exists() + assert (root / "entry_0" / "guards.txt").exists() + finally: + shutil.rmtree(tmpdir, ignore_errors=True) + + def test_write_graph_break_has_resume_dirs(self): + _reset() + tmpdir = tempfile.mkdtemp(prefix="magi_test_") + try: + + def fn(x, y): + z = x + y + print("[test]", z.shape) + return z * 2 + + compiled = torch.compile(fn, backend="eager") + compiled(torch.randn(3), torch.randn(3)) + info = Introspector.build_function_info(fn) + root = write_function(info, tmpdir) + + rfns_dir = root / "entry_0" / "resume_fns" + assert rfns_dir.exists(), "Expected resume_fns/ directory" + resume_dirs = [d for d in rfns_dir.iterdir() if d.is_dir() and d.name.startswith("__resume")] + assert len(resume_dirs) >= 1, "Expected resume function subdirectories" + + for rd in resume_dirs: + assert (rd / "entry_0" / "decompiled_code.py").exists() + finally: + shutil.rmtree(tmpdir, ignore_errors=True) diff --git a/pkgs/MagiCompiler/tests/magi_depyf/inspect/test_introspect.py b/pkgs/MagiCompiler/tests/magi_depyf/inspect/test_introspect.py new file mode 100644 index 0000000000000000000000000000000000000000..9270092cb9102b82ae0c0c2653eea9ee6f595e6f --- /dev/null +++ b/pkgs/MagiCompiler/tests/magi_depyf/inspect/test_introspect.py @@ -0,0 +1,123 @@ +# Copyright (c) 2026 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. + +"""Introspector API tests: build_function_info, guard tree, writer, debug_compiled.""" + +from __future__ import annotations + +import os +import shutil +import tempfile + +import pytest + +torch = pytest.importorskip("torch") + +from magi_compiler.magi_depyf.inspect import debug_compiled, write_function +from magi_compiler.magi_depyf.inspect.introspect import Introspector + + +def _reset(): + torch._dynamo.reset() + + +class TestBuildFunctionInfo: + def test_basic(self): + _reset() + + def fn(x): + return x + 1 + + compiled = torch.compile(fn, backend="eager") + compiled(torch.randn(4)) + + info = Introspector.build_function_info(fn) + assert info.name == "fn" + assert len(info.entries) >= 1 + assert info.entries[0].decompiled_source + assert "def" in info.entries[0].decompiled_source + + def test_guard_tree(self): + _reset() + + def fn(x): + return x * 2 + + compiled = torch.compile(fn, backend="eager") + compiled(torch.randn(3)) + + info = Introspector.build_function_info(fn) + entry = info.entries[0] + assert entry.guard is not None + assert entry.guard.tree is not None + assert entry.guard.tree.type_name == "RootGuardManager" + + def test_format_output(self): + _reset() + + def fn(x): + return x + 1 + + compiled = torch.compile(fn, backend="eager") + compiled(torch.randn(3)) + + info = Introspector.build_function_info(fn) + text = info.format() + assert "fn" in text + assert "entry" in text + + +class TestWriter: + def test_write_files(self): + _reset() + tmpdir = tempfile.mkdtemp(prefix="magi_depyf_test_") + try: + + def fn(x, y): + return x + y + + compiled = torch.compile(fn, backend="eager") + compiled(torch.randn(3), torch.randn(3)) + + info = Introspector.build_function_info(fn) + root = write_function(info, tmpdir) + + assert root.exists() + assert (root / "overview.md").exists() + assert (root / "entry_0").exists() + assert (root / "entry_0" / "decompiled_code.py").exists() + assert (root / "entry_0" / "guards.txt").exists() + + overview = (root / "overview.md").read_text() + assert "fn" in overview + assert "entry\\[0\\]" in overview + finally: + shutil.rmtree(tmpdir, ignore_errors=True) + + def test_debug_compiled_convenience(self): + _reset() + tmpdir = tempfile.mkdtemp(prefix="magi_depyf_test_") + try: + + def fn(x): + return x.sum() + + compiled = torch.compile(fn, backend="eager") + compiled(torch.randn(5)) + + info = debug_compiled(fn, output_dir=tmpdir) + assert info.name == "fn" + assert os.path.exists(os.path.join(tmpdir, "fn")) + finally: + shutil.rmtree(tmpdir, ignore_errors=True) diff --git a/pkgs/MagiCompiler/tests/magi_depyf/inspect/test_magi_backend.py b/pkgs/MagiCompiler/tests/magi_depyf/inspect/test_magi_backend.py new file mode 100644 index 0000000000000000000000000000000000000000..53aa3996da2c895141b6db58030f3d24db9d8fb7 --- /dev/null +++ b/pkgs/MagiCompiler/tests/magi_depyf/inspect/test_magi_backend.py @@ -0,0 +1,319 @@ +# Copyright (c) 2026 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. + +"""Test magi_compile backend introspection: verify inductor artifact extraction. + +This test uses a minimal model compiled through MagiBackend (with Inductor) +and verifies that the introspection pipeline can: + 1. Detect the magi_compile backend + 2. Extract full-graph and split-graph info + 3. Read Inductor-generated kernel source from saved artifacts + 4. Write inductor_output.py files for each compiled subgraph +""" + +from __future__ import annotations + +import shutil +import sys +import tempfile +from pathlib import Path +from unittest.mock import patch + +import pytest + +torch = pytest.importorskip("torch") + +pytestmark = pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required for Inductor compilation") + + +def _make_config(cache_root_dir: str, cudagraph_mode: str = "NONE"): + """Create a minimal CompileConfig suitable for testing.""" + from magi_compiler.config import CompileConfig, CompileMode, CudaGraphMode + + mode = CudaGraphMode[cudagraph_mode] + with patch.object(sys, "argv", ["test"]): + return CompileConfig( + compile_mode=CompileMode.MAGI_COMPILE, + backend="inductor", + cache_root_dir=cache_root_dir, + splitting_ops=[], + cudagraph_mode=mode, + compile_sizes=[], + ) + + +def _cudagraph_passthrough(self, func, *args, layer_number=None, **kwargs): + """Replacement for CudaGraphMgr.run that skips capture/replay.""" + return func(*args, **kwargs) + + +def _compile_simple_model(tmpdir: str, cudagraph_mode: str = "NONE"): + """Compile a simple model via MagiBackend and return the original function. + + The model is: Linear(32→64) → ReLU → Linear(64→16). + + For cudagraph modes (PIECEWISE/FULL), we mock ``CudaGraphMgr.run`` to + skip the actual CUDA graph capture/replay while preserving the wrapping + structure — this is sufficient to verify introspection correctness. + """ + from magi_compiler.cuda_graph_mgr import CudaGraphMgr + from magi_compiler.magi_backend import MagiBackend + + w1 = torch.randn(64, 32, device="cuda") + b1 = torch.randn(64, device="cuda") + w2 = torch.randn(16, 64, device="cuda") + b2 = torch.randn(16, device="cuda") + + def fn(x): + h = torch.nn.functional.linear(x, w1, b1) + h = torch.nn.functional.relu(h) + return torch.nn.functional.linear(h, w2, b2) + + config = _make_config(tmpdir, cudagraph_mode=cudagraph_mode) + backend = MagiBackend(config) + compiled = torch.compile(fn, backend=backend) + + ctx = patch.object(CudaGraphMgr, "run", _cudagraph_passthrough) if cudagraph_mode != "NONE" else _nullcontext() + with torch.no_grad(), ctx: + compiled(torch.randn(4, 32, device="cuda")) + + return fn + + +class _nullcontext: + """Minimal no-op context manager (avoid importing contextlib).""" + + def __enter__(self): + return self + + def __exit__(self, *exc): + return False + + +def _assert_magi_backend_detected(fn, expected_cudagraph_mode: str = "NONE"): + from magi_compiler.magi_depyf.inspect.introspect import Introspector + + info = Introspector.build_function_info(fn) + assert len(info.entries) >= 1 + entry = info.entries[0] + magi_fns = [cf for cf in entry.compiled_fns if cf.backend == "magi_compile"] + assert len(magi_fns) >= 1, ( + f"Should detect magi_compile backend, got backends: " f"{[cf.backend for cf in entry.compiled_fns]}" + ) + cf = magi_fns[0] + assert ( + cf.cudagraph_mode == expected_cudagraph_mode + ), f"Expected cudagraph_mode={expected_cudagraph_mode}, got {cf.cudagraph_mode}" + return cf + + +def _assert_inductor_source(cf): + compiled_sgs = [sg for sg in cf.subgraph_infos if not sg.is_splitting_graph] + assert len(compiled_sgs) > 0, "Should have compiled (non-splitting) subgraphs" + for sg in compiled_sgs: + assert sg.inductor_code is not None, f"Subgraph {sg.name} should have inductor source code" + assert len(sg.inductor_code) > 100, f"Subgraph {sg.name} inductor code seems too short ({len(sg.inductor_code)} chars)" + + +class TestMagiBackendInductorSource: + """Verify Inductor kernel source extraction — no cudagraph.""" + + def setup_method(self): + torch._dynamo.reset() + self.tmpdir = tempfile.mkdtemp(prefix="magi_depyf_test_magi_") + + def teardown_method(self): + shutil.rmtree(self.tmpdir, ignore_errors=True) + torch._dynamo.reset() + + def test_detects_magi_compile_backend(self): + fn = _compile_simple_model(self.tmpdir) + _assert_magi_backend_detected(fn) + + def test_has_full_graph_and_split_graph(self): + fn = _compile_simple_model(self.tmpdir) + cf = _assert_magi_backend_detected(fn) + assert cf.readable_code is not None, "Should have full graph readable code" + assert len(cf.subgraph_infos) > 0, "Should have subgraph infos" + + def test_inductor_source_extracted(self): + fn = _compile_simple_model(self.tmpdir) + cf = _assert_magi_backend_detected(fn) + _assert_inductor_source(cf) + + def test_write_inductor_output_files(self): + from magi_compiler.magi_depyf.inspect import write_function + from magi_compiler.magi_depyf.inspect.introspect import Introspector + + fn = _compile_simple_model(self.tmpdir) + info = Introspector.build_function_info(fn) + output_dir = Path(self.tmpdir) / "write_output" + root = write_function(info, output_dir) + + assert (root / "overview.md").exists() + + inductor_files = list(root.rglob("inductor_output.py")) + assert len(inductor_files) > 0, f"Should generate inductor_output.py files. " f"Files found: {list(root.rglob('*'))}" + for f in inductor_files: + content = f.read_text() + assert len(content) > 100, f"{f} seems too short" + + +class TestMagiBackendPiecewiseCudaGraph: + """Verify introspection works with PIECEWISE cudagraph mode. + + In PIECEWISE mode, each PiecewiseBackend submodule is wrapped by + gen_wrap_func_for_cudagraph. The wrapper copies __dict__ from + PiecewiseBackend, so attribute-based detection should still work. + """ + + def setup_method(self): + torch._dynamo.reset() + self.tmpdir = tempfile.mkdtemp(prefix="magi_depyf_test_pw_cg_") + + def teardown_method(self): + shutil.rmtree(self.tmpdir, ignore_errors=True) + torch._dynamo.reset() + + def test_detects_backend(self): + fn = _compile_simple_model(self.tmpdir, cudagraph_mode="PIECEWISE") + _assert_magi_backend_detected(fn, expected_cudagraph_mode="PIECEWISE") + + def test_has_subgraph_info(self): + fn = _compile_simple_model(self.tmpdir, cudagraph_mode="PIECEWISE") + cf = _assert_magi_backend_detected(fn, expected_cudagraph_mode="PIECEWISE") + assert cf.readable_code is not None + assert len(cf.subgraph_infos) > 0 + + def test_inductor_source_extracted(self): + fn = _compile_simple_model(self.tmpdir, cudagraph_mode="PIECEWISE") + cf = _assert_magi_backend_detected(fn, expected_cudagraph_mode="PIECEWISE") + _assert_inductor_source(cf) + + +class TestMagiBackendFullCudaGraph: + """Verify introspection works with FULL cudagraph mode. + + In FULL mode, the entire split_gm is wrapped by + gen_wrap_func_for_cudagraph, so MSF.optimized_call is a function + rather than a GraphModule. The introspector must unwrap the closure + chain to find the actual GraphModule. + """ + + def setup_method(self): + torch._dynamo.reset() + self.tmpdir = tempfile.mkdtemp(prefix="magi_depyf_test_full_cg_") + + def teardown_method(self): + shutil.rmtree(self.tmpdir, ignore_errors=True) + torch._dynamo.reset() + + def test_detects_backend(self): + fn = _compile_simple_model(self.tmpdir, cudagraph_mode="FULL") + _assert_magi_backend_detected(fn, expected_cudagraph_mode="FULL") + + def test_has_subgraph_info(self): + fn = _compile_simple_model(self.tmpdir, cudagraph_mode="FULL") + cf = _assert_magi_backend_detected(fn, expected_cudagraph_mode="FULL") + assert cf.readable_code is not None + assert len(cf.subgraph_infos) > 0 + + def test_inductor_source_extracted(self): + fn = _compile_simple_model(self.tmpdir, cudagraph_mode="FULL") + cf = _assert_magi_backend_detected(fn, expected_cudagraph_mode="FULL") + _assert_inductor_source(cf) + + +class TestDumpSrcEndToEnd: + """Integration test: dump_src context-manager across all cudagraph modes. + + Verifies the full pipeline (CaptureSession → Introspector → writer) produces + consistent output regardless of cudagraph wrapping. + """ + + def setup_method(self): + torch._dynamo.reset() + self.tmpdir = tempfile.mkdtemp(prefix="magi_depyf_dump_src_") + self.cache_dir = tempfile.mkdtemp(prefix="magi_depyf_dump_cache_") + + def teardown_method(self): + shutil.rmtree(self.tmpdir, ignore_errors=True) + shutil.rmtree(self.cache_dir, ignore_errors=True) + torch._dynamo.reset() + + def _run_dump_src(self, cudagraph_mode: str): + from magi_compiler.cuda_graph_mgr import CudaGraphMgr + from magi_compiler.magi_backend import MagiBackend + from magi_compiler.magi_depyf.inspect import dump_src + + w = torch.randn(16, 8, device="cuda") + b = torch.randn(16, device="cuda") + + def fn(x): + return torch.nn.functional.relu(torch.nn.functional.linear(x, w, b)) + + cache = tempfile.mkdtemp(dir=self.cache_dir) + config = _make_config(cache, cudagraph_mode=cudagraph_mode) + backend = MagiBackend(config) + compiled = torch.compile(fn, backend=backend) + + output_dir = Path(self.tmpdir) / cudagraph_mode + ctx = patch.object(CudaGraphMgr, "run", _cudagraph_passthrough) if cudagraph_mode != "NONE" else _nullcontext() + with torch.no_grad(), ctx: + with dump_src(str(output_dir)): + compiled(torch.randn(2, 8, device="cuda")) + + return output_dir + + def _assert_output(self, output_dir: Path): + fn_dir = list(output_dir.iterdir()) + assert len(fn_dir) == 1, f"Expected one function dir, got {fn_dir}" + root = fn_dir[0] + assert (root / "overview.md").exists() + assert (root / "decompiled_code.py").exists() + assert (root / "entry_0" / "decompiled_code.py").exists() + + inductor_files = list(root.rglob("inductor_output.py")) + assert len(inductor_files) > 0, f"Should have inductor_output.py, files: {list(root.rglob('*'))}" + for f in inductor_files: + assert f.stat().st_size > 100 + + def test_dump_src_none(self): + out = self._run_dump_src("NONE") + self._assert_output(out) + + def test_dump_src_piecewise(self): + out = self._run_dump_src("PIECEWISE") + self._assert_output(out) + + def test_dump_src_full(self): + out = self._run_dump_src("FULL") + self._assert_output(out) + + def test_structure_identical_across_modes(self): + """All three modes should produce the same file tree (ignoring hashed names).""" + import re + + trees = {} + for mode in ("NONE", "PIECEWISE", "FULL"): + out = self._run_dump_src(mode) + torch._dynamo.reset() + files = sorted(str(p.relative_to(out)) for p in out.rglob("*") if p.is_file()) + normalized = [re.sub(r"__compiled_fn_\d+_[0-9a-f_]+", "COMPILED_FN", f) for f in files] + trees[mode] = normalized + + assert ( + trees["NONE"] == trees["PIECEWISE"] + ), f"NONE vs PIECEWISE structure differs:\n{trees['NONE']}\nvs\n{trees['PIECEWISE']}" + assert trees["NONE"] == trees["FULL"], f"NONE vs FULL structure differs:\n{trees['NONE']}\nvs\n{trees['FULL']}" diff --git a/pkgs/MagiCompiler/tests/magi_depyf/inspect/test_session.py b/pkgs/MagiCompiler/tests/magi_depyf/inspect/test_session.py new file mode 100644 index 0000000000000000000000000000000000000000..5043f6dad1799f52675868661530ca7b4eaa5144 --- /dev/null +++ b/pkgs/MagiCompiler/tests/magi_depyf/inspect/test_session.py @@ -0,0 +1,167 @@ +# Copyright (c) 2026 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. + +"""CaptureSession and CaptureResult tests: dataclass fields, lifecycle, basic capture.""" + +from __future__ import annotations + +import time +import types + +import pytest +from magi_compiler.magi_depyf.inspect.result import CaptureResult +from magi_compiler.magi_depyf.inspect.session import CaptureSession + + +def _make_dummy_code(name="test_fn"): + src = f"def {name}(x): return x + 1" + code = compile(src, "", "exec") + for c in code.co_consts: + if isinstance(c, types.CodeType) and c.co_name == name: + return c + raise RuntimeError(f"No code object named {name}") + + +# ── CaptureResult ───────────────────────────────────────────────────────── + + +class TestCaptureResult: + def test_fields_populated(self): + code = _make_dummy_code() + r = CaptureResult( + function_name="my_func", + original_code=code, + dynamo_code=code, + decompiled_source="def my_func(x):\n return x + 1\n", + guards=["guard1", "guard2"], + graph_source="graph code here", + ) + assert r.function_name == "my_func" + assert r.original_code is code + assert r.dynamo_code is code + assert "my_func" in r.decompiled_source + assert r.guards == ["guard1", "guard2"] + assert r.graph_source == "graph code here" + assert isinstance(r.timestamp, float) + assert r.timestamp > 0 + + def test_summary(self): + code = _make_dummy_code("sample") + r = CaptureResult( + function_name="sample", + original_code=code, + dynamo_code=code, + decompiled_source="...", + guards=["g1", "g2", "g3"], + graph_source="some graph", + ) + s = r.summary() + assert "sample" in s + assert "guards=3" in s + assert "graph=yes" in s + + def test_defaults(self): + code = _make_dummy_code() + r = CaptureResult(function_name="fn", original_code=code, dynamo_code=code, decompiled_source="...") + assert "graph=no" in r.summary() + assert r.guards == [] + assert r.fn_globals is None + + def test_timestamp_auto(self): + before = time.time() + code = _make_dummy_code() + r = CaptureResult(function_name="fn", original_code=code, dynamo_code=code, decompiled_source="...") + after = time.time() + assert before <= r.timestamp <= after + + +# ── CaptureSession lifecycle (no torch) ─────────────────────────────────── + + +class TestCaptureSessionLifecycle: + def test_init_state(self): + s = CaptureSession() + assert s.results == [] + assert s._hook_handle is None + + def test_clear(self): + s = CaptureSession() + s._results.append("fake") + assert len(s.results) == 1 + s.clear() + assert s.results == [] + + def test_results_returns_copy(self): + s = CaptureSession() + s._results.append("item") + r = s.results + r.append("should not affect internal") + assert len(s._results) == 1 + + +# ── CaptureSession with torch ──────────────────────────────────────────── + +torch_available = False +try: + import torch + + torch_available = True +except ImportError: + pass + + +@pytest.mark.skipif(not torch_available, reason="torch not installed") +class TestCaptureSessionWithTorch: + def test_capture_simple_compile(self): + """Compile a simple function and verify CaptureResult contents.""" + + def fn(x): + return x + 1 + + torch._dynamo.reset() + with CaptureSession() as session: + compiled = torch.compile(fn, backend="eager") + compiled(torch.tensor([1.0, 2.0, 3.0])) + + assert len(session.results) >= 1 + r = session.results[0] + assert isinstance(r, CaptureResult) + assert isinstance(r.function_name, str) + assert isinstance(r.original_code, types.CodeType) + assert isinstance(r.dynamo_code, types.CodeType) + assert isinstance(r.decompiled_source, str) + assert "def" in r.decompiled_source or "Failed" in r.decompiled_source + + def test_hook_removed_after_exit(self): + session = CaptureSession() + session.__enter__() + assert session._hook_handle is not None + session.__exit__(None, None, None) + assert session._hook_handle is None + + def test_multiple_captures(self): + def fn_a(x): + return x * 2 + + def fn_b(x): + return x - 1 + + torch._dynamo.reset() + with CaptureSession() as session: + ca = torch.compile(fn_a, backend="eager") + cb = torch.compile(fn_b, backend="eager") + ca(torch.tensor([1.0])) + cb(torch.tensor([2.0])) + + assert len(session.results) >= 2 diff --git a/pkgs/MagiCompiler/tests/model_definition.py b/pkgs/MagiCompiler/tests/model_definition.py new file mode 100644 index 0000000000000000000000000000000000000000..67f3d82f70b9b30f977070f08c012945f5112933 --- /dev/null +++ b/pkgs/MagiCompiler/tests/model_definition.py @@ -0,0 +1,113 @@ +# 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. + +from dataclasses import dataclass + +import torch +import torch.nn as nn +import torch.nn.functional as F +from magi_compiler import magi_compile + + +@dataclass +class MLPConfig: + """Configuration for the MLP module""" + + hidden_size: int + intermediate_size: int + params_dtype: torch.dtype = torch.bfloat16 + + +class RMSNorm(nn.Module): + """Simple RMSNorm implementation""" + + def __init__(self, dim: int, eps: float = 1e-6): + super().__init__() + self.dim = dim + self.eps = eps + self.weight = nn.Parameter(torch.ones(dim)) + + def forward(self, x: torch.Tensor) -> torch.Tensor: + input_dtype = x.dtype + variance = x.to(torch.float32).pow(2).mean(dim=-1, keepdim=True) + x = x * torch.rsqrt(variance + self.eps) + x = x.to(self.weight.dtype) * self.weight + return x.to(input_dtype) + + +@magi_compile(dynamic_arg_dims={"x": 0}) +class MLP(torch.nn.Module): + """MLP module with traditional architecture (up-projection, activation, and down-projection)""" + + config: MLPConfig + + def __init__(self, config: MLPConfig): + super().__init__() + self.config = config + self.pre_norm = RMSNorm(config.hidden_size) + self.up_proj = nn.Linear(config.hidden_size, config.intermediate_size, bias=False, dtype=config.params_dtype) + self.down_proj = nn.Linear(config.intermediate_size, config.hidden_size, bias=False, dtype=config.params_dtype) + + def forward(self, x: torch.Tensor) -> torch.Tensor: + """Forward pass of the MLP module. + + Args: + x (torch.Tensor): Input tensor + + Returns: + output (torch.Tensor): Output tensor + + Shape: + - x: (num_tokens, hidden_size) + - output: (num_tokens, hidden_size) + """ + # Pre-normalization + x = self.pre_norm(x).to(torch.bfloat16) + # Up-projection + x = self.up_proj(x).to(torch.float32) + # Activation (SiLU) + x = F.silu(x).to(torch.bfloat16) + # Down-projection + x = self.down_proj(x).to(torch.float32) + return x + + +def create_mlp_model(config: MLPConfig, device: torch.device) -> MLP: + """Create MLP model + + Args: + config: MLP configuration + device: Target device + + Returns: + model: Created MLP model + """ + model = MLP(config).to(device) + return model + + +def create_mlp_model_with_initial_params(config: MLPConfig, device: torch.device) -> tuple[MLP, list[torch.Tensor]]: + """Create MLP model and return model with initial parameter snapshot + + Args: + config: MLP configuration + device: Target device + + Returns: + model: Created MLP model + initial_params: Initial snapshot of model parameters for verifying parameter updates + """ + model = MLP(config).to(device) + initial_params = [p.clone().detach() for p in model.parameters()] + return model, initial_params diff --git a/pkgs/MagiCompiler/tests/test_capture_scalar.py b/pkgs/MagiCompiler/tests/test_capture_scalar.py new file mode 100644 index 0000000000000000000000000000000000000000..c69171cae7156dea5348562a32a39f15305b842d --- /dev/null +++ b/pkgs/MagiCompiler/tests/test_capture_scalar.py @@ -0,0 +1,75 @@ +# Copyright (c) 2026 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. + + +# NOTE: Currentlly, magi compiler's dynamo config is independent from the outer dynamo config. + + +import pytest +import torch +from magi_compiler import magi_compile +from magi_compiler.api import _DEFAULT_DYNAMO_CONFIG + + +@pytest.fixture(autouse=True) +def _enable_capture_scalar_outputs(): + """Enable capture_scalar_outputs for magi_compile's dynamo config in this module.""" + old_value = _DEFAULT_DYNAMO_CONFIG["capture_scalar_outputs"] + _DEFAULT_DYNAMO_CONFIG["capture_scalar_outputs"] = True + yield + _DEFAULT_DYNAMO_CONFIG["capture_scalar_outputs"] = old_value + + +SEQ_LEN = 512 +HIDDEN_SIZE = 1024 + + +@magi_compile(dynamic_arg_dims={"x": 0}) +class MulItemModel(torch.nn.Module): + def __init__(self): + super().__init__() + self.linear = torch.nn.Linear(HIDDEN_SIZE, HIDDEN_SIZE, bias=False, dtype=torch.bfloat16) + + @torch.no_grad() + def forward(self, x: torch.Tensor, y: torch.Tensor) -> torch.Tensor: + return self.linear(x) * y[0].item() + + +def test_mul_item_model_torch_compile(_enable_capture_scalar_outputs): + from unittest.mock import patch + + from magi_compiler.config import CompileMode, get_compile_config + + with patch.object(get_compile_config(), "compile_mode", CompileMode.TORCH_COMPILE): + model = MulItemModel() + x = torch.randn(SEQ_LEN, HIDDEN_SIZE, dtype=torch.bfloat16) + y = torch.randn(HIDDEN_SIZE, dtype=torch.bfloat16) + output = model(x, y) + assert output.shape == (SEQ_LEN, HIDDEN_SIZE) + + +# FIXME: Support item() with MAGI_COMPILE +def test_mul_item_model_magi_compile(_enable_capture_scalar_outputs): + try: + model = MulItemModel() + x = torch.randn(SEQ_LEN, HIDDEN_SIZE, dtype=torch.bfloat16) + y = torch.randn(HIDDEN_SIZE, dtype=torch.bfloat16) + output = model(x, y) + assert output.shape == (SEQ_LEN, HIDDEN_SIZE) + except Exception as e: + print(f"Error: {e}") + + +if __name__ == "__main__": + pytest.main([__file__, "-v"]) diff --git a/pkgs/MagiCompiler/tests/test_cpu_offload.py b/pkgs/MagiCompiler/tests/test_cpu_offload.py new file mode 100644 index 0000000000000000000000000000000000000000..3fda1ef852b516921d7cb80797e7bba444389be8 --- /dev/null +++ b/pkgs/MagiCompiler/tests/test_cpu_offload.py @@ -0,0 +1,234 @@ +# 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. + +from contextlib import contextmanager +from typing import Type + +import pytest +import torch +import torch.nn as nn +import torch.nn.functional as F +from magi_compiler import magi_compile +from magi_compiler.config import OffloadPolicy, get_compile_config + +from .model_definition import MLPConfig, RMSNorm + + +class TransformerWrapper(nn.Module): + """ + A wrapper class simulating a Transformer Block. + Accepts mlp_cls to support injecting dynamically defined classes. + """ + + def __init__(self, config: MLPConfig, mlp_cls: Type[nn.Module]): + super().__init__() + # Standard layer (should move to GPU) + self.attention_proj = nn.Linear(config.hidden_size, config.hidden_size, dtype=config.params_dtype) + + # Compiled layer (should stay on CPU if offload is enabled) + self.mlp = mlp_cls(config) + + def forward(self, x): + x = self.mlp(x) + x = my_attention(x, x, x) + x = self.attention_proj(x) + return x + + +@contextmanager +def set_cpu_offload(enable: bool, offload_policy: OffloadPolicy = OffloadPolicy.COST_EFFECTIVE): + """ + Context manager to temporarily override the cpu_offload setting in global config. + """ + config = get_compile_config() + original_value = config.offload_config.model_cpu_offload + config.offload_config.model_cpu_offload = enable + original_offload_policy = config.offload_config.offload_policy + config.offload_config.offload_policy = offload_policy + try: + yield + finally: + config.offload_config.model_cpu_offload = original_value + config.offload_config.offload_policy = original_offload_policy + + +def create_offload_mlp_class(): + """ + Create MLP class at runtime so that @magi_compile decorator captures the *current* config state. + + This is necessary because the decorator runs at class definition time. + By defining the class inside a function called within `set_cpu_offload(True)` context, + we ensure the decorator sees `model_cpu_offload=True`. + """ + + @magi_compile(dynamic_arg_dims={"x": 0}) + class OffloadMLP(torch.nn.Module): + config: MLPConfig + + def __init__(self, config: MLPConfig): + super().__init__() + self.config = config + self.pre_norm = RMSNorm(config.hidden_size) + self.up_proj = nn.Linear(config.hidden_size, config.intermediate_size, bias=False, dtype=config.params_dtype) + self.down_proj = nn.Linear(config.intermediate_size, config.hidden_size, bias=False, dtype=config.params_dtype) + + def forward(self, x: torch.Tensor) -> torch.Tensor: + x = self.pre_norm(x).to(torch.bfloat16) + x = self.up_proj(x).to(torch.float32) + x = F.silu(x).to(torch.bfloat16) + x = self.down_proj(x) + return x + + return OffloadMLP + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="Requires CUDA support") +def test_cpu_offload_placement(device, mlp_config): + """ + Test that the decorated module stays on CPU when .cuda() is called on parent, + while other modules move correctly. + """ + # Use the context manager to enable CPU offload + with set_cpu_offload(True): + # 1. Initialize the parent model + OffloadMLP = create_offload_mlp_class() + + model = TransformerWrapper(mlp_config, mlp_cls=OffloadMLP) + + # Verify initial state (everything on CPU by default in PyTorch) + assert model.attention_proj.weight.device.type == "cpu" + assert model.mlp.up_proj.weight.device.type == "cpu" + + # 2. Move the model to GPU + # This triggers the _apply hook in _magi_compile + model.cuda() + + # 3. Verify devices + # The standard layer should be on GPU + assert model.attention_proj.weight.device.type == "cuda", "Standard layers should move to CUDA" + + # The compiled/offloaded layer should stay on CPU + assert ( + model.mlp.up_proj.weight.device.type == "cpu" + ), "Compiled MLP layer should remain on CPU due to offload configuration" + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="Requires CUDA support") +def test_cpu_offload_manual_move(device, mlp_config): + """ + Test that the offload hook only blocks the move ONCE. + Subsequent calls to .to(device) on the specific module should allow movement. + """ + with set_cpu_offload(True): + OffloadMLP = create_offload_mlp_class() + + model = TransformerWrapper(mlp_config, mlp_cls=OffloadMLP) + + # 1. First move (Should trigger offload logic) + model.cuda() + assert model.mlp.up_proj.weight.device.type == "cpu" + assert model.attention_proj.weight.device.type == "cuda" + + # 2. Check if the internal flag is set (optional debugging check) + # Note: This relies on the implementation detail _magi_offloaded_once + if hasattr(model.mlp, "_magi_offloaded_once"): + assert model.mlp._magi_offloaded_once is True + + # 3. Second move (Should bypass hook and actually move to GPU) + # Manually force the submodule to GPU + model.mlp.to(device) + + assert model.mlp.up_proj.weight.device.type == "cuda", "Subsequent .to() calls should allow moving the module to GPU" + + +@torch.library.custom_op("athena::my_attention", mutates_args=()) +def my_attention(q: torch.Tensor, k: torch.Tensor, v: torch.Tensor) -> torch.Tensor: + return q + k + v + + +@my_attention.register_fake +def _(q: torch.Tensor, k: torch.Tensor, v: torch.Tensor) -> torch.Tensor: + return torch.empty_like(q) + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="Requires CUDA support") +def test_cpu_offload_inference(device, mlp_config): + """ + Test that the offload hook only blocks the move ONCE. + Subsequent calls to .to(device) on the specific module should allow movement. + """ + + test_shapes = [ + (32, mlp_config.hidden_size), # Small batch + (128, mlp_config.hidden_size), # Medium batch + (512, mlp_config.hidden_size), # Large batch + # NOTE: compiler will specialize for single token, so we move it to the last + (1, mlp_config.hidden_size), # Single token + ] + with set_cpu_offload(True): + get_compile_config().splitting_ops.extend(["athena::my_attention"]) + + OffloadMLP = create_offload_mlp_class() + + model = TransformerWrapper(mlp_config, mlp_cls=OffloadMLP) + + # 1. First move (Should trigger offload logic) + model.cuda() + assert model.mlp.up_proj.weight.device.type == "cpu" + assert model.attention_proj.weight.device.type == "cuda" + + with torch.no_grad(): + for num_tokens, hidden_size in test_shapes: + input_tensor = torch.randn(num_tokens, hidden_size, device=device, dtype=mlp_config.params_dtype) + output = model(input_tensor) + + assert output.shape == ( + num_tokens, + hidden_size, + ), f"For input shape ({num_tokens}, {hidden_size}), output shape should be ({num_tokens}, {hidden_size}), but got {output.shape}" + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="Requires CUDA support") +def test_cpu_offload_heuristic(device, mlp_config): + """ + Test that the heuristic scheduler is working correctly. + """ + test_shapes = [ + (32, mlp_config.hidden_size), # Small batch + (128, mlp_config.hidden_size), # Medium batch + (512, mlp_config.hidden_size), # Large batch + # NOTE: compiler will specialize for single token, so we move it to the last + (1, mlp_config.hidden_size), # Single token + ] + with set_cpu_offload(True, OffloadPolicy.HEURISTIC): + get_compile_config().splitting_ops.extend(["athena::my_attention"]) + OffloadMLP = create_offload_mlp_class() + model = TransformerWrapper(mlp_config, mlp_cls=OffloadMLP) + model.cuda() + assert model.mlp.up_proj.weight.device.type == "cpu" + assert model.attention_proj.weight.device.type == "cuda" + + with torch.no_grad(): + for num_tokens, hidden_size in test_shapes: + input_tensor = torch.randn(num_tokens, hidden_size, device=device, dtype=mlp_config.params_dtype) + output = model(input_tensor) + + assert output.shape == ( + num_tokens, + hidden_size, + ), f"For input shape ({num_tokens}, {hidden_size}), output shape should be ({num_tokens}, {hidden_size}), but got {output.shape}" + + +if __name__ == "__main__": + pytest.main([__file__, "-v"]) diff --git a/pkgs/MagiCompiler/tests/test_cuda_graph.py b/pkgs/MagiCompiler/tests/test_cuda_graph.py new file mode 100644 index 0000000000000000000000000000000000000000..a63850debfa809d383ca0a4752a1218bb45cb52c --- /dev/null +++ b/pkgs/MagiCompiler/tests/test_cuda_graph.py @@ -0,0 +1,298 @@ +# 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 magi_compiler.api import magi_compile +from magi_compiler.config import CompileConfig, CompileMode, CudaGraphMode +from magi_compiler.cuda_graph_mgr import CudaGraphMgr +from torch.testing import assert_close + +pytestmark = pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA not available") + + +def _register_custom_split_op(): + """Register a custom split op for piecewise graph testing.""" + try: + torch.ops.athena.my_split_op + return + except AttributeError: + pass + + @torch.library.custom_op("athena::my_split_op", mutates_args=()) + def my_split_op(x: torch.Tensor) -> torch.Tensor: + return x.clone() + + @my_split_op.register_fake + def _(x: torch.Tensor) -> torch.Tensor: + return torch.empty_like(x) + + +@pytest.fixture(autouse=True) +def cuda_graph_test_env(tmp_path): + """Fixture to set up a clean CUDA graph test environment.""" + _register_custom_split_op() + + # Create an isolated CudaGraphMgr instance for testing + test_mgr = CudaGraphMgr() + test_mgr.cache = dict() + + with patch("magi_compiler.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, splitting_ops: list = None) -> CompileConfig: + """Create a compile configuration for CUDA graph testing.""" + return CompileConfig( + compile_mode=CompileMode.MAGI_COMPILE, + backend="inductor", + cudagraph_mode=cudagraph_mode, + cudagraph_copy_inputs=True, + splitting_ops=splitting_ops or [], + cache_root_dir=cache_dir, + dynamic_sources="", + traced_files=set(), + ) + + +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.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, splitting_ops=["athena::my_split_op"]) + + 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.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.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.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"]) diff --git a/pkgs/MagiCompiler/tests/test_magi_compile.py b/pkgs/MagiCompiler/tests/test_magi_compile.py new file mode 100644 index 0000000000000000000000000000000000000000..2e8d398a16e7d65d4471d4dc4274a4f2c41b9649 --- /dev/null +++ b/pkgs/MagiCompiler/tests/test_magi_compile.py @@ -0,0 +1,258 @@ +# 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 @magi_compile decorator functionality. + +This module tests: +- Automatic inference of dynamic_arg_dims +- Negative index handling in dynamic_arg_dims +- Compilation correctness (compiled output matches native forward) +- Nested function calls within compiled models +- Multiple output compilation +- Source code change detection and recompilation +""" + +import tempfile +from typing import Tuple +from unittest.mock import MagicMock, patch + +import pytest +import torch +from magi_compiler.api import magi_compile +from magi_compiler.config import CompileConfig, CompileMode +from torch import nn +from torch.testing import assert_close + + +@pytest.fixture(autouse=True) +def compile_config_fixture(): + """Fixture to set up a clean compile configuration for each test.""" + compile_config = CompileConfig( + compile_mode=CompileMode.TORCH_COMPILE, cache_root_dir=tempfile.mkdtemp(), dynamic_sources="", traced_files=set() + ) + + with patch("magi_compiler.api.get_compile_config") as mock_get_config, patch("torch.distributed.get_rank") as mock_rank: + mock_get_config.return_value = compile_config + mock_rank.return_value = 0 + yield compile_config + + import shutil + + shutil.rmtree(compile_config.cache_root_dir, ignore_errors=True) + + +class TestDynamicArgDims: + """Tests for dynamic_arg_dims inference and handling.""" + + def test_automatic_inference(self): + """Test that dynamic_arg_dims is automatically inferred when not specified.""" + model_config = MagicMock() + + @magi_compile + class InferredModel(nn.Module): + def __init__(self, *, model_config): + super().__init__() + self.linear = nn.Linear(10, 5) + + def forward(self, x: torch.Tensor, y: torch.Tensor) -> torch.Tensor: + return self.linear(x + y) + + model = InferredModel(model_config=model_config) + + # Test with different batch sizes + x1 = torch.randn(4, 10) + y1 = torch.randn(4, 10) + output1 = model(x1, y1) + assert output1.shape == (4, 5) + + x2 = torch.randn(8, 10) + y2 = torch.randn(8, 10) + output2 = model(x2, y2) + assert output2.shape == (8, 5) + + def test_negative_index(self): + """Test that negative indices in dynamic_arg_dims are handled correctly.""" + model_config = MagicMock() + + @magi_compile(dynamic_arg_dims={"x": -1}) + class DynamicLastDimModel(nn.Module): + def __init__(self, *, model_config): + super().__init__() + self.out_features = 5 + self._weight = None + self._bias = None + + def forward(self, x: torch.Tensor) -> torch.Tensor: + in_features = x.size(-1) + # Dynamically generate weights based on input size + if self._weight is None or self._weight.size(1) != in_features: + self._weight = torch.randn(self.out_features, in_features, device=x.device, dtype=x.dtype) + self._bias = torch.randn(self.out_features, device=x.device, dtype=x.dtype) + return torch.matmul(x, self._weight.t()) + self._bias + + model = DynamicLastDimModel(model_config=model_config) + + # Test inputs with different last dimension sizes + x1 = torch.randn(4, 10) + x2 = torch.randn(4, 15) + + output1 = model(x1) + output2 = model(x2) + + assert output1.shape == (4, 5) + assert output2.shape == (4, 5) + assert model._weight.size(1) == 15 # Weight updated for last input + + +class TestCompilationCorrectness: + """Tests verifying that compiled models produce correct outputs.""" + + def test_compiled_matches_native_forward(self): + """Test that compiled model produces same output as native forward.""" + model_config = MagicMock() + + @magi_compile(dynamic_arg_dims={"x": 0, "y": 0}) + class AddModel(nn.Module): + def __init__(self, *, model_config): + super().__init__() + self.linear = nn.Linear(10, 10) + + def forward(self, x: torch.Tensor, y: torch.Tensor) -> torch.Tensor: + return self.linear(x) + self.linear(y) + + model = AddModel(model_config=model_config) + x = torch.randn(4, 10) + y = torch.randn(4, 10) + + native_output = model.forward(x, y) + compiled_output = model(x, y) + + assert_close(compiled_output, native_output, rtol=1e-5, atol=1e-5) + + def test_nested_function_calls(self): + """Test compilation of model with nested function calls.""" + model_config = MagicMock() + + @magi_compile(dynamic_arg_dims={"x": 0}) + class NestedModel(nn.Module): + def __init__(self, *, model_config): + super().__init__() + self.conv = nn.Conv2d(3, 8, kernel_size=1) + + def forward(self, x: torch.Tensor) -> torch.Tensor: + x = self._preprocess(x) + return self.conv(x) + + def _preprocess(self, x: torch.Tensor) -> torch.Tensor: + return x * 2 + 1 + + model = NestedModel(model_config=model_config) + x = torch.randn(2, 3, 8, 8) + + native_output = model.forward(x) + compiled_output = model(x) + + assert_close(compiled_output, native_output, rtol=1e-5, atol=1e-5) + + def test_multiple_outputs(self): + """Test compilation of model with multiple outputs.""" + model_config = MagicMock() + + @magi_compile(dynamic_arg_dims={"x": 0}) + class MultiOutputModel(nn.Module): + def __init__(self, *, model_config): + super().__init__() + self.linear1 = nn.Linear(10, 5) + self.linear2 = nn.Linear(10, 5) + + def forward(self, x: torch.Tensor) -> Tuple[torch.Tensor, torch.Tensor]: + return self.linear1(x), self.linear2(x) + + model = MultiOutputModel(model_config=model_config) + x = torch.randn(4, 10) + + native_out1, native_out2 = model.forward(x) + compiled_out1, compiled_out2 = model(x) + + assert_close(compiled_out1, native_out1, rtol=1e-5, atol=1e-5) + assert_close(compiled_out2, native_out2, rtol=1e-5, atol=1e-5) + + +class TestRecompilation: + """Tests for source code change detection and recompilation.""" + + def test_source_change_triggers_recompile(self, tmpdir): + """Test that modifying source code triggers recompilation.""" + import importlib + import sys + + model_config = MagicMock() + + # Create a temporary source file for the model + src_file = tmpdir.join("test_model.py") + src_content = """ +import torch +from torch import nn + +class TempModel(nn.Module): + def __init__(self, *, model_config): + super().__init__() + self.linear = nn.Linear(10, 5) + + def forward(self, x: torch.Tensor) -> torch.Tensor: + return self.linear(x) +""" + src_file.write(src_content) + sys.path.insert(0, str(tmpdir)) + + import test_model + from test_model import TempModel + + CompiledTempModel = magi_compile(dynamic_arg_dims={"x": 0})(TempModel) + model1 = CompiledTempModel(model_config=model_config) + x = torch.randn(4, 10) + output1 = model1(x) + + # Modify the source file to change the forward logic + modified_src = """ +import torch +from torch import nn + +class TempModel(nn.Module): + def __init__(self, *, model_config): + super().__init__() + self.linear = nn.Linear(10, 5) + + def forward(self, x: torch.Tensor) -> torch.Tensor: + return self.linear(x) * 2 +""" + src_file.write(modified_src) + importlib.reload(test_model) + from test_model import TempModel + + CompiledTempModel2 = magi_compile(dynamic_arg_dims={"x": 0})(TempModel) + model2 = CompiledTempModel2(model_config=model_config) + output2 = model2(x) + + # Outputs should be different due to the *2 multiplication + assert not torch.allclose(output1, output2, rtol=1e-5, atol=1e-5) + + # Verify compiled output matches native forward + native_output = model2.forward(x) + assert_close(output2, native_output, rtol=1e-5, atol=1e-5) + + +if __name__ == "__main__": + pytest.main([__file__, "-v"]) diff --git a/pkgs/MagiCompiler/tests/test_mlp_infer.py b/pkgs/MagiCompiler/tests/test_mlp_infer.py new file mode 100644 index 0000000000000000000000000000000000000000..789bf0da7f0d6c634e4e16500cd674e6ea527034 --- /dev/null +++ b/pkgs/MagiCompiler/tests/test_mlp_infer.py @@ -0,0 +1,111 @@ +# 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. + +import pytest +import torch + +from .model_definition import create_mlp_model + + +@pytest.fixture(scope="function") +def mlp_model(device, mlp_config): + """MLP model fixture""" + model = create_mlp_model(mlp_config, device) + model.eval() + return model + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="Requires CUDA support") +def test_mlp_basic_inference(device, mlp_config, mlp_model): + """Test basic inference functionality""" + num_tokens = 128 + input_tensor = torch.randn(num_tokens, mlp_config.hidden_size, device=device, dtype=torch.bfloat16) + + with torch.no_grad(): + output = mlp_model(input_tensor) + + # Verify output shape + assert output.shape == ( + num_tokens, + mlp_config.hidden_size, + ), f"Output shape should be ({num_tokens}, {mlp_config.hidden_size}), but got {output.shape}" + + # Verify output data type + assert output.dtype == torch.float32, f"Output data type should be torch.float32, but got {output.dtype}" + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="Requires CUDA support") +def test_mlp_different_input_shapes(device, mlp_config, mlp_model): + """Test different input shapes""" + test_shapes = [ + (32, mlp_config.hidden_size), # Small batch + (128, mlp_config.hidden_size), # Medium batch + (512, mlp_config.hidden_size), # Large batch + # NOTE: compiler will specialize for single token, so we move it to the last + (1, mlp_config.hidden_size), # Single token + ] + + with torch.no_grad(): + for num_tokens, hidden_size in test_shapes: + input_tensor = torch.randn(num_tokens, hidden_size, device=device, dtype=torch.bfloat16) + output = mlp_model(input_tensor) + + assert output.shape == ( + num_tokens, + hidden_size, + ), f"For input shape ({num_tokens}, {hidden_size}), output shape should be ({num_tokens}, {hidden_size}), but got {output.shape}" + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="Requires CUDA support") +def test_mlp_inference_consistency(device, mlp_config, mlp_model): + """Test inference consistency (multiple runs with same input should produce same output)""" + num_tokens = 64 + torch.manual_seed(42) + input_tensor = torch.randn(num_tokens, mlp_config.hidden_size, device=device, dtype=torch.bfloat16) + + outputs = [] + with torch.no_grad(): + for _ in range(3): + output = mlp_model(input_tensor) + outputs.append(output) + + # Verify all outputs are the same + for i in range(1, len(outputs)): + assert torch.allclose(outputs[0], outputs[i], atol=1e-5), f"Output from run {i+1} is inconsistent with the first run" + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="Requires CUDA support") +def test_mlp_compiled_consistency(device, mlp_config, mlp_model): + """Test inference using magi_compiler (verify compiled code consistency)""" + num_tokens = 128 + input_tensor = torch.randn(num_tokens, mlp_config.hidden_size, device=device, dtype=torch.bfloat16) + + with torch.no_grad(): + # First run (may trigger compilation) + output1 = mlp_model(input_tensor) + + # Second run (should use compiled code) + output2 = mlp_model(input_tensor) + + # Verify output shape and consistency + assert output1.shape == ( + num_tokens, + mlp_config.hidden_size, + ), f"Output shape should be ({num_tokens}, {mlp_config.hidden_size}), but got {output1.shape}" + + assert torch.allclose(output1, output2, atol=1e-5), "Outputs from two inference runs should be consistent" + + +if __name__ == "__main__": + pytest.main([__file__, "-v"]) diff --git a/pkgs/MagiCompiler/tests/test_mlp_training.py b/pkgs/MagiCompiler/tests/test_mlp_training.py new file mode 100644 index 0000000000000000000000000000000000000000..80e31bad525eafd087ff5fde349d307ac2c29d40 --- /dev/null +++ b/pkgs/MagiCompiler/tests/test_mlp_training.py @@ -0,0 +1,159 @@ +# 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. + +import pytest +import torch +import torch.nn as nn +from tests.model_definition import MLP, MLPConfig, create_mlp_model_with_initial_params +from tests.utils import CleanupCacheContext, enable_remote_debug + + +def train_mlp_model( + model: MLP, + optimizer: torch.optim.Optimizer, + device: torch.device, + num_tokens: int, + hidden_size: int, + num_epochs: int, + batches_per_epoch: int, + gradient_accumulation_steps: int = 1, +) -> list[float]: + """Execute training loop for MLP model (supports gradient accumulation) + + Args: + model: MLP model to train + optimizer: Optimizer + device: Training device + num_tokens: Number of tokens per batch + hidden_size: Hidden layer dimension + num_epochs: Number of training epochs + batches_per_epoch: Number of batches per epoch + gradient_accumulation_steps: Gradient accumulation steps, default is 1 (no accumulation) + + Returns: + epoch_losses: List of average losses per epoch + """ + epoch_losses = [] + + print(f"Starting training: {num_epochs} epochs, {batches_per_epoch} batches per epoch") + if gradient_accumulation_steps > 1: + print(f"Using gradient accumulation, accumulation steps: {gradient_accumulation_steps}") + + for epoch in range(num_epochs): + epoch_loss_sum = 0.0 + + for batch_idx in range(batches_per_epoch): + # Zero gradients at the start of each accumulation cycle + if batch_idx % gradient_accumulation_steps == 0: + optimizer.zero_grad() + + # Generate random input and target data + input_tensor = torch.randn(num_tokens, hidden_size, device=device, dtype=torch.bfloat16) + target_tensor = torch.ones(num_tokens, hidden_size, device=device, dtype=torch.float32) + + # Forward pass + output = model(input_tensor) + + # Compute loss, divided by accumulation steps to maintain effective batch size consistency + loss = nn.functional.mse_loss(output, target_tensor) / gradient_accumulation_steps + + # Backward pass (gradients are automatically accumulated) + loss.backward() + + # Accumulate loss for logging (multiply by accumulation steps to restore original value) + epoch_loss_sum += loss.item() * gradient_accumulation_steps + + # Update parameters after accumulating gradient_accumulation_steps batches + if (batch_idx + 1) % gradient_accumulation_steps == 0: + optimizer.step() + + # Handle the last incomplete accumulation batch + if batches_per_epoch % gradient_accumulation_steps != 0: + optimizer.step() + optimizer.zero_grad() + + avg_loss = epoch_loss_sum / batches_per_epoch + epoch_losses.append(avg_loss) + print(f"Epoch {epoch + 1}/{num_epochs}, Average Loss: {avg_loss:.6f}") + + print("Training completed!") + return epoch_losses + + +def verify_model_parameters_updated( + initial_params: list[torch.Tensor], current_params: list[torch.Tensor], tolerance: float = 1e-6 +) -> bool: + """Verify whether model parameters have been updated after training + + Args: + initial_params: Parameter snapshot before training + current_params: Current parameters after training + tolerance: Tolerance for determining if parameters are the same + + Returns: + True if parameters have been updated, False otherwise + """ + for initial_param, current_param in zip(initial_params, current_params): + if not torch.allclose(initial_param, current_param, atol=tolerance): + return True + return False + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA is not available, skipping test") +def test_mlp_training_with_magi_compiler(): + """Test MLP training with magi_compiler in training scenario""" + + # Set device + device = torch.device("cuda") + + # Create MLP configuration + mlp_config = MLPConfig(hidden_size=512, intermediate_size=2048, params_dtype=torch.bfloat16) + + # Create model and save initial parameters + model, initial_params = create_mlp_model_with_initial_params(mlp_config, device) + + # Create optimizer + optimizer = torch.optim.AdamW(model.parameters(), lr=1e-4) + + # Training parameters + num_tokens = 128 + hidden_size = mlp_config.hidden_size + num_epochs = 10 + batches_per_epoch = 4 + + # Execute training + epoch_losses = train_mlp_model( + model=model, + optimizer=optimizer, + device=device, + num_tokens=num_tokens, + hidden_size=hidden_size, + num_epochs=num_epochs, + batches_per_epoch=batches_per_epoch, + ) + + # Verify model parameters have been updated + params_updated = verify_model_parameters_updated(initial_params=initial_params, current_params=list(model.parameters())) + + assert params_updated, "Model parameters should change after training" + + print("Test passed: Model successfully completed multiple training epochs, parameters have been updated") + + +if __name__ == "__main__": + # Usage: + # ENABLE_REMOTE_DEBUG=true MAGI_ENABLE_FX_GRAPH_VIZ=true TORCH_LOGS=aot CUDA_VISIBLE_DEVICES=1 python pkgs/MagiCompiler/tests/test_mlp_training.py + with CleanupCacheContext(): + enable_remote_debug() + test_mlp_training_with_magi_compiler() diff --git a/pkgs/MagiCompiler/tests/test_nested_compile.py b/pkgs/MagiCompiler/tests/test_nested_compile.py new file mode 100644 index 0000000000000000000000000000000000000000..678eaecc2c37948c328d340520e7fcdc8cd79e79 --- /dev/null +++ b/pkgs/MagiCompiler/tests/test_nested_compile.py @@ -0,0 +1,501 @@ +# 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. + +"""测试嵌套 compile 场景：torch.compile 与 magi_compile 的各种组合""" + +import os + +import pytest +import torch +import torch.nn as nn +from magi_compiler import magi_compile +from magi_compiler.config import CompileMode, get_compile_config + +DEVICE = "cuda" +HIDDEN_SIZE = 64 +TOLERANCE = 1e-3 + + +# ============ 辅助函数 ============ + + +def is_torch_compiled(module: nn.Module) -> bool: + """ + 检查模块是否被 torch.compile 编译 + + 两种方式： + 1. torch.compile(instance) -> OptimizedModule + 2. @torch.compile def forward -> forward 有 _torchdynamo_orig_callable + + 注意：@torch.compiler.disable 也设置 _torchdynamo_orig_callable， + 但会额外设置 _torchdynamo_disable=True，需排除 + """ + if type(module).__name__ == "OptimizedModule": + return True + forward_method = type(module).forward + if hasattr(forward_method, "_torchdynamo_orig_callable"): + if not getattr(forward_method, "_torchdynamo_disable", False): + return True + return False + + +def is_torch_disabled(module: nn.Module) -> bool: + """检查 forward 是否被 @torch.compiler.disable 装饰""" + return getattr(type(module).forward, "_torchdynamo_disable", False) + + +def assert_torch_compiled(module: nn.Module, msg: str = ""): + assert is_torch_compiled(module), ( + f"Expected torch.compile'd. type={type(module).__name__}, " + f"has _torchdynamo_orig_callable={hasattr(type(module).forward, '_torchdynamo_orig_callable')}. {msg}" + ) + + +def assert_not_torch_compiled_or_disabled(module: nn.Module, msg: str = ""): + assert not is_torch_compiled(module), ( + f"Expected NOT torch.compile'd. type={type(module).__name__}, " + f"has _torchdynamo_orig_callable={hasattr(type(module).forward, '_torchdynamo_orig_callable')}. {msg}" + ) + + +def assert_magi_compiled(module: nn.Module, msg: str = ""): + assert hasattr(module, "compiled_code"), f"Missing compiled_code. {msg}" + assert module.compiled_code is not None, f"compiled_code is None. {msg}" + + +def assert_not_magi_compiled(module: nn.Module, msg: str = ""): + if hasattr(module, "compiled_code"): + assert module.compiled_code is None, f"compiled_code should be None. {msg}" + + +def assert_torch_disabled(module: nn.Module, msg: str = ""): + assert is_torch_disabled(module), ( + f"Expected @torch.compiler.disable. " + f"_torchdynamo_disable={getattr(type(module).forward, '_torchdynamo_disable', False)}. {msg}" + ) + + +# ============ Fixtures ============ + + +@pytest.fixture(autouse=True) +def set_magi_compile_mode(): + """测试期间 compile_mode=MAGI_COMPILE""" + config = get_compile_config() + old_value = config.compile_mode + config.compile_mode = CompileMode.MAGI_COMPILE + config.cache_root_dir = os.environ.get("MAGI_COMPILE_CACHE_ROOT_DIR", config.cache_root_dir) + print(f"set magi compile mode: {config.compile_mode}, cache root dir: {config.cache_root_dir}") + yield + config.compile_mode = old_value + + +# ============ torch.compile 嵌套行为 ============ + + +def test_torch_compile_nested(): + """torch.compile 嵌套：内层已编译的 OptimizedModule 作为 opaque 节点""" + + class InnerBlock(nn.Module): + def __init__(self, hidden_size): + super().__init__() + self.linear = nn.Linear(hidden_size, hidden_size) + + def forward(self, x: torch.Tensor) -> torch.Tensor: + return torch.relu(self.linear(x)) + + class OuterModel(nn.Module): + def __init__(self, hidden_size): + super().__init__() + self.inner = InnerBlock(hidden_size) + self.output = nn.Linear(hidden_size, hidden_size) + + def forward(self, x: torch.Tensor) -> torch.Tensor: + x = self.inner(x) + return self.output(x) + + model = OuterModel(HIDDEN_SIZE).to(DEVICE) + x = torch.randn(4, 16, HIDDEN_SIZE, device=DEVICE) + + with torch.no_grad(): + baseline = model(x) + + model.inner = torch.compile(model.inner, fullgraph=False, dynamic=True) + assert_torch_compiled(model.inner) + with torch.no_grad(): + inner_compiled_out = model(x) + assert torch.allclose(baseline, inner_compiled_out, atol=TOLERANCE, rtol=TOLERANCE) + + compiled_model = torch.compile(model, fullgraph=False, dynamic=True) + assert_torch_compiled(compiled_model) + assert_torch_compiled(compiled_model.inner) + with torch.no_grad(): + nested_out = compiled_model(x) + + assert torch.allclose(baseline, nested_out, atol=TOLERANCE, rtol=TOLERANCE) + + +def test_torch_compile_with_disable_inner(): + """torch.compile + @torch.compiler.disable：disable 的函数产生 graph break""" + + class InnerBlock(nn.Module): + def __init__(self, hidden_size): + super().__init__() + self.linear = nn.Linear(hidden_size, hidden_size) + + @torch.compiler.disable + def forward(self, x: torch.Tensor) -> torch.Tensor: + return torch.relu(self.linear(x)) + + class OuterModel(nn.Module): + def __init__(self, hidden_size): + super().__init__() + self.inner = InnerBlock(hidden_size) + self.output = nn.Linear(hidden_size, hidden_size) + + def forward(self, x: torch.Tensor) -> torch.Tensor: + x = self.inner(x) + return self.output(x) + + model = OuterModel(HIDDEN_SIZE).to(DEVICE) + x = torch.randn(4, 16, HIDDEN_SIZE, device=DEVICE) + + with torch.no_grad(): + baseline = model(x) + + compiled_model = torch.compile(model, fullgraph=False, dynamic=True) + assert_torch_compiled(compiled_model) + assert_not_torch_compiled_or_disabled(compiled_model.inner) + with torch.no_grad(): + compiled_out = compiled_model(x) + + assert torch.allclose(baseline, compiled_out, atol=TOLERANCE, rtol=TOLERANCE) + + +# ============ torch.compile + magi_compile 嵌套 ============ + + +def test_nested_torch_compile_magi_compile(): + """外层 torch.compile + 内层 magi_compile""" + + @magi_compile() + class InnerMagiBlock(nn.Module): + def __init__(self, hidden_size: int): + super().__init__() + self.linear1 = nn.Linear(hidden_size, hidden_size * 4) + self.linear2 = nn.Linear(hidden_size * 4, hidden_size) + self.norm = nn.LayerNorm(hidden_size) + + def forward(self, x: torch.Tensor) -> torch.Tensor: + residual = x + x = self.norm(x) + x = self.linear1(x) + x = torch.nn.functional.gelu(x) + x = self.linear2(x) + return x + residual + + class OuterModel(nn.Module): + def __init__(self, hidden_size: int, num_layers: int = 2): + super().__init__() + self.embed = nn.Linear(hidden_size, hidden_size) + self.blocks = nn.ModuleList([InnerMagiBlock(hidden_size) for _ in range(num_layers)]) + self.output = nn.Linear(hidden_size, hidden_size) + + def forward(self, x: torch.Tensor) -> torch.Tensor: + x = self.embed(x) + for block in self.blocks: + x = block(x) + return self.output(x) + + num_layers = 2 + model = OuterModel(HIDDEN_SIZE, num_layers=num_layers).to(DEVICE) + x = torch.randn(4, 16, HIDDEN_SIZE, device=DEVICE) + + for i, block in enumerate(model.blocks): + assert hasattr(block, "enable_compile") + assert block.enable_compile is True + + with torch.no_grad(): + baseline = model(x) + + for i, block in enumerate(model.blocks): + assert_magi_compiled(block) + assert_not_torch_compiled_or_disabled(block) + + compiled_model = torch.compile(model, fullgraph=False, dynamic=True) + assert_torch_compiled(compiled_model) + assert compiled_model._orig_mod is model + + with torch.no_grad(): + compiled_out = compiled_model(x) + + for i, block in enumerate(model.blocks): + assert block.enable_compile is True + assert_magi_compiled(block) + + assert torch.allclose(baseline, compiled_out, atol=TOLERANCE, rtol=TOLERANCE) + + +def test_nested_torch_compile_multiple_magi_compile(): + """外层 torch.compile 包含多个 magi_compile 模块""" + + @magi_compile() + class MagiBlock1(nn.Module): + def __init__(self, hidden_size): + super().__init__() + self.linear = nn.Linear(hidden_size, hidden_size) + + def forward(self, x: torch.Tensor) -> torch.Tensor: + return torch.relu(self.linear(x)) + + @magi_compile() + class MagiBlock2(nn.Module): + def __init__(self, hidden_size): + super().__init__() + self.linear = nn.Linear(hidden_size, hidden_size) + + def forward(self, x: torch.Tensor) -> torch.Tensor: + return torch.relu(self.linear(x)) + + class OuterModel(nn.Module): + def __init__(self, hidden_size): + super().__init__() + self.block1 = MagiBlock1(hidden_size) + self.block2 = MagiBlock2(hidden_size) + self.output = nn.Linear(hidden_size, hidden_size) + + def forward(self, x: torch.Tensor) -> torch.Tensor: + x = self.block1(x) + x = self.block2(x) + return self.output(x) + + model = OuterModel(HIDDEN_SIZE).to(DEVICE) + x = torch.randn(4, 16, HIDDEN_SIZE, device=DEVICE) + + with torch.no_grad(): + baseline = model(x) + + assert_magi_compiled(model.block1) + assert_magi_compiled(model.block2) + assert_not_torch_compiled_or_disabled(model.block1) + assert_not_torch_compiled_or_disabled(model.block2) + + compiled_model = torch.compile(model, fullgraph=False, dynamic=True) + assert_torch_compiled(compiled_model) + assert_not_torch_compiled_or_disabled(model.block1) + assert_not_torch_compiled_or_disabled(model.block2) + with torch.no_grad(): + compiled_out = compiled_model(x) + + assert torch.allclose(baseline, compiled_out, atol=TOLERANCE, rtol=TOLERANCE) + + +# ============ torch.compile 使用装饰器 + magi_compile 嵌套 ============ + + +def test_decorator_torch_compile_on_forward(): + """@torch.compile 装饰 forward：模块类型不变，但 is_torch_compiled 返回 True""" + + class MyModel(nn.Module): + def __init__(self, hidden_size): + super().__init__() + self.linear = nn.Linear(hidden_size, hidden_size) + + def forward(self, x: torch.Tensor) -> torch.Tensor: + return torch.relu(self.linear(x)) + + model = MyModel(HIDDEN_SIZE).to(DEVICE) + x = torch.randn(4, 16, HIDDEN_SIZE, device=DEVICE) + + # eager baseline + with torch.no_grad(): + baseline = model(x) + + # 创建带 @torch.compile forward 的版本 + class CompiledModel(nn.Module): + def __init__(self, hidden_size): + super().__init__() + self.linear = nn.Linear(hidden_size, hidden_size) + + @torch.compile(fullgraph=False, dynamic=True) + def forward(self, x: torch.Tensor) -> torch.Tensor: + return torch.relu(self.linear(x)) + + compiled_model = CompiledModel(HIDDEN_SIZE).to(DEVICE) + compiled_model.load_state_dict(model.state_dict()) + + assert type(compiled_model).__name__ == "CompiledModel" + assert_torch_compiled(compiled_model) + + with torch.no_grad(): + out = compiled_model(x) + + assert torch.allclose(baseline, out, atol=TOLERANCE, rtol=TOLERANCE) + + +def test_decorator_nested_torch_compile_forward_magi_inner(): + """外层 forward @torch.compile + 内层 @magi_compile""" + + class InnerBlock(nn.Module): + def __init__(self, hidden_size): + super().__init__() + self.linear = nn.Linear(hidden_size, hidden_size) + + def forward(self, x: torch.Tensor) -> torch.Tensor: + return torch.relu(self.linear(x)) + + class OuterModel(nn.Module): + def __init__(self, hidden_size): + super().__init__() + self.inner = InnerBlock(hidden_size) + self.output = nn.Linear(hidden_size, hidden_size) + + def forward(self, x: torch.Tensor) -> torch.Tensor: + x = self.inner(x) + return self.output(x) + + model = OuterModel(HIDDEN_SIZE).to(DEVICE) + x = torch.randn(4, 16, HIDDEN_SIZE, device=DEVICE) + + # eager baseline + with torch.no_grad(): + baseline = model(x) + + # 创建 magi inner + torch.compile forward outer 版本 + MagiInnerBlock = magi_compile()(InnerBlock) + + class CompiledOuterModel(nn.Module): + def __init__(self, hidden_size): + super().__init__() + self.inner = MagiInnerBlock(hidden_size) + self.output = nn.Linear(hidden_size, hidden_size) + + @torch.compile(fullgraph=False, dynamic=True) + def forward(self, x: torch.Tensor) -> torch.Tensor: + x = self.inner(x) + return self.output(x) + + compiled_model = CompiledOuterModel(HIDDEN_SIZE).to(DEVICE) + compiled_model.load_state_dict(model.state_dict()) + + assert_torch_compiled(compiled_model) + + with torch.no_grad(): + out = compiled_model(x) + + assert_magi_compiled(compiled_model.inner) + assert_not_torch_compiled_or_disabled(compiled_model.inner) + assert torch.allclose(baseline, out, atol=TOLERANCE, rtol=TOLERANCE) + + +# ============ torch._dynamo.config 正确性验证 ============ + + +def test_dynamo_config_nested_patch_restore(): + """验证 config.patch() 嵌套时能正确恢复到上一层的值""" + import torch._dynamo.config as config + + # 记录初始值 + initial_value = config.assume_static_by_default + + # 模拟外层 compile 设置 dynamic=True (assume_static_by_default=False) + with config.patch(assume_static_by_default=False): + assert config.assume_static_by_default is False, "外层 patch 应将值设为 False" + + # 模拟内层 magi_compile 恢复默认 (assume_static_by_default=True) + with config.patch(assume_static_by_default=True): + assert config.assume_static_by_default is True, "内层 patch 应将值设为 True" + + # 内层退出后，应该恢复到外层的值 + assert config.assume_static_by_default is False, "内层退出后应恢复到外层值 False" + + # 外层退出后，应该恢复到初始值 + assert config.assume_static_by_default == initial_value, f"外层退出后应恢复到初始值 {initial_value}" + + +def test_dynamo_config_multiple_options_patch(): + """验证同时 patch 多个配置项时的正确性""" + import torch._dynamo.config as config + + # 记录初始值 + initial_assume_static = config.assume_static_by_default + initial_suppress_errors = config.suppress_errors + initial_verbose = config.verbose + + # 同时 patch 多个配置项 + with config.patch( + assume_static_by_default=not initial_assume_static, + suppress_errors=not initial_suppress_errors, + verbose=not initial_verbose, + ): + # 验证所有配置项都已修改 + assert config.assume_static_by_default == (not initial_assume_static), "assume_static_by_default 应被修改" + assert config.suppress_errors == (not initial_suppress_errors), "suppress_errors 应被修改" + assert config.verbose == (not initial_verbose), "verbose 应被修改" + + # 嵌套 patch 部分配置项 + with config.patch(assume_static_by_default=initial_assume_static): + assert config.assume_static_by_default == initial_assume_static, "内层应恢复 assume_static_by_default" + # 其他配置项应保持外层 patch 的值 + assert config.suppress_errors == (not initial_suppress_errors), "suppress_errors 应保持外层值" + assert config.verbose == (not initial_verbose), "verbose 应保持外层值" + + # 内层退出后，assume_static_by_default 应恢复到外层 patch 的值 + assert config.assume_static_by_default == (not initial_assume_static), "内层退出后应恢复到外层 patch 值" + + # 外层退出后，所有配置项都应恢复到初始值 + assert config.assume_static_by_default == initial_assume_static, "外层退出后 assume_static_by_default 应恢复" + assert config.suppress_errors == initial_suppress_errors, "外层退出后 suppress_errors 应恢复" + assert config.verbose == initial_verbose, "外层退出后 verbose 应恢复" + + +def test_dynamo_config_restore_on_exception(): + """验证在 with 块内抛出异常时配置能正确恢复""" + import torch._dynamo.config as config + + # 记录初始值 + initial_value = config.assume_static_by_default + + # 测试单层 patch 在异常时的恢复 + try: + with config.patch(assume_static_by_default=not initial_value): + assert config.assume_static_by_default == (not initial_value), "patch 应生效" + raise RuntimeError("测试异常") + except RuntimeError: + pass + + # 异常后配置应恢复 + assert config.assume_static_by_default == initial_value, "单层异常后应恢复到初始值" + + # 测试嵌套 patch 在内层异常时的恢复 + try: + with config.patch(assume_static_by_default=False): + assert config.assume_static_by_default is False, "外层 patch 应生效" + try: + with config.patch(assume_static_by_default=True): + assert config.assume_static_by_default is True, "内层 patch 应生效" + raise ValueError("内层测试异常") + except ValueError: + pass + # 内层异常捕获后，应恢复到外层值 + assert config.assume_static_by_default is False, "内层异常后应恢复到外层值" + except Exception: + pytest.fail("外层不应捕获到异常") + + # 最终应恢复到初始值 + assert config.assume_static_by_default == initial_value, "嵌套异常后应恢复到初始值" + + +if __name__ == "__main__": + pytest.main([__file__, "-v", "-s"]) diff --git a/pkgs/MagiCompiler/tests/test_register_custom_op.py b/pkgs/MagiCompiler/tests/test_register_custom_op.py new file mode 100644 index 0000000000000000000000000000000000000000..dc94f44d6bdc5160fbbd8ce02fb1dff11977f9d3 --- /dev/null +++ b/pkgs/MagiCompiler/tests/test_register_custom_op.py @@ -0,0 +1,660 @@ +# 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 @magi_register_custom_op decorator functionality. + +This module tests: +- Basic custom op registration (forward only) +- Custom op with infer_output_meta_fn for torch.compile tracing +- Custom op with autograd support (setup_context + backward) +- Full custom op with all components +- Multiple outputs support +- Integration with magi_compile decorator +""" + +import tempfile +from unittest.mock import patch + +import pytest +import torch +from magi_compiler.api import magi_compile, magi_register_custom_op +from magi_compiler.config import CompileConfig, CompileMode +from torch import nn +from torch.testing import assert_close + + +class TestBasicRegistration: + """Tests for basic custom op registration without autograd.""" + + def test_forward_only(self): + """Test registering a custom op with only forward implementation.""" + + @magi_register_custom_op(name="test::forward_only_op", mutates_args=()) + def _forward_only_op(x: torch.Tensor) -> torch.Tensor: + return x * 2 + 1 + + x = torch.randn(4, 8) + output = _forward_only_op(x) + expected = x * 2 + 1 + + assert_close(output, expected) + + def test_multiple_inputs(self): + """Test custom op with multiple input tensors.""" + + @magi_register_custom_op(name="test::multi_input_op", mutates_args=()) + def _multi_input_op(a: torch.Tensor, b: torch.Tensor, scale: float) -> torch.Tensor: + return (a + b) * scale + + a = torch.randn(4, 8) + b = torch.randn(4, 8) + scale = 2.5 + output = _multi_input_op(a, b, scale) + expected = (a + b) * scale + + assert_close(output, expected) + + +class TestInferOutputMeta: + """Tests for custom op with infer_output_meta_fn.""" + + def test_with_infer_output_meta(self): + """Test that infer_output_meta_fn is correctly registered for tracing.""" + + def _scaled_add_infer_output_meta(x: torch.Tensor, y: torch.Tensor, scale: float) -> torch.Tensor: + return torch.empty_like(x) + + @magi_register_custom_op( + name="test::scaled_add_op", mutates_args=(), infer_output_meta_fn=_scaled_add_infer_output_meta + ) + def _scaled_add_op(x: torch.Tensor, y: torch.Tensor, scale: float) -> torch.Tensor: + return (x + y) * scale + + x = torch.randn(4, 8) + y = torch.randn(4, 8) + scale = 3.0 + output = _scaled_add_op(x, y, scale) + expected = (x + y) * scale + + assert_close(output, expected) + + def test_multiple_outputs_infer_meta(self): + """Test infer_output_meta_fn with multiple outputs.""" + + def _split_op_infer_output_meta(x: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]: + half_size = x.shape[-1] // 2 + return (x.new_empty((*x.shape[:-1], half_size)), x.new_empty((*x.shape[:-1], half_size))) + + @magi_register_custom_op(name="test::split_op", mutates_args=(), infer_output_meta_fn=_split_op_infer_output_meta) + def _split_op(x: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]: + half_size = x.shape[-1] // 2 + # NOTE: Output cannot share the same memory with input + return torch.clone(x[..., :half_size]), torch.clone(x[..., half_size:]) + + x = torch.randn(4, 8) + out1, out2 = _split_op(x) + + assert out1.shape == (4, 4) + assert out2.shape == (4, 4) + assert_close(out1, x[..., :4]) + assert_close(out2, x[..., 4:]) + + +class TestAutograd: + """Tests for custom op with autograd support.""" + + def test_with_autograd(self): + """Test custom op with setup_context and backward functions.""" + + def _square_infer_output_meta(x: torch.Tensor) -> torch.Tensor: + return torch.empty_like(x) + + def _square_setup_context(ctx, inputs, output): + (x,) = inputs + ctx.save_for_backward(x) + + def _square_backward(ctx, grad_output): + (x,) = ctx.saved_tensors + return grad_output * 2 * x + + @magi_register_custom_op( + name="test::square_op", + mutates_args=(), + infer_output_meta_fn=_square_infer_output_meta, + setup_context_fn=_square_setup_context, + backward_fn=_square_backward, + ) + def _square_op(x: torch.Tensor) -> torch.Tensor: + return x * x + + x = torch.randn(4, 8, requires_grad=True) + output = _square_op(x) + loss = output.sum() + loss.backward() + + # Gradient of x^2 is 2x + expected_grad = 2 * x + assert_close(x.grad, expected_grad) + + def test_autograd_multiple_inputs(self): + """Test autograd with multiple input tensors.""" + + def _weighted_sum_infer_output_meta(a: torch.Tensor, b: torch.Tensor, weight: float) -> torch.Tensor: + return torch.empty_like(a) + + def _weighted_sum_setup_context(ctx, inputs, output): + a, b, weight = inputs + ctx.save_for_backward(a, b) + ctx.weight = weight + + def _weighted_sum_backward(ctx, grad_output): + a, b = ctx.saved_tensors + weight = ctx.weight + grad_a = grad_output * weight + grad_b = grad_output * (1 - weight) + return grad_a, grad_b, None # None for non-tensor input + + @magi_register_custom_op( + name="test::weighted_sum_op", + mutates_args=(), + infer_output_meta_fn=_weighted_sum_infer_output_meta, + setup_context_fn=_weighted_sum_setup_context, + backward_fn=_weighted_sum_backward, + ) + def _weighted_sum_op(a: torch.Tensor, b: torch.Tensor, weight: float) -> torch.Tensor: + return a * weight + b * (1 - weight) + + a = torch.randn(4, 8, requires_grad=True) + b = torch.randn(4, 8, requires_grad=True) + weight = 0.7 + + output = _weighted_sum_op(a, b, weight) + loss = output.sum() + loss.backward() + + expected_grad_a = torch.ones_like(a) * weight + expected_grad_b = torch.ones_like(b) * (1 - weight) + + assert_close(a.grad, expected_grad_a) + assert_close(b.grad, expected_grad_b) + + def test_autograd_multiple_outputs(self): + """Test autograd with multiple output tensors.""" + + def _split_scale_infer_output_meta(x: torch.Tensor, scale: float) -> tuple[torch.Tensor, torch.Tensor]: + half = x.shape[-1] // 2 + return (x.new_empty((*x.shape[:-1], half)), x.new_empty((*x.shape[:-1], half))) + + def _split_scale_setup_context(ctx, inputs, output): + x, scale = inputs + ctx.save_for_backward(x) + ctx.scale = scale + ctx.half = x.shape[-1] // 2 + + def _split_scale_backward(ctx, grad_out1, grad_out2): + (x,) = ctx.saved_tensors + scale = ctx.scale + # Reconstruct gradient for x + grad_x = torch.cat([grad_out1 * scale, grad_out2 * scale], dim=-1) + return grad_x, None + + @magi_register_custom_op( + name="test::split_scale_op", + mutates_args=(), + infer_output_meta_fn=_split_scale_infer_output_meta, + setup_context_fn=_split_scale_setup_context, + backward_fn=_split_scale_backward, + ) + def _split_scale_op(x: torch.Tensor, scale: float) -> tuple[torch.Tensor, torch.Tensor]: + half = x.shape[-1] // 2 + return x[..., :half] * scale, x[..., half:] * scale + + x = torch.randn(4, 8, requires_grad=True) + scale = 2.0 + + out1, out2 = _split_scale_op(x, scale) + loss = out1.sum() + out2.sum() + loss.backward() + + expected_grad = torch.ones_like(x) * scale + assert_close(x.grad, expected_grad) + + +class TestAutoGeneratedName: + """Tests for auto-generated operator name when name is not provided.""" + + def test_auto_name_single_output(self): + """Test auto-generated name with single tensor output.""" + + @magi_register_custom_op() + def _auto_name_single_op(x: torch.Tensor) -> torch.Tensor: + return x * 2 + + def fn(x): + return _auto_name_single_op(x) + + compiled_fn = torch.compile(fn, backend="eager") + + x = torch.randn(4, 8) + output = compiled_fn(x) + expected = x * 2 + + assert_close(output, expected) + + def test_auto_name_multiple_outputs(self): + """Test auto-generated name with multiple tensor outputs.""" + + @magi_register_custom_op() + def _auto_name_multi_out_op(a: torch.Tensor, b: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]: + return torch.clone(a + 1), torch.clone(b + 2) + + def fn(a, b): + return _auto_name_multi_out_op(a, b) + + compiled_fn = torch.compile(fn, backend="eager") + + a = torch.randn(3, 5) + b = torch.randn(3, 5) + out1, out2 = compiled_fn(a, b) + + assert_close(out1, a + 1) + assert_close(out2, b + 2) + + def test_auto_name_with_autograd(self): + """Test auto-generated name with autograd support.""" + + def _auto_grad_setup_context(ctx, inputs, output): + (x,) = inputs + ctx.save_for_backward(x) + + def _auto_grad_backward(ctx, grad_output): + (x,) = ctx.saved_tensors + return grad_output * 2 * x + + @magi_register_custom_op(setup_context_fn=_auto_grad_setup_context, backward_fn=_auto_grad_backward) + def _auto_name_square_op(x: torch.Tensor) -> torch.Tensor: + return x * x + + x = torch.randn(4, 8, requires_grad=True) + output = _auto_name_square_op(x) + loss = output.sum() + loss.backward() + + expected_grad = 2 * x + assert_close(x.grad, expected_grad) + + +class TestDefaultIdentityMetaFn: + """Tests for the default identity meta function when infer_output_meta_fn is not provided.""" + + def test_single_output_default_meta(self): + """Test default meta function with single tensor output.""" + + @magi_register_custom_op(name="test::default_meta_single") + def _default_meta_single_op(x: torch.Tensor) -> torch.Tensor: + return x * 2 + + def fn(x): + return _default_meta_single_op(x) + + compiled_fn = torch.compile(fn, backend="eager") + + x = torch.randn(4, 8) + output = compiled_fn(x) + expected = x * 2 + + assert_close(output, expected) + + def test_single_output_multiple_inputs_default_meta(self): + """Test default meta function with multiple inputs but single tensor output.""" + + @magi_register_custom_op(name="test::default_meta_multi_in") + def _default_meta_multi_in_op(a: torch.Tensor, b: torch.Tensor, scale: float) -> torch.Tensor: + return (a + b) * scale + + def fn(a, b, scale): + return _default_meta_multi_in_op(a, b, scale) + + compiled_fn = torch.compile(fn, backend="eager") + + a = torch.randn(4, 8) + b = torch.randn(4, 8) + scale = 2.5 + output = compiled_fn(a, b, scale) + expected = (a + b) * scale + + assert_close(output, expected) + + def test_multiple_outputs_default_meta(self): + """Test default meta function with multiple tensor outputs.""" + + @magi_register_custom_op(name="test::default_meta_multi_out") + def _default_meta_multi_out_op(x: torch.Tensor, y: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]: + # Clone to avoid aliasing issues + return torch.clone(x * 2), torch.clone(y * 3) + + def fn(x, y): + return _default_meta_multi_out_op(x, y) + + compiled_fn = torch.compile(fn, backend="eager") + + x = torch.randn(4, 8) + y = torch.randn(4, 8) + out1, out2 = compiled_fn(x, y) + + assert_close(out1, x * 2) + assert_close(out2, y * 3) + + def test_three_outputs_default_meta(self): + """Test default meta function with three tensor outputs.""" + + @magi_register_custom_op(name="test::default_meta_three_out", mutates_args=()) + def _default_meta_three_out_op( + a: torch.Tensor, b: torch.Tensor, c: torch.Tensor + ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + return torch.clone(a + 1), torch.clone(b + 2), torch.clone(c + 3) + + def fn(a, b, c): + return _default_meta_three_out_op(a, b, c) + + compiled_fn = torch.compile(fn, backend="eager") + + a = torch.randn(2, 4) + b = torch.randn(2, 4) + c = torch.randn(2, 4) + out1, out2, out3 = compiled_fn(a, b, c) + + assert_close(out1, a + 1) + assert_close(out2, b + 2) + assert_close(out3, c + 3) + + def test_default_meta_with_non_tensor_args(self): + """Test default meta function correctly skips non-tensor arguments.""" + + @magi_register_custom_op(name="test::default_meta_mixed_args") + def _default_meta_mixed_args_op( + scale: float, x: torch.Tensor, offset: int, y: torch.Tensor + ) -> tuple[torch.Tensor, torch.Tensor]: + return torch.clone(x * scale + offset), torch.clone(y * scale + offset) + + def fn(scale, x, offset, y): + return _default_meta_mixed_args_op(scale, x, offset, y) + + compiled_fn = torch.compile(fn, backend="eager") + + scale = 2.0 + x = torch.randn(3, 5) + offset = 10 + y = torch.randn(3, 5) + out1, out2 = compiled_fn(scale, x, offset, y) + + assert_close(out1, x * scale + offset) + assert_close(out2, y * scale + offset) + + +class TestTorchCompileIntegration: + """Tests for integration with torch.compile.""" + + def test_custom_op_in_compiled_function(self): + """Test that custom op works inside a torch.compile'd function.""" + + def _double_infer_output_meta(x: torch.Tensor) -> torch.Tensor: + return torch.empty_like(x) + + @magi_register_custom_op(name="test::double_op", mutates_args=(), infer_output_meta_fn=_double_infer_output_meta) + def _double_op(x: torch.Tensor) -> torch.Tensor: + return x * 2 + + def fn(x): + y = _double_op(x) + return y + 1 + + compiled_fn = torch.compile(fn, backend="eager") + + x = torch.randn(4, 8) + output = compiled_fn(x) + expected = x * 2 + 1 + + assert_close(output, expected) + + def test_custom_op_with_autograd_in_compiled_function(self): + """Test custom op with autograd inside torch.compile'd function.""" + + def _cube_infer_output_meta(x: torch.Tensor) -> torch.Tensor: + return torch.empty_like(x) + + def _cube_setup_context(ctx, inputs, output): + (x,) = inputs + ctx.save_for_backward(x) + + def _cube_backward(ctx, grad_output): + (x,) = ctx.saved_tensors + return grad_output * 3 * x * x + + @magi_register_custom_op( + name="test::cube_op", + mutates_args=(), + infer_output_meta_fn=_cube_infer_output_meta, + setup_context_fn=_cube_setup_context, + backward_fn=_cube_backward, + ) + def _cube_op(x: torch.Tensor) -> torch.Tensor: + return x * x * x + + def fn(x): + return _cube_op(x) + + compiled_fn = torch.compile(fn, backend="eager") + + x = torch.randn(4, 8, requires_grad=True) + output = compiled_fn(x) + loss = output.sum() + loss.backward() + + # Gradient of x^3 is 3x^2 + expected_grad = 3 * x * x + assert_close(x.grad, expected_grad) + + +@pytest.fixture() +def magi_compile_config(): + """Fixture to set up a clean compile configuration for magi_compile tests.""" + compile_config = CompileConfig( + compile_mode=CompileMode.TORCH_COMPILE, cache_root_dir=tempfile.mkdtemp(), dynamic_sources="", traced_files=set() + ) + + with patch("magi_compiler.api.get_compile_config") as mock_get_config, patch("torch.distributed.get_rank") as mock_rank: + mock_get_config.return_value = compile_config + mock_rank.return_value = 0 + yield compile_config + + import shutil + + shutil.rmtree(compile_config.cache_root_dir, ignore_errors=True) + + +class TestMagiCompileIntegration: + """Tests for integration with magi_compile decorator.""" + + def test_custom_op_in_magi_compiled_module(self, magi_compile_config): + """Test that custom op works inside a magi_compile'd nn.Module.""" + + def _triple_infer_output_meta(x: torch.Tensor) -> torch.Tensor: + return torch.empty_like(x) + + @magi_register_custom_op(name="test::triple_op", mutates_args=(), infer_output_meta_fn=_triple_infer_output_meta) + def _triple_op(x: torch.Tensor) -> torch.Tensor: + return x * 3 + + @magi_compile(dynamic_arg_dims={"x": 0}) + class TripleModule(nn.Module): + def __init__(self): + super().__init__() + + def forward(self, x: torch.Tensor) -> torch.Tensor: + return _triple_op(x) + 1 + + model = TripleModule() + + x = torch.randn(4, 8) + output = model(x) + expected = x * 3 + 1 + assert_close(output, expected) + + # Test with different batch size to exercise dynamic shapes + x2 = torch.randn(8, 8) + output2 = model(x2) + expected2 = x2 * 3 + 1 + assert_close(output2, expected2) + + def test_custom_op_with_autograd_in_magi_compiled_module(self, magi_compile_config): + """Test custom op with autograd inside a magi_compile'd nn.Module.""" + + def _square_v2_infer_output_meta(x: torch.Tensor) -> torch.Tensor: + return torch.empty_like(x) + + def _square_v2_setup_context(ctx, inputs, output): + (x,) = inputs + ctx.save_for_backward(x) + + def _square_v2_backward(ctx, grad_output): + (x,) = ctx.saved_tensors + return grad_output * 2 * x + + @magi_register_custom_op( + name="test::square_v2_op", + mutates_args=(), + infer_output_meta_fn=_square_v2_infer_output_meta, + setup_context_fn=_square_v2_setup_context, + backward_fn=_square_v2_backward, + ) + def _square_v2_op(x: torch.Tensor) -> torch.Tensor: + return x * x + + @magi_compile(dynamic_arg_dims={"x": 0}) + class SquareModule(nn.Module): + def __init__(self): + super().__init__() + + def forward(self, x: torch.Tensor) -> torch.Tensor: + return _square_v2_op(x) + + model = SquareModule() + + x = torch.randn(4, 8, requires_grad=True) + output = model(x) + loss = output.sum() + loss.backward() + + # Gradient of x^2 is 2x + expected_grad = 2 * x + assert_close(x.grad, expected_grad) + + def test_custom_op_with_linear_in_magi_compiled_module(self, magi_compile_config): + """Test custom op combined with nn.Linear inside a magi_compile'd module.""" + + def _relu_custom_infer_output_meta(x: torch.Tensor) -> torch.Tensor: + return torch.empty_like(x) + + @magi_register_custom_op( + name="test::relu_custom_op", mutates_args=(), infer_output_meta_fn=_relu_custom_infer_output_meta + ) + def _relu_custom_op(x: torch.Tensor) -> torch.Tensor: + return torch.relu(x) + + @magi_compile(dynamic_arg_dims={"x": 0}) + class LinearReluModule(nn.Module): + def __init__(self): + super().__init__() + self.linear = nn.Linear(8, 8) + + @torch.no_grad() + def forward(self, x: torch.Tensor) -> torch.Tensor: + return _relu_custom_op(self.linear(x)) + + model = LinearReluModule() + + x = torch.randn(4, 8) + output = model(x) + expected = torch.relu(model.linear(x)) + assert_close(output, expected) + + def test_multiple_custom_ops_in_magi_compiled_module(self, magi_compile_config): + """Test multiple custom ops used together inside a magi_compile'd module.""" + + def _add_one_infer_output_meta(x: torch.Tensor) -> torch.Tensor: + return torch.empty_like(x) + + def _mul_two_infer_output_meta(x: torch.Tensor) -> torch.Tensor: + return torch.empty_like(x) + + @magi_register_custom_op(name="test::add_one_op", mutates_args=(), infer_output_meta_fn=_add_one_infer_output_meta) + def _add_one_op(x: torch.Tensor) -> torch.Tensor: + return x + 1 + + @magi_register_custom_op(name="test::mul_two_op", mutates_args=(), infer_output_meta_fn=_mul_two_infer_output_meta) + def _mul_two_op(x: torch.Tensor) -> torch.Tensor: + return x * 2 + + @magi_compile(dynamic_arg_dims={"x": 0}) + class ChainedOpsModule(nn.Module): + def __init__(self): + super().__init__() + + def forward(self, x: torch.Tensor) -> torch.Tensor: + # (x + 1) * 2 + return _mul_two_op(_add_one_op(x)) + + model = ChainedOpsModule() + + x = torch.randn(4, 8) + output = model(x) + expected = (x + 1) * 2 + assert_close(output, expected) + + def test_custom_op_multiple_outputs_in_magi_compiled_module(self, magi_compile_config): + """Test custom op with multiple outputs inside a magi_compile'd module.""" + + def _split_v2_infer_output_meta(x: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]: + half = x.shape[-1] // 2 + return (x.new_empty((*x.shape[:-1], half)), x.new_empty((*x.shape[:-1], half))) + + @magi_register_custom_op(name="test::split_v2_op", mutates_args=(), infer_output_meta_fn=_split_v2_infer_output_meta) + def _split_v2_op(x: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]: + half = x.shape[-1] // 2 + return torch.clone(x[..., :half]), torch.clone(x[..., half:]) + + @magi_compile(dynamic_arg_dims={"x": 0}) + class SplitModule(nn.Module): + def __init__(self): + super().__init__() + + def forward(self, x: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]: + a, b = _split_v2_op(x) + return a + 1, b * 2 + + model = SplitModule() + + x = torch.randn(4, 8) + out1, out2 = model(x) + + assert out1.shape == (4, 4) + assert out2.shape == (4, 4) + assert_close(out1, x[..., :4] + 1) + assert_close(out2, x[..., 4:] * 2) + + +if __name__ == "__main__": + pytest.main([__file__, "-v"]) diff --git a/pkgs/MagiCompiler/tests/test_symbolic_dim_ability.py b/pkgs/MagiCompiler/tests/test_symbolic_dim_ability.py new file mode 100644 index 0000000000000000000000000000000000000000..c7e49292ce06f32c180bd0cbb6f7ec24002cbfc8 --- /dev/null +++ b/pkgs/MagiCompiler/tests/test_symbolic_dim_ability.py @@ -0,0 +1,97 @@ +# Copyright (c) 2026 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. + +"""Basic ability test for symbolic tensors.""" + +import pytest +import sympy +import torch +from torch.fx.experimental.symbolic_shapes import ShapeEnv, SymNode +from torch.utils._sympy.functions import IntTrueDiv + + +@pytest.fixture +def shape_env(): + """Create a fresh ShapeEnv for each test.""" + return ShapeEnv() + + +def create_symint(env: ShapeEnv, name: str): + """Create a symbolic integer with the given name.""" + sympy_symbol = sympy.Symbol(name, integer=True) + sym_node = SymNode(expr=sympy_symbol, shape_env=env, pytype=int, hint=32) + return torch.SymInt(sym_node) + + +class TestSymbolicDimAbility: + """Test suite for symbolic dimension operations.""" + + def test_symint_division(self, shape_env): + """Test symbolic integer division creates IntTrueDiv expression.""" + s1 = create_symint(shape_env, "s1") + s2 = create_symint(shape_env, "s2") + + div = s1 / s2 + + assert str(div) == "IntTrueDiv(s1, s2)" + + def test_symint_addition(self, shape_env): + """Test symbolic integer addition with constant.""" + s1 = create_symint(shape_env, "s1") + + add = s1 + 5 + + assert str(add) == "s1 + 5" + + def test_symint_multiplication(self, shape_env): + """Test symbolic integer multiplication.""" + s1 = create_symint(shape_env, "s1") + k = create_symint(shape_env, "k") + n = create_symint(shape_env, "n") + + mm_flops = s1 * n * k + + assert str(mm_flops) == "k*n*s1" + + def test_sympy_simplification(self, shape_env): + """Test that sympy can simplify symbolic expressions. + + This test verifies that (k*n*s1) / (n*s1) simplifies to k. + """ + s1 = create_symint(shape_env, "s1") + k = create_symint(shape_env, "k") + n = create_symint(shape_env, "n") + + mm_flops = s1 * n * k + assert str(mm_flops) == "k*n*s1" + + mm_memory = s1 * n + assert str(mm_memory) == "n*s1" + + # IntTrueDiv(k*n*s1, n*s1) + mm_f_m = mm_flops / mm_memory + assert str(mm_f_m) == "IntTrueDiv(k*n*s1, n*s1)" + + # Get the raw expression and simplify + raw_expr = mm_f_m.node.expr + simplified_expr_raw = sympy.simplify(raw_expr) + assert str(simplified_expr_raw) == "IntTrueDiv(k*n*s1, n*s1)" + + simplified_expr_sympy = sympy.simplify(raw_expr) + # The expression (k*n*s1) / (n*s1) should simplify to k + assert str(simplified_expr_sympy) == "IntTrueDiv(k*n*s1, n*s1)" + + # Replace IntTrueDiv with standard division for sympy simplification + simplified_expr_replace = raw_expr.replace(IntTrueDiv, lambda x, y: x / y) + assert str(simplified_expr_replace) == "k" diff --git a/pkgs/MagiCompiler/tests/tokenflow/test_graph_executor.py b/pkgs/MagiCompiler/tests/tokenflow/test_graph_executor.py new file mode 100644 index 0000000000000000000000000000000000000000..456fff7e96721da7683d053e0eab2522db6308f3 --- /dev/null +++ b/pkgs/MagiCompiler/tests/tokenflow/test_graph_executor.py @@ -0,0 +1,344 @@ +# Copyright (c) 2026 SandAI. All Rights Reserved. +from typing import Dict, Tuple + +import torch +from magi_compiler.tokenflow.graph_executor import ( + GraphNormalExecutor, + GraphOptimizer, + GraphRawExecutor, + GraphStageExecutor, + LaneType, +) +from magi_compiler.tokenflow.green_ctx import GreenCtxManager +from magi_compiler.tokenflow.sampler import exponential_aligned_sampler +from magi_compiler.tokenflow.utils import ModelConfig, TransformerModel, benchmark_func +from torch import fx + +DEVICE = torch.device("cuda:0" if torch.cuda.is_available() else "cpu") +ATOL = 1e-5 + + +def _setup_executors(model) -> Tuple[Dict[str, object], torch.nn.Module, fx.GraphModule]: + """内部辅助函数：初始化模型和所有执行器""" + + model.eval() + gm = fx.symbolic_trace(model) + + # 生成阶段配置 + optimizer = GraphOptimizer() + stage_configs = optimizer.generate_stages_per_op(gm.graph) + + # 初始化所有执行器 + executors = {} + + executors["model"] = model # 原始模型 + executors["fx"] = gm # FX GraphModule + + executors["raw"] = GraphRawExecutor(gm, DEVICE) # 原始执行器（未优化） + + # 普通执行器 + executors["normal_default"] = GraphNormalExecutor(gm, DEVICE) + executors["normal_multi"] = GraphNormalExecutor(gm, DEVICE) + executors["normal_green"] = GraphNormalExecutor(gm, DEVICE) + + for node_name in executors["normal_default"].stream_map.keys(): + executors["normal_default"].stream_map[node_name] = torch.cuda.default_stream(DEVICE) + for node_name in executors["normal_multi"].stream_map.keys(): + executors["normal_multi"].stream_map[node_name] = torch.cuda.Stream(DEVICE) + for node_name in executors["normal_green"].stream_map.keys(): + gmgr = GreenCtxManager(DEVICE.index) + executors["normal_green"].stream_map[node_name] = gmgr.create_stream(sm_count=gmgr.max_sm) + + # 阶段化执行器 + executors["stage_default"] = GraphStageExecutor(gm, stage_configs, DEVICE) + executors["stage_multi"] = GraphStageExecutor(gm, stage_configs, DEVICE) + executors["stage_green"] = GraphStageExecutor(gm, stage_configs, DEVICE) + + # 配置阶段化执行器Stream + for stage_cfg in stage_configs: + stage_name = stage_cfg.name + gmgr = GreenCtxManager(DEVICE.index) + for lane_type in stage_cfg.lane_node_dict.keys(): + executors["stage_default"].stage_lane_stream[stage_name][lane_type] = torch.cuda.default_stream(DEVICE) + gmgr = GreenCtxManager(DEVICE.index) + for lane_type in stage_cfg.lane_node_dict.keys(): + executors["stage_multi"].stage_lane_stream[stage_name][lane_type] = torch.cuda.Stream(DEVICE) + gmgr = GreenCtxManager(DEVICE.index) + for lane_type in stage_cfg.lane_node_dict.keys(): + if lane_type == LaneType.COMPUTE: + executors["stage_green"].stage_lane_stream[stage_name][lane_type] = gmgr.create_stream(sm_count=gmgr.max_sm) + + return executors + + +def test_executor_correctness_basic(): + """测试基础序列长度下所有执行器的正确性""" + + model_config = ModelConfig( + hidden_size=4096, + num_layers=1, + num_heads_q=32, + num_heads_kv=8, + head_dim=128, + intermediate_size=16384, + activation_type="gelu", + ) + model = TransformerModel(model_config).to(DEVICE) + model.eval() + + test_seq_lengths = exponential_aligned_sampler(min_val=16, max_val=2048, num_samples=10, align=7) + + executors = _setup_executors(model) + + test_input = None + + def run_orig(): + with torch.no_grad(): + res = executors["model"](test_input) + torch.cuda.synchronize(DEVICE) + return res + + def run_fx(): + with torch.no_grad(): + res = executors["fx"](test_input) + torch.cuda.synchronize(DEVICE) + return res + + def run_raw(): + with torch.no_grad(): + res = executors["raw"].execute(test_input) + torch.cuda.synchronize(DEVICE) + return res + + def run_normal_default(): + with torch.no_grad(): + res = executors["normal_default"].execute(test_input) + torch.cuda.synchronize(DEVICE) + return res + + def run_normal_multi(): + with torch.no_grad(): + res = executors["normal_multi"].execute(test_input) + torch.cuda.synchronize(DEVICE) + return res + + def run_normal_green(): + with torch.no_grad(): + res = executors["normal_green"].execute(test_input) + torch.cuda.synchronize(DEVICE) + return res + + def run_stage_default(): + with torch.no_grad(): + res = executors["stage_default"].execute(test_input) + torch.cuda.synchronize(DEVICE) + return res + + def run_stage_multi(): + with torch.no_grad(): + res = executors["stage_multi"].execute(test_input) + torch.cuda.synchronize(DEVICE) + return res + + def run_stage_green(): + with torch.no_grad(): + res = executors["stage_green"].execute(test_input) + torch.cuda.synchronize(DEVICE) + return res + + for seq_len in test_seq_lengths: + # 生成测试输入 + test_input = torch.randn(seq_len, model_config.hidden_size).to(DEVICE) + + print(f"\n--- 正确性验证，序列长度={seq_len} ---") + + out_orig = run_orig() + out_fx = run_fx() + out_raw = run_raw() + out_normal_default = run_normal_default() + out_normal_multi = run_normal_multi() + out_normal_green = run_normal_green() + out_stage_default = run_stage_default() + out_stage_multi = run_stage_multi() + out_stage_green = run_stage_green() + try: + torch.testing.assert_close(out_fx, out_orig, rtol=1e-5, atol=1e-5) + torch.testing.assert_close(out_raw, out_orig, rtol=1e-5, atol=1e-5) + torch.testing.assert_close(out_normal_default, out_orig, rtol=1e-5, atol=1e-5) + torch.testing.assert_close(out_normal_multi, out_orig, rtol=1e-5, atol=1e-5) + torch.testing.assert_close(out_normal_green, out_orig, rtol=1e-5, atol=1e-5) + torch.testing.assert_close(out_stage_default, out_orig, rtol=1e-5, atol=1e-5) + torch.testing.assert_close(out_stage_multi, out_orig, rtol=1e-5, atol=1e-5) + torch.testing.assert_close(out_stage_green, out_orig, rtol=1e-5, atol=1e-5) + print(f"序列长度={seq_len} 正确性验证通过！") + except Exception as e: + print(f"序列长度={seq_len} 正确性验证失败！错误信息: {e}") + raise e + + +def test_executor_correctness_large_sequence(): + """测试大序列长度下的执行器正确性""" + + # 大序列配置（减小hidden_size避免OOM） + model_config = ModelConfig( + hidden_size=1024, + num_layers=1, + num_heads_q=8, + num_heads_kv=4, + head_dim=128, + intermediate_size=4096, + activation_type="gelu", + ) + + model = TransformerModel(model_config).to(DEVICE) + model.eval() + + executors = _setup_executors(model) + + # 生成测试输入 + seq_len = 8192 + test_input = torch.randn(seq_len, model_config.hidden_size).to(DEVICE) + + # 基准结果 + with torch.no_grad(): + baseline = model(test_input) + + # 验证阶段化绿色Stream（重点验证最优性能执行器） + with torch.no_grad(): + stage_green_result = executors["stage_green"].execute(test_input) + executors["stage_green"].synchronize() + + assert torch.allclose(baseline, stage_green_result, atol=ATOL), f"大序列长度 {seq_len} 阶段化绿色Stream结果不匹配" + + +def test_executor_efficiency(): + """测试所有执行器的效率（输出耗时和加速比）""" + DEVICE = torch.device("cuda:0" if torch.cuda.is_available() else "cpu") + BASE_CONFIG = ModelConfig( + hidden_size=4096, + num_layers=1, ### fx382->raw858 ### fx3104->raw7409 + num_heads_q=32, + num_heads_kv=8, + head_dim=128, + intermediate_size=16384, + activation_type="gelu", + ) + + seq_len = 1 + warmup_steps = 10 + run_steps = 10 + + model = TransformerModel(BASE_CONFIG).to(DEVICE) + # model = MiniMLP(BASE_CONFIG).to(DEVICE) + executors = _setup_executors(model) + + test_input = torch.randn(seq_len, BASE_CONFIG.hidden_size).to(DEVICE) + + # 定义各执行函数 + def run_original(): + with torch.no_grad(): + executors["model"](test_input) + torch.cuda.synchronize(DEVICE) + + def run_fx(): + with torch.no_grad(): + executors["fx"](test_input) + torch.cuda.synchronize(DEVICE) + + def run_raw(): + with torch.no_grad(): + executors["raw"].execute(test_input) + torch.cuda.synchronize(DEVICE) + + def run_normal_default(): + with torch.no_grad(): + executors["normal_default"].execute(test_input) + torch.cuda.synchronize(DEVICE) + + def run_normal_multi(): + with torch.no_grad(): + executors["normal_multi"].execute(test_input) + torch.cuda.synchronize(DEVICE) + + def run_normal_green(): + with torch.no_grad(): + executors["normal_green"].execute(test_input) + torch.cuda.synchronize(DEVICE) + + def run_stage_default(): + with torch.no_grad(): + executors["stage_default"].execute(test_input) + torch.cuda.synchronize(DEVICE) + + def run_stage_multi(): + with torch.no_grad(): + executors["stage_multi"].execute(test_input) + torch.cuda.synchronize(DEVICE) + + def run_stage_green(): + with torch.no_grad(): + executors["stage_green"].execute(test_input) + torch.cuda.synchronize(DEVICE) + + # 执行基准测试 + times = { + "original": benchmark_func(run_original, warmup_steps, run_steps), + "fx": benchmark_func(run_fx, warmup_steps, run_steps), + "raw": benchmark_func(run_raw, warmup_steps, run_steps), + "normal_default": benchmark_func(run_normal_default, warmup_steps, run_steps), + "normal_multi": benchmark_func(run_normal_multi, warmup_steps, run_steps), + "normal_green": benchmark_func(run_normal_green, warmup_steps, run_steps), + "stage_default": benchmark_func(run_stage_default, warmup_steps, run_steps), + "stage_multi": benchmark_func(run_stage_multi, warmup_steps, run_steps), + "stage_green": benchmark_func(run_stage_green, warmup_steps, run_steps), + } + + # 计算加速比 + speedups = {k: times["original"] / v for k, v in times.items()} + + # 输出结果（pytest会捕获print输出） + print(f"\n=== 执行器效率测试结果（{seq_len=}） ===") + for name, t in times.items(): + print(f"{name:15s}: {t:.6f} 秒/次 (加速比: {speedups[name]:.2f}x)") + + +def test_executor_edge_cases(): + """测试边界情况""" + + # 1. 极小序列长度 + model_config = ModelConfig( + hidden_size=128, + num_layers=1, + num_heads_q=4, + num_heads_kv=2, + head_dim=32, + intermediate_size=512, + activation_type="gelu", + ) + model = TransformerModel(model_config).to(DEVICE) + model.eval() + executors = _setup_executors(model) + + test_input = torch.randn(1, model_config.hidden_size).to(DEVICE) + + with torch.no_grad(): + baseline = model(test_input) + result = executors["normal_default"].execute(test_input) + executors["normal_default"].synchronize() + + assert torch.allclose(baseline, result, atol=ATOL), "极小序列长度执行失败" + + # 2. 空依赖模型测试 + simple_model = torch.nn.Sequential(torch.nn.Linear(128, 256), torch.nn.GELU(), torch.nn.Linear(256, 128)).to(DEVICE) + simple_model.eval() + gm_simple = fx.symbolic_trace(simple_model) + + executor = GraphNormalExecutor(gm_simple, DEVICE) + test_input = torch.randn(32, 128).to(DEVICE) + + with torch.no_grad(): + baseline = simple_model(test_input) + executor_result = executor.execute(test_input) + executor.synchronize() + + assert torch.allclose(baseline, executor_result, atol=ATOL), "空依赖节点执行失败" diff --git a/pkgs/MagiCompiler/tests/tokenflow/test_graph_profile.py b/pkgs/MagiCompiler/tests/tokenflow/test_graph_profile.py new file mode 100644 index 0000000000000000000000000000000000000000..59a5fa3a7d451cfeffde320355bf84b62e28e8fa --- /dev/null +++ b/pkgs/MagiCompiler/tests/tokenflow/test_graph_profile.py @@ -0,0 +1,131 @@ +# Copyright (c) 2026 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. + +from unittest.mock import patch + +import pytest +import torch +import torch.fx as fx +from magi_compiler.config import get_compile_config +from magi_compiler.tokenflow.graph_profile import GraphProfileWrapper +from magi_compiler.tokenflow.utils import CompiledTransformerModel, ModelConfig +from magi_compiler.utils import envs + + +@pytest.fixture(scope="function") +def simple_graph_profile_wrapper() -> GraphProfileWrapper: + class SimpleModel(torch.nn.Module): + def __init__(self): + super().__init__() + self.linear = torch.nn.Linear(128, 128) + + def forward(self, x): + y = self.linear(x) + z = torch.relu(y) + return z + + model = SimpleModel() + graph_module = fx.symbolic_trace(model) + wrapper = GraphProfileWrapper(graph_module) + return wrapper + + +def test_resolve_symint_expression(simple_graph_profile_wrapper): + wrapper = simple_graph_profile_wrapper + seq_len = 64 + + assert wrapper._resolve_symint_expression(128, seq_len) == 128 + + class FakeSymInt: + def __init__(self, expr_str): + self.expr_str = expr_str + + def __str__(self): + return self.expr_str + + with patch("torch.SymInt", FakeSymInt): + sym_simple = torch.SymInt("s0") + res_simple = wrapper._resolve_symint_expression(sym_simple, seq_len) + assert res_simple == 64 + + sym_complex = torch.SymInt("s0 * 2 + 10") + res_complex = wrapper._resolve_symint_expression(sym_complex, seq_len) + assert res_complex == 138 + + sym_multi = torch.SymInt("s0 + s1") + res_multi = wrapper._resolve_symint_expression(sym_multi, seq_len) + assert res_multi == 128 + + +def test_generate_real_tensor(simple_graph_profile_wrapper): + seq_len = 1 + + class FakeSymInt: + def __init__(self, *args, **kwargs): + pass + + def __str__(self): + return "s0 * 64" + + def __int__(self): + return seq_len * 64 + + with patch("torch.SymInt", FakeSymInt): + wrapper = simple_graph_profile_wrapper + sym_dim = torch.SymInt() + shape = (sym_dim, 128) + stride = (128, 1) + dtype = torch.float32 + device = torch.device("cpu") + + tensor = wrapper._generate_real_tensor(shape, stride, dtype, device, seq_len) + + assert tensor.shape == (seq_len * 64, 128) + assert tensor.stride() == (128, 1) + assert tensor.dtype == dtype + assert not torch.allclose(tensor, torch.zeros_like(tensor)) + + +def test_e2e_correctness(): + envs.MAGI_ENABLE_PROFILE = True + # envs.MAGI_ENABLE_FX_GRAPH_VIZ = True + + get_compile_config().splitting_ops.extend(["athena::my_attention"]) + + performer_config = ModelConfig( + hidden_size=4096, + num_layers=1, + num_heads_q=32, + num_heads_kv=8, + head_dim=128, + intermediate_size=16384, + activation_type="gelu", + ) + + device = "cuda" if torch.cuda.is_available() else "cpu" + class_constructor = CompiledTransformerModel + model = class_constructor(performer_config).to(device).to(performer_config.params_dtype) + uncompiled_model = model.mod + + test_seq_lens = [4096, 1014, 512, 101, 64, 7, 1] + for seq_len in test_seq_lens: + x = torch.randn(seq_len, performer_config.hidden_size, device=device, dtype=performer_config.params_dtype) + with torch.no_grad(): + output = model(x) + uncompiled_output = uncompiled_model(x) + assert torch.allclose(output, uncompiled_output, atol=1e-3) + + +if __name__ == "__main__": + pytest.main(["-v", __file__]) diff --git a/pkgs/MagiCompiler/tests/tokenflow/test_green_ctx.py b/pkgs/MagiCompiler/tests/tokenflow/test_green_ctx.py new file mode 100644 index 0000000000000000000000000000000000000000..a6983a7558571a35085977599f5873cbeec29d0e --- /dev/null +++ b/pkgs/MagiCompiler/tests/tokenflow/test_green_ctx.py @@ -0,0 +1,284 @@ +# Copyright (c) 2026 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. + +import time + +import torch +from magi_compiler.tokenflow.green_ctx import GreenCtxManager +from magi_compiler.utils import nvtx + + +# ==================== 正确性测试 ==================== +def test_stream_creation_basic(): + gm = GreenCtxManager(0) + + stream = gm.create_stream(8) + assert stream is not None, "Stream创建失败" + assert isinstance(stream, torch.cuda.Stream), "返回类型非CUDA Stream" + + assert gm.used_sm == 8, f"SM计数错误，预期8实际{gm.used_sm}" + assert len(gm.green_streams) == 1, f"Stream列表长度错误" + + gm.cleanup() + assert gm.used_sm == 0, "清理后SM计数未归零" + assert gm.remaining_resource is not None, "清理后剩余资源为空" + + +def test_batch_stream_creation(): + gm = GreenCtxManager(0) + + streams = gm.batch_create_streams([32, 64]) + assert len(streams) == 2, "批量创建返回数量错误" + assert gm.used_sm == 96, f"批量SM计数错误，预期96实际{gm.used_sm}" + + remaining_sm = gm.remaining_resource.sm.smCount if gm.remaining_resource else 0 + assert remaining_sm == gm.max_sm - 96, f"剩余SM计算错误" + + gm.cleanup() + + +def test_sm_alignment(): + gm = GreenCtxManager(0) + gm.cleanup() + + stream = gm.create_stream(5) + + if gm.min_sm == 4 and gm.align_sm == 2: + assert gm.used_sm == 6, f"SM对齐错误，预期6实际{gm.used_sm}" + elif gm.min_sm == 8 and gm.align_sm == 8: + assert gm.used_sm == 8, f"SM对齐错误，预期8实际{gm.used_sm}" + + gm.cleanup() + + +# ==================== SM索引测试 ==================== +def get_sm_probe(): + from torch.utils.cpp_extension import load_inline + + cuda_source = ''' + #include + #include + #include + + __global__ void get_sm_ids_kernel(int* out_sm_ids, int n) { + int idx = blockIdx.x * blockDim.x + threadIdx.x; + if (idx < n) { + unsigned int smid; + asm volatile("mov.u32 %0, %%smid;" : "=r"(smid)); + out_sm_ids[idx] = (int)smid; + } + } + + torch::Tensor launch_get_sm_ids(torch::Tensor out, int64_t stream_ptr) { + int n = out.size(0); + int threads = 256; + int blocks = (n + threads - 1) / threads; + + cudaStream_t cu_stream = (cudaStream_t)stream_ptr; + + get_sm_ids_kernel<<>>( + out.data_ptr(), n); + + return out; + } + ''' + + cpp_source = "torch::Tensor launch_get_sm_ids(torch::Tensor out, int64_t stream_ptr);" + + sm_probe = load_inline( + name='sm_probe_fixed', + cpp_sources=cpp_source, + cuda_sources=cuda_source, + functions=['launch_get_sm_ids'], + with_cuda=True, + extra_cflags=['-O3'], + ) + return sm_probe + + +def get_actual_sm_indices(stream, device): + num_samples = 8192 * 32 + ids = torch.zeros(num_samples, dtype=torch.int32, device=device) + + sm_probe = get_sm_probe() + sm_probe.launch_get_sm_ids(ids, stream.cuda_stream) + + torch.cuda.synchronize(device) + + unique_ids = torch.unique(ids).cpu().numpy().tolist() + return sorted([int(x) for x in unique_ids]) + + +def test_sm_index_allocation(): + gm = GreenCtxManager(0) + gm.cleanup() + + test_cases = [[8], [64], [32, 32], [16] * 4] + for sm_counts in test_cases: + streams = gm.batch_create_streams(sm_counts) + total_allocated = 0 + + for i, (stream, sm_count) in enumerate(zip(streams, sm_counts)): + actual_sms = get_actual_sm_indices(stream, gm.device) + actual_count = len(actual_sms) + + assert actual_count == sm_count, f"SM数量不匹配：请求{sm_count}实际{actual_count}" + total_allocated += actual_count + + assert len(set(actual_sms)) == actual_count, "SM ID存在重复" + + assert total_allocated == sum(sm_counts), f"总SM分配错误：预期{sum(sm_counts)}实际{total_allocated}" + gm.cleanup() + + +# ==================== 剩余资源测试 ==================== +def test_remaining_resource_creation(): + gm = GreenCtxManager(0) + gm.cleanup() + + stream1 = gm.create_stream(64) + remaining_sm = gm.max_sm - 64 + + stream2 = gm.create_stream(remaining_sm) + + assert stream2 is not None, "剩余资源Stream创建失败" + assert gm.remaining_resource is None, "剩余资源未正确置空" + + with torch.cuda.stream(stream2): + a = torch.randn(128, 128, device=gm.device, dtype=torch.float16) + b = torch.randn(128, 128, device=gm.device, dtype=torch.float16) + c = torch.matmul(a, b) + torch.cuda.synchronize() + + actual_sms = get_actual_sm_indices(stream2, gm.device) + assert len(actual_sms) == remaining_sm, f"剩余资源SM数量错误：预期{remaining_sm}实际{len(actual_sms)}" + + gm.cleanup() + + +# ==================== 效率测试 ==================== +@nvtx.instrument_nvtx +def measure_green_overlap_shard(M, K, N, streams, control_stream, warmup=100, runs=100): + num_shards = len(streams) + device = torch.device("cuda:0") + + A = torch.randn(M, K, device=device, dtype=torch.float16) + B = torch.randn(K, N, device=device, dtype=torch.float16) + shard_m = M // num_shards + a_shards = [A[i * shard_m : (i + 1) * shard_m, :].contiguous() for i in range(num_shards)] + + @nvtx.instrument_nvtx + def warmup_phase(): + for i in range(num_shards): + with torch.cuda.stream(streams[i]): + for _ in range(warmup): + torch.mm(a_shards[i], B) + torch.cuda.synchronize() + + warmup_phase() + + _tmp_start_event = torch.cuda.Event(enable_timing=True) + _tmp_end_event = torch.cuda.Event(enable_timing=True) + + @nvtx.instrument_nvtx + def blocking_task(): + tmp = torch.randn(2048 * 16, 2048 * 16, device=device) + _tmp_start_event.record(control_stream) + with torch.cuda.stream(control_stream): + for _ in range(10): + tmp = tmp * tmp + _tmp_end_event.record(control_stream) + + blocking_task() + + @nvtx.instrument_nvtx + def main_measurement(): + nonlocal _cpu_issue_start, _cpu_issue_end + _cpu_issue_start = time.time() + + start_event.record(control_stream) + for i in range(num_shards): + with torch.cuda.stream(streams[i]): + streams[i].wait_event(start_event) + for _ in range(runs): + torch.mm(a_shards[i], B) + streams[i].record_event(stream_end_events[i]) + + _cpu_issue_end = time.time() + + for i in range(num_shards): + control_stream.wait_event(stream_end_events[i]) + end_event.record(control_stream) + torch.cuda.synchronize() + + _cpu_issue_start = None + _cpu_issue_end = None + start_event = torch.cuda.Event(enable_timing=True) + end_event = torch.cuda.Event(enable_timing=True) + stream_end_events = [torch.cuda.Event() for _ in range(num_shards)] + main_measurement() + + _cpu_elapsed = _cpu_issue_end - _cpu_issue_start + _tmp_elapsed_ms = _tmp_start_event.elapsed_time(_tmp_end_event) + + assert _cpu_elapsed < _tmp_elapsed_ms / 1000.0, "CPU submission too slow, measurement invalid!" + + avg_time = start_event.elapsed_time(end_event) / runs + return avg_time + + +def test_green_stream_performance(): + gm = GreenCtxManager(0) + + M, K, N = 8192, 4096, 4096 + num_shards = 2 + + control_stream = torch.cuda.Stream() + + default_streams = [torch.cuda.default_stream() for _ in range(num_shards)] + t_default = measure_green_overlap_shard(M, K, N, default_streams, control_stream) + + multi_streams = [torch.cuda.Stream() for _ in range(num_shards)] + t_multi = measure_green_overlap_shard(M, K, N, multi_streams, control_stream) + + sm_count = (gm.max_sm // num_shards) // gm.align_sm * gm.align_sm + green_streams = gm.batch_create_streams([sm_count] * num_shards) + t_green = measure_green_overlap_shard(M, K, N, green_streams, control_stream) + + assert t_multi <= t_default * (num_shards - 0.5), f"多流性能不应劣化过多！{t_multi=} vs {t_default=}" + assert t_green >= t_default, f"Green Stream 性能不应优于默认流！{t_green=} vs {t_default=}" + assert t_green <= t_default * (num_shards + 0.5), f"Green Stream 性能不应劣化过多！{t_green=} vs {t_default=}" + + +def test_remaining_resource_stream(): + gm = GreenCtxManager(0) + total_sm = gm.max_sm + + stream0 = gm.create_stream(64) + stream1 = gm.create_stream(total_sm - 64) + + assert stream1 is not None, "剩余资源Stream创建失败" + + with torch.cuda.stream(stream1): + a = torch.randn(128, 128, device=gm.device, dtype=torch.float16) + b = torch.randn(128, 128, device=gm.device, dtype=torch.float16) + c = torch.matmul(a, b) + torch.cuda.synchronize() + + actual_sms = get_actual_sm_indices(stream1, gm.device) + assert len(actual_sms) == total_sm - 64, f"剩余资源SM数量错误" + + assert gm.remaining_resource is None, "剩余资源未正确置空！" + + gm.cleanup() diff --git a/pkgs/MagiCompiler/tests/tokenflow/test_sampler.py b/pkgs/MagiCompiler/tests/tokenflow/test_sampler.py new file mode 100644 index 0000000000000000000000000000000000000000..fcf46e0c5f77a433e67290ac2465ff33386e996b --- /dev/null +++ b/pkgs/MagiCompiler/tests/tokenflow/test_sampler.py @@ -0,0 +1,72 @@ +# Copyright (c) 2026 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. + +import numpy as np +import pytest +from magi_compiler.tokenflow.sampler import exponential_aligned_sampler + + +def test_sm_sampling(): + result = exponential_aligned_sampler(min_val=8, max_val=132, num_samples=5, align=8) + assert isinstance(result, list), "返回值必须是列表" + assert len(result) == 5, f"返回长度应为5，实际为{len(result)}" + assert all(isinstance(x, int) for x in result), "所有元素必须是整数" + assert all(x % 8 == 0 for x in result), "所有元素必须对齐到8的倍数" + assert result[0] == 8, f"首元素应为8，实际为{result[0]}" + assert result[-1] == 128, f"尾元素应为128（132对齐后），实际为{result[-1]}" + assert all(result[i] < result[i + 1] for i in range(len(result) - 1)), "采样结果应严格递增" + + mid_value = (8 + 128) / 2 + mid_value_found = np.median(result) + assert mid_value_found < mid_value, "中间值应小于范围中点，符合指数分布特性" + + +def test_seqlen_sampling(): + result = exponential_aligned_sampler(min_val=1, max_val=65536, num_samples=5, align=32) + assert len(result) == 5 + assert all(x % 32 == 0 for x in result) + assert result[0] == 32 + assert result[-1] == 65536 + + mid_value = (32 + 65536) / 2 + mid_value_found = np.median(result) + assert mid_value_found < mid_value, "中间值应小于范围中点，符合指数分布特性" + + +def test_min_max_aligned_exactly(): + result = exponential_aligned_sampler(min_val=16, max_val=64, num_samples=4, align=16) + assert len(result) == 4 + assert result[0] == 16 + assert result[-1] == 64 + assert all(x in [16, 32, 48, 64] for x in result) + + +def test_output_len_2(): + result = exponential_aligned_sampler(min_val=8, max_val=132, num_samples=2, align=8) + assert len(result) == 2 + assert result == [8, 128] + + +def test_large_range_sampling(): + result = exponential_aligned_sampler(min_val=1, max_val=1000001, num_samples=10, align=64) + assert len(result) == 10 + assert all(x % 64 == 0 for x in result) + assert result[0] == 64 + assert result[-1] == 1000000 + + +def test_large_output_len(): + with pytest.raises(ValueError) as excinfo: + result = exponential_aligned_sampler(min_val=8, max_val=64, num_samples=10, align=8) + assert excinfo is not None diff --git a/pkgs/MagiCompiler/tests/utils.py b/pkgs/MagiCompiler/tests/utils.py new file mode 100644 index 0000000000000000000000000000000000000000..a5fa647165c216d05e38b7cf13723fec833e9634 --- /dev/null +++ b/pkgs/MagiCompiler/tests/utils.py @@ -0,0 +1,54 @@ +# 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. + +import shutil + +from magi_compiler.config import get_compile_config + + +class CleanupCacheContext: + """Context manager for cleaning cache before and after execution""" + + def __enter__(self): + shutil.rmtree(get_compile_config().cache_root_dir, ignore_errors=True) + return self + + def __exit__(self, exc_type, exc_val, exc_tb): + shutil.rmtree(get_compile_config().cache_root_dir, ignore_errors=True) + + +def enable_remote_debug(): + import os + + import debugpy + + ENABLE_MAGI_REMOTE_DEBUG = os.environ.get("ENABLE_MAGI_REMOTE_DEBUG", "false").lower() + if ENABLE_MAGI_REMOTE_DEBUG == "false": + return + + debug_ranks = [] + if ENABLE_MAGI_REMOTE_DEBUG == "true": + debug_ranks = [0] + elif ENABLE_MAGI_REMOTE_DEBUG == "all": + debug_ranks = [i for i in range(1)] + else: + debug_ranks = [int(i) for i in ENABLE_MAGI_REMOTE_DEBUG.split(",")] + + rank = 0 + if rank in debug_ranks: + debug_port = 5678 + int(rank) + print(f"[rank {rank}] Starting remote debug on port {debug_port}") + debugpy.listen(("0.0.0.0", debug_port)) + debugpy.wait_for_client() + print(f"[rank {rank}] Remote debug attached") diff --git a/pkgs/magife_stable_audio_open-1.0.0+mav.1-py3-none-any.whl b/pkgs/magife_stable_audio_open-1.0.0+mav.1-py3-none-any.whl new file mode 100644 index 0000000000000000000000000000000000000000..5a65dd1e0b4510073c3c06c3f33c21960bbc1f0f --- /dev/null +++ b/pkgs/magife_stable_audio_open-1.0.0+mav.1-py3-none-any.whl @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:3535ef0239ba5014dce5a66210e2a8542fd8d5bab78007f9f7236b193da2b1c7 +size 30557 diff --git a/requirements.txt b/requirements.txt index 6323be56ddd0b6a256fdc528e681ebd7c756928f..05e2af50098eda5bbf2677c8eb199214a75857f6 100644 --- a/requirements.txt +++ b/requirements.txt @@ -1,12 +1,3 @@ -# -- HF Spaces / Gradio SDK additions -- -spaces -gradio -huggingface_hub -soundfile -flash-attn -magi_compiler @ git+https://github.com/sandai/MagiCompiler.git - -# -- Original project deps -- accelerate==1.10.1 av==15.1.0 beautifulsoup4