├── .github └── workflows │ ├── build-and-test.yml │ ├── format.yml │ └── regression-benchmark.yml ├── .gitignore ├── .gitmodules ├── CMakeLists.txt ├── LICENSE ├── README.md ├── benchmark ├── CMakeLists.txt ├── README.md └── python │ ├── benchmarks │ ├── gcn_benchmark.py │ ├── kernels_benchmark.py │ ├── lif_benchmark.py │ ├── regression_benchmark.py │ └── resnet_benchmark.py │ ├── manual │ ├── README.md │ └── sum_of_sq.py │ └── utils │ ├── benchmark_utils.py │ ├── profiler.py │ └── tensor_generator.py ├── include ├── CMakeLists.txt ├── mpact-c │ └── Registration.h └── mpact │ ├── CMakeLists.txt │ └── Transforms │ ├── CMakeLists.txt │ ├── Passes.h │ ├── Passes.td │ └── Sparsity │ └── SparseEncodingPropagate.h ├── lib ├── CAPI │ ├── CMakeLists.txt │ └── Registration.cpp ├── CMakeLists.txt └── Transforms │ ├── CMakeLists.txt │ ├── Passes.cpp │ └── Sparsity │ ├── CMakeLists.txt │ └── SparseEncodingPropagate.cpp ├── python ├── CMakeLists.txt ├── MPACTModule.cpp └── mpact │ ├── models │ ├── gat.py │ ├── gcn.py │ ├── kernels.py │ ├── lif.py │ ├── resnet.py │ └── train.py │ └── mpactbackend.py ├── setup.py ├── test ├── CMakeLists.txt ├── lit.cfg.py ├── lit.site.cfg.py.in └── python │ ├── add.py │ ├── counteq.py │ ├── file_formats.py │ ├── gat.py │ ├── gcn.py │ ├── lif.py │ ├── mm.py │ ├── mm_print.py │ ├── mul.py │ ├── norm.py │ ├── resnet.py │ ├── scale.py │ ├── sddmm.py │ ├── spmv.py │ ├── sqsum.py │ └── train_simple.py └── tools ├── CMakeLists.txt └── mpact-opt ├── CMakeLists.txt └── mpact_opt.cpp /.github/workflows/build-and-test.yml: -------------------------------------------------------------------------------- 1 | name: Build and test 2 | 3 | on: 4 | push: 5 | branches: [ "main" ] 6 | pull_request: 7 | branches: [ "main" ] 8 | 9 | concurrency: 10 | # Use PR number as key for a pull request or the commit hash otherwise. This cancels 11 | # queued and in-progress runs for the same PR (presubmit) or commit 12 | # (postsubmit). 13 | group: ci-build-test-cpp-linux-${{ github.event.number || github.sha }} 14 | cancel-in-progress: true 15 | 16 | jobs: 17 | build: 18 | runs-on: ubuntu-latest 19 | env: 20 | CACHE_DIR: ${{ github.workspace }}/.ccache 21 | steps: 22 | - uses: actions/checkout@v4 23 | with: 24 | submodules: recursive 25 | token: ${{ secrets.GITHUB_TOKEN }} 26 | 27 | - name: Setup Python Version 28 | uses: actions/setup-python@v5 29 | with: 30 | python-version: 3.11 # Install the python version needed 31 | 32 | - name: Set PYTHONPATH 33 | run: export PYTHONPATH=build/tools/mpact/python_packages/mpact 34 | shell: bash 35 | 36 | - name: Set up ccache 37 | uses: hendrikmuhs/ccache-action@v1.2 38 | 39 | - name: Install requirements 40 | run: | 41 | export CCACHE_DIR=${{ env.CACHE_DIR }} 42 | python -m pip install --upgrade pip 43 | python -m pip install -r externals/torch-mlir/requirements.txt 44 | python -m pip install -r externals/torch-mlir/torchvision-requirements.txt 45 | 46 | - name: Create build directory 47 | run: mkdir build 48 | 49 | - name: Configure CMake 50 | run: > 51 | cmake -GNinja -Bbuild 52 | -DCMAKE_BUILD_TYPE=Release 53 | -DLLVM_ENABLE_PROJECTS=mlir 54 | -DLLVM_ENABLE_ASSERTIONS=ON 55 | -DLLVM_EXTERNAL_PROJECTS="torch-mlir;mpact" 56 | -DLLVM_EXTERNAL_TORCH_MLIR_SOURCE_DIR="${PWD}/externals/torch-mlir" 57 | -DLLVM_EXTERNAL_MPACT_SOURCE_DIR="${PWD}" 58 | -DLLVM_TARGETS_TO_BUILD=host 59 | -DMLIR_ENABLE_BINDINGS_PYTHON=ON 60 | -DCMAKE_C_COMPILER_LAUNCHER=ccache 61 | -DCMAKE_CXX_COMPILER_LAUNCHER=ccache 62 | -DCMAKE_C_COMPILER=clang 63 | -DCMAKE_CXX_COMPILER=clang++ 64 | "externals/torch-mlir/externals/llvm-project/llvm" 65 | 66 | - name: Build 67 | run: cmake --build build --target check-mpact 68 | 69 | -------------------------------------------------------------------------------- /.github/workflows/format.yml: -------------------------------------------------------------------------------- 1 | name: "Check code formatting" 2 | 3 | permissions: 4 | contents: read 5 | 6 | on: 7 | pull_request: 8 | branches: 9 | - main 10 | 11 | jobs: 12 | code_formatter: 13 | runs-on: ubuntu-latest 14 | steps: 15 | - name: Fetch sources 16 | uses: actions/checkout@v4 17 | with: 18 | ref: ${{ github.event.pull_request.head.sha }} 19 | 20 | - name: Checkout through merge base 21 | uses: rmacklin/fetch-through-merge-base@v0 22 | with: 23 | base_ref: ${{ github.event.pull_request.base.ref }} 24 | head_ref: ${{ github.event.pull_request.head.sha }} 25 | deepen_length: 500 26 | 27 | - name: Get changed files 28 | id: changed-files 29 | uses: tj-actions/changed-files@v41 30 | with: 31 | separator: "," 32 | skip_initial_fetch: true 33 | 34 | # We need to pull the script from the main branch, so that we ensure 35 | # we get the latest version of this script. 36 | - name: Fetch code formatting utils 37 | uses: actions/checkout@v4 38 | with: 39 | repository: llvm/llvm-project 40 | ref: refs/heads/main 41 | sparse-checkout: | 42 | llvm/utils/git/requirements_formatting.txt 43 | llvm/utils/git/code-format-helper.py 44 | sparse-checkout-cone-mode: false 45 | path: code-format-tools 46 | 47 | - name: "Listed files" 48 | env: 49 | CHANGED_FILES: ${{ steps.changed-files.outputs.all_changed_files }} 50 | run: | 51 | echo "Formatting files:" 52 | echo "$CHANGED_FILES" 53 | - name: Install clang-format 54 | uses: aminya/setup-cpp@v1 55 | with: 56 | clangformat: 18.1.1 57 | 58 | - name: Setup Python env 59 | uses: actions/setup-python@v4 60 | with: 61 | python-version: '3.11' 62 | cache: 'pip' 63 | cache-dependency-path: 'code-format-tools/llvm/utils/git/requirements_formatting.txt' 64 | 65 | - name: Install python dependencies 66 | run: pip install -r code-format-tools/llvm/utils/git/requirements_formatting.txt 67 | 68 | - name: Run code formatter 69 | env: 70 | GITHUB_PR_NUMBER: ${{ github.event.pull_request.number }} 71 | START_REV: ${{ github.event.pull_request.base.sha }} 72 | END_REV: ${{ github.event.pull_request.head.sha }} 73 | CHANGED_FILES: ${{ steps.changed-files.outputs.all_changed_files }} 74 | # TODO: Once clang v18 is released, we should be able 75 | # to take advantage of the new --diff_from_common_commit option 76 | # explicitly in code-format-helper.py and not have to diff starting at 77 | # the merge base. 78 | # Create an empty comments file so the pr-write job doesn't fail. 79 | run: | 80 | echo "[]" > comments && 81 | python ./code-format-tools/llvm/utils/git/code-format-helper.py \ 82 | --write-comment-to-file \ 83 | --token ${{ secrets.GITHUB_TOKEN }} \ 84 | --issue-number $GITHUB_PR_NUMBER \ 85 | --start-rev $(git merge-base $START_REV $END_REV) \ 86 | --end-rev $END_REV \ 87 | --changed-files "$CHANGED_FILES" 88 | - uses: actions/upload-artifact@26f96dfa697d77e81fd5907df203aa23a56210a8 #v4.3.0 89 | if: always() 90 | with: 91 | name: workflow-args 92 | path: | 93 | comments 94 | -------------------------------------------------------------------------------- /.github/workflows/regression-benchmark.yml: -------------------------------------------------------------------------------- 1 | name: Regression benchmark 2 | 3 | on: 4 | push: 5 | branches: [ "main" ] 6 | 7 | permissions: 8 | contents: write 9 | deployments: write 10 | pull-requests: write 11 | repository-projects: write 12 | 13 | jobs: 14 | benchmark: 15 | name: Performance regression check 16 | runs-on: ubuntu-latest 17 | env: 18 | CACHE_DIR: ${{ github.workspace }}/.ccache 19 | PYTHONPATH: ${{ github.workspace }}/build/tools/mpact/python_packages/mpact 20 | steps: 21 | - uses: actions/checkout@v4 22 | with: 23 | submodules: recursive 24 | token: ${{ secrets.GITHUB_TOKEN }} 25 | 26 | - name: Setup Python Version 27 | uses: actions/setup-python@v5 28 | with: 29 | python-version: 3.11 # Install the python version needed 30 | 31 | - name: Set up ccache 32 | uses: hendrikmuhs/ccache-action@v1.2 33 | 34 | - name: Install requirements 35 | run: | 36 | export CCACHE_DIR=${{ env.CACHE_DIR }} 37 | python -m pip install --upgrade pip 38 | python -m pip install --upgrade pip 39 | python -m pip install pytest pytest-benchmark 40 | python -m pip install -r externals/torch-mlir/requirements.txt 41 | python -m pip install -r externals/torch-mlir/torchvision-requirements.txt 42 | 43 | - name: Create build directory 44 | run: mkdir build 45 | 46 | - name: Configure CMake 47 | run: > 48 | cmake -GNinja -Bbuild 49 | -DCMAKE_BUILD_TYPE=Release 50 | -DLLVM_ENABLE_PROJECTS=mlir 51 | -DLLVM_ENABLE_ASSERTIONS=ON 52 | -DLLVM_EXTERNAL_PROJECTS="torch-mlir;mpact" 53 | -DLLVM_EXTERNAL_TORCH_MLIR_SOURCE_DIR="${PWD}/externals/torch-mlir" 54 | -DLLVM_EXTERNAL_MPACT_SOURCE_DIR="${PWD}" 55 | -DLLVM_TARGETS_TO_BUILD=host 56 | -DMLIR_ENABLE_BINDINGS_PYTHON=ON 57 | -DCMAKE_C_COMPILER_LAUNCHER=ccache 58 | -DCMAKE_CXX_COMPILER_LAUNCHER=ccache 59 | -DCMAKE_C_COMPILER=clang 60 | -DCMAKE_CXX_COMPILER=clang++ 61 | "externals/torch-mlir/externals/llvm-project/llvm" 62 | 63 | - name: Build 64 | run: cmake --build build --target build-benchmark-mpact 65 | 66 | - name: Run benchmark 67 | run: pytest benchmark/python/benchmarks/regression_benchmark.py --benchmark-json output.json 68 | 69 | - name: Store benchmark result 70 | uses: benchmark-action/github-action-benchmark@v1 71 | with: 72 | tool: 'pytest' 73 | output-file-path: output.json 74 | fail-on-alert: true 75 | # GitHub API token to make a commit comment 76 | github-token: ${{ secrets.GITHUB_TOKEN }} 77 | # Enable alert commit comment 78 | comment-on-alert: true 79 | # Mention @reidtatge in the commit comment 80 | alert-comment-cc-users: '@reidtatge' 81 | # Push and deploy GitHub pages branch automatically 82 | auto-push: true 83 | alert-threshold: 120% 84 | -------------------------------------------------------------------------------- /.gitignore: -------------------------------------------------------------------------------- 1 | *_venv/ 2 | __pycache__ 3 | /build*/ 4 | 5 | # lsp files 6 | .cache/ 7 | compile_commands.json 8 | -------------------------------------------------------------------------------- /.gitmodules: -------------------------------------------------------------------------------- 1 | [submodule "torch-mlir"] 2 | path = externals/torch-mlir 3 | url = https://github.com/llvm/torch-mlir.git 4 | [submodule "externals/Enzyme"] 5 | path = externals/Enzyme 6 | url = https://github.com/EnzymeAD/Enzyme.git 7 | -------------------------------------------------------------------------------- /CMakeLists.txt: -------------------------------------------------------------------------------- 1 | #------------------------------------------------------------------------------- 2 | # The MPACT Compiler 3 | #------------------------------------------------------------------------------- 4 | 5 | cmake_minimum_required(VERSION 3.12) 6 | 7 | project(mpact VERSION 1.0 LANGUAGES CXX C) 8 | 9 | set(CMAKE_C_STANDARD 11) 10 | set(CMAKE_CXX_STANDARD 17) 11 | 12 | set(MPACT_SOURCE_DIR "${CMAKE_CURRENT_SOURCE_DIR}") 13 | set(MPACT_BINARY_DIR "${CMAKE_CURRENT_BINARY_DIR}") 14 | message(STATUS "Building the MPACT compiler at ${MPACT_SOURCE_DIR} (into ${MPACT_BINARY_DIR})") 15 | 16 | set(MPACT_PYTHON_PACKAGES_DIR "${MPACT_BINARY_DIR}/python_packages") 17 | 18 | #------------------------------------------------------------------------------- 19 | # Configure out-of-tree vs in-tree build 20 | #------------------------------------------------------------------------------- 21 | 22 | if(CMAKE_SOURCE_DIR STREQUAL CMAKE_CURRENT_SOURCE_DIR) 23 | message(STATUS "MPACT out-of-tree build.") 24 | message(FATAL_ERROR "TODO") 25 | else() 26 | message(STATUS "MPACT in-tree build.") 27 | # In-tree build with LLVM_EXTERNAL_PROJECTS=mpact 28 | option(MLIR_ENABLE_BINDINGS_PYTHON "Enables MLIR Python Bindings" OFF) 29 | set(MLIR_MAIN_SRC_DIR ${LLVM_MAIN_SRC_DIR}/../mlir) 30 | set(MLIR_INCLUDE_DIR ${LLVM_MAIN_SRC_DIR}/../mlir/include) 31 | set(MLIR_GENERATED_INCLUDE_DIR ${LLVM_BINARY_DIR}/tools/mlir/include) 32 | set(MLIR_INCLUDE_DIRS "${MLIR_INCLUDE_DIR};${MLIR_GENERATED_INCLUDE_DIR}") 33 | endif() 34 | 35 | include_directories(${LLVM_INCLUDE_DIRS}) 36 | include_directories(${MLIR_INCLUDE_DIRS}) 37 | include_directories(${CMAKE_CURRENT_SOURCE_DIR}/include) 38 | include_directories(${CMAKE_CURRENT_BINARY_DIR}/include) 39 | 40 | # Needed to build TorchMLIRExtensions. 41 | include_directories(${CMAKE_CURRENT_SOURCE_DIR}/externals/torch-mlir/include) 42 | 43 | function(mpact_target_includes target) 44 | set(_dirs 45 | $ 46 | $ 47 | $ 48 | ) 49 | target_include_directories(${target} PUBLIC ${_dirs}) 50 | if(TARGET obj.${target}) 51 | target_include_directories(obj.${target} PRIVATE ${_dirs}) 52 | endif() 53 | endfunction() 54 | 55 | list(APPEND CMAKE_MODULE_PATH ${MLIR_MAIN_SRC_DIR}/cmake/modules) 56 | list(APPEND CMAKE_MODULE_PATH ${LLVM_MAIN_SRC_DIR}/cmake) 57 | list(APPEND CMAKE_MODULE_PATH ${CMAKE_CURRENT_SOURCE_DIR}/build_tools/cmake) 58 | 59 | include(TableGen) 60 | include(AddLLVM) 61 | include(AddMLIR) 62 | include(AddMLIRPython) 63 | 64 | include(MLIRDetectPythonEnv) 65 | mlir_configure_python_dev_packages() 66 | 67 | add_subdirectory(include) 68 | add_subdirectory(lib) 69 | add_subdirectory(tools) 70 | 71 | add_subdirectory(benchmark) 72 | add_subdirectory(python) 73 | add_subdirectory(test) 74 | -------------------------------------------------------------------------------- /LICENSE: -------------------------------------------------------------------------------- 1 | Apache License 2 | Version 2.0, January 2004 3 | http://www.apache.org/licenses/ 4 | 5 | TERMS AND CONDITIONS FOR USE, REPRODUCTION, AND DISTRIBUTION 6 | 7 | 1. Definitions. 8 | 9 | "License" shall mean the terms and conditions for use, reproduction, 10 | and distribution as defined by Sections 1 through 9 of this document. 11 | 12 | "Licensor" shall mean the copyright owner or entity authorized by 13 | the copyright owner that is granting the License. 14 | 15 | "Legal Entity" shall mean the union of the acting entity and all 16 | other entities that control, are controlled by, or are under common 17 | control with that entity. For the purposes of this definition, 18 | "control" means (i) the power, direct or indirect, to cause the 19 | direction or management of such entity, whether by contract or 20 | otherwise, or (ii) ownership of fifty percent (50%) or more of the 21 | outstanding shares, or (iii) beneficial ownership of such entity. 22 | 23 | "You" (or "Your") shall mean an individual or Legal Entity 24 | exercising permissions granted by this License. 25 | 26 | "Source" form shall mean the preferred form for making modifications, 27 | including but not limited to software source code, documentation 28 | source, and configuration files. 29 | 30 | "Object" form shall mean any form resulting from mechanical 31 | transformation or translation of a Source form, including but 32 | not limited to compiled object code, generated documentation, 33 | and conversions to other media types. 34 | 35 | "Work" shall mean the work of authorship, whether in Source or 36 | Object form, made available under the License, as indicated by a 37 | copyright notice that is included in or attached to the work 38 | (an example is provided in the Appendix below). 39 | 40 | "Derivative Works" shall mean any work, whether in Source or Object 41 | form, that is based on (or derived from) the Work and for which the 42 | editorial revisions, annotations, elaborations, or other modifications 43 | represent, as a whole, an original work of authorship. For the purposes 44 | of this License, Derivative Works shall not include works that remain 45 | separable from, or merely link (or bind by name) to the interfaces of, 46 | the Work and Derivative Works thereof. 47 | 48 | "Contribution" shall mean any work of authorship, including 49 | the original version of the Work and any modifications or additions 50 | to that Work or Derivative Works thereof, that is intentionally 51 | submitted to Licensor for inclusion in the Work by the copyright owner 52 | or by an individual or Legal Entity authorized to submit on behalf of 53 | the copyright owner. For the purposes of this definition, "submitted" 54 | means any form of electronic, verbal, or written communication sent 55 | to the Licensor or its representatives, including but not limited to 56 | communication on electronic mailing lists, source code control systems, 57 | and issue tracking systems that are managed by, or on behalf of, the 58 | Licensor for the purpose of discussing and improving the Work, but 59 | excluding communication that is conspicuously marked or otherwise 60 | designated in writing by the copyright owner as "Not a Contribution." 61 | 62 | "Contributor" shall mean Licensor and any individual or Legal Entity 63 | on behalf of whom a Contribution has been received by Licensor and 64 | subsequently incorporated within the Work. 65 | 66 | 2. Grant of Copyright License. Subject to the terms and conditions of 67 | this License, each Contributor hereby grants to You a perpetual, 68 | worldwide, non-exclusive, no-charge, royalty-free, irrevocable 69 | copyright license to reproduce, prepare Derivative Works of, 70 | publicly display, publicly perform, sublicense, and distribute the 71 | Work and such Derivative Works in Source or Object form. 72 | 73 | 3. Grant of Patent License. Subject to the terms and conditions of 74 | this License, each Contributor hereby grants to You a perpetual, 75 | worldwide, non-exclusive, no-charge, royalty-free, irrevocable 76 | (except as stated in this section) patent license to make, have made, 77 | use, offer to sell, sell, import, and otherwise transfer the Work, 78 | where such license applies only to those patent claims licensable 79 | by such Contributor that are necessarily infringed by their 80 | Contribution(s) alone or by combination of their Contribution(s) 81 | with the Work to which such Contribution(s) was submitted. If You 82 | institute patent litigation against any entity (including a 83 | cross-claim or counterclaim in a lawsuit) alleging that the Work 84 | or a Contribution incorporated within the Work constitutes direct 85 | or contributory patent infringement, then any patent licenses 86 | granted to You under this License for that Work shall terminate 87 | as of the date such litigation is filed. 88 | 89 | 4. Redistribution. You may reproduce and distribute copies of the 90 | Work or Derivative Works thereof in any medium, with or without 91 | modifications, and in Source or Object form, provided that You 92 | meet the following conditions: 93 | 94 | (a) You must give any other recipients of the Work or 95 | Derivative Works a copy of this License; and 96 | 97 | (b) You must cause any modified files to carry prominent notices 98 | stating that You changed the files; and 99 | 100 | (c) You must retain, in the Source form of any Derivative Works 101 | that You distribute, all copyright, patent, trademark, and 102 | attribution notices from the Source form of the Work, 103 | excluding those notices that do not pertain to any part of 104 | the Derivative Works; and 105 | 106 | (d) If the Work includes a "NOTICE" text file as part of its 107 | distribution, then any Derivative Works that You distribute must 108 | include a readable copy of the attribution notices contained 109 | within such NOTICE file, excluding those notices that do not 110 | pertain to any part of the Derivative Works, in at least one 111 | of the following places: within a NOTICE text file distributed 112 | as part of the Derivative Works; within the Source form or 113 | documentation, if provided along with the Derivative Works; or, 114 | within a display generated by the Derivative Works, if and 115 | wherever such third-party notices normally appear. The contents 116 | of the NOTICE file are for informational purposes only and 117 | do not modify the License. You may add Your own attribution 118 | notices within Derivative Works that You distribute, alongside 119 | or as an addendum to the NOTICE text from the Work, provided 120 | that such additional attribution notices cannot be construed 121 | as modifying the License. 122 | 123 | You may add Your own copyright statement to Your modifications and 124 | may provide additional or different license terms and conditions 125 | for use, reproduction, or distribution of Your modifications, or 126 | for any such Derivative Works as a whole, provided Your use, 127 | reproduction, and distribution of the Work otherwise complies with 128 | the conditions stated in this License. 129 | 130 | 5. Submission of Contributions. Unless You explicitly state otherwise, 131 | any Contribution intentionally submitted for inclusion in the Work 132 | by You to the Licensor shall be under the terms and conditions of 133 | this License, without any additional terms or conditions. 134 | Notwithstanding the above, nothing herein shall supersede or modify 135 | the terms of any separate license agreement you may have executed 136 | with Licensor regarding such Contributions. 137 | 138 | 6. Trademarks. This License does not grant permission to use the trade 139 | names, trademarks, service marks, or product names of the Licensor, 140 | except as required for reasonable and customary use in describing the 141 | origin of the Work and reproducing the content of the NOTICE file. 142 | 143 | 7. Disclaimer of Warranty. Unless required by applicable law or 144 | agreed to in writing, Licensor provides the Work (and each 145 | Contributor provides its Contributions) on an "AS IS" BASIS, 146 | WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or 147 | implied, including, without limitation, any warranties or conditions 148 | of TITLE, NON-INFRINGEMENT, MERCHANTABILITY, or FITNESS FOR A 149 | PARTICULAR PURPOSE. You are solely responsible for determining the 150 | appropriateness of using or redistributing the Work and assume any 151 | risks associated with Your exercise of permissions under this License. 152 | 153 | 8. Limitation of Liability. In no event and under no legal theory, 154 | whether in tort (including negligence), contract, or otherwise, 155 | unless required by applicable law (such as deliberate and grossly 156 | negligent acts) or agreed to in writing, shall any Contributor be 157 | liable to You for damages, including any direct, indirect, special, 158 | incidental, or consequential damages of any character arising as a 159 | result of this License or out of the use or inability to use the 160 | Work (including but not limited to damages for loss of goodwill, 161 | work stoppage, computer failure or malfunction, or any and all 162 | other commercial damages or losses), even if such Contributor 163 | has been advised of the possibility of such damages. 164 | 165 | 9. Accepting Warranty or Additional Liability. While redistributing 166 | the Work or Derivative Works thereof, You may choose to offer, 167 | and charge a fee for, acceptance of support, warranty, indemnity, 168 | or other liability obligations and/or rights consistent with this 169 | License. However, in accepting such obligations, You may act only 170 | on Your own behalf and on Your sole responsibility, not on behalf 171 | of any other Contributor, and only if You agree to indemnify, 172 | defend, and hold each Contributor harmless for any liability 173 | incurred by, or claims asserted against, such Contributor by reason 174 | of your accepting any such warranty or additional liability. 175 | 176 | END OF TERMS AND CONDITIONS 177 | 178 | APPENDIX: How to apply the Apache License to your work. 179 | 180 | To apply the Apache License to your work, attach the following 181 | boilerplate notice, with the fields enclosed by brackets "[]" 182 | replaced with your own identifying information. (Don't include 183 | the brackets!) The text should be enclosed in the appropriate 184 | comment syntax for the file format. We also recommend that a 185 | file or class name and description of purpose be included on the 186 | same "printed page" as the copyright notice for easier 187 | identification within third-party archives. 188 | 189 | Copyright [yyyy] [name of copyright owner] 190 | 191 | Licensed under the Apache License, Version 2.0 (the "License"); 192 | you may not use this file except in compliance with the License. 193 | You may obtain a copy of the License at 194 | 195 | http://www.apache.org/licenses/LICENSE-2.0 196 | 197 | Unless required by applicable law or agreed to in writing, software 198 | distributed under the License is distributed on an "AS IS" BASIS, 199 | WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 200 | See the License for the specific language governing permissions and 201 | limitations under the License. 202 | -------------------------------------------------------------------------------- /README.md: -------------------------------------------------------------------------------- 1 | # The MPACT Project 2 | 3 | ## Introduction 4 | 5 | The MPACT project's main objective is to dramatically reduce the effort 6 | required to create highly optimizing HPC and ML compilers for a large 7 | class of architectures using LLVM and MLIR. We do this by providing 8 | a declarative language-based mechanism for collecting and expressing 9 | critical aspects of a target architecture in a way that can be reasoned 10 | about and leveraged by all passes in both MLIR and LLVM. 11 | 12 | ## Installing the MPACT compiler 13 | 14 | To install the MPACT compiler through [PyPI](https://pypi.org/project/mpact/), please use the following command: 15 | 16 | ```shell 17 | pip install mpact 18 | ``` 19 | 20 | ## Building the MPACT compiler 21 | 22 | To build and run the MPACT compiler from source (for developers), 23 | please follow the steps below. 24 | 25 | ### Check out code and sync submodules 26 | 27 | Use the following commands to clone the MPACT compiler repository. 28 | 29 | ```shell 30 | git clone https://github.com/MPACT-ORG/mpact-compiler.git 31 | cd mpact-compiler 32 | git submodule update --init --recursive --progress 33 | ``` 34 | 35 | To always get updated submodules through `git pull`, set the following flag: 36 | 37 | ```shell 38 | git config --global submodule.recurse true 39 | ``` 40 | 41 | NOTE: All following commands assume you remain in the `mpact-compiler` directory. 42 | 43 | ### Setup Python virtual environment 44 | 45 | The following commands initialize a virtual environment under bash/sh/etc. For other shells, see Note 1, [below](README.md#notes). 46 | 47 | ```shell 48 | python3.11 -m venv mpact_venv # one time set up 49 | source mpact_venv/bin/activate # MUST BE REPEATED FOR EVERY SESSION 50 | ``` 51 | 52 | Next, set the Python paths as follows; for shells not in the bash/sh family, see Note 2, [below](README.md#notes). 53 | ```shell 54 | export PYTHONPATH=`pwd`/build/tools/mpact/python_packages/mpact 55 | ``` 56 | 57 | ### Install build requirements 58 | 59 | Note that currently we rely on `torch-mlir` requirements defined in that 60 | submodule to ensure all the build requirements are consistent. 61 | 62 | ```shell 63 | python -m pip install --upgrade pip 64 | python -m pip install -r externals/torch-mlir/requirements.txt 65 | python -m pip install -r externals/torch-mlir/torchvision-requirements.txt 66 | ``` 67 | For shells not in the bash/sh family, see Note 3, [below](README.md#notes). 68 | 69 | ### Building the MPACT compiler in-tree 70 | 71 | The following command generates configuration files to build the MPACT compiler 72 | project completely *in-tree*, which means that both LLVM as well as torch-mlir 73 | are built from source. 74 | 75 | ```shell 76 | cmake -GNinja -Bbuild \ 77 | -DCMAKE_BUILD_TYPE=Release \ 78 | -DPython3_FIND_VIRTUALENV=ONLY \ 79 | -DLLVM_ENABLE_PROJECTS=mlir \ 80 | -DLLVM_EXTERNAL_PROJECTS="torch-mlir;mpact" \ 81 | -DLLVM_EXTERNAL_TORCH_MLIR_SOURCE_DIR="${PWD}/externals/torch-mlir" \ 82 | -DLLVM_EXTERNAL_MPACT_SOURCE_DIR="${PWD}" \ 83 | -DLLVM_TARGETS_TO_BUILD=host \ 84 | -DMLIR_ENABLE_BINDINGS_PYTHON=ON \ 85 | externals/torch-mlir/externals/llvm-project/llvm 86 | ``` 87 | 88 | To speed up the build process, you can set up [ccache](https://ccache.dev/download.html) and add the following flags to the command above: 89 | 90 | ```shell 91 | -DCMAKE_C_COMPILER_LAUNCHER=ccache -DCMAKE_CXX_COMPILER_LAUNCHER=ccache 92 | ``` 93 | 94 | Run the following to ensure the MPACT compiler builds and runs correctly. 95 | 96 | ```shell 97 | cmake --build build --target check-mpact 98 | ``` 99 | 100 | And the following to run all benchmarks 101 | (see [Benchmarks](benchmark/README.md) for more details). 102 | 103 | ```shell 104 | cmake --build build --target benchmark-mpact 105 | ``` 106 | 107 | ## Performance 108 | 109 | Performance of critical kernels for ML models are tracked for each commit. Any regression > 120% will be notified. Graphs could be found in the [github page](https://mpact-org.github.io/mpact-compiler/dev/bench/). 110 | 111 | ## Notes 112 | 113 | 1. Shells other than bash/sh/etc. require a different `activate` script, as shown. Because the python environment has to be set up for every session, we recommend putting it in your .*sh startup file. 114 | - For csh/tcsh/etc.: 115 | ```shell 116 | source `pwd`/mpact_venv/bin/activate.csh 117 | ``` 118 | - For fish/etc.: 119 | ```shell 120 | source /mpact_venv/bin/activate.fish 121 | ``` 122 | 2. Shells other than bash/sh/etc. set their environment variables differently: 123 | - For csh/tcsh/etc.: 124 | ```shell 125 | setenv PYTHONPATH `pwd`/build/tools/mpact/python_packages/mpact 126 | ``` 127 | 3. If using csh/tcsh/etc., run the following command before trying to build the compiler: 128 | ```shell 129 | rehash 130 | ``` 131 | -------------------------------------------------------------------------------- /benchmark/CMakeLists.txt: -------------------------------------------------------------------------------- 1 | #------------------------------------------------------------------------------- 2 | # The MPACT Compiler Python Benchmarks 3 | #------------------------------------------------------------------------------- 4 | 5 | declare_mlir_python_sources(MPACTBenchmarkPythonSources) 6 | 7 | declare_mlir_python_sources(MPACTBenchmarkPythonSources.BenchmarkSuite 8 | ROOT_DIR "${CMAKE_CURRENT_SOURCE_DIR}/python" 9 | ADD_TO_PARENT MPACTBenchmarkPythonSources 10 | SOURCES_GLOB 11 | benchmarks/*.py 12 | ) 13 | 14 | declare_mlir_python_sources(MPACTBenchmarkPythonSources.BenchmarkUtils 15 | ROOT_DIR "${CMAKE_CURRENT_SOURCE_DIR}/python" 16 | ADD_TO_PARENT MPACTBenchmarkPythonSources 17 | SOURCES_GLOB 18 | utils/*.py 19 | ) 20 | 21 | #------------------------------------------------------------------------------- 22 | # Python Modules 23 | #------------------------------------------------------------------------------- 24 | 25 | add_mlir_python_modules(MPACTBenchmarkPythonModules 26 | ROOT_PREFIX "${MPACT_PYTHON_PACKAGES_DIR}/mpact/mpact_benchmark" 27 | INSTALL_PREFIX "python_packages/mpact/mpact_benchmark" 28 | DECLARED_SOURCES MPACTBenchmarkPythonSources 29 | ) 30 | 31 | add_custom_target(build-benchmark-mpact) 32 | add_dependencies(build-benchmark-mpact MPACTPythonModules MPACTBenchmarkPythonModules) 33 | 34 | add_custom_target(benchmark-mpact) 35 | add_dependencies(benchmark-mpact build-benchmark-mpact) 36 | file(GLOB PYTHON_FILES "${CMAKE_CURRENT_SOURCE_DIR}/python/benchmarks/*.py") 37 | 38 | # Loop over each matched .py file and create a custom command to run it. 39 | foreach(PY_FILE IN LISTS PYTHON_FILES) 40 | add_custom_command( 41 | TARGET benchmark-mpact 42 | COMMAND cmake -E echo "Running ${PY_FILE}" 43 | COMMAND python ${PY_FILE} 44 | DEPENDS ${PY_FILE} 45 | USES_TERMINAL 46 | ) 47 | endforeach() 48 | -------------------------------------------------------------------------------- /benchmark/README.md: -------------------------------------------------------------------------------- 1 | ### Run benchmarks 2 | 3 | To run all benchmarks: 4 | 5 | ```shell 6 | cmake --build build --target benchmark-mpact 7 | ``` 8 | 9 | To run selected benchmarks, build the benchmark modules first: 10 | 11 | ```shell 12 | cmake --build build --target build-benchmark-mpact 13 | ``` 14 | 15 | And then run the benchmark file: 16 | 17 | ```shell 18 | python path/to/the/_benchmark.py 19 | ``` 20 | 21 | If you would like to run selected kernels in kernels_benchmark.py, 22 | you can use `--benchmark-filter` flag like the following example: 23 | 24 | ```shell 25 | python path/to/the/kernels_benchmark.py --benchmark-filter=add 26 | ``` 27 | 28 | ### Profiler 29 | 30 | Utils for profiling python scripts and pytorch models could be found in 31 | `benchmark/python/utils/profiler.py`. 32 | 33 | To run the profiling example, use the following command: 34 | 35 | ```shell 36 | python benchmark/python/utils/profiler.py 37 | ``` 38 | -------------------------------------------------------------------------------- /benchmark/python/benchmarks/gcn_benchmark.py: -------------------------------------------------------------------------------- 1 | import torch 2 | import numpy as np 3 | from mpact.models.gcn import GraphConv 4 | from mpact_benchmark.utils.benchmark_utils import benchmark, Backends 5 | 6 | 7 | @benchmark( 8 | [ 9 | { 10 | "name": f"{fmt}_{shape}_{dtype.__name__}", 11 | "shape": shape, 12 | "formats": fmt, 13 | "dtype": dtype, 14 | "drange": (1, 100), 15 | "sparsity": [0, 0.5, 0.9, 0.99], 16 | "backends": [b for b in Backends], 17 | } 18 | for shape in [ 19 | [[128, 128], [128, 128]], 20 | [[512, 512], [512, 512]], 21 | [[1024, 1024], [1024, 1024]], 22 | ] 23 | for fmt in [["dense", "csr"]] 24 | for dtype in [np.float32] 25 | ] 26 | ) 27 | def GCN() -> torch.nn.Module: 28 | """Graph Convolution Network.""" 29 | return GraphConv 30 | 31 | 32 | if __name__ == "__main__": 33 | GCN() 34 | -------------------------------------------------------------------------------- /benchmark/python/benchmarks/kernels_benchmark.py: -------------------------------------------------------------------------------- 1 | import torch 2 | import argparse 3 | import numpy as np 4 | from mpact.models.kernels import * 5 | from mpact_benchmark.utils.benchmark_utils import benchmark, Backends 6 | 7 | 8 | @benchmark( 9 | [ 10 | { 11 | "name": f"{lhs_fmt}_{rhs_fmt}_{shape}_{dtype.__name__}", 12 | "shape": shape, 13 | "formats": (lhs_fmt, rhs_fmt), 14 | "dtype": dtype, 15 | "backends": [b for b in Backends], 16 | "drange": (1, 100), 17 | "sparsity": [0, 0.5, 0.9, 0.99], 18 | } 19 | for shape in [([2**i, 2**i], [2**i, 2**i]) for i in range(5, 8)] 20 | for lhs_fmt in ["dense", "csr"] 21 | for rhs_fmt in ["dense", "csr"] 22 | for dtype in [np.float64] 23 | ] 24 | ) 25 | def matmul() -> torch.nn.Module: 26 | """Matrix multiplication.""" 27 | return MMNet() 28 | 29 | 30 | @benchmark( 31 | [ 32 | { 33 | "name": f"{lhs_fmt}_{rhs_fmt}_{shape}_{dtype.__name__}", 34 | "shape": shape, 35 | "formats": (lhs_fmt, rhs_fmt), 36 | "dtype": dtype, 37 | "backends": [b for b in Backends], 38 | "drange": (1, 100), 39 | "sparsity": [0, 0.5, 0.9, 0.99], 40 | } 41 | for shape in [([2**i, 2**i], [2**i]) for i in range(5, 8)] 42 | for lhs_fmt in ["dense", "csr"] 43 | for rhs_fmt in ["dense"] # torch.mv only supports dense vector for now. 44 | for dtype in [np.float64] 45 | ] 46 | ) 47 | def matvec() -> torch.nn.Module: 48 | """Matrix-vector multiplication.""" 49 | return MVNet() 50 | 51 | 52 | @benchmark( 53 | [ 54 | { 55 | "name": f"{lhs_fmt}_{rhs_fmt}_{shape}_{dtype.__name__}", 56 | "shape": shape, 57 | "formats": (lhs_fmt, rhs_fmt), 58 | "dtype": dtype, 59 | "backends": [b for b in Backends], 60 | "drange": (1, 100), 61 | "sparsity": [0, 0.5, 0.9, 0.99], 62 | } 63 | for shape in [ 64 | ([2**i, 2**i], [2**i, 2**i]) for i in range(5, 8) 65 | ] # 512x512 crashes runtime. 66 | for lhs_fmt in ["dense", "csr"] 67 | for rhs_fmt in ["dense", "csr"] 68 | for dtype in [np.float64] 69 | ] 70 | ) 71 | def add() -> torch.nn.Module: 72 | """Element-wise addition.""" 73 | return AddNet() 74 | 75 | 76 | @benchmark( 77 | [ 78 | { 79 | "name": f"{lhs_fmt}_{rhs_fmt}_{shape}_{dtype.__name__}", 80 | "shape": shape, 81 | "formats": (lhs_fmt, rhs_fmt), 82 | "dtype": dtype, 83 | "backends": [b for b in Backends], 84 | "drange": (1, 100), 85 | "sparsity": [0, 0.5, 0.9, 0.99], 86 | } 87 | for shape in [([2**i, 2**i], [2**i, 2**i]) for i in range(5, 8)] 88 | for lhs_fmt in ["dense", "csr"] 89 | for rhs_fmt in ["dense", "csr"] 90 | for dtype in [np.float64] 91 | ] 92 | ) 93 | def elt_mul() -> torch.nn.Module: 94 | """Element-wise addition.""" 95 | return MulNet() 96 | 97 | 98 | @benchmark( 99 | [ 100 | { 101 | "name": f"{fmt}_{shape}_{dtype.__name__}", 102 | "shape": shape, 103 | "formats": (fmt,), 104 | "dtype": dtype, 105 | "backends": [b for b in Backends], 106 | "drange": (1, 100), 107 | "sparsity": [0, 0.5, 0.9, 0.99], 108 | } 109 | for shape in [([2**i, 2**i],) for i in range(2, 3)] 110 | for fmt in ["dense", "csr"] 111 | for dtype in [np.float64] 112 | ] 113 | ) 114 | def nop() -> torch.nn.Module: 115 | """Returns matrix unmodified (speed of light).""" 116 | return SelfNet() 117 | 118 | 119 | @benchmark( 120 | [ 121 | { 122 | "name": f"{sample_fmt}_sample_{shape}_{dtype.__name__}", 123 | "shape": shape, 124 | "formats": (sample_fmt, "dense", "dense"), 125 | "dtype": dtype, 126 | "backends": [b for b in Backends], 127 | "drange": (1, 100), 128 | "sparsity": [0, 0.5, 0.9, 0.99], 129 | } 130 | for shape in [ 131 | ([2**i, 2**i], [2**i, 2**i], [2**i, 2**i]) for i in range(5, 8) 132 | ] 133 | for sample_fmt in ["dense", "csr"] 134 | for dtype in [np.float64] 135 | ] 136 | ) 137 | def sddmm() -> torch.nn.Module: 138 | """SDDMM: C = S ○ (A X B) Sampled dense-dense matrix-matrix multiplication.""" 139 | return SDDMMNet() 140 | 141 | 142 | @benchmark( 143 | [ 144 | { 145 | "name": f"{fmt}_{shape}_{dtype.__name__}", 146 | "shape": shape, 147 | "formats": (fmt,), 148 | "dtype": dtype, 149 | # TODO: add mpact and torch inductor once they work. 150 | "backends": [ 151 | b 152 | for b in Backends 153 | if b.value 154 | in ( 155 | Backends.TORCH_SPARSE_EAGER.value, 156 | Backends.TORCH_DENSE_EAGER.value, 157 | ) 158 | ], 159 | "drange": (1, 100), 160 | "sparsity": [0, 0.5, 0.9, 0.99], 161 | } 162 | for shape in [([2**i, 2**i],) for i in range(5, 8)] 163 | for fmt in ["dense"] 164 | for dtype in [np.float64] 165 | ] 166 | ) 167 | def feature_scale() -> torch.nn.Module: 168 | """Scales feature matrix in GNN.""" 169 | return FeatureScale() 170 | 171 | 172 | @benchmark( 173 | [ 174 | { 175 | "name": f"{fmt}_{shape}_{dtype.__name__}", 176 | "shape": shape, 177 | "formats": (fmt,), 178 | "dtype": dtype, 179 | # TODO: add mpact and torch inductor once they work. 180 | "backends": [ 181 | b 182 | for b in Backends 183 | if b.value 184 | in ( 185 | Backends.TORCH_SPARSE_EAGER.value, 186 | Backends.TORCH_DENSE_EAGER.value, 187 | ) 188 | ], 189 | "drange": (1, 100), 190 | "sparsity": [0, 0.5, 0.9, 0.99], 191 | } 192 | for shape in [([2**i, 2**i],) for i in range(5, 8)] 193 | for fmt in ["dense"] 194 | for dtype in [np.float64] 195 | ] 196 | ) 197 | def normalization() -> torch.nn.Module: 198 | """Normalizes adjacency matrix in GNN.""" 199 | return Normalization() 200 | 201 | 202 | if __name__ == "__main__": 203 | parser = argparse.ArgumentParser( 204 | prog="pytorch_kernel_benchmarks", 205 | description="Run a set of given PyTorch kernel benchmarks", 206 | ) 207 | parser.add_argument("--benchmark-filter", type=str, default="", required=False) 208 | arguments = parser.parse_args() 209 | 210 | benchmark_list = [ 211 | "nop", 212 | "add", 213 | "matmul", 214 | "matvec", 215 | "elt_mul", 216 | "sddmm", 217 | "feature_scale", 218 | "normalization", 219 | ] 220 | if arguments.benchmark_filter: 221 | benchmark_list = arguments.benchmark_filter.split(",") 222 | 223 | # Run selected benchmarks. 224 | for benchmark_name in benchmark_list: 225 | globals()[benchmark_name]() 226 | -------------------------------------------------------------------------------- /benchmark/python/benchmarks/lif_benchmark.py: -------------------------------------------------------------------------------- 1 | import torch 2 | import numpy as np 3 | from mpact.models.lif import LIFSumOfSq 4 | from mpact_benchmark.utils.benchmark_utils import benchmark, Backends 5 | 6 | 7 | @benchmark( 8 | [ 9 | { 10 | "name": f"{fmt}_{shape}_{dtype.__name__}", 11 | "shape": shape, 12 | "formats": fmt, 13 | "dtype": dtype, 14 | # Simulate batch normalization. 15 | "drange": (-1, 1), 16 | "sparsity": [0, 0.5, 0.9, 0.99], 17 | # to_dense() in LIF prop hack is not supported in torch inductor. 18 | # TODO: add torch inductor once prop hack is no longer needed. 19 | "backends": [ 20 | b 21 | for b in Backends 22 | if b.value 23 | not in ( 24 | Backends.TORCH_SPARSE_INDUCTOR.value, 25 | Backends.TORCH_DENSE_INDUCTOR.value, 26 | ) 27 | ], 28 | } 29 | for shape in [ 30 | [[64, 3, 32, 32, 1]], 31 | [[32, 3, 64, 64, 1]], 32 | [[16, 3, 224, 224, 1]], 33 | ] 34 | for fmt in [["dense"]] 35 | for dtype in [np.float64] 36 | ] 37 | ) 38 | def LifSumOfSq() -> torch.nn.Module: 39 | """LIF feeding into sum of squares.""" 40 | return LIFSumOfSq() 41 | 42 | 43 | if __name__ == "__main__": 44 | LifSumOfSq() 45 | -------------------------------------------------------------------------------- /benchmark/python/benchmarks/regression_benchmark.py: -------------------------------------------------------------------------------- 1 | import pytest 2 | from mpact.models.kernels import * 3 | from mpact_benchmark.utils.tensor_generator import generate_tensor 4 | 5 | SHAPE = (1024, 1024) 6 | SPARSITY = 0.8 7 | 8 | dense_tensor1 = generate_tensor(0, SHAPE, SPARSITY) 9 | dense_tensor2 = generate_tensor(1, SHAPE, SPARSITY) 10 | dense_tensor3 = generate_tensor(2, SHAPE, SPARSITY) 11 | dense_vector = generate_tensor(1, (SHAPE[0],), SPARSITY) 12 | 13 | sparse_tensor1 = dense_tensor1.to_sparse_csr() 14 | sparse_tensor2 = dense_tensor2.to_sparse_csr() 15 | sparse_tensor3 = dense_tensor3.to_sparse_csr() 16 | 17 | 18 | def test_mv_dense(benchmark): 19 | benchmark(MVNet(), dense_tensor1, dense_vector) 20 | 21 | 22 | def test_mm_dense(benchmark): 23 | benchmark(MMNet(), dense_tensor1, dense_tensor2) 24 | 25 | 26 | def test_add_dense(benchmark): 27 | benchmark(AddNet(), dense_tensor1, dense_tensor2) 28 | 29 | 30 | def test_mul_dense(benchmark): 31 | benchmark(MulNet(), dense_tensor1, dense_tensor2) 32 | 33 | 34 | def test_nop_dense(benchmark): 35 | benchmark(SelfNet(), dense_tensor1) 36 | 37 | 38 | def test_sddmm_dense(benchmark): 39 | benchmark(SDDMMNet(), dense_tensor1, dense_tensor2, dense_tensor3) 40 | 41 | 42 | def test_mv_sparse(benchmark): 43 | benchmark(MVNet(), sparse_tensor1, dense_vector) 44 | 45 | 46 | def test_mm_sparse(benchmark): 47 | benchmark(MMNet(), sparse_tensor1, sparse_tensor2) 48 | 49 | 50 | def test_add_sparse(benchmark): 51 | benchmark(AddNet(), sparse_tensor1, sparse_tensor2) 52 | 53 | 54 | def test_mul_sparse(benchmark): 55 | benchmark(MulNet(), sparse_tensor1, sparse_tensor2) 56 | 57 | 58 | def test_nop_sparse(benchmark): 59 | benchmark(SelfNet(), sparse_tensor1) 60 | 61 | 62 | def test_sddmm_sparse(benchmark): 63 | benchmark(SDDMMNet(), sparse_tensor1, dense_tensor2, dense_tensor3) 64 | -------------------------------------------------------------------------------- /benchmark/python/benchmarks/resnet_benchmark.py: -------------------------------------------------------------------------------- 1 | import torch 2 | import numpy as np 3 | from mpact.models.resnet import resnet_20 4 | from mpact_benchmark.utils.benchmark_utils import benchmark, Backends 5 | 6 | 7 | @benchmark( 8 | [ 9 | { 10 | "name": f"{fmt}_{shape}_{dtype.__name__}", 11 | "shape": shape, 12 | "formats": fmt, 13 | "dtype": dtype, 14 | "drange": (1, 100), 15 | "sparsity": [0.5, 0.9], 16 | # TODO: Torch inductor requires lower precision with larger input size, 17 | # such as [8, 3, 32, 32]. 18 | "precision": 1e-3, 19 | "backends": [b for b in Backends], 20 | } 21 | for shape in [ 22 | [[1, 3, 16, 16]], 23 | ] 24 | for fmt in [["dense"]] 25 | for dtype in [np.float32] 26 | ] 27 | ) 28 | def resnet() -> torch.nn.Module: 29 | """Restnet20 model.""" 30 | resnet_model = resnet_20() 31 | resnet_model.train(False) 32 | return resnet_model 33 | 34 | 35 | if __name__ == "__main__": 36 | resnet() 37 | -------------------------------------------------------------------------------- /benchmark/python/manual/README.md: -------------------------------------------------------------------------------- 1 | ### Benchmarks run by hand 2 | 3 | These benchmarks are not run as part of MPACT's regular testing or benchmarking. 4 | To run an individual test, build the MPACT compiler, cd into this directory, 5 | and then simply run a benchmark as follows: 6 | 7 | ```shell 8 | python .py 9 | ``` 10 | -------------------------------------------------------------------------------- /benchmark/python/manual/sum_of_sq.py: -------------------------------------------------------------------------------- 1 | import torch 2 | import numpy as np 3 | import time 4 | 5 | from mpact.mpactbackend import mpact_jit_compile, mpact_jit_run 6 | from mpact_benchmark.utils.tensor_generator import generate_tensor 7 | 8 | 9 | def runbench_eager(tag, sp, net, x, num_iters=1000): 10 | net(x) # warmup 11 | checksum = 0 12 | start = time.time() 13 | for i in range(num_iters): 14 | res = net(x).item() 15 | checksum = checksum + res 16 | end = time.time() 17 | time_ms = (end - start) * 1000 / num_iters 18 | print("%s : %.2f : %8.4f ms. : checksum=%d" % (tag, sp, time_ms, checksum)) 19 | 20 | 21 | def runbench_mpact(tag, sp, net, x, num_iters=1000): 22 | invoker, fn = mpact_jit_compile(net, x) 23 | mpact_jit_run(invoker, fn, x) # warmup 24 | checksum = 0 25 | start = time.time() 26 | for i in range(num_iters): 27 | res = mpact_jit_run(invoker, fn, x) 28 | checksum = checksum + res 29 | end = time.time() 30 | time_ms = (end - start) * 1000 / num_iters 31 | print("%s : %.2f : %8.4f ms. : checksum=%d" % (tag, sp, time_ms, checksum)) 32 | 33 | 34 | class SqSumNet(torch.nn.Module): 35 | def forward(self, x): 36 | # TODO: make this work too: return (x ** 2).sum() 37 | return (x * x).sum() 38 | 39 | 40 | net = SqSumNet() 41 | h = 1024 * 4 42 | w = 1024 * 4 43 | 44 | for d in range(0, 101, 10): 45 | sparsity = 1.0 - (d / 100.0) 46 | x = generate_tensor( 47 | seed=0, shape=(h, w), sparsity=sparsity, dtype=np.float32, drange=(1.0, 1.0) 48 | ) 49 | 50 | # Note, we don't have binary-valued sparse tensors in PyTorch 51 | # so we are using csr. For now, we have to hack the 52 | # "explicitVal=1.0:f32" 53 | # into the MLIR sparse tensor type to make optimize it fully. 54 | s = x.to_sparse_csr() 55 | 56 | runbench_eager("PyTorch (dense) ", sparsity, net, x) 57 | runbench_mpact("MPACT (dense) ", sparsity, net, x) 58 | runbench_eager("PyTorch (sparse)", sparsity, net, s) 59 | runbench_mpact("MPACT (sparse)", sparsity, net, s) 60 | -------------------------------------------------------------------------------- /benchmark/python/utils/benchmark_utils.py: -------------------------------------------------------------------------------- 1 | import functools 2 | import torch 3 | from enum import Enum 4 | from typing import Any, Callable 5 | from torch.utils import benchmark as torch_benchmark 6 | from mpact.mpactbackend import mpact_jit, mpact_jit_compile, mpact_jit_run 7 | from mpact_benchmark.utils.tensor_generator import generate_inputs 8 | 9 | 10 | class Backends(Enum): 11 | TORCH_SPARSE_EAGER = 1 12 | TORCH_DENSE_EAGER = 2 13 | TORCH_SPARSE_INDUCTOR = 3 14 | TORCH_DENSE_INDUCTOR = 4 15 | MPACT_SPARSE = 5 16 | MPACT_DENSE = 6 17 | 18 | 19 | def timer(stmt: str, description: str, setup: str = "", **kwargs: Any) -> Any: 20 | """Timer for benchmark.""" 21 | return torch_benchmark.Timer( 22 | stmt=stmt, 23 | globals=kwargs["variables"], 24 | setup=setup, 25 | num_threads=1, 26 | label=kwargs["variables"]["label"], 27 | sub_label=kwargs["variables"]["sub_label"], 28 | description=description, 29 | ).adaptive_autorange() 30 | 31 | 32 | def get_dynamo_compile_time(sub_label: str, label: str, description: str) -> Any: 33 | """Get compile time from dynamo and create a benchmark measurement object.""" 34 | try: 35 | compile_time = torch_benchmark.Measurement( 36 | 1, 37 | [ 38 | float( 39 | torch._dynamo.utils.compile_times(repr="csv")[1][0] 40 | .split(",")[-1] 41 | .strip() 42 | ) 43 | ], 44 | torch_benchmark.TaskSpec( 45 | sub_label, 46 | None, 47 | description=description, 48 | label=label, 49 | ), 50 | ) 51 | return compile_time 52 | except ValueError: 53 | print(f"No compilation happened for {description}: {sub_label}.") 54 | return None 55 | 56 | 57 | def run_benchmark( 58 | sparse_inputs: tuple[torch.Tensor, ...], 59 | dense_inputs: tuple[torch.Tensor, ...], 60 | torch_net: torch.nn.Module, 61 | variables: dict[str, Any], 62 | backends: tuple[Backends, ...], 63 | runtime_results: list[torch_benchmark.Measurement], 64 | compile_time_results: list[torch_benchmark.Measurement], 65 | ): 66 | """Run benchmark with specified backends.""" 67 | output = [] 68 | output_type = None 69 | 70 | with torch.no_grad(): 71 | for backend in backends: 72 | match backend: 73 | case Backends.TORCH_SPARSE_EAGER: 74 | sparse_out = torch_net(*sparse_inputs) 75 | output_type = sparse_out.layout 76 | output.append(sparse_out) 77 | runtime_results.append( 78 | timer( 79 | "torch_net(*sparse_inputs)", 80 | "torch-sparse-eager", 81 | variables=variables, 82 | ) 83 | ) 84 | case Backends.TORCH_DENSE_EAGER: 85 | output.append(torch_net(*dense_inputs)) 86 | runtime_results.append( 87 | timer( 88 | "torch_net(*dense_inputs)", 89 | "torch-dense-eager", 90 | variables=variables, 91 | ) 92 | ) 93 | case Backends.TORCH_SPARSE_INDUCTOR: 94 | torch_inductor = torch.compile(torch_net) 95 | torch_out = torch_inductor(*sparse_inputs) 96 | output.append(torch_out) 97 | compile_time = get_dynamo_compile_time( 98 | variables["sub_label"], 99 | variables["label"], 100 | "torch-sparse-inductor-compile", 101 | ) 102 | if compile_time: 103 | compile_time_results.append(compile_time) 104 | runtime_results.append( 105 | timer( 106 | "torch_inductor(*sparse_inputs)", 107 | "torch-sparse-inductor-runtime", 108 | variables=dict(variables, **locals()), 109 | ) 110 | ) 111 | case Backends.TORCH_DENSE_INDUCTOR: 112 | torch_inductor = torch.compile(torch_net) 113 | output.append(torch_inductor(*dense_inputs)) 114 | compile_time = get_dynamo_compile_time( 115 | variables["sub_label"], 116 | variables["label"], 117 | "torch-dense-inductor-compile", 118 | ) 119 | if compile_time: 120 | compile_time_results.append(compile_time) 121 | runtime_results.append( 122 | timer( 123 | "torch_inductor(*dense_inputs)", 124 | "torch-dense-inductor-runtime", 125 | variables=dict(variables, **locals()), 126 | ) 127 | ) 128 | case Backends.MPACT_SPARSE: 129 | sp_out = mpact_jit(torch_net, *sparse_inputs) 130 | # Construct sparse csr tensor if the output type is csr. 131 | # TODO: return sparse tensor directly instead of a tuple of arrays. 132 | if type(sp_out) is tuple: 133 | # torch.sparse_csr_tensor could deduce the size incorrectly, 134 | # so pass the dense_out's shape explicitly. 135 | dense_out = mpact_jit(torch_net, *dense_inputs) 136 | output.append( 137 | torch.sparse_csr_tensor(*sp_out, size=dense_out.shape) 138 | ) 139 | # Check MPACT and torch eager both return sparse csr output 140 | # only when torch sparse eager has been run. 141 | if output_type: 142 | assert output_type == torch.sparse_csr 143 | else: 144 | output.append(torch.from_numpy(sp_out)) 145 | # Check MPACT and torch eager both return dense output 146 | # only when torch sparse eager has been run. 147 | if output_type: 148 | assert output_type == torch.strided 149 | invoker, f = mpact_jit_compile(torch_net, *sparse_inputs) 150 | compile_time_results.append( 151 | timer( 152 | "mpact_jit_compile(torch_net, *sparse_inputs)", 153 | "mpact-sparse-compile", 154 | "from mpact.mpactbackend import mpact_jit_compile", 155 | variables=dict(variables, **locals()), 156 | ) 157 | ) 158 | runtime_results.append( 159 | timer( 160 | "mpact_jit_run(invoker, f, *sparse_inputs)", 161 | "mpact-sparse-runtime", 162 | "from mpact.mpactbackend import mpact_jit_run", 163 | variables=dict(variables, **locals()), 164 | ) 165 | ) 166 | case Backends.MPACT_DENSE: 167 | output.append(torch.from_numpy(mpact_jit(torch_net, *dense_inputs))) 168 | invoker, f = mpact_jit_compile(torch_net, *dense_inputs) 169 | compile_time_results.append( 170 | timer( 171 | "mpact_jit_compile(torch_net, *dense_inputs)", 172 | "mpact-dense-compile", 173 | "from mpact.mpactbackend import mpact_jit_compile", 174 | variables=dict(variables, **locals()), 175 | ) 176 | ) 177 | runtime_results.append( 178 | timer( 179 | "mpact_jit_run(invoker, f, *dense_inputs)", 180 | "mpact-dense-runtime", 181 | "from mpact.mpactbackend import mpact_jit_run", 182 | variables=dict(variables, **locals()), 183 | ) 184 | ) 185 | case _: 186 | print(f"{backend} is not supported yet.") 187 | 188 | # Sanity check. 189 | if output: 190 | rtol = variables["precision"] if "precision" in variables else 1e-5 191 | assert all( 192 | torch.allclose(output[0].to_dense(), out.to_dense(), rtol=rtol) 193 | for out in output 194 | ) 195 | 196 | 197 | def benchmark(*args: Any) -> Callable: 198 | """Wrapper for benchmark.""" 199 | 200 | def decorator(func): 201 | @functools.wraps(func) 202 | def wrapper(test_cases=args[0]): 203 | runtime_results = [] 204 | compile_time_results = [] 205 | torch_net = net = func() 206 | for test_case in test_cases: 207 | label = func.__name__ 208 | for sparsity in test_case["sparsity"]: 209 | sub_label = f"{test_case['name']}_{sparsity}" 210 | dense_inputs, sparse_inputs = generate_inputs( 211 | test_case["shape"], 212 | sparsity, 213 | test_case["formats"], 214 | test_case["dtype"], 215 | test_case["drange"], 216 | ) 217 | 218 | if "GCN" in label: 219 | torch_net = net(*test_case["shape"][0]) 220 | if "precision" in test_case: 221 | precision = test_case["precision"] 222 | 223 | run_benchmark( 224 | sparse_inputs, 225 | dense_inputs, 226 | torch_net, 227 | locals(), 228 | test_case["backends"], 229 | runtime_results, 230 | compile_time_results, 231 | ) 232 | 233 | compare1 = torch_benchmark.Compare(runtime_results) 234 | compare1.print() 235 | compare2 = torch_benchmark.Compare(compile_time_results) 236 | compare2.print() 237 | 238 | return func 239 | 240 | return wrapper 241 | 242 | return decorator 243 | -------------------------------------------------------------------------------- /benchmark/python/utils/profiler.py: -------------------------------------------------------------------------------- 1 | import torch 2 | import cProfile 3 | from pstats import Stats 4 | 5 | 6 | def profile_torch( 7 | func, args, row_limit=10, save_output=False, func_name=None, file_name="trace" 8 | ): 9 | """Use PyTorch's profiler to profile torch ops. 10 | 11 | To see the graph: upload trace.json to chrome://tracing 12 | 13 | More details about PyTorch profiler: 14 | https://pytorch.org/tutorials/recipes/recipes/profiler_recipe.html 15 | """ 16 | func_name = func_name if func_name else func.__name__ 17 | with torch.profiler.profile() as prof: 18 | with torch.profiler.record_function(func_name): 19 | func(*args) 20 | print(prof.key_averages().table(sort_by="cpu_time_total", row_limit=row_limit)) 21 | if save_output: 22 | prof.export_chrome_trace(f"{file_name}.json") 23 | 24 | 25 | def profile_python(func, args, row_limit=10, save_output=False, file_name="stats"): 26 | """Use cProfile to profile python function calls. 27 | 28 | To see the graph, run the following commands: 29 | 1. python -m pip install snakeviz 30 | 2. snakeviz stats.prof 31 | """ 32 | pr = cProfile.Profile() 33 | pr.enable() 34 | func(*args) 35 | pr.disable() 36 | stats = Stats(pr) 37 | stats.sort_stats("tottime").print_stats(row_limit) 38 | if save_output: 39 | pr.dump_stats(f"{file_name}.prof") 40 | 41 | 42 | if __name__ == "__main__": 43 | # Example usage of the profiler. 44 | from mpact.models.kernels import MMNet 45 | from mpact_benchmark.utils.tensor_generator import generate_tensor 46 | from mpact.mpactbackend import mpact_jit 47 | 48 | # Generate input tensors. 49 | dense_tensor1 = generate_tensor(seed=0, shape=(32, 32), sparsity=0.8) 50 | dense_tensor2 = generate_tensor(seed=1, shape=(32, 32), sparsity=0.8) 51 | sparse_tensor1 = dense_tensor1.to_sparse_csr() 52 | sparse_tensor2 = dense_tensor2.to_sparse_csr() 53 | 54 | # Profile with PyTorch profiler for torch operators. 55 | # MPACT sparse. 56 | profile_torch(mpact_jit, (MMNet(), sparse_tensor1, sparse_tensor2)) 57 | # Torch sparse. 58 | profile_torch( 59 | MMNet(), (sparse_tensor1, sparse_tensor2), func_name="sparsexsparse matmul" 60 | ) 61 | 62 | # Profile with cProfile for Python function calls. 63 | # MPACT sparse. 64 | profile_python(mpact_jit, (MMNet(), sparse_tensor1, sparse_tensor2)) 65 | # Torch sparse. 66 | profile_python(MMNet(), (sparse_tensor1, sparse_tensor2)) 67 | -------------------------------------------------------------------------------- /benchmark/python/utils/tensor_generator.py: -------------------------------------------------------------------------------- 1 | import torch 2 | import math 3 | import numpy as np 4 | from typing import Any 5 | 6 | 7 | def generate_inputs( 8 | shapes: tuple[Any, ...], 9 | sparsity: float, 10 | formats: tuple[str, ...], 11 | dtype: Any = np.float64, 12 | drange: tuple[Any, ...] = (1, 100), 13 | ) -> tuple[tuple[torch.Tensor, ...], tuple[torch.Tensor, ...]]: 14 | """Generates dense and sparse tensor inputs. 15 | 16 | Args: 17 | shapes: Shape for each input. 18 | sparsity: Sparsity level for the inputs. 19 | formats: Sparsity format for each input. 20 | dtype: Data type of the generated inputs. Default is np.float64. 21 | drange: Data range of the non-zero values. Default is (1, 100). 22 | 23 | Returns: 24 | dense_inputs: all dense tensors. 25 | sparse_inputs: inputs are of the specified sparsity format, such as CSR. 26 | """ 27 | dense_inputs = [] 28 | sparse_inputs = [] 29 | # Each input has a different seed. 30 | for seed, shape in enumerate(shapes): 31 | dense_inputs.append(generate_tensor(seed, shape, sparsity, dtype, drange)) 32 | for idx, dense_input in enumerate(dense_inputs): 33 | if formats[idx] == "dense": 34 | sparse_inputs.append(dense_input) 35 | else: 36 | # TODO: support more sparsity formats. 37 | sparse_inputs.append(dense_input.to_sparse_csr()) 38 | return dense_inputs, sparse_inputs 39 | 40 | 41 | def generate_tensor( 42 | seed: int, 43 | shape: tuple[Any, ...], 44 | sparsity: float, 45 | dtype: Any = np.float64, 46 | drange: tuple[Any, ...] = (1, 100), 47 | ) -> torch.Tensor: 48 | """Generates a tensor given sparsity level, shape and data type. 49 | 50 | Args: 51 | seed: Seed value for np.random. 52 | shape: A tuple for the shape of tensor. 53 | sparsity: Sparsity level in the range of [0, 1], viz. 0=dense and 1=all-zeros 54 | dtype: Data type of the generated tensor. Default is np.float64. 55 | drange: Data range of the non-zero values (inclusive). Default is (1, 100). 56 | 57 | Returns: 58 | A dense torch tensor with the specified shape, sparsity level and type. 59 | 60 | Note: the tensor generated doesn't guarantee each batch will have the same 61 | number of specified elements. Therefore, for batched CSR, torch.cat can be 62 | used to concatenate generated tensors in the specified dimension. 63 | """ 64 | if sparsity < 0.0 or sparsity > 1.0: 65 | raise ValueError("Invalid sparsity level: %f" % sparsity) 66 | 67 | np.random.seed(seed) 68 | size = math.prod(shape) 69 | nse = size - int(math.ceil(sparsity * size)) 70 | 71 | flat_output = np.zeros(size) 72 | indices = np.random.choice(size, nse, replace=False) 73 | values = np.random.uniform(drange[0], drange[1], nse) 74 | flat_output[indices] = values 75 | 76 | result = np.reshape(flat_output, shape).astype(dtype) 77 | return torch.from_numpy(result) 78 | 79 | 80 | def print_matrix_market_format(tensor: torch.Tensor): 81 | """Prints the matrix market format for a sparse matrix. 82 | 83 | Args: 84 | tensor: sparse matrix (real type) 85 | """ 86 | if len(tensor.shape) != 2: 87 | raise ValueError("Unexpected rank : %d (matrices only)" % len(tensor.shape)) 88 | if tensor.dtype != torch.float32 and tensor.dtype != torch.float64: 89 | raise ValueError("Unexpected type : %s" % tensor.dtype) 90 | 91 | h = tensor.shape[0] 92 | w = tensor.shape[1] 93 | nnz = sum([1 if tensor[i, j] != 0 else 0 for i in range(h) for j in range(w)]) 94 | density = (100.0 * nnz) / tensor.numel() 95 | print("%%MatrixMarket matrix coordinate real general") 96 | print("% https://math.nist.gov/MatrixMarket/formats.html") 97 | print("%") 98 | print("%% density = %4.2f%%" % density) 99 | print("%") 100 | print(h, w, nnz) 101 | for i in range(h): 102 | for j in range(w): 103 | if tensor[i, j] != 0: 104 | print(i + 1, j + 1, tensor[i, j].item()) 105 | 106 | 107 | def print_extended_frostt_format(tensor: torch.Tensor): 108 | """Prints the Extended FROSTT format for a sparse tensor. 109 | 110 | Args: 111 | tensor: sparse tensor 112 | """ 113 | a = tensor.numpy() 114 | nnz = sum([1 if x != 0 else 0 for x in np.nditer(a)]) 115 | density = (100.0 * nnz) / tensor.numel() 116 | print("# Tensor in Extended FROSTT file format") 117 | print("# http://frostt.io/tensors/file-formats.html") 118 | print("# extended with two metadata lines:") 119 | print("# rank nnz") 120 | print("# dims (one per rank)") 121 | print("#") 122 | print("# density = %4.2f%%" % density) 123 | print("#") 124 | print(len(tensor.shape), nnz) 125 | print(*tensor.shape, sep=" ") 126 | it = np.nditer(a, flags=["multi_index"]) 127 | for x in it: 128 | if x != 0: 129 | print(*[i + 1 for i in it.multi_index], sep=" ", end=" ") 130 | print(x) 131 | -------------------------------------------------------------------------------- /include/CMakeLists.txt: -------------------------------------------------------------------------------- 1 | add_subdirectory(mpact) 2 | -------------------------------------------------------------------------------- /include/mpact-c/Registration.h: -------------------------------------------------------------------------------- 1 | /*===-- mpact-c/Registration.h - Registration functions -----*- C -*-===*\ 2 | |* *| 3 | |* Part of the MPACT Project, under the Apache License v2.0 with LLVM *| 4 | |* Exceptions. *| 5 | |* See https://llvm.org/LICENSE.txt for license information. *| 6 | |* SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception *| 7 | |* *| 8 | \*===----------------------------------------------------------------------===*/ 9 | 10 | #ifndef MPACT_C_REGISTRATION_H 11 | #define MPACT_C_REGISTRATION_H 12 | 13 | #include "mlir-c/IR.h" 14 | #include "mlir-c/Support.h" 15 | 16 | #ifdef __cplusplus 17 | extern "C" { 18 | #endif 19 | 20 | /** Registers all passes for symbolic access with the global registry. */ 21 | MLIR_CAPI_EXPORTED void mpactRegisterAllPasses(void); 22 | 23 | #ifdef __cplusplus 24 | } 25 | #endif 26 | 27 | #endif // MPACT_C_REGISTRATION_H 28 | -------------------------------------------------------------------------------- /include/mpact/CMakeLists.txt: -------------------------------------------------------------------------------- 1 | add_subdirectory(Transforms) 2 | -------------------------------------------------------------------------------- /include/mpact/Transforms/CMakeLists.txt: -------------------------------------------------------------------------------- 1 | set(LLVM_TARGET_DEFINITIONS Passes.td) 2 | mlir_tablegen(Passes.h.inc -gen-pass-decls) 3 | add_public_tablegen_target(MPACTTransformsPassIncGen) 4 | 5 | add_mlir_doc(Passes MPACTTransformsPass ./ -gen-pass-doc) 6 | -------------------------------------------------------------------------------- /include/mpact/Transforms/Passes.h: -------------------------------------------------------------------------------- 1 | //===------------------------------------------------------------*- C++ -*-===// 2 | // 3 | // Part of the MPACT Project, under the Apache License v2.0 with LLVM 4 | // Exceptions. See https://llvm.org/LICENSE.txt for license information. 5 | // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception 6 | // Also available under a BSD-style license. See LICENSE. 7 | // 8 | //===----------------------------------------------------------------------===// 9 | 10 | #ifndef MPACT_TRANSFORMS_PASSES_H 11 | #define MPACT_TRANSFORMS_PASSES_H 12 | 13 | namespace mlir { 14 | namespace mpact { 15 | 16 | /// Registers all mpact transform passes. 17 | void registerTransformPasses(); 18 | 19 | } // namespace mpact 20 | } // namespace mlir 21 | 22 | #endif // MPACT_TRANSFORMS_PASSES_H 23 | -------------------------------------------------------------------------------- /include/mpact/Transforms/Passes.td: -------------------------------------------------------------------------------- 1 | //===-- Passes.td - Transforms pass definition file --------*- tablegen -*-===// 2 | // 3 | // Part of the MPACT Project, under the Apache License v2.0 with LLVM Exceptions. 4 | // See https://llvm.org/LICENSE.txt for license information. 5 | // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception 6 | // 7 | //===----------------------------------------------------------------------===// 8 | // 9 | // This file contains definitions for passes within the Transforms/ directory. 10 | // 11 | //===----------------------------------------------------------------------===// 12 | 13 | #ifndef MPACT_TRANSFORMS_PASSES 14 | #define MPACT_TRANSFORMS_PASSES 15 | 16 | include "mlir/Pass/PassBase.td" 17 | 18 | def SparseEncodingPropagation : Pass<"sparse-encoding-propagation", "func::FuncOp"> { 19 | let summary = "Propagate sparse tensor encodings"; 20 | let description = [{ 21 | A pass that propagates sparse tensor encodings. 22 | 23 | Background: To avoid introducing repetitive operations, sparse tensors 24 | in MLIR try to reuse tensor operations whenever available. However, most 25 | tensor operations are canonicalized/transformed without the knowledge 26 | of sparsity. The pass tries to propagate missing sparse encodings. 27 | 28 | For example: 29 | ```mlir 30 | %s = tensor.extract_slice %input[0, 0,] [2, 1] [1, 1] 31 | : tensor<2x3xf32, #sparse> to tensor<2x1xf32, #sparse> 32 | 33 | // After rank reducing (by tensor dialect transformation) 34 | %t = tensor.extract_slice %input[0, 0,] [2, 1] [1, 1] 35 | : tensor<2x3xf32, #sparse> to tensor<2xf32> 36 | %s = tensor.expand_shape [[0, 1]] %t 37 | : tensor<2xf32> to tensor<2x1xf32, #sparse> 38 | 39 | // After sparsity propagation 40 | %t = tensor.extract_slice %input[0, 0,] [2, 1] [1, 1] 41 | : tensor<2x3xf32, #sparse> to tensor<2xf32, #sparse1> 42 | %s = tensor.expand_shape [[0, 1]] %t 43 | : tensor<2xf32, #sparse1> to tensor<2x1xf32, #sparse> 44 | ``` 45 | }]; 46 | 47 | let constructor = "mlir::mpact::createSparseEncodingPropagationPass()"; 48 | let dependentDialects = []; 49 | } 50 | 51 | #endif // MPACT_TRANSFORMS_PASSES 52 | -------------------------------------------------------------------------------- /include/mpact/Transforms/Sparsity/SparseEncodingPropagate.h: -------------------------------------------------------------------------------- 1 | //===------------------------------------------------------------*- C++ -*-===// 2 | // 3 | // Part of the MPACT Project, under the Apache License v2.0 with LLVM 4 | // Exceptions. See https://llvm.org/LICENSE.txt for license information. 5 | // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception 6 | // Also available under a BSD-style license. See LICENSE. 7 | // 8 | //===----------------------------------------------------------------------===// 9 | 10 | #ifndef MPACT_TRANSFORMS_SPARSITY_SPARSEENCODINGPROPAGATE_H 11 | #define MPACT_TRANSFORMS_SPARSITY_SPARSEENCODINGPROPAGATE_H 12 | 13 | #include "mlir/Dialect/Func/IR/FuncOps.h" 14 | #include "mlir/IR/BuiltinOps.h" 15 | #include "mlir/Pass/Pass.h" 16 | 17 | namespace mlir { 18 | namespace mpact { 19 | std::unique_ptr> 20 | createSparseEncodingPropagationPass(); 21 | } 22 | } // namespace mlir 23 | 24 | #endif // MPACT_TRANSFORMS_SPARSITY_SPARSEENCODINGPROPAGATE_H 25 | -------------------------------------------------------------------------------- /lib/CAPI/CMakeLists.txt: -------------------------------------------------------------------------------- 1 | add_mlir_public_c_api_library(MPACTCAPI 2 | Registration.cpp 3 | 4 | ENABLE_AGGREGATION 5 | 6 | LINK_LIBS PUBLIC 7 | MLIRIR 8 | MLIRSupport 9 | MPACTTransformPasses 10 | ) 11 | 12 | mpact_target_includes(MPACTCAPI) 13 | -------------------------------------------------------------------------------- /lib/CAPI/Registration.cpp: -------------------------------------------------------------------------------- 1 | //===- Registration.cpp - C Interface for MLIR Registration ---------------===// 2 | // 3 | // Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions. 4 | // See https://llvm.org/LICENSE.txt for license information. 5 | // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception 6 | // Also available under a BSD-style license. See LICENSE. 7 | // 8 | //===----------------------------------------------------------------------===// 9 | 10 | #include "mpact-c/Registration.h" 11 | 12 | #include "mlir/CAPI/IR.h" 13 | #include "mlir/Conversion/Passes.h" 14 | #include "mlir/Dialect/Linalg/Passes.h" 15 | #include "mlir/Transforms/Passes.h" 16 | #include "mpact/Transforms/Passes.h" 17 | 18 | MLIR_CAPI_EXPORTED void mpactRegisterAllPasses() { 19 | mlir::mpact::registerTransformPasses(); 20 | } 21 | -------------------------------------------------------------------------------- /lib/CMakeLists.txt: -------------------------------------------------------------------------------- 1 | add_subdirectory(CAPI) 2 | add_subdirectory(Transforms) 3 | -------------------------------------------------------------------------------- /lib/Transforms/CMakeLists.txt: -------------------------------------------------------------------------------- 1 | add_subdirectory(Sparsity) 2 | 3 | set(linked_libs MPACTSparsityPropagation) 4 | 5 | add_mlir_library(MPACTTransformPasses 6 | Passes.cpp 7 | 8 | DEPENDS 9 | MPACTTransformsPassIncGen 10 | 11 | LINK_LIBS PUBLIC 12 | ${linked_libs} 13 | ) 14 | -------------------------------------------------------------------------------- /lib/Transforms/Passes.cpp: -------------------------------------------------------------------------------- 1 | //===----------------------------------------------------------------------===// 2 | // 3 | // Part of the MPACT Project, under the Apache License v2.0 with LLVM 4 | // Exceptions. See https://llvm.org/LICENSE.txt for license information. 5 | // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception 6 | // Also available under a BSD-style license. See LICENSE. 7 | // 8 | //===----------------------------------------------------------------------===// 9 | 10 | #include "mpact/Transforms/Passes.h" 11 | #include "mpact/Transforms/Sparsity/SparseEncodingPropagate.h" 12 | 13 | //===----------------------------------------------------------------------===// 14 | // Pass registration 15 | //===----------------------------------------------------------------------===// 16 | 17 | namespace { 18 | #define GEN_PASS_REGISTRATION 19 | #include "mpact/Transforms/Passes.h.inc" 20 | } // end namespace 21 | 22 | void mlir::mpact::registerTransformPasses() { ::registerPasses(); } 23 | -------------------------------------------------------------------------------- /lib/Transforms/Sparsity/CMakeLists.txt: -------------------------------------------------------------------------------- 1 | add_mlir_conversion_library(MPACTSparsityPropagation 2 | SparseEncodingPropagate.cpp 3 | 4 | ADDITIONAL_HEADER_DIRS 5 | ${PROJECT_SOURCE_DIR}/include/mpact/Transforms/Sparsity 6 | 7 | DEPENDS 8 | MPACTTransformsPassIncGen 9 | 10 | LINK_LIBS PUBLIC 11 | MLIRIR 12 | MLIRPass 13 | ) 14 | 15 | mpact_target_includes(MPACTSparsityPropagation) 16 | -------------------------------------------------------------------------------- /lib/Transforms/Sparsity/SparseEncodingPropagate.cpp: -------------------------------------------------------------------------------- 1 | //===- SparseEncodingPropagate.cpp ---------------------------------------===// 2 | // 3 | // Part of the MPACT Project, under the Apache License v2.0 with LLVM 4 | // Exceptions. See https://llvm.org/LICENSE.txt for license information. 5 | // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception 6 | // 7 | //===----------------------------------------------------------------------===// 8 | 9 | #include "mpact/Transforms/Sparsity/SparseEncodingPropagate.h" 10 | 11 | namespace mlir { 12 | #define GEN_PASS_DEF_SPARSEENCODINGPROPAGATION 13 | #include "mpact/Transforms/Passes.h.inc" 14 | } // namespace mlir 15 | 16 | using namespace mlir; 17 | 18 | // ----------------------------------------------------------------------------- 19 | // The pass 20 | // ----------------------------------------------------------------------------- 21 | 22 | namespace { 23 | struct SparseEncodingPropagation 24 | : public impl::SparseEncodingPropagationBase { 25 | SparseEncodingPropagation() = default; 26 | SparseEncodingPropagation(const SparseEncodingPropagation &pass) = default; 27 | 28 | void runOnOperation() override {} 29 | }; 30 | } // namespace 31 | 32 | std::unique_ptr> 33 | mlir::mpact::createSparseEncodingPropagationPass() { 34 | return std::make_unique(); 35 | } 36 | -------------------------------------------------------------------------------- /python/CMakeLists.txt: -------------------------------------------------------------------------------- 1 | #------------------------------------------------------------------------------- 2 | # The MPACT Compiler Python Modules 3 | #------------------------------------------------------------------------------- 4 | 5 | # Disables generation of "version soname" (i.e. libFoo.so.). 6 | set(CMAKE_PLATFORM_NO_VERSIONED_SONAME ON) 7 | 8 | # The directory at which the Python import tree begins. 9 | set(MPACT_PYTHON_ROOT_DIR "${CMAKE_CURRENT_SOURCE_DIR}/mpact") 10 | 11 | # We vendor our own MLIR instance in the `mpact` namespace. 12 | add_compile_definitions("MLIR_PYTHON_PACKAGE_PREFIX=mpact.") 13 | 14 | declare_mlir_python_sources(MPACTPythonSources) 15 | declare_mlir_python_sources(MPACTPythonExtensions) 16 | 17 | declare_mlir_python_sources(MPACTPythonSources.PublicAPI 18 | ROOT_DIR "${MPACT_PYTHON_ROOT_DIR}" 19 | ADD_TO_PARENT MPACTPythonSources 20 | SOURCES 21 | mpactbackend.py 22 | ) 23 | 24 | declare_mlir_python_sources(MPACTPythonSources.SampleModels 25 | ROOT_DIR "${MPACT_PYTHON_ROOT_DIR}" 26 | ADD_TO_PARENT MPACTPythonSources 27 | SOURCES_GLOB 28 | models/*.py 29 | ) 30 | 31 | #------------------------------------------------------------------------------- 32 | # Extensions 33 | #------------------------------------------------------------------------------- 34 | 35 | declare_mlir_python_extension(MPACTPythonExtensions.Main 36 | MODULE_NAME _mpact 37 | ADD_TO_PARENT MPACTPythonExtensions 38 | SOURCES 39 | MPACTModule.cpp 40 | EMBED_CAPI_LINK_LIBS 41 | MPACTCAPI 42 | PRIVATE_LINK_LIBS 43 | LLVMSupport 44 | ) 45 | 46 | #------------------------------------------------------------------------------- 47 | # Python Modules 48 | #------------------------------------------------------------------------------- 49 | 50 | set(_source_components 51 | MLIRPythonSources 52 | MLIRPythonExtension.Core 53 | MLIRPythonExtension.RegisterEverything 54 | 55 | # We need various modules from torch-mlir. 56 | TorchMLIRPythonSources.Importers 57 | TorchMLIRPythonSources.Dialects 58 | TorchMLIRPythonSources.PublicAPI 59 | TorchMLIRPythonExtensions 60 | 61 | MPACTPythonSources 62 | MPACTPythonExtensions 63 | ) 64 | 65 | add_mlir_python_common_capi_library(MPACTAggregateCAPI 66 | INSTALL_COMPONENT MPACTPythonModules 67 | INSTALL_DESTINATION python_packages/mpact/mpact/_mlir_libs 68 | OUTPUT_DIRECTORY "${MPACT_PYTHON_PACKAGES_DIR}/mpact/mpact/_mlir_libs" 69 | RELATIVE_INSTALL_ROOT ".." 70 | DECLARED_SOURCES ${_source_components} 71 | ) 72 | 73 | add_mlir_python_modules(MPACTPythonModules 74 | ROOT_PREFIX "${MPACT_PYTHON_PACKAGES_DIR}/mpact/mpact" 75 | INSTALL_PREFIX "python_packages/mpact/mpact" 76 | DECLARED_SOURCES ${_source_components} 77 | COMMON_CAPI_LINK_LIBS 78 | MPACTAggregateCAPI 79 | ) 80 | -------------------------------------------------------------------------------- /python/MPACTModule.cpp: -------------------------------------------------------------------------------- 1 | //===-- MPACTModule.cpp ------------------------------------*- cpp -*-===// 2 | // 3 | // Part of the MPACT Project, under the Apache License v2.0 with LLVM 4 | // Exceptions. See https://llvm.org/LICENSE.txt for license information. 5 | // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception 6 | // Also available under a BSD-style license. See LICENSE. 7 | // 8 | //===----------------------------------------------------------------------===// 9 | 10 | #include "mlir/Bindings/Python/PybindAdaptors.h" 11 | #include "mpact-c/Registration.h" 12 | 13 | PYBIND11_MODULE(_mpact, m) { 14 | mpactRegisterAllPasses(); 15 | 16 | m.doc() = "mpact main python extension"; 17 | } 18 | -------------------------------------------------------------------------------- /python/mpact/models/gat.py: -------------------------------------------------------------------------------- 1 | import torch 2 | import torch.nn.functional as F 3 | 4 | 5 | class GraphAttentionLayer(torch.nn.Module): 6 | def __init__( 7 | self, 8 | in_features: int, 9 | out_features: int, 10 | n_heads: int, 11 | dropout: float = 0.4, 12 | leaky_relu_slope: float = 0.2, 13 | ): 14 | super(GraphAttentionLayer, self).__init__() 15 | self.n_heads = n_heads 16 | self.dropout = dropout 17 | self.n_hidden = out_features 18 | self.W = torch.nn.Parameter( 19 | torch.empty(size=(in_features, self.n_hidden * n_heads)) 20 | ) 21 | self.a = torch.nn.Parameter(torch.empty(size=(n_heads, 2 * self.n_hidden, 1))) 22 | self.leakyrelu = torch.nn.LeakyReLU(leaky_relu_slope) 23 | self.softmax = torch.nn.Softmax(dim=1) 24 | torch.nn.init.ones_(self.W) 25 | torch.nn.init.ones_(self.a) 26 | 27 | def forward(self, h: torch.Tensor, adj_mat: torch.Tensor): 28 | n_nodes = h.shape[0] 29 | h_transformed = torch.mm(h, self.W) 30 | h_transformed = F.dropout(h_transformed, self.dropout, training=self.training) 31 | h_transformed = h_transformed.view( 32 | n_nodes, self.n_heads, self.n_hidden 33 | ).permute(1, 0, 2) 34 | e = self._get_attention_scores(h_transformed) 35 | connectivity_mask = -9e16 * torch.ones_like(e) 36 | e = torch.where(adj_mat > 0, e, connectivity_mask) 37 | attention = F.softmax(e, dim=-1) 38 | attention = F.dropout(attention, self.dropout, training=self.training) 39 | h_prime = torch.matmul(attention, h_transformed) 40 | return h_prime.mean(dim=0) 41 | 42 | def _get_attention_scores(self, h_transformed: torch.Tensor): 43 | source_scores = torch.matmul(h_transformed, self.a[:, : self.n_hidden, :]) 44 | target_scores = torch.matmul(h_transformed, self.a[:, self.n_hidden :, :]) 45 | e = source_scores + target_scores.mT 46 | return self.leakyrelu(e) 47 | 48 | 49 | class GAT(torch.nn.Module): 50 | """ 51 | Graph Attention Network (GAT) inspired by . 52 | """ 53 | 54 | def __init__( 55 | self, 56 | in_features, 57 | n_hidden, 58 | n_heads, 59 | num_classes, 60 | dropout=0.4, 61 | leaky_relu_slope=0.2, 62 | ): 63 | super(GAT, self).__init__() 64 | self.gat1 = GraphAttentionLayer( 65 | in_features=in_features, 66 | out_features=n_hidden, 67 | n_heads=n_heads, 68 | dropout=dropout, 69 | leaky_relu_slope=leaky_relu_slope, 70 | ) 71 | self.gat2 = GraphAttentionLayer( 72 | in_features=n_hidden, 73 | out_features=num_classes, 74 | n_heads=1, 75 | dropout=dropout, 76 | leaky_relu_slope=leaky_relu_slope, 77 | ) 78 | 79 | def forward(self, input_tensor: torch.Tensor, adj_mat: torch.Tensor): 80 | x = self.gat1(input_tensor, adj_mat) 81 | x = F.elu(x) 82 | x = self.gat2(x, adj_mat) 83 | return F.log_softmax(x, dim=1) 84 | 85 | 86 | def gat_4_64_8_3(): 87 | return GAT(in_features=4, n_hidden=64, n_heads=8, num_classes=3) 88 | -------------------------------------------------------------------------------- /python/mpact/models/gcn.py: -------------------------------------------------------------------------------- 1 | import torch 2 | import torch.nn.functional as F 3 | 4 | 5 | class GraphConv(torch.nn.Module): 6 | def __init__(self, input_dim, output_dim): 7 | super(GraphConv, self).__init__() 8 | self.kernel = torch.nn.Parameter(torch.Tensor(input_dim, output_dim)) 9 | torch.nn.init.ones_(self.kernel) 10 | self.bias = torch.nn.Parameter(torch.Tensor(output_dim)) 11 | torch.nn.init.ones_(self.bias) 12 | 13 | def forward(self, inp, adj_mat): 14 | # Input matrix times weight matrix. 15 | support = torch.mm(inp, self.kernel) 16 | # Sparse adjacency matrix times support matrix. 17 | output = torch.spmm(adj_mat, support) 18 | # Add bias. 19 | output = output + self.bias 20 | return output 21 | 22 | 23 | class GCN(torch.nn.Module): 24 | """ 25 | Graph Convolutional Network (GCN) inspired by . 26 | """ 27 | 28 | def __init__(self, input_dim, hidden_dim, output_dim, dropout_p=0.1): 29 | super(GCN, self).__init__() 30 | self.gc1 = GraphConv(input_dim, hidden_dim) 31 | self.gc2 = GraphConv(hidden_dim, output_dim) 32 | self.dropout = torch.nn.Dropout(dropout_p) 33 | 34 | def forward(self, input_tensor, adj_mat): 35 | x = self.gc1(input_tensor, adj_mat) 36 | x = F.relu(x) 37 | x = self.dropout(x) 38 | x = self.gc2(x, adj_mat) 39 | return F.log_softmax(x, dim=1) 40 | 41 | 42 | def graphconv_4_4(): 43 | return GraphConv(input_dim=4, output_dim=4) 44 | 45 | 46 | def gcn_4_16_4(): 47 | return GCN(input_dim=4, hidden_dim=16, output_dim=4) 48 | -------------------------------------------------------------------------------- /python/mpact/models/kernels.py: -------------------------------------------------------------------------------- 1 | import torch 2 | 3 | 4 | class MVNet(torch.nn.Module): 5 | def forward(self, x, v): 6 | return torch.mv(x, v) 7 | 8 | 9 | class MMNet(torch.nn.Module): 10 | def forward(self, x, y): 11 | return torch.mm(x, y) 12 | 13 | 14 | class AddNet(torch.nn.Module): 15 | def forward(self, x, y): 16 | return torch.add(x, y) 17 | 18 | 19 | class MulNet(torch.nn.Module): 20 | def forward(self, x, y): 21 | return torch.mul(x, y) 22 | 23 | 24 | class SelfNet(torch.nn.Module): 25 | def forward(self, x): 26 | return x 27 | 28 | 29 | class SDDMMNet(torch.nn.Module): 30 | def forward(self, x, y, z): 31 | return torch.mul(x, torch.mm(y, z)) 32 | 33 | 34 | class SqSum(torch.nn.Module): 35 | def forward(self, x): 36 | return (x * x).sum() 37 | 38 | 39 | class CountEq(torch.nn.Module): 40 | def forward(self, x, s): 41 | nums = (x == s).sum() 42 | return nums 43 | 44 | 45 | class FeatureScale(torch.nn.Module): 46 | def forward(self, F): 47 | sum_vector = torch.sum(F, dim=1) 48 | reciprocal_vector = 1 / sum_vector 49 | reciprocal_vector[reciprocal_vector == float("inf")] = 0 50 | scaling_diagonal = torch.diag(reciprocal_vector).to_sparse() 51 | return scaling_diagonal @ F 52 | 53 | 54 | class Normalization(torch.nn.Module): 55 | def forward(self, A): 56 | sum_vector = torch.sum(A, dim=1) 57 | reciprocal_vector = 1 / sum_vector 58 | reciprocal_vector[reciprocal_vector == float("inf")] = 0 59 | scaling_diagonal = torch.diag(reciprocal_vector).to_sparse() 60 | return scaling_diagonal @ A @ scaling_diagonal 61 | 62 | 63 | class SimpleNet(torch.nn.Module): 64 | def __init__(self): 65 | super(SimpleNet, self).__init__() 66 | # Model parameters (weights and biases of linear layers). 67 | self.fc1 = torch.nn.Linear(16, 8) 68 | self.fc2 = torch.nn.Linear(8, 4) 69 | self.fc3 = torch.nn.Linear(4, 2) 70 | 71 | def forward(self, x): 72 | x = x.view(-1, 16) 73 | x = torch.nn.functional.relu(self.fc1(x)) 74 | x = torch.nn.functional.relu(self.fc2(x)) 75 | return self.fc3(x) # assumes: softmax in loss function 76 | -------------------------------------------------------------------------------- /python/mpact/models/lif.py: -------------------------------------------------------------------------------- 1 | import torch 2 | 3 | 4 | def spike(input): 5 | return (input >= 0).float() 6 | 7 | 8 | def sqSum(input): 9 | return (input * input).sum() 10 | 11 | 12 | class LIF(torch.nn.Module): 13 | def __init__(self): 14 | super(LIF, self).__init__() 15 | self.thresh = 1.0 16 | self.decay = 0.5 17 | self.act = spike 18 | 19 | def forward(self, X): 20 | """A filter that yields a binary-valued sparse tensor.""" 21 | mem = 0 22 | spike_pot = [] 23 | T = X.size(-1) 24 | for t in range(T): 25 | mem = mem * self.decay + X[..., t] 26 | spike = self.act(mem - self.thresh) 27 | spike = spike.to_sparse().to_dense() # prop hack 28 | mem = mem * (1.0 - spike) 29 | spike_pot.append(spike) 30 | spike_pot = torch.stack(spike_pot, dim=-1) 31 | return spike_pot 32 | 33 | 34 | class tdLayer(torch.nn.Module): 35 | def __init__(self, layer): 36 | super(tdLayer, self).__init__() 37 | self.layer = layer 38 | 39 | def forward(self, X): 40 | T = X.size(-1) 41 | out = [] 42 | for t in range(T): 43 | m = self.layer(X[..., t]) 44 | out.append(m) 45 | out = torch.stack(out, dim=-1) 46 | return out 47 | 48 | 49 | class LIFSumOfSq(torch.nn.Module): 50 | def __init__(self): 51 | super(LIFSumOfSq, self).__init__() 52 | self.spike = LIF() 53 | self.layer = tdLayer(sqSum) 54 | 55 | def forward(self, X): 56 | out = self.spike(X) 57 | out = self.layer(out) 58 | return out 59 | -------------------------------------------------------------------------------- /python/mpact/models/resnet.py: -------------------------------------------------------------------------------- 1 | import torch 2 | import numpy as np 3 | 4 | 5 | def spike(input): 6 | return (input >= 0).float() 7 | 8 | 9 | class Straight(torch.nn.Module): 10 | def forward(self, input): 11 | return input 12 | 13 | 14 | class tdLayer(torch.nn.Module): 15 | def __init__(self, layer, bn=None): 16 | super(tdLayer, self).__init__() 17 | self.layer = layer 18 | self.bn = bn if bn is not None else Straight() 19 | 20 | def forward(self, X): 21 | T = X.size(-1) 22 | out = [] 23 | for t in range(T): 24 | m = self.layer(X[..., t]) 25 | out.append(m) 26 | out = torch.stack(out, dim=-1) 27 | out = self.bn(out) 28 | return out 29 | 30 | 31 | class LIF(torch.nn.Module): 32 | def __init__(self): 33 | super(LIF, self).__init__() 34 | self.thresh = 1.0 35 | self.decay = 0.5 36 | self.act = spike 37 | self.gama = 1.0 38 | 39 | def forward(self, X, gama=1): 40 | mem = 0 41 | spike_pot = [] 42 | T = X.size(-1) 43 | for t in range(T): 44 | mem = mem * self.decay + X[..., t] 45 | spike = self.act(mem - self.thresh) 46 | mem = mem * (1.0 - spike) 47 | spike_pot.append(spike) 48 | spike_pot = torch.stack(spike_pot, dim=-1) 49 | return spike_pot 50 | 51 | 52 | class tdBatchNorm(torch.nn.BatchNorm2d): 53 | def __init__( 54 | self, 55 | num_features, 56 | eps=1e-05, 57 | momentum=0.1, 58 | alpha=1, 59 | affine=True, 60 | track_running_stats=True, 61 | ): 62 | super(tdBatchNorm, self).__init__( 63 | num_features, eps, momentum, affine, track_running_stats 64 | ) 65 | self.alpha = alpha 66 | 67 | def forward(self, input): 68 | exponential_average_factor = 0.0 69 | mean = self.running_mean 70 | var = self.running_var 71 | input = ( 72 | self.alpha 73 | * (input - mean[None, :, None, None, None]) 74 | / (torch.sqrt(var[None, :, None, None, None] + self.eps)) 75 | ) 76 | if self.affine: 77 | input = ( 78 | input * self.weight[None, :, None, None, None] 79 | + self.bias[None, :, None, None, None] 80 | ) 81 | return input 82 | 83 | 84 | def conv3x3(in_planes, out_planes, stride=1, groups=1, dilation=1): 85 | return torch.nn.Conv2d( 86 | in_planes, 87 | out_planes, 88 | kernel_size=3, 89 | stride=stride, 90 | padding=dilation, 91 | groups=groups, 92 | bias=False, 93 | dilation=dilation, 94 | ) 95 | 96 | 97 | def conv1x1(in_planes, out_planes, stride=1): 98 | return torch.nn.Conv2d( 99 | in_planes, out_planes, kernel_size=1, stride=stride, bias=False 100 | ) 101 | 102 | 103 | class BasicBlock(torch.nn.Module): 104 | expansion = 1 105 | 106 | def __init__( 107 | self, 108 | inplanes, 109 | planes, 110 | stride=1, 111 | downsample=None, 112 | groups=1, 113 | base_width=64, 114 | dilation=1, 115 | norm_layer=None, 116 | ): 117 | super(BasicBlock, self).__init__() 118 | if norm_layer is None: 119 | norm_layer = tdBatchNorm 120 | # norm_layer = nn.BatchNorm2d 121 | if groups != 1 or base_width != 64: 122 | raise ValueError("BasicBlock only supports groups=1 and base_width=64") 123 | if dilation > 1: 124 | raise NotImplementedError("Dilation > 1 not supported in BasicBlock") 125 | # Both self.conv1 and self.downsample layers downsample the input when stride != 1 126 | self.conv1 = conv3x3(inplanes, planes, stride) 127 | self.bn1 = norm_layer(planes) 128 | self.conv2 = conv3x3(planes, planes) 129 | self.bn2 = norm_layer(planes) 130 | self.downsample = downsample 131 | self.stride = stride 132 | self.conv1_s = tdLayer(self.conv1, self.bn1) 133 | self.conv2_s = tdLayer(self.conv2, self.bn2) 134 | self.spike1 = LIF() 135 | self.spike2 = LIF() 136 | 137 | def forward(self, x): 138 | identity = x 139 | 140 | out = self.conv1_s(x) 141 | out = self.spike1(out) 142 | out = self.conv2_s(out) 143 | 144 | if self.downsample is not None: 145 | identity = self.downsample(x) 146 | 147 | out += identity 148 | out = self.spike2(out) 149 | 150 | return out 151 | 152 | 153 | class ResNety(torch.nn.Module): 154 | def __init__( 155 | self, 156 | block, 157 | layers, 158 | num_classes=10, 159 | zero_init_residual=False, 160 | groups=1, 161 | width_per_group=64, 162 | replace_stride_with_dilation=None, 163 | norm_layer=None, 164 | ): 165 | super(ResNety, self).__init__() 166 | if norm_layer is None: 167 | norm_layer = tdBatchNorm 168 | # norm_layer = nn.BatchNorm2d 169 | self._norm_layer = norm_layer 170 | self.inplanes = 64 171 | self.dilation = 1 172 | self.groups = groups 173 | self.base_width = width_per_group 174 | self.pre = torch.nn.Sequential( 175 | tdLayer( 176 | layer=torch.nn.Conv2d( 177 | 3, self.inplanes, kernel_size=(3, 3), stride=(1, 1) 178 | ), 179 | bn=self._norm_layer(self.inplanes), 180 | ), 181 | LIF(), 182 | ) 183 | self.layer1 = self._make_layer(block, 64, layers[0], stride=2) 184 | self.layer2 = self._make_layer(block, 128, layers[1], stride=2) 185 | self.layer3 = self._make_layer(block, 256, layers[2], stride=2) 186 | self.avgpool = tdLayer(torch.nn.AdaptiveAvgPool2d((1, 1))) 187 | self.fc = tdLayer(torch.nn.Linear(256, num_classes)) 188 | self.T = 6 189 | for m in self.modules(): 190 | if isinstance(m, torch.nn.Conv2d): 191 | torch.nn.init.kaiming_normal_( 192 | m.weight, mode="fan_out", nonlinearity="relu" 193 | ) 194 | 195 | def _make_layer(self, block, planes, blocks, stride=1, dilate=False): 196 | norm_layer = self._norm_layer 197 | downsample = None 198 | previous_dilation = self.dilation 199 | if dilate: 200 | self.dilation *= stride 201 | stride = 1 202 | if stride != 1 or self.inplanes != planes * block.expansion: 203 | downsample = tdLayer( 204 | conv1x1(self.inplanes, planes * block.expansion, stride), 205 | norm_layer(planes * block.expansion), 206 | ) 207 | 208 | layers = [] 209 | layers.append( 210 | block( 211 | self.inplanes, 212 | planes, 213 | stride, 214 | downsample, 215 | self.groups, 216 | self.base_width, 217 | previous_dilation, 218 | norm_layer, 219 | ) 220 | ) 221 | self.inplanes = planes * block.expansion 222 | for _ in range(1, blocks): 223 | layers.append( 224 | block( 225 | self.inplanes, 226 | planes, 227 | groups=self.groups, 228 | base_width=self.base_width, 229 | dilation=self.dilation, 230 | norm_layer=norm_layer, 231 | ) 232 | ) 233 | 234 | return torch.nn.Sequential(*layers) 235 | 236 | def _forward_impl(self, input): 237 | out = [] 238 | input = input.unsqueeze(-1).repeat(1, 1, 1, 1, self.T) 239 | x = self.pre(input) 240 | x = self.layer1(x) 241 | x = self.layer2(x) 242 | x = self.layer3(x) 243 | x = self.avgpool(x) 244 | x = x.view(x.size(0), -1, x.size(-1)) 245 | x = self.fc(x) 246 | for t in range(self.T): 247 | out.append(x[..., t]) 248 | return torch.stack(out, dim=1) 249 | 250 | def forward(self, x): 251 | return self._forward_impl(x) 252 | 253 | 254 | def resnet_20(): 255 | return ResNety(block=BasicBlock, layers=[2, 2, 2], num_classes=10) 256 | -------------------------------------------------------------------------------- /python/mpact/models/train.py: -------------------------------------------------------------------------------- 1 | import torch 2 | import torch.nn.functional as F 3 | 4 | 5 | def num_all_parameters(model): 6 | """Returns the number of all parameters in a model.""" 7 | return sum(p.numel() for p in model.parameters()) 8 | 9 | 10 | def num_parameters(model): 11 | """Returns the number of trainable parameters in a model.""" 12 | return sum(p.numel() for p in model.parameters() if p.requires_grad) 13 | 14 | 15 | def training_loop(model, optimizer, loss_function, train, validation, epochs=10): 16 | """A rudimentary PyTorch training loop for classification with training and validation data.""" 17 | for epoch in range(epochs): 18 | # Switch to training mode. 19 | model.train() 20 | tloss = 0.0 21 | num_train = len(train) # in batches 22 | for inp, target in train: # batch loop (training) 23 | optimizer.zero_grad() 24 | output = model(inp) 25 | loss = loss_function(output, target) 26 | loss.backward() 27 | optimizer.step() 28 | tloss += loss.data.item() 29 | 30 | # Switch to inference mode. 31 | model.eval() # disables e.g. model drop-out 32 | with torch.no_grad(): # disables gradient computations 33 | vloss = 0.0 34 | num_validation = len(validation) # in batches 35 | num_correct = 0 36 | num_total = 0 37 | for inp, target in validation: # batch loop (validation) 38 | output = model(inp) 39 | loss = loss_function(output, target) 40 | vloss += loss.data.item() 41 | correct = torch.eq( 42 | torch.max(F.softmax(output, dim=1), dim=1)[1], target 43 | ).view(-1) 44 | num_correct += torch.sum(correct).item() 45 | num_total += correct.shape[0] 46 | 47 | # Report stats. 48 | print( 49 | "Epoch {:d}, Training loss = {:.2f} #{:d}, Validation loss = {:.2f} #{:d}, Accuracy = {:.2f} #{:d}".format( 50 | epoch, 51 | (tloss / num_train) if num_train != 0 else 0, 52 | num_train, 53 | (vloss / num_validation) if num_validation != 0 else 0, 54 | num_validation, 55 | (num_correct / num_total) if num_total != 0 else 0, 56 | num_total, 57 | ) 58 | ) 59 | -------------------------------------------------------------------------------- /python/mpact/mpactbackend.py: -------------------------------------------------------------------------------- 1 | # Initialize mpact python extension. 2 | import mpact._mlir_libs._mpact 3 | 4 | import ctypes 5 | from io import StringIO 6 | import numpy as np 7 | import os 8 | import sys 9 | import tempfile 10 | import torch 11 | 12 | from typing import Any, Callable, Optional, Tuple, Dict, TypeVar, Union 13 | 14 | from mpact import ir 15 | from mpact.ir import Module 16 | from mpact.dialects import torch as torch_d 17 | from mpact.execution_engine import * 18 | from mpact.extras.fx_decomp_util import get_decomposition_table 19 | from mpact.extras.fx_importer import FxImporter 20 | from mpact.ir import * 21 | from mpact.passmanager import * 22 | from mpact.runtime import * 23 | 24 | # One time set up of support library. 25 | SUPPORT_LIB = os.getenv("SUPPORT_LIB", default=None) 26 | SHARED_LIBS = [] if SUPPORT_LIB is None else [SUPPORT_LIB] 27 | 28 | # The result of MPACT compile() and input to load(). 29 | MpactCompiledArtifact = TypeVar("MpactCompiledArtifact") 30 | 31 | 32 | def get_module_name_for_debug_dump(module): 33 | """Gets a name suitable for a debug dump. 34 | 35 | The name is not guaranteed to be unique. 36 | """ 37 | if not "torch.debug_module_name" in module.operation.attributes: 38 | return "UnnammedModule" 39 | return StringAttr(module.operation.attributes["torch.debug_module_name"]).value 40 | 41 | 42 | class MPACTCompilerError(Exception): 43 | pass 44 | 45 | 46 | def run_pipeline_with_repro_report( 47 | module, pipeline: str, description: str, enable_ir_printing: bool = False 48 | ): 49 | """Runs `pipeline` on `module`, with a nice repro report if it fails.""" 50 | module_name = get_module_name_for_debug_dump(module) 51 | original_stderr = sys.stderr 52 | try: 53 | sys.stderr = StringIO() 54 | asm_for_error_report = module.operation.get_asm( 55 | large_elements_limit=10, enable_debug_info=True 56 | ) 57 | # Lower module in place to make it ready for compiler backends. 58 | with module.context as ctx: 59 | pm = PassManager.parse(pipeline) 60 | if enable_ir_printing: 61 | ctx.enable_multithreading(False) 62 | pm.enable_ir_printing() 63 | pm.run(module.operation) 64 | except Exception as e: 65 | filename = os.path.join(tempfile.gettempdir(), module_name + ".mlir") 66 | with open(filename, "w") as f: 67 | f.write(asm_for_error_report) 68 | debug_options = "-mlir-print-ir-after-all -mlir-disable-threading" 69 | # Put something descriptive here even if description is empty. 70 | description = description or f"{module_name} compile" 71 | 72 | message = f"""\ 73 | {description} failed with the following diagnostics: 74 | {sys.stderr.getvalue()} 75 | 76 | python exception: {e} 77 | 78 | The error can be reproduced with: 79 | $ mpact-opt -pass-pipeline='{pipeline}' {filename} 80 | Add '{debug_options}' to get the IR dump for debugging purpose. 81 | """ 82 | trimmed_message = "\n".join([m.lstrip() for m in message.split("\n")]) 83 | raise MPACTCompilerError(trimmed_message) from None 84 | finally: 85 | sys.stderr = original_stderr 86 | 87 | 88 | def assert_arg_type_is_supported(ty): 89 | SUPPORTED = [ 90 | np.float16, 91 | np.float32, 92 | np.float64, 93 | np.uint8, 94 | np.int8, 95 | np.int32, 96 | np.int64, 97 | np.bool_, 98 | np.complex64, 99 | np.complex128, 100 | ] 101 | assert ( 102 | ty in SUPPORTED 103 | ), f"Only numpy arrays with dtypes in {SUPPORTED} are supported, but got {ty}" 104 | 105 | 106 | memref_type_to_np_dtype = { 107 | "mrf16": np.float16, 108 | "mrf32": np.float32, 109 | "mrf64": np.float64, 110 | "mri1": np.bool_, 111 | "mri8": np.int8, 112 | "mri32": np.int32, 113 | "mri64": np.int64, 114 | "mrc32": np.complex64, 115 | "mrc64": np.complex128, 116 | } 117 | elemental_type_to_ctype = { 118 | "i1": ctypes.c_bool, 119 | "i8": ctypes.c_byte, 120 | "i64": ctypes.c_int, 121 | "f32": ctypes.c_float, 122 | "f64": ctypes.c_double, 123 | } 124 | 125 | CONSUME_RETURN_FUNC_PREFIX = "refbackend_consume_func_return_" 126 | 127 | 128 | def get_return_funcs(module): 129 | return_prefix_len = len(CONSUME_RETURN_FUNC_PREFIX) 130 | return_funcs = [] 131 | with module.context: 132 | for func in module.body: 133 | # Returns strings of the form `"refbackend.."` so `"` is deleted. 134 | func_name = str(func.attributes["sym_name"]).replace('"', "") 135 | if func_name[:return_prefix_len] == CONSUME_RETURN_FUNC_PREFIX: 136 | return_funcs.append(func_name) 137 | 138 | return return_funcs 139 | 140 | 141 | def get_ctype_func(func_name): 142 | return_prefix_len = len(CONSUME_RETURN_FUNC_PREFIX) 143 | ret_types = func_name[return_prefix_len:].split("_") 144 | ctypes_arg = [None] 145 | for type in ret_types: 146 | if type in elemental_type_to_ctype: 147 | ctypes_arg.append(elemental_type_to_ctype[type]) 148 | elif type in memref_type_to_np_dtype: 149 | ctypes_arg.append(ctypes.POINTER(UnrankedMemRefDescriptor)) 150 | else: 151 | assert False, f"Not supported type: {type}" 152 | 153 | return ctypes.CFUNCTYPE(*ctypes_arg), ret_types 154 | 155 | 156 | class MpactBackendInvoker: 157 | def __init__(self, module, opt_level): 158 | self.ee = ExecutionEngine(module, opt_level=opt_level, shared_libs=SHARED_LIBS) 159 | self.result = None 160 | 161 | return_funcs = get_return_funcs(module) 162 | 163 | for ret_func in return_funcs: 164 | ctype_wrapper, ret_types = get_ctype_func(ret_func) 165 | 166 | def consume_return_funcs(*args): 167 | self.result = tuple( 168 | [ 169 | ( 170 | arg 171 | if type in elemental_type_to_ctype 172 | else unranked_memref_to_numpy( 173 | arg, memref_type_to_np_dtype[type] 174 | ) 175 | ) 176 | for arg, type in zip(args, ret_types) 177 | ] 178 | ) 179 | if len(self.result) == 1: 180 | self.result = self.result[0] 181 | 182 | self.ee.register_runtime(ret_func, ctype_wrapper(consume_return_funcs)) 183 | 184 | def __getattr__(self, function_name: str): 185 | def invoke(*args): 186 | ffi_args = [] 187 | for arg in args: 188 | assert_arg_type_is_supported(arg.dtype) 189 | ffi_args.append( 190 | ctypes.pointer(ctypes.pointer(get_unranked_memref_descriptor(arg))) 191 | ) 192 | 193 | self.ee.invoke(function_name, *ffi_args) 194 | result = self.result 195 | assert result is not None, "Invocation didn't produce a result" 196 | self.result = None 197 | return result 198 | 199 | return invoke 200 | 201 | 202 | LOWERING_PIPELINE_TEMPLATE = ( 203 | "builtin.module(" 204 | + ",".join( 205 | [ 206 | "func.func(linalg-generalize-named-ops)", 207 | # Run pre-sparsification pass to fuse convert/cast op into 208 | # producer as they might hinder kernel fusions. 209 | "pre-sparsification-rewrite", 210 | "func.func(linalg-fuse-elementwise-ops)", 211 | "convert-shape-to-std", 212 | # Propagate sparse encodings before sparsifier mini-pipeline. 213 | # TODO: the following pass currently contains no pattern. Will be 214 | # added as needed. 215 | "func.func(sparse-encoding-propagation)", 216 | # MLIR Sparsifier mini-pipeline: 217 | # use the PyTorch assembler conventions 218 | # enable vectorization with VL=16 (more or less assumes AVX512 for float) 219 | # allow 32-bit index optimizations (unsafe for very large dimensions) 220 | "sparse-assembler{{direct-out}}", 221 | "sparsification-and-bufferization{{{sp_options}}}", 222 | "sparse-storage-specifier-to-llvm", 223 | # Buffer deallocation pass does not know how to handle realloc. 224 | "func.func(expand-realloc)", 225 | # Generalize pad and concat after sparse compiler, as they are handled 226 | # differently when the operations involve sparse operands. 227 | "func.func(refback-generalize-tensor-pad)", 228 | "func.func(refback-generalize-tensor-concat)", 229 | # Bufferize. 230 | "func.func(tm-tensor-bufferize)", 231 | "one-shot-bufferize{{copy-before-write bufferize-function-boundaries function-boundary-type-conversion=identity-layout-map}}", 232 | "refback-mlprogram-bufferize", 233 | "func.func(finalizing-bufferize)", 234 | "func.func(buffer-deallocation)", 235 | # Inline sparse helper methods where useful (but after dealloc). 236 | "inline", 237 | "refback-munge-calling-conventions", 238 | "func.func(tm-tensor-to-loops)", 239 | "func.func(refback-munge-memref-copy)", 240 | "func.func(convert-linalg-to-loops)", 241 | "func.func(lower-affine)", 242 | "convert-scf-to-cf", 243 | "func.func(refback-expand-ops-for-llvm)", 244 | "func.func(arith-expand)", 245 | "func.func(convert-math-to-llvm)", 246 | "convert-math-to-libm", 247 | "expand-strided-metadata", 248 | "finalize-memref-to-llvm", 249 | "lower-affine", 250 | "convert-bufferization-to-memref", 251 | "finalize-memref-to-llvm", 252 | "func.func(convert-arith-to-llvm)", 253 | # Vector code (SIMD): 254 | # allow fp reductions to reassociate 255 | # allow 32-bit index optimizations (unsafe for very large dimensions) 256 | # assume we are running on a good ol' Intel X86 (disable for ARM/other) 257 | "convert-vector-to-llvm{{reassociate-fp-reductions force-32bit-vector-indices enable-x86vector}}", 258 | "convert-func-to-llvm", 259 | "convert-cf-to-llvm", 260 | "convert-complex-to-llvm", 261 | "reconcile-unrealized-casts", 262 | ] 263 | ) 264 | + ")" 265 | ) 266 | 267 | 268 | class MpactBackendCompiler: 269 | """Main entry-point for the MPACT backend compiler.""" 270 | 271 | def __init__(self, opt_level, use_sp_it): 272 | self.opt_level = opt_level 273 | self.use_sp_it = use_sp_it 274 | 275 | def compile(self, imported_module: Module) -> MpactCompiledArtifact: 276 | sp_options = ( 277 | "sparse-emit-strategy=sparse-iterator" 278 | if self.use_sp_it 279 | else "vl=16 enable-simd-index32" 280 | ) 281 | LOWERING_PIPELINE = LOWERING_PIPELINE_TEMPLATE.format(sp_options=sp_options) 282 | """Compiles an imported module, with a flat list of functions. 283 | The module is expected to be in linalg-on-tensors + scalar code form. 284 | 285 | Args: 286 | imported_module: The MLIR module in the torch dialect. 287 | Returns: 288 | An opaque artifact that can be passed to `load`. 289 | """ 290 | run_pipeline_with_repro_report( 291 | imported_module, 292 | LOWERING_PIPELINE, 293 | "Lowering Linalg-on-Tensors IR to LLVM with MpactBackendCompiler", 294 | enable_ir_printing=False, 295 | ) 296 | return imported_module 297 | 298 | def load(self, module: MpactCompiledArtifact) -> MpactBackendInvoker: 299 | """Loads a compiled artifact into the runtime. 300 | 301 | Args: 302 | module: The result of a previous call to `compile`. 303 | Returns: 304 | MPactInvoker to call a compiled method (viz `invoker.foo(...)`). 305 | """ 306 | return MpactBackendInvoker(module, self.opt_level) 307 | 308 | 309 | def export_and_import(f, *args, **kwargs): 310 | """A FX graph importer, stripped down to essentials.""" 311 | context = ir.Context() 312 | torch_d.register_dialect(context) 313 | fx_importer = FxImporter(context=context) 314 | prog = torch.export.export(f, args, kwargs) 315 | decomposition_table = get_decomposition_table() 316 | if decomposition_table: 317 | prog = prog.run_decompositions(decomposition_table) 318 | fx_importer.import_frozen_program(prog) 319 | return fx_importer.module 320 | 321 | 322 | def mpact_linalg(f, *args, **kwargs): 323 | """Imports a callable as module and lowers it into Linalg IR.""" 324 | module = export_and_import(f, *args, **kwargs) 325 | run_pipeline_with_repro_report( 326 | module, 327 | ( 328 | "builtin.module(" 329 | "func.func(torch-decompose-complex-ops)," 330 | "torch-backend-to-linalg-on-tensors-backend-pipeline)" 331 | ), 332 | "Lowering TorchFX IR -> Linalg IR", 333 | enable_ir_printing=False, 334 | ) 335 | return module 336 | 337 | 338 | def mpact_stablehlo(f, *args, **kwargs): 339 | """Imports a callable as module and lowers it into StableHLO IR.""" 340 | module = export_and_import(f, *args, **kwargs) 341 | run_pipeline_with_repro_report( 342 | module, 343 | ( 344 | "builtin.module(" 345 | "func.func(torch-decompose-complex-ops)," 346 | "torch-backend-to-stablehlo-backend-pipeline)" 347 | ), 348 | "Lowering TorchFX IR -> StableHLO IR", 349 | enable_ir_printing=False, 350 | ) 351 | return module 352 | 353 | 354 | def mpact_jit_compile(f, *args, opt_level=2, use_sp_it=False, **kwargs): 355 | """This method compiles the given callable using the MPACT backend.""" 356 | module = mpact_linalg(f, *args, **kwargs) 357 | backend = MpactBackendCompiler(opt_level=opt_level, use_sp_it=use_sp_it) 358 | compiled = backend.compile(module) 359 | invoker = backend.load(compiled) 360 | return invoker, f 361 | 362 | 363 | def mpact_jit_run(invoker, f, *args, **kwargs): 364 | """This method runs the given callable using the given MPACT invoker.""" 365 | xargs = [] 366 | # Prepare all the named buffer parameters (assume all dense). 367 | # All scalar arguments are filtered out since they appear inline. 368 | params = dict(f.named_buffers(remove_duplicate=True)) 369 | params_flat, params_spec = torch.utils._pytree.tree_flatten(params) 370 | for p in params_flat: 371 | if len(p.shape) > 0: 372 | xargs.append(p.numpy()) 373 | # Prepare input parameters. Sparse input tensors are split into 374 | # their composite tensors. All PyTorch tensors are converted 375 | # to their backing numpy arrays. Note that the output consists 376 | # of numpy arrays as well, which can trivially be reconstructed 377 | # into PyTorch tensors (dense and sparse). 378 | for a in args: 379 | if a.layout is torch.sparse_coo: 380 | # Construct the additional position array required by MLIR with data 381 | # array([0, nnz]). The COO format always uses int64 indices. 382 | xargs.append(np.array([0, a._nnz()], dtype=np.int64)) 383 | # Transform a tensor into ndim x tensor to conform 384 | # to the MLIR SoA COO representation. 385 | for idx in a._indices(): 386 | xargs.append(idx.numpy()) 387 | xargs.append(a._values().numpy()) 388 | elif a.layout is torch.sparse_csr or a.layout is torch.sparse_bsr: 389 | xargs.append(a.crow_indices().numpy()) 390 | xargs.append(a.col_indices().numpy()) 391 | xargs.append(a.values().numpy()) 392 | elif a.layout is torch.sparse_csc or a.layout is torch.sparse_bsc: 393 | xargs.append(a.ccol_indices().numpy()) 394 | xargs.append(a.row_indices().numpy()) 395 | xargs.append(a.values().numpy()) 396 | else: 397 | xargs.append(a.numpy()) 398 | # Invoke. 399 | return invoker.main(*xargs) 400 | 401 | 402 | # Convenience wrapper. 403 | def mpact_jit(f, *args, **kwargs): 404 | """This method compiles and runs the given callable using the MPACT backend.""" 405 | invoker, fn = mpact_jit_compile(f, *args, **kwargs) 406 | return mpact_jit_run(invoker, fn, *args, **kwargs) 407 | -------------------------------------------------------------------------------- /setup.py: -------------------------------------------------------------------------------- 1 | # Script for generating the mpact wheel. 2 | # ``` 3 | # $ python setup.py bdist_wheel 4 | # ``` 5 | # Environment variables you are probably interested in: 6 | # 7 | # CMAKE_BUILD_TYPE: 8 | # specify the build type: DEBUG/RelWithDebInfo/Release 9 | # 10 | # MPACT_CMAKE_ALREADY_BUILT: 11 | # the `MPACT_CMAKE_BUILD_DIR` directory has already been compiled, 12 | # and the CMake compilation process will not be executed again. 13 | # On CIs, it is often advantageous to re-use/control the CMake build directory. 14 | # 15 | # It is recommended to build with Ninja and ccache. To do so, set environment 16 | # variables by prefixing to above invocations: 17 | # ``` 18 | # CMAKE_GENERATOR=Ninja CMAKE_C_COMPILER_LAUNCHER=ccache CMAKE_CXX_COMPILER_LAUNCHER=ccache 19 | # ``` 20 | # 21 | # Implementation notes: 22 | # The contents of the wheel is just the contents of the `python_packages` 23 | # directory that our CMake build produces. We go through quite a bit of effort 24 | # on the CMake side to organize that directory already, so we avoid duplicating 25 | # that here, and just package up its contents. 26 | 27 | import os 28 | import pathlib 29 | import shutil 30 | import subprocess 31 | import sys 32 | 33 | from datetime import date 34 | from distutils.command.build import build as _build 35 | from setuptools import setup, Extension 36 | from setuptools.command.build_ext import build_ext 37 | from setuptools.command.build_py import build_py 38 | 39 | 40 | def _check_env_flag(name: str, default=None) -> bool: 41 | return str(os.getenv(name, default)).upper() in ["ON", "1", "YES", "TRUE", "Y"] 42 | 43 | 44 | PACKAGE_VERSION = "".join(str(date.today()).split("-")) 45 | SRC_DIR = pathlib.Path(__file__).parent.absolute() 46 | CMAKE_BUILD_TYPE = os.getenv("CMAKE_BUILD_TYPE", "Release") 47 | MPACT_CMAKE_ALREADY_BUILT = _check_env_flag("MPACT_CMAKE_ALREADY_BUILT", False) 48 | MPACT_CMAKE_BUILD_DIR = os.path.join(SRC_DIR, "build") 49 | 50 | 51 | # Build phase discovery is unreliable. Just tell it what phases to run. 52 | class CustomBuild(_build): 53 | def initialize_options(self): 54 | _build.initialize_options(self) 55 | # Make setuptools not steal the build directory name, 56 | # because the mlir c++ developers are quite 57 | # used to having build/ be for cmake 58 | self.build_base = "setup_build" 59 | 60 | def run(self): 61 | self.run_command("build_py") 62 | self.run_command("build_ext") 63 | self.run_command("build_scripts") 64 | 65 | 66 | class CMakeBuild(build_py): 67 | def cmake_build(self, cmake_build_dir): 68 | llvm_dir = str( 69 | SRC_DIR / "externals" / "torch-mlir" / "externals" / "llvm-project" / "llvm" 70 | ) 71 | cmake_config_args = [ 72 | f"cmake", 73 | f"-GNinja", 74 | f"-DCMAKE_BUILD_TYPE=Release", 75 | f"-DPython3_FIND_VIRTUALENV=ONLY", 76 | f"-DLLVM_ENABLE_PROJECTS=mlir", 77 | f"-DLLVM_EXTERNAL_PROJECTS='torch-mlir;mpact'", 78 | f"-DLLVM_EXTERNAL_TORCH_MLIR_SOURCE_DIR='{SRC_DIR}/externals/torch-mlir'", 79 | f"-DLLVM_EXTERNAL_MPACT_SOURCE_DIR='{SRC_DIR}'", 80 | f"-DLLVM_TARGETS_TO_BUILD=host", 81 | f"-DMLIR_ENABLE_BINDINGS_PYTHON=ON", 82 | # Optimization options for building wheels. 83 | f"-DCMAKE_VISIBILITY_INLINES_HIDDEN=ON", 84 | f"-DCMAKE_C_VISIBILITY_PRESET=hidden", 85 | f"-DCMAKE_CXX_VISIBILITY_PRESET=hidden", 86 | f"{llvm_dir}", 87 | ] 88 | 89 | cmake_build_args = [ 90 | f"cmake", 91 | f"--build", 92 | f".", 93 | f"--target", 94 | f"MPACTPythonModules", 95 | f"MPACTBenchmarkPythonModules", 96 | ] 97 | 98 | try: 99 | subprocess.check_call(cmake_config_args, cwd=cmake_build_dir) 100 | subprocess.check_call(cmake_build_args, cwd=cmake_build_dir) 101 | except subprocess.CalledProcessError as e: 102 | print("cmake build failed with\n", e) 103 | print("debug by follow cmake command:") 104 | sys.exit(e.returncode) 105 | finally: 106 | print(f"cmake config: {' '.join(cmake_config_args)}") 107 | print(f"cmake build: {' '.join(cmake_build_args)}") 108 | print(f"cmake workspace: {cmake_build_dir}") 109 | print(SRC_DIR) 110 | 111 | def run(self): 112 | target_dir = self.build_lib 113 | cmake_build_dir = MPACT_CMAKE_BUILD_DIR 114 | if not cmake_build_dir: 115 | cmake_build_dir = os.path.abspath(os.path.join(target_dir, "..", "build")) 116 | 117 | python_package_dir = os.path.join( 118 | cmake_build_dir, "tools", "mpact", "python_packages", "mpact" 119 | ) 120 | if not MPACT_CMAKE_ALREADY_BUILT: 121 | os.makedirs(cmake_build_dir, exist_ok=True) 122 | cmake_cache_file = os.path.join(cmake_build_dir, "CMakeCache.txt") 123 | if os.path.exists(cmake_cache_file): 124 | os.remove(cmake_cache_file) 125 | # NOTE: With repeated builds for different Python versions, the 126 | # prior version binaries will continue to accumulate. Here we just 127 | # delete the directory where we build native extensions to keep 128 | # this from happening but still take advantage of most of the 129 | # build cache. 130 | mlir_libs_dir = os.path.join(python_package_dir, "mpact", "_mlir_libs") 131 | if os.path.exists(mlir_libs_dir): 132 | print(f"Removing _mlir_mlibs dir to force rebuild: {mlir_libs_dir}") 133 | shutil.rmtree(mlir_libs_dir) 134 | else: 135 | print(f"Not removing _mlir_libs dir (does not exist): {mlir_libs_dir}") 136 | self.cmake_build(cmake_build_dir) 137 | 138 | if os.path.exists(target_dir): 139 | shutil.rmtree(target_dir, ignore_errors=False, onerror=None) 140 | 141 | shutil.copytree(python_package_dir, target_dir, symlinks=False) 142 | 143 | 144 | class CMakeExtension(Extension): 145 | def __init__(self, name, sourcedir=""): 146 | Extension.__init__(self, name, sources=[]) 147 | self.sourcedir = os.path.abspath(sourcedir) 148 | 149 | 150 | class NoopBuildExtension(build_ext): 151 | def build_extension(self, ext): 152 | pass 153 | 154 | 155 | with open("README.md", "r", encoding="utf-8") as fh: 156 | long_description = fh.read() 157 | 158 | 159 | # Requires and extension modules depend on whether building PyTorch 160 | # extensions. 161 | INSTALL_REQUIRES = [ 162 | "numpy", 163 | "packaging", 164 | ] 165 | EXT_MODULES = [ 166 | CMakeExtension("mpact._mlir_libs._mpact"), 167 | ] 168 | 169 | setup( 170 | name="mpact", 171 | version=f"{PACKAGE_VERSION}", 172 | author="Reid Tatge", 173 | author_email="tatge@google.com", 174 | description="MPACT retargetable ML compiler", 175 | long_description=long_description, 176 | long_description_content_type="text/markdown", 177 | include_package_data=True, 178 | cmdclass={ 179 | "build": CustomBuild, 180 | "built_ext": NoopBuildExtension, 181 | "build_py": CMakeBuild, 182 | }, 183 | ext_modules=EXT_MODULES, 184 | python_requires=">=3.8", 185 | install_requires=INSTALL_REQUIRES, 186 | zip_safe=False, 187 | ) 188 | -------------------------------------------------------------------------------- /test/CMakeLists.txt: -------------------------------------------------------------------------------- 1 | #------------------------------------------------------------------------------- 2 | # The MPACT Compiler Tests 3 | #------------------------------------------------------------------------------- 4 | 5 | configure_lit_site_cfg( 6 | ${CMAKE_CURRENT_SOURCE_DIR}/lit.site.cfg.py.in 7 | ${CMAKE_CURRENT_BINARY_DIR}/lit.site.cfg.py 8 | MAIN_CONFIG 9 | ${CMAKE_CURRENT_SOURCE_DIR}/lit.cfg.py 10 | ) 11 | 12 | set(MPACT_TEST_DEPENDS 13 | FileCheck count not 14 | MPACTBenchmarkPythonModules 15 | MPACTPythonModules 16 | TorchMLIRPythonModules 17 | torch-mlir-opt 18 | ) 19 | 20 | add_lit_testsuite(check-mpact "Running the MPACT regression tests" 21 | ${CMAKE_CURRENT_BINARY_DIR} 22 | DEPENDS ${MPACT_TEST_DEPENDS} 23 | ) 24 | set_target_properties(check-mpact PROPERTIES FOLDER "Tests") 25 | 26 | add_lit_testsuites(MPACT ${CMAKE_CURRENT_SOURCE_DIR} DEPENDS ${TORCH_MLIR_TEST_DEPENDS}) 27 | -------------------------------------------------------------------------------- /test/lit.cfg.py: -------------------------------------------------------------------------------- 1 | #------------------------------------------------------------------------------- 2 | # The MPACT Compiler LIT Configuration 3 | #------------------------------------------------------------------------------- 4 | 5 | import os 6 | import platform 7 | import re 8 | import subprocess 9 | import tempfile 10 | 11 | import lit.formats 12 | import lit.util 13 | 14 | from lit.llvm import llvm_config 15 | from lit.llvm.subst import ToolSubst 16 | from lit.llvm.subst import FindTool 17 | 18 | # The name of this test suite. 19 | config.name = "MPACT" 20 | 21 | # The test format. 22 | config.test_format = lit.formats.ShTest(not llvm_config.use_lit_shell) 23 | 24 | # A list of file extensions to treat as test files. 25 | config.suffixes = [".py"] 26 | 27 | # A list of files to exclude from the test suite. 28 | config.excludes = [ 29 | "CMakeLists.txt", 30 | "README.txt", 31 | "LICENSE.txt", 32 | "lit.cfg.py", 33 | "lit.site.cfg.py", 34 | ] 35 | 36 | # The root path where tests are located. 37 | config.test_source_root = os.path.dirname(__file__) 38 | 39 | # The root path where tests should be run. 40 | config.test_exec_root = os.path.join(config.mpact_obj_root, "test") 41 | config.standalone_tools_dir = os.path.join(config.mpact_obj_root, "bin") 42 | 43 | # Substitutions. 44 | config.substitutions.append(("%PATH%", config.environment["PATH"])) 45 | config.substitutions.append(("%shlibext", config.llvm_shlib_ext)) 46 | 47 | # Tweak the PATH to include the tools dir. 48 | llvm_config.with_environment("PATH", config.llvm_tools_dir, append_path=True) 49 | llvm_config.with_environment( 50 | "PATH", os.path.join(config.llvm_build_dir, "bin"), append_path=True 51 | ) 52 | llvm_config.with_system_environment(["HOME", "INCLUDE", "LIB", "TMP", "TEMP"]) 53 | 54 | # On Windows the path to python could contain spaces in which case it needs to 55 | # be provided in quotes. This is the equivalent of how %python is setup in 56 | # llvm/utils/lit/lit/llvm/config.py. 57 | if "Windows" in config.host_os: 58 | config.python_executable = '"%s"' % (config.python_executable) 59 | 60 | # Tools. 61 | tool_dirs = [ 62 | config.standalone_tools_dir, 63 | config.llvm_tools_dir, 64 | config.mpact_obj_root, 65 | ] 66 | tools = [ 67 | "mpact-opt", 68 | ToolSubst("%PYTHON", config.python_executable, unresolved="ignore"), 69 | ] 70 | 71 | llvm_config.add_tool_substitutions(tools, tool_dirs) 72 | 73 | llvm_config.with_environment( 74 | "PYTHONPATH", 75 | [ 76 | os.path.join(config.mpact_obj_root, "python_packages/mpact"), 77 | ], 78 | append_path=True, 79 | ) 80 | -------------------------------------------------------------------------------- /test/lit.site.cfg.py.in: -------------------------------------------------------------------------------- 1 | @LIT_SITE_CFG_IN_HEADER@ 2 | 3 | import sys 4 | 5 | config.host_os = "@HOST_OS@" 6 | config.mpact_src_root = "@MPACT_SOURCE_DIR@" 7 | config.mpact_obj_root = "@MPACT_BINARY_DIR@" 8 | config.torch_mlir_obj_root = "@LLVM_BINARY_DIR@/tools/torch-mlir" 9 | config.llvm_src_root = "@LLVM_SOURCE_DIR@" 10 | config.llvm_obj_root = "@LLVM_BINARY_DIR@" 11 | config.llvm_tools_dir = "@LLVM_TOOLS_DIR@" 12 | config.llvm_build_dir = "@CMAKE_BINARY_DIR@" 13 | config.llvm_lib_dir = "@LLVM_LIBS_DIR@" 14 | config.llvm_shlib_dir = "@SHLIBDIR@" 15 | config.llvm_shlib_ext = "@SHLIBEXT@" 16 | config.llvm_exe_ext = "@EXEEXT@" 17 | config.python_executable = "@Python3_EXECUTABLE@" 18 | 19 | import lit.llvm 20 | lit.llvm.initialize(lit_config, config) 21 | 22 | # Let the main config do the real work. 23 | lit_config.load_config(config, "@MPACT_SOURCE_DIR@/test/lit.cfg.py") 24 | -------------------------------------------------------------------------------- /test/python/add.py: -------------------------------------------------------------------------------- 1 | # RUN: %PYTHON %s | FileCheck %s 2 | 3 | import torch 4 | import numpy as np 5 | 6 | from mpact.mpactbackend import mpact_jit 7 | 8 | from mpact.models.kernels import AddNet 9 | 10 | 11 | def print_sparse(res): 12 | print(res[0]) 13 | print(res[1]) 14 | print(res[2]) 15 | 16 | 17 | net = AddNet() 18 | 19 | # Construct dense and sparse matrices. 20 | X = torch.arange(0, 16, dtype=torch.float32).view(4, 4) 21 | Y = torch.arange(16, 32, dtype=torch.float32).view(4, 4) 22 | A = torch.tensor( 23 | [ 24 | [0.0, 1.0, 0.0, 0.0], 25 | [0.0, 0.0, 0.0, 2.0], 26 | [0.0, 0.0, 0.0, 0.0], 27 | [3.0, 0.0, 0.0, 0.0], 28 | ], 29 | dtype=torch.float32, 30 | ) 31 | S = A.to_sparse_csr() 32 | 33 | # 34 | # CHECK: pytorch 35 | # CHECK: tensor({{\[}}[16., 18., 20., 22.], 36 | # CHECK: [24., 26., 28., 30.], 37 | # CHECK: [32., 34., 36., 38.], 38 | # CHECK: [40., 42., 44., 46.]{{\]}}) 39 | # CHECK: tensor({{\[}}[16., 18., 18., 19.], 40 | # CHECK: [20., 21., 22., 25.], 41 | # CHECK: [24., 25., 26., 27.], 42 | # CHECK: [31., 29., 30., 31.]{{\]}}) 43 | # CHECK: tensor({{\[}}[ 0., 2., 2., 3.], 44 | # CHECK: [ 4., 5., 6., 9.], 45 | # CHECK: [ 8., 9., 10., 11.], 46 | # CHECK: [15., 13., 14., 15.]{{\]}}) 47 | # CHECK: tensor(crow_indices=tensor([0, 1, 2, 2, 3]), 48 | # CHECK: col_indices=tensor([1, 3, 0]), 49 | # CHECK: values=tensor([2., 4., 6.]), size=(4, 4), nnz=3, 50 | # CHECK: layout=torch.sparse_csr) 51 | # CHECK: mpact 52 | # CHECK: {{\[}}[16. 18. 20. 22.] 53 | # CHECK: [24. 26. 28. 30.] 54 | # CHECK: [32. 34. 36. 38.] 55 | # CHECK: [40. 42. 44. 46.]{{\]}} 56 | # CHECK: {{\[}}[16. 18. 18. 19.] 57 | # CHECK: [20. 21. 22. 25.] 58 | # CHECK: [24. 25. 26. 27.] 59 | # CHECK: [31. 29. 30. 31.]{{\]}} 60 | # CHECK: {{\[}}[ 0. 2. 2. 3.] 61 | # CHECK: [ 4. 5. 6. 9.] 62 | # CHECK: [ 8. 9. 10. 11.] 63 | # CHECK: [15. 13. 14. 15.]{{\]}} 64 | # CHECK: [0 1 2 2 3] 65 | # CHECK: [1 3 0] 66 | # CHECK: [2. 4. 6.] 67 | # 68 | 69 | # Run it with PyTorch. 70 | print("pytorch", torch.__version__) 71 | res = net(X, Y) 72 | print(res) 73 | res = net(S, Y) 74 | print(res) 75 | res = net(X, S) 76 | print(res) 77 | res = net(S, S) 78 | print(res) 79 | 80 | # Run it with MPACT. 81 | print("mpact") 82 | res = mpact_jit(net, X, Y) 83 | print(res) 84 | res = mpact_jit(net, S, Y) 85 | print(res) 86 | res = mpact_jit(net, X, S) 87 | print(res) 88 | res = mpact_jit(net, S, S) 89 | print_sparse(res) 90 | -------------------------------------------------------------------------------- /test/python/counteq.py: -------------------------------------------------------------------------------- 1 | # RUN: %PYTHON %s | FileCheck %s 2 | 3 | import torch 4 | import numpy as np 5 | 6 | from mpact.mpactbackend import mpact_jit 7 | 8 | from mpact.models.kernels import CountEq 9 | 10 | 11 | net = CountEq() 12 | 13 | # Construct dense and sparse matrices. 14 | A = torch.tensor( 15 | [ 16 | [0.0, 1.0, 0.0, 0.0], 17 | [0.0, 0.0, 0.0, 2.0], 18 | [0.0, 0.0, 1.0, 1.0], 19 | [3.0, 0.0, 3.0, 0.0], 20 | ], 21 | dtype=torch.float32, 22 | ) 23 | 24 | # TODO: very interesting idiom to sparsify (collapse the sum 25 | # into the eq for full sparsity), but needs PyTorch support 26 | S = A 27 | # S = A.to_sparse() 28 | # S = A.to_sparse_csr() 29 | 30 | # 31 | # CHECK: pytorch 32 | # CHECK: 10 33 | # CHECK: 3 34 | # CHECK: 1 35 | # CHECK: 2 36 | # CHECK: 0 37 | # CHECK: mpact 38 | # CHECK: 10 39 | # CHECK: 3 40 | # CHECK: 1 41 | # CHECK: 2 42 | # CHECK: 0 43 | # 44 | 45 | # Run it with PyTorch. 46 | print("pytorch") 47 | for i in range(5): 48 | target = torch.tensor(i) 49 | res = net(S, target).item() 50 | print(res) 51 | 52 | print("mpact") 53 | for i in range(5): 54 | target = torch.tensor(i) 55 | res = mpact_jit(net, S, target) 56 | print(res) 57 | -------------------------------------------------------------------------------- /test/python/file_formats.py: -------------------------------------------------------------------------------- 1 | # RUN: %PYTHON %s | FileCheck %s 2 | 3 | import numpy as np 4 | 5 | from mpact_benchmark.utils.tensor_generator import ( 6 | generate_tensor, 7 | print_matrix_market_format, 8 | print_extended_frostt_format, 9 | ) 10 | 11 | x = generate_tensor( 12 | seed=0, shape=(4, 7), sparsity=0.5, dtype=np.float32, drange=(4.0, 4.0) 13 | ) 14 | 15 | # CHECK: %%MatrixMarket matrix coordinate real general 16 | # CHECK: % https://math.nist.gov/MatrixMarket/formats.html 17 | # CHECK: % 18 | # CHECK: % density = 50.00% 19 | # CHECK: % 20 | # CHECK: 4 7 14 21 | # CHECK: 1 2 4.0 22 | # CHECK: 1 3 4.0 23 | # CHECK: 1 6 4.0 24 | # CHECK: 2 4 4.0 25 | # CHECK: 2 5 4.0 26 | # CHECK: 2 7 4.0 27 | # CHECK: 3 1 4.0 28 | # CHECK: 3 3 4.0 29 | # CHECK: 3 4 4.0 30 | # CHECK: 3 7 4.0 31 | # CHECK: 4 2 4.0 32 | # CHECK: 4 4 4.0 33 | # CHECK: 4 5 4.0 34 | # CHECK: 4 7 4.0 35 | print_matrix_market_format(x) 36 | 37 | # CHECK: # Tensor in Extended FROSTT file format 38 | # CHECK: # http://frostt.io/tensors/file-formats.html 39 | # CHECK: # extended with two metadata lines: 40 | # CHECK: # rank nnz 41 | # CHECK: # dims (one per rank) 42 | # CHECK: # 43 | # CHECK: # density = 50.00% 44 | # CHECK: # 45 | # CHECK: 2 14 46 | # CHECK: 4 7 47 | # CHECK: 1 2 4.0 48 | # CHECK: 1 3 4.0 49 | # CHECK: 1 6 4.0 50 | # CHECK: 2 4 4.0 51 | # CHECK: 2 5 4.0 52 | # CHECK: 2 7 4.0 53 | # CHECK: 3 1 4.0 54 | # CHECK: 3 3 4.0 55 | # CHECK: 3 4 4.0 56 | # CHECK: 3 7 4.0 57 | # CHECK: 4 2 4.0 58 | # CHECK: 4 4 4.0 59 | # CHECK: 4 5 4.0 60 | # CHECK: 4 7 4.0 61 | print_extended_frostt_format(x) 62 | -------------------------------------------------------------------------------- /test/python/gat.py: -------------------------------------------------------------------------------- 1 | # RUN: %PYTHON %s | FileCheck %s 2 | 3 | import torch 4 | import numpy as np 5 | 6 | from mpact.mpactbackend import mpact_jit 7 | 8 | from mpact.models.gat import gat_4_64_8_3 9 | 10 | net = gat_4_64_8_3() 11 | net.eval() # Switch to inference. 12 | 13 | # Sparse input. 14 | idx = torch.tensor([[0, 0, 1, 2], [0, 2, 3, 1]], dtype=torch.int64) 15 | val = torch.tensor([14.0, 3.0, -8.0, 11.0], dtype=torch.float32) 16 | S = torch.sparse_coo_tensor(idx, val, size=[4, 4]) 17 | 18 | # Construct adjacency matrix. 19 | V = 4 20 | edges = np.array([[0, 1], [0, 2], [1, 2], [1, 3], [2, 3]], dtype=np.int32) 21 | E = edges.shape[0] 22 | adj_mat = torch.sparse_coo_tensor(edges.T, torch.ones(E), (V, V), dtype=torch.int64) 23 | adj_mat = ( 24 | torch.eye(V) + adj_mat 25 | ) # Add self-loops to the adjacency matrix (becomes dense) 26 | 27 | 28 | # 29 | # CHECK: pytorch gat 30 | # CHECK: tensor({{\[}}[-1.0986, -1.0986, -1.0986], 31 | # CHECK: [-1.0986, -1.0986, -1.0986], 32 | # CHECK: [-1.0986, -1.0986, -1.0986], 33 | # CHECK: [-1.0986, -1.0986, -1.0986]{{\]}} 34 | # CHECK: mpact gat 35 | # CHECK: {{\[}}[-1.0986123 -1.0986123 -1.0986123] 36 | # CHECK: [-1.0986123 -1.0986123 -1.0986123] 37 | # CHECK: [-1.0986123 -1.0986123 -1.0986123] 38 | # CHECK: [-1.0986123 -1.0986123 -1.0986123]{{\]}} 39 | # 40 | with torch.no_grad(): 41 | # Run it with PyTorch. 42 | print("pytorch gat") 43 | res = net(S, adj_mat) 44 | print(res) 45 | 46 | # Run it with MPACT. 47 | print("mpact gat") 48 | res = mpact_jit(net, S, adj_mat) 49 | print(res) 50 | -------------------------------------------------------------------------------- /test/python/gcn.py: -------------------------------------------------------------------------------- 1 | # RUN: %PYTHON %s | FileCheck %s 2 | 3 | import torch 4 | 5 | from mpact.mpactbackend import mpact_jit, mpact_jit_compile, mpact_jit_run 6 | 7 | from mpact.models.gcn import graphconv_4_4, gcn_4_16_4 8 | 9 | net = graphconv_4_4() 10 | net.eval() # Switch to inference. 11 | 12 | # Get random (but reproducible) matrices. 13 | torch.manual_seed(0) 14 | inp = torch.rand(4, 4) 15 | adj_mat = torch.rand(4, 4).to_sparse() 16 | 17 | # 18 | # CHECK: pytorch 19 | # CHECK: tensor({{\[}}[4.4778, 4.4778, 4.4778, 4.4778], 20 | # CHECK: [5.7502, 5.7502, 5.7502, 5.7502], 21 | # CHECK: [4.6980, 4.6980, 4.6980, 4.6980], 22 | # CHECK: [3.6407, 3.6407, 3.6407, 3.6407]{{\]}}) 23 | # CHECK: mpact compile and run 24 | # CHECK: {{\[}}[4.477828 4.477828 4.477828 4.477828 ] 25 | # CHECK: [5.7501717 5.7501717 5.7501717 5.7501717] 26 | # CHECK: [4.697952 4.697952 4.697952 4.697952 ] 27 | # CHECK: [3.640687 3.640687 3.640687 3.640687 ]{{\]}} 28 | # CHECK: mpact compile 29 | # CHECK: mpact run 30 | # CHECK: {{\[}}[4.477828 4.477828 4.477828 4.477828 ] 31 | # CHECK: [5.7501717 5.7501717 5.7501717 5.7501717] 32 | # CHECK: [4.697952 4.697952 4.697952 4.697952 ] 33 | # CHECK: [3.640687 3.640687 3.640687 3.640687 ]{{\]}} 34 | # CHECK: mpact compile opt=3 35 | # CHECK: mpact run 36 | # CHECK: {{\[}}[4.477828 4.477828 4.477828 4.477828 ] 37 | # CHECK: [5.7501717 5.7501717 5.7501717 5.7501717] 38 | # CHECK: [4.697952 4.697952 4.697952 4.697952 ] 39 | # 40 | with torch.no_grad(): 41 | # Run it with PyTorch. 42 | print("pytorch") 43 | res = net(inp, adj_mat) 44 | print(res) 45 | 46 | # Run it with MPACT (compile and run at once). 47 | print("mpact compile and run") 48 | res = mpact_jit(net, inp, adj_mat) 49 | print(res) 50 | 51 | # Run it with MPACT (with separate compile and run steps). 52 | print("mpact compile") 53 | invoker, fn = mpact_jit_compile(net, inp, adj_mat) 54 | print("mpact run") 55 | res = mpact_jit_run(invoker, fn, inp, adj_mat) 56 | print(res) 57 | 58 | # Run it with MPACT (with separate compile and run steps, given opt_level). 59 | print("mpact compile opt=3") 60 | invoker, fn = mpact_jit_compile(net, inp, adj_mat, opt_level=3) 61 | print("mpact run") 62 | res = mpact_jit_run(invoker, fn, inp, adj_mat) 63 | print(res) 64 | 65 | net = gcn_4_16_4() 66 | net.eval() # Switch to inference. 67 | 68 | 69 | # Sparse input. 70 | idx = torch.tensor([[0, 0, 1, 2], [0, 2, 3, 1]], dtype=torch.int64) 71 | val = torch.tensor([14.0, 3.0, -8.0, 11.0], dtype=torch.float32) 72 | S = torch.sparse_coo_tensor(idx, val, size=[4, 4]) 73 | 74 | # 75 | # CHECK: pytorch gcn 76 | # CHECK: tensor({{\[}}[-1.3863, -1.3863, -1.3863, -1.3863], 77 | # CHECK: [-1.3863, -1.3863, -1.3863, -1.3863], 78 | # CHECK: [-1.3863, -1.3863, -1.3863, -1.3863], 79 | # CHECK: [-1.3863, -1.3863, -1.3863, -1.3863]]) 80 | # CHECK: mpact gcn 81 | # CHECK: {{\[}}[-1.3862944 -1.3862944 -1.3862944 -1.3862944] 82 | # CHECK: [-1.3862944 -1.3862944 -1.3862944 -1.3862944] 83 | # CHECK: [-1.3862944 -1.3862944 -1.3862944 -1.3862944] 84 | # CHECK: [-1.3862944 -1.3862944 -1.3862944 -1.3862944]{{\]}} 85 | # 86 | with torch.no_grad(): 87 | # Run it with PyTorch. 88 | print("pytorch gcn") 89 | res = net(S, adj_mat) 90 | print(res) 91 | 92 | # Run it with MPACT. 93 | print("mpact gcn") 94 | res = mpact_jit(net, S, adj_mat) 95 | print(res) 96 | -------------------------------------------------------------------------------- /test/python/lif.py: -------------------------------------------------------------------------------- 1 | # RUN: %PYTHON %s | FileCheck %s 2 | 3 | import torch 4 | 5 | from mpact.mpactbackend import mpact_jit 6 | 7 | from mpact.models.lif import LIFSumOfSq 8 | 9 | net = LIFSumOfSq() 10 | 11 | # Get a random (but reproducible) input, so that a 12 | # general sparse tensor appears after LIF. 13 | torch.manual_seed(0) 14 | x = torch.rand(2, 3, 8, 8) 15 | 16 | # 17 | # CHECK: pytorch 18 | # CHECK: tensor([ 0., 11., 9., 11., 13., 11., 10., 12.]) 19 | # CHECK: mpact 20 | # CHECK: [ 0. 11. 9. 11. 13. 11. 10. 12.] 21 | # 22 | 23 | # Run it with PyTorch. 24 | print("pytorch") 25 | res = net(x) 26 | print(res) 27 | 28 | # Run it with MPACT. 29 | print("mpact") 30 | res = mpact_jit(net, x) 31 | print(res) 32 | -------------------------------------------------------------------------------- /test/python/mm.py: -------------------------------------------------------------------------------- 1 | # RUN: %PYTHON %s | FileCheck %s 2 | 3 | import torch 4 | import numpy as np 5 | 6 | from mpact.mpactbackend import mpact_jit 7 | 8 | from mpact.models.kernels import MMNet 9 | 10 | 11 | def print_sparse(res): 12 | print(res[0]) 13 | print(res[1]) 14 | print(res[2]) 15 | 16 | 17 | net = MMNet() 18 | 19 | # Construct dense and sparse matrices. 20 | X = torch.arange(0, 16, dtype=torch.float32).view(4, 4) 21 | Y = torch.arange(16, 32, dtype=torch.float32).view(4, 4) 22 | A = torch.tensor( 23 | [ 24 | [0.0, 1.0, 0.0, 0.0], 25 | [0.0, 0.0, 0.0, 2.0], 26 | [0.0, 0.0, 0.0, 0.0], 27 | [3.0, 0.0, 0.0, 0.0], 28 | ], 29 | dtype=torch.float32, 30 | ) 31 | S = A.to_sparse_csr() 32 | 33 | # 34 | # CHECK: pytorch 35 | # CHECK: tensor({{\[}}[ 152., 158., 164., 170.], 36 | # CHECK: [ 504., 526., 548., 570.], 37 | # CHECK: [ 856., 894., 932., 970.], 38 | # CHECK: [1208., 1262., 1316., 1370.]{{\]}}) 39 | # CHECK: tensor({{\[}}[20., 21., 22., 23.], 40 | # CHECK: [56., 58., 60., 62.], 41 | # CHECK: [ 0., 0., 0., 0.], 42 | # CHECK: [48., 51., 54., 57.]{{\]}}) 43 | # CHECK: tensor({{\[}}[ 9., 0., 0., 2.], 44 | # CHECK: [21., 4., 0., 10.], 45 | # CHECK: [33., 8., 0., 18.], 46 | # CHECK: [45., 12., 0., 26.]{{\]}}) 47 | # CHECK: tensor(crow_indices=tensor([0, 1, 2, 2, 3]), 48 | # CHECK: col_indices=tensor([3, 0, 1]), 49 | # CHECK: values=tensor([2., 6., 3.]), size=(4, 4), nnz=3, 50 | # CHECK: layout=torch.sparse_csr) 51 | # CHECK: mpact 52 | # CHECK: {{\[}}[ 152. 158. 164. 170.] 53 | # CHECK: [ 504. 526. 548. 570.] 54 | # CHECK: [ 856. 894. 932. 970.] 55 | # CHECK: [1208. 1262. 1316. 1370.]{{\]}} 56 | # CHECK: {{\[}}[20. 21. 22. 23.] 57 | # CHECK: [56. 58. 60. 62.] 58 | # CHECK: [ 0. 0. 0. 0.] 59 | # CHECK: [48. 51. 54. 57.]{{\]}} 60 | # CHECK: {{\[}}[ 9. 0. 0. 2.] 61 | # CHECK: [21. 4. 0. 10.] 62 | # CHECK: [33. 8. 0. 18.] 63 | # CHECK: [45. 12. 0. 26.]{{\]}} 64 | # CHECK: [0 1 2 2 3] 65 | # CHECK: [3 0 1] 66 | # CHECK: [2. 6. 3.] 67 | # 68 | 69 | # Run it with PyTorch. 70 | print("pytorch") 71 | res = net(X, Y) 72 | print(res) 73 | res = net(S, Y) 74 | print(res) 75 | res = net(X, S) 76 | print(res) 77 | res = net(S, S) 78 | print(res) 79 | 80 | # Run it with MPACT. 81 | print("mpact") 82 | res = mpact_jit(net, X, Y) 83 | print(res) 84 | res = mpact_jit(net, S, Y) 85 | print(res) 86 | res = mpact_jit(net, X, S) 87 | print(res) 88 | res = mpact_jit(net, S, S) 89 | print_sparse(res) 90 | -------------------------------------------------------------------------------- /test/python/mm_print.py: -------------------------------------------------------------------------------- 1 | # RUN: %PYTHON %s | FileCheck %s 2 | 3 | import torch 4 | import numpy as np 5 | 6 | from mpact.mpactbackend import mpact_linalg, mpact_stablehlo 7 | 8 | from mpact.models.kernels import MMNet 9 | 10 | 11 | net = MMNet() 12 | 13 | X = torch.arange(0, 16, dtype=torch.float32).view(4, 4) 14 | Y = torch.arange(16, 32, dtype=torch.float32).view(4, 4) 15 | 16 | # 17 | # CHECK: module { 18 | # CHECK: func.func @main(%[[A0:.*]]: tensor<4x4xf32>, %[[A1:.*]]: tensor<4x4xf32>) -> tensor<4x4xf32> { 19 | # CHECK: %[[C0:.*]] = arith.constant 0.000000e+00 : f32 20 | # CHECK: %[[T0:.*]] = tensor.empty() : tensor<4x4xf32> 21 | # CHECK: %[[T1:.*]] = linalg.fill ins(%[[C0]] : f32) outs(%[[T0]] : tensor<4x4xf32>) -> tensor<4x4xf32> 22 | # CHECK: %[[T2:.*]] = linalg.matmul 23 | # CHECK-SAME: ins(%[[A0]], %[[A1]] : tensor<4x4xf32>, tensor<4x4xf32>) 24 | # CHECK-SAME: outs(%[[T1]] : tensor<4x4xf32>) -> tensor<4x4xf32> 25 | # CHECK: return %2 : tensor<4x4xf32> 26 | # CHECK: } 27 | # CHECK: } 28 | # 29 | 30 | linalg = mpact_linalg(net, X, Y) 31 | print(linalg) 32 | 33 | # 34 | # CHECK: module { 35 | # CHECK: func.func @main(%[[A0:.*]]: tensor<4x4xf32>, %[[A1:.*]]: tensor<4x4xf32>) -> tensor<4x4xf32> { 36 | # CHECK: %[[T0:.*]] = stablehlo.dot_general %[[A0]], %[[A1]], contracting_dims = [1] x [0] : (tensor<4x4xf32>, tensor<4x4xf32>) -> tensor<4x4xf32> 37 | # CHECK: return %[[T0]] : tensor<4x4xf32> 38 | # CHECK: } 39 | # CHECK: } 40 | 41 | stablehlo = mpact_stablehlo(net, X, Y) 42 | print(stablehlo) 43 | -------------------------------------------------------------------------------- /test/python/mul.py: -------------------------------------------------------------------------------- 1 | # RUN: %PYTHON %s | FileCheck %s 2 | 3 | import torch 4 | import numpy as np 5 | 6 | from mpact.mpactbackend import mpact_jit 7 | 8 | from mpact.models.kernels import MulNet 9 | 10 | 11 | def print_sparse(res): 12 | print(res[0]) 13 | print(res[1]) 14 | print(res[2]) 15 | 16 | 17 | net = MulNet() 18 | 19 | # Construct dense and sparse matrices. 20 | X = torch.arange(0, 16, dtype=torch.float32).view(4, 4) 21 | Y = torch.arange(16, 32, dtype=torch.float32).view(4, 4) 22 | A = torch.tensor( 23 | [ 24 | [0.0, 1.0, 0.0, 0.0], 25 | [0.0, 0.0, 0.0, 2.0], 26 | [0.0, 0.0, 0.0, 0.0], 27 | [3.0, 0.0, 0.0, 0.0], 28 | ], 29 | dtype=torch.float32, 30 | ) 31 | S = A.to_sparse_csr() 32 | 33 | # 34 | # CHECK: pytorch 35 | # CHECK: tensor({{\[}}[ 0., 17., 36., 57.], 36 | # CHECK: [ 80., 105., 132., 161.], 37 | # CHECK: [192., 225., 260., 297.], 38 | # CHECK: [336., 377., 420., 465.]{{\]}}) 39 | # CHECK: tensor(crow_indices=tensor([0, 1, 2, 2, 3]), 40 | # CHECK: col_indices=tensor([1, 3, 0]), 41 | # CHECK: values=tensor([17., 46., 84.]), size=(4, 4), nnz=3, 42 | # CHECK: layout=torch.sparse_csr) 43 | # CHECK: tensor(crow_indices=tensor([0, 1, 2, 2, 3]), 44 | # CHECK: col_indices=tensor([1, 3, 0]), 45 | # CHECK: values=tensor([ 1., 14., 36.]), size=(4, 4), nnz=3, 46 | # CHECK: layout=torch.sparse_csr) 47 | # CHECK: tensor(crow_indices=tensor([0, 1, 2, 2, 3]), 48 | # CHECK: col_indices=tensor([1, 3, 0]), 49 | # CHECK: values=tensor([1., 4., 9.]), size=(4, 4), nnz=3, 50 | # CHECK: layout=torch.sparse_csr) 51 | # CHECK: mpact 52 | # CHECK: {{\[}}[ 0. 17. 36. 57.] 53 | # CHECK: [ 80. 105. 132. 161.] 54 | # CHECK: [192. 225. 260. 297.] 55 | # CHECK: [336. 377. 420. 465.]{{\]}} 56 | # CHECK: [0 1 2 2 3] 57 | # CHECK: [1 3 0] 58 | # CHECK: [17. 46. 84.] 59 | # CHECK: [0 1 2 2 3] 60 | # CHECK: [1 3 0] 61 | # CHECK: [ 1. 14. 36.] 62 | # CHECK: [0 1 2 2 3] 63 | # CHECK: [1 3 0] 64 | # CHECK: [1. 4. 9.] 65 | # 66 | 67 | # Run it with PyTorch. 68 | print("pytorch") 69 | res = net(X, Y) 70 | print(res) 71 | res = net(S, Y) 72 | print(res) 73 | res = net(X, S) 74 | print(res) 75 | res = net(S, S) 76 | print(res) 77 | 78 | # Run it with MPACT. 79 | print("mpact") 80 | res = mpact_jit(net, X, Y) 81 | print(res) 82 | res = mpact_jit(net, S, Y) 83 | print_sparse(res) 84 | res = mpact_jit(net, X, S) 85 | print_sparse(res) 86 | res = mpact_jit(net, S, S) 87 | print_sparse(res) 88 | -------------------------------------------------------------------------------- /test/python/norm.py: -------------------------------------------------------------------------------- 1 | # RUN: %PYTHON %s | FileCheck %s 2 | 3 | import torch 4 | import numpy as np 5 | 6 | from mpact.mpactbackend import mpact_jit 7 | 8 | from mpact.models.kernels import Normalization 9 | 10 | net = Normalization() 11 | 12 | # Construct adjacency matrix. 13 | V = 8 14 | edges = np.array([[0, 1], [0, 4], [1, 4], [3, 4], [4, 3]], dtype=np.int32) 15 | E = edges.shape[0] 16 | adj_mat = torch.sparse_coo_tensor(edges.T, torch.ones(E), (V, V), dtype=torch.int64) 17 | adj_mat = ( 18 | torch.eye(V) + adj_mat 19 | ) # Add self-loops to the adjacency matrix (becomes dense) 20 | 21 | # 22 | # CHECK: pytorch 23 | # CHECK: tensor({{\[}}[0.1111, 0.1667, 0.0000, 0.0000, 0.1667, 0.0000, 0.0000, 0.0000], 24 | # CHECK: [0.0000, 0.2500, 0.0000, 0.0000, 0.2500, 0.0000, 0.0000, 0.0000], 25 | # CHECK: [0.0000, 0.0000, 1.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000], 26 | # CHECK: [0.0000, 0.0000, 0.0000, 0.2500, 0.2500, 0.0000, 0.0000, 0.0000], 27 | # CHECK: [0.0000, 0.0000, 0.0000, 0.2500, 0.2500, 0.0000, 0.0000, 0.0000], 28 | # CHECK: [0.0000, 0.0000, 0.0000, 0.0000, 0.0000, 1.0000, 0.0000, 0.0000], 29 | # CHECK: [0.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000, 1.0000, 0.0000], 30 | # CHECK: [0.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000, 1.0000]{{\]}}) 31 | # 32 | # TODO: first row? 33 | # 34 | # CHECK: mpact 35 | # CHECK: {{\[}}[0. 0. 0. 0. 0. 0. 0. 0. ] 36 | # CHECK: [0. 0.25 0. 0. 0.25 0. 0. 0. ] 37 | # CHECK: [0. 0. 1. 0. 0. 0. 0. 0. ] 38 | # CHECK: [0. 0. 0. 0.25 0.25 0. 0. 0. ] 39 | # CHECK: [0. 0. 0. 0.25 0.25 0. 0. 0. ] 40 | # CHECK: [0. 0. 0. 0. 0. 1. 0. 0. ] 41 | # CHECK: [0. 0. 0. 0. 0. 0. 1. 0. ] 42 | # CHECK: [0. 0. 0. 0. 0. 0. 0. 1. ]{{\]}} 43 | # 44 | 45 | # Run it with PyTorch. 46 | print("pytorch") 47 | res = net(adj_mat) 48 | print(res) 49 | 50 | # Run it with MPACT. 51 | print("mpact") 52 | res = mpact_jit(net, adj_mat) 53 | print(res) 54 | -------------------------------------------------------------------------------- /test/python/resnet.py: -------------------------------------------------------------------------------- 1 | # RUN: %PYTHON %s | FileCheck %s 2 | 3 | import torch 4 | import numpy as np 5 | 6 | from mpact.mpactbackend import mpact_jit 7 | 8 | from mpact.models.resnet import resnet_20 9 | 10 | resnet = resnet_20() 11 | resnet.eval() # Switch to inference. 12 | 13 | # Get a random input. 14 | # B x RGB x H x W 15 | x = torch.rand(1, 3, 16, 16) 16 | 17 | # 18 | # CHECK: pytorch 19 | # CHECK: mpact 20 | # CHECK: passed 21 | # 22 | 23 | with torch.no_grad(): 24 | # Run it with PyTorch. 25 | print("pytorch") 26 | res1 = resnet(x) 27 | print(res1) 28 | 29 | # Run it with MPACT. 30 | print("mpact") 31 | res2 = mpact_jit(resnet, x) 32 | print(res2) 33 | 34 | # Completely different inputs and weights for each run, 35 | # so we simply verify the two results are the same. 36 | np.testing.assert_allclose(res1.numpy(), res2, rtol=1e-5, atol=0) 37 | print("passed") 38 | -------------------------------------------------------------------------------- /test/python/scale.py: -------------------------------------------------------------------------------- 1 | # RUN: %PYTHON %s | FileCheck %s 2 | 3 | import torch 4 | import numpy as np 5 | 6 | from mpact.mpactbackend import mpact_jit 7 | 8 | from mpact.models.kernels import FeatureScale 9 | 10 | net = FeatureScale() 11 | 12 | # Get random (but reproducible) matrices. 13 | torch.manual_seed(0) 14 | features = torch.rand(7, 7) 15 | 16 | # 17 | # CHECK: pytorch 18 | # CHECK: tensor({{\[}}[0.1702, 0.2634, 0.0303, 0.0453, 0.1054, 0.2174, 0.1680], 19 | # CHECK: [0.3064, 0.1557, 0.2161, 0.1192, 0.1373, 0.0076, 0.0577], 20 | # CHECK: [0.0856, 0.1510, 0.2031, 0.2329, 0.0469, 0.0822, 0.1984], 21 | # CHECK: [0.2207, 0.0957, 0.2108, 0.1011, 0.1333, 0.2297, 0.0087], 22 | # CHECK: [0.0774, 0.1561, 0.1275, 0.3896, 0.0735, 0.1128, 0.0630], 23 | # CHECK: [0.0093, 0.0611, 0.2731, 0.2124, 0.2180, 0.1546, 0.0716], 24 | # CHECK: [0.2026, 0.0115, 0.0481, 0.0839, 0.2826, 0.2749, 0.0964]{{\]}}) 25 | # 26 | # TODO: first row? 27 | # 28 | # CHECK: mpact 29 | # CHECK: {{\[}}[0. 0. 0. 0. 0. 0. 30 | # CHECK: 0. ] 31 | # CHECK: [0.30635384 0.15570773 0.21608633 0.11923195 0.13728413 0.00762967 32 | # CHECK: 0.05770639] 33 | # CHECK: [0.08555716 0.15095268 0.20310582 0.23290026 0.04687909 0.08217437 34 | # CHECK: 0.19843069] 35 | # CHECK: [0.22065267 0.09574053 0.2107584 0.10111907 0.13330552 0.22970453 36 | # CHECK: 0.00871931] 37 | # CHECK: [0.07743214 0.15609969 0.12754099 0.3896042 0.07353575 0.11279855 38 | # CHECK: 0.06298868] 39 | # CHECK: [0.00931544 0.06112389 0.2730649 0.2123639 0.21801054 0.15456341 40 | # CHECK: 0.07155795] 41 | # CHECK: [0.20259099 0.01148908 0.04807246 0.08394676 0.28260148 0.2748705 42 | # CHECK: 0.09642864]] 43 | # 44 | 45 | # Run it with PyTorch. 46 | print("pytorch") 47 | res = net(features) 48 | print(res) 49 | 50 | # Run it with MPACT. 51 | print("mpact") 52 | res = mpact_jit(net, features) 53 | print(res) 54 | -------------------------------------------------------------------------------- /test/python/sddmm.py: -------------------------------------------------------------------------------- 1 | # RUN: %PYTHON %s | FileCheck %s 2 | 3 | import torch 4 | import numpy as np 5 | 6 | from mpact.mpactbackend import mpact_jit, mpact_jit_compile, mpact_jit_run 7 | 8 | from mpact.models.kernels import MMNet, SDDMMNet 9 | 10 | 11 | def print_sparse(res): 12 | print(res[0]) 13 | print(res[1]) 14 | print(res[2]) 15 | print(res[3]) 16 | 17 | 18 | mmnet = MMNet() 19 | sddmmnet = SDDMMNet() 20 | 21 | # Construct very sparse matrix. 22 | idx = torch.tensor([[0, 4], [0, 4]], dtype=torch.int64) 23 | val = torch.tensor([2.0, 3.0], dtype=torch.float64) 24 | S = torch.sparse_coo_tensor(idx, val, size=[5, 5]) 25 | 26 | # Trivial dense inputs. 27 | A = torch.arange(0, 25, dtype=torch.float32).view(5, 5) 28 | B = torch.arange(25, 50, dtype=torch.float32).view(5, 5) 29 | 30 | # 31 | # CHECK: pytorch 32 | # CHECK: tensor({{\[}}[ 400., 410., 420., 430., 440.], 33 | # CHECK: [1275., 1310., 1345., 1380., 1415.], 34 | # CHECK: [2150., 2210., 2270., 2330., 2390.], 35 | # CHECK: [3025., 3110., 3195., 3280., 3365.], 36 | # CHECK: [3900., 4010., 4120., 4230., 4340.]{{\]}}) 37 | # CHECK: tensor(indices=tensor({{\[}}[0, 4], 38 | # CHECK: [0, 4]{{\]}}), 39 | # CHECK: values=tensor([ 800., 13020.]), 40 | # CHECK: size=(5, 5), nnz=2, dtype=torch.float64, layout=torch.sparse_coo) 41 | # CHECK: mpact 42 | # CHECK: {{\[}}[ 400. 410. 420. 430. 440.] 43 | # CHECK: [1275. 1310. 1345. 1380. 1415.] 44 | # CHECK: [2150. 2210. 2270. 2330. 2390.] 45 | # CHECK: [3025. 3110. 3195. 3280. 3365.] 46 | # CHECK: [3900. 4010. 4120. 4230. 4340.]{{\]}} 47 | # CHECK: [0 2] 48 | # CHECK: [0 4] 49 | # CHECK: [0 4] 50 | # CHECK: [ 800. 13020.] 51 | # 52 | 53 | # Run it with PyTorch. 54 | print("pytorch") 55 | dense = mmnet(A, B) 56 | print(dense) 57 | res = sddmmnet(S, A, B) 58 | print(res) 59 | 60 | # Run it with MPACT. 61 | print("mpact") 62 | dense = mpact_jit(mmnet, A, B) 63 | print(dense) 64 | res = mpact_jit(sddmmnet, S, A, B) 65 | print_sparse(res) 66 | -------------------------------------------------------------------------------- /test/python/spmv.py: -------------------------------------------------------------------------------- 1 | # RUN: %PYTHON %s | FileCheck %s 2 | 3 | import torch 4 | 5 | from mpact.mpactbackend import mpact_jit 6 | 7 | from mpact.models.kernels import MVNet 8 | 9 | net = MVNet() 10 | 11 | # Get a fixed vector and matrix (which we make 2x2 block "sparse"). 12 | dense_vector = torch.arange(1, 11, dtype=torch.float32) 13 | dense_input = torch.arange(1, 101, dtype=torch.float32).view(10, 10) 14 | sparse_matrix = dense_input.to_sparse_bsr(blocksize=(2, 2)) 15 | 16 | # 17 | # CHECK: pytorch 18 | # CHECK: tensor([ 385., 935., 1485., 2035., 2585., 3135., 3685., 4235., 4785., 5335.]) 19 | # CHECK: mpact 20 | # CHECK: [ 385. 935. 1485. 2035. 2585. 3135. 3685. 4235. 4785. 5335.] 21 | # 22 | 23 | # Run it with PyTorch. 24 | print("pytorch") 25 | res = net(sparse_matrix, dense_vector) 26 | print(res) 27 | 28 | # Run it with MPACT. 29 | print("mpact") 30 | res = mpact_jit(net, sparse_matrix, dense_vector) 31 | print(res) 32 | -------------------------------------------------------------------------------- /test/python/sqsum.py: -------------------------------------------------------------------------------- 1 | # RUN: %PYTHON %s | FileCheck %s 2 | 3 | import torch 4 | import numpy as np 5 | 6 | from mpact.mpactbackend import mpact_jit 7 | 8 | from mpact.models.kernels import SqSum 9 | 10 | net = SqSum() 11 | 12 | # Construct adjacency matrix. 13 | V = 8 14 | edges = np.array([[0, 1], [0, 4], [1, 4], [3, 4], [4, 3]], dtype=np.int32) 15 | E = edges.shape[0] 16 | adj_mat = torch.sparse_coo_tensor(edges.T, torch.ones(E), (V, V), dtype=torch.int64) 17 | 18 | # 19 | # CHECK: pytorch 20 | # CHECK: tensor(5) 21 | # CHECK: mpact 22 | # CHECK: 5 23 | 24 | # Run it with PyTorch. 25 | print("pytorch") 26 | res = net(adj_mat) 27 | print(res) 28 | 29 | # Run it with MPACT. 30 | print("mpact") 31 | # Try sparsification with sparse iterator 32 | # TODO: will work after explicit value is specified in the encoding. 33 | # res = mpact_jit(net, adj_mat, use_sp_it=True) 34 | # print(res) 35 | # Try sparsification directly with scf.for/while 36 | res = mpact_jit(net, adj_mat) 37 | print(res) 38 | -------------------------------------------------------------------------------- /test/python/train_simple.py: -------------------------------------------------------------------------------- 1 | # RUN: %PYTHON %s | FileCheck %s 2 | 3 | import torch 4 | import numpy as np 5 | 6 | from torch.utils.data import Dataset, DataLoader 7 | 8 | from mpact.mpactbackend import mpact_jit 9 | from mpact.models.kernels import SimpleNet 10 | from mpact.models.train import training_loop, num_all_parameters, num_parameters 11 | 12 | 13 | A = torch.tensor( 14 | [ 15 | [ 16 | [1.0, 1.0, 1.0, 1.0], 17 | [0.0, 1.0, 1.0, 1.0], 18 | [1.0, 1.0, 1.0, 1.0], 19 | [1.0, 1.0, 1.0, 1.0], 20 | ], 21 | [ 22 | [0.0, 0.0, 0.0, 0.0], 23 | [0.0, 0.0, 0.0, 0.0], 24 | [0.0, 0.0, 0.0, 0.0], 25 | [1.0, 0.0, 0.0, 0.0], 26 | ], 27 | [ 28 | [1.0, 1.0, 1.0, 1.0], 29 | [1.0, 1.0, 1.0, 1.0], 30 | [1.0, 1.0, 1.0, 1.0], 31 | [1.0, 1.0, 1.0, 0.0], 32 | ], 33 | [ 34 | [0.0, 0.0, 0.0, 1.0], 35 | [0.0, 0.0, 0.0, 0.0], 36 | [0.0, 0.0, 0.0, 0.0], 37 | [0.0, 0.0, 0.0, 0.0], 38 | ], 39 | ], 40 | dtype=torch.float32, 41 | ) 42 | 43 | 44 | B = torch.tensor( 45 | [ 46 | [ 47 | [0.0, 0.0, 0.0, 0.0], 48 | [0.0, 0.0, 0.0, 0.0], 49 | [0.0, 0.0, 0.0, 0.0], 50 | [0.0, 0.0, 0.0, 0.0], 51 | ], 52 | [ 53 | [1.0, 1.0, 1.0, 1.0], 54 | [1.0, 1.0, 1.0, 1.0], 55 | [1.0, 1.0, 1.0, 1.0], 56 | [1.0, 1.0, 1.0, 1.0], 57 | ], 58 | ], 59 | dtype=torch.float32, 60 | ) 61 | 62 | # Labels 0:sparse 1:dense 63 | 64 | labA = torch.tensor([1, 0, 1, 0]) 65 | 66 | labB = torch.tensor([0, 1]) 67 | 68 | # A toy training and validation data set consisting of dense/sparse tensors. 69 | 70 | 71 | class TrainData(Dataset): 72 | def __len__(self): 73 | return A.shape[0] 74 | 75 | def __getitem__(self, index): 76 | return A[index], labA[index] 77 | 78 | 79 | class ValidationData(Dataset): 80 | def __len__(self): 81 | return B.shape[0] 82 | 83 | def __getitem__(self, index): 84 | return B[index], labB[index] 85 | 86 | 87 | train_data = TrainData() 88 | validation_data = ValidationData() 89 | 90 | net = SimpleNet() 91 | optimizer = torch.optim.Adam(net.parameters(), lr=0.001) 92 | loss_function = torch.nn.CrossEntropyLoss() 93 | train = DataLoader(train_data, batch_size=2) 94 | validation = DataLoader(validation_data, batch_size=2) 95 | 96 | 97 | # CHECK-LABEL: parameters 98 | # CHECK-COUNT-2: 182 99 | print("parameters") 100 | print(num_all_parameters(net)) 101 | print(num_parameters(net)) 102 | 103 | # Run it with PyTorch. 104 | # CHECK-LABEL: pytorch 105 | # CHECK: Epoch 9 106 | print("pytorch") 107 | training_loop(net, optimizer, loss_function, train, validation, epochs=10) 108 | 109 | # Run it with MPACT. 110 | # CHECK-LABEL: mpact 111 | print("mpact") 112 | # TODO: teach MPACT about autograd 113 | -------------------------------------------------------------------------------- /tools/CMakeLists.txt: -------------------------------------------------------------------------------- 1 | add_subdirectory(mpact-opt) 2 | -------------------------------------------------------------------------------- /tools/mpact-opt/CMakeLists.txt: -------------------------------------------------------------------------------- 1 | add_llvm_executable(mpact-opt mpact_opt.cpp) 2 | 3 | set(dependency_libraries) 4 | if(TORCH_MLIR_ENABLE_STABLEHLO) 5 | list(APPEND dependency_libraries StablehloRegister) 6 | endif() 7 | 8 | target_link_libraries(mpact-opt PRIVATE 9 | MLIROptLib 10 | MLIRTransforms 11 | TorchMLIRInitAll 12 | TorchMLIRTorchDialect 13 | TorchMLIRTorchPasses 14 | 15 | MPACTTransformPasses 16 | ${dependency_libraries} 17 | ) 18 | -------------------------------------------------------------------------------- /tools/mpact-opt/mpact_opt.cpp: -------------------------------------------------------------------------------- 1 | //===- mpact-opt.cpp - MLIR Optimizer Driver -------------------------===// 2 | // 3 | // Part of the MPACT Project, under the Apache License v2.0 with LLVM 4 | // Exceptions. 5 | // See https://llvm.org/LICENSE.txt for license information. 6 | // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception 7 | // Also available under a BSD-style license. See LICENSE. 8 | // 9 | //===----------------------------------------------------------------------===// 10 | 11 | #include "mlir/Dialect/SCF/IR/SCF.h" 12 | #include "mlir/Dialect/Tensor/IR/Tensor.h" 13 | #include "mlir/Tools/mlir-opt/MlirOptMain.h" 14 | #include "mlir/Transforms/Passes.h" 15 | #include "mpact/Transforms/Passes.h" 16 | #include "torch-mlir/InitAll.h" 17 | 18 | #ifdef TORCH_MLIR_ENABLE_STABLEHLO 19 | #include "stablehlo/dialect/Register.h" 20 | #endif 21 | 22 | using namespace mlir; 23 | 24 | int main(int argc, char **argv) { 25 | mlir::mpact::registerTransformPasses(); 26 | 27 | mlir::torch::registerAllPasses(); 28 | 29 | // Core Transforms 30 | registerCanonicalizerPass(); 31 | registerCSEPass(); 32 | registerInlinerPass(); 33 | registerLocationSnapshotPass(); 34 | registerLoopInvariantCodeMotionPass(); 35 | registerPrintOpStatsPass(); 36 | registerViewOpGraphPass(); 37 | registerStripDebugInfoPass(); 38 | registerSymbolDCEPass(); 39 | 40 | DialectRegistry registry; 41 | mlir::torch::registerAllDialects(registry); 42 | mlir::torch::registerOptionalInputDialects(registry); 43 | 44 | #ifdef TORCH_MLIR_ENABLE_STABLEHLO 45 | mlir::stablehlo::registerAllDialects(registry); 46 | #endif 47 | return mlir::asMainReturnCode(mlir::MlirOptMain( 48 | argc, argv, "MLIR modular optimizer driver\n", registry)); 49 | } 50 | --------------------------------------------------------------------------------