initial commit
diff --git a/.github/workflows/build-and-test.yml b/.github/workflows/build-and-test.yml new file mode 100644 index 0000000..423d705 --- /dev/null +++ b/.github/workflows/build-and-test.yml
@@ -0,0 +1,66 @@ +name: Build and test + +on: + pull_request: + branches: [ "main" ] + +concurrency: + # Use PR number as key for a pull request or the commit hash otherwise. This cancels + # queued and in-progress runs for the same PR (presubmit) or commit + # (postsubmit). + group: ci-build-test-cpp-linux-${{ github.event.number || github.sha }} + cancel-in-progress: true + +jobs: + build: + runs-on: ubuntu-latest + env: + CACHE_DIR: ${{ github.workspace }}/.ccache + steps: + - uses: actions/checkout@v4 + with: + submodules: recursive + token: ${{ secrets.GITHUB_TOKEN }} + + - name: Setup Python Version + uses: actions/setup-python@v5 + with: + python-version: 3.11 # Install the python version needed + + - name: Set PYTHONPATH + run: export PYTHONPATH=build/tools/mpact/python_packages/mpact + shell: bash + + - name: Set up ccache + uses: hendrikmuhs/ccache-action@v1.2 + + - name: Install requirements + run: | + export CCACHE_DIR=${{ env.CACHE_DIR }} + python -m pip install --upgrade pip + python -m pip install -r externals/torch-mlir/requirements.txt + python -m pip install -r externals/torch-mlir/torchvision-requirements.txt + + - name: Create build directory + run: mkdir build + + - name: Configure CMake + run: > + cmake -GNinja -Bbuild + -DCMAKE_BUILD_TYPE=Release + -DLLVM_ENABLE_PROJECTS=mlir + -DLLVM_ENABLE_ASSERTIONS=ON + -DLLVM_EXTERNAL_PROJECTS="torch-mlir;mpact" + -DLLVM_EXTERNAL_TORCH_MLIR_SOURCE_DIR="${PWD}/externals/torch-mlir" + -DLLVM_EXTERNAL_MPACT_SOURCE_DIR="${PWD}" + -DLLVM_TARGETS_TO_BUILD=host + -DMLIR_ENABLE_BINDINGS_PYTHON=ON + -DCMAKE_C_COMPILER_LAUNCHER=ccache + -DCMAKE_CXX_COMPILER_LAUNCHER=ccache + -DCMAKE_C_COMPILER=clang + -DCMAKE_CXX_COMPILER=clang++ + "externals/torch-mlir/externals/llvm-project/llvm" + + - name: Build + run: cmake --build build --target check-mpact +
diff --git a/.github/workflows/format.yml b/.github/workflows/format.yml new file mode 100644 index 0000000..8b1a835 --- /dev/null +++ b/.github/workflows/format.yml
@@ -0,0 +1,93 @@ +name: "Check code formatting" + +permissions: + contents: read + +on: + pull_request: + branches: + - main + +jobs: + code_formatter: + runs-on: ubuntu-latest + steps: + - name: Fetch sources + uses: actions/checkout@v4 + with: + ref: ${{ github.event.pull_request.head.sha }} + + - name: Checkout through merge base + uses: rmacklin/fetch-through-merge-base@v0 + with: + base_ref: ${{ github.event.pull_request.base.ref }} + head_ref: ${{ github.event.pull_request.head.sha }} + deepen_length: 500 + + - name: Get changed files + id: changed-files + uses: tj-actions/changed-files@v41 + with: + separator: "," + skip_initial_fetch: true + + # We need to pull the script from the main branch, so that we ensure + # we get the latest version of this script. + - name: Fetch code formatting utils + uses: actions/checkout@v4 + with: + repository: llvm/llvm-project + ref: refs/heads/main + sparse-checkout: | + llvm/utils/git/requirements_formatting.txt + llvm/utils/git/code-format-helper.py + sparse-checkout-cone-mode: false + path: code-format-tools + + - name: "Listed files" + env: + CHANGED_FILES: ${{ steps.changed-files.outputs.all_changed_files }} + run: | + echo "Formatting files:" + echo "$CHANGED_FILES" + - name: Install clang-format + uses: aminya/setup-cpp@v1 + with: + clangformat: 18.1.1 + + - name: Setup Python env + uses: actions/setup-python@v4 + with: + python-version: '3.11' + cache: 'pip' + cache-dependency-path: 'code-format-tools/llvm/utils/git/requirements_formatting.txt' + + - name: Install python dependencies + run: pip install -r code-format-tools/llvm/utils/git/requirements_formatting.txt + + - name: Run code formatter + env: + GITHUB_PR_NUMBER: ${{ github.event.pull_request.number }} + START_REV: ${{ github.event.pull_request.base.sha }} + END_REV: ${{ github.event.pull_request.head.sha }} + CHANGED_FILES: ${{ steps.changed-files.outputs.all_changed_files }} + # TODO: Once clang v18 is released, we should be able + # to take advantage of the new --diff_from_common_commit option + # explicitly in code-format-helper.py and not have to diff starting at + # the merge base. + # Create an empty comments file so the pr-write job doesn't fail. + run: | + echo "[]" > comments && + python ./code-format-tools/llvm/utils/git/code-format-helper.py \ + --write-comment-to-file \ + --token ${{ secrets.GITHUB_TOKEN }} \ + --issue-number $GITHUB_PR_NUMBER \ + --start-rev $(git merge-base $START_REV $END_REV) \ + --end-rev $END_REV \ + --changed-files "$CHANGED_FILES" + - uses: actions/upload-artifact@26f96dfa697d77e81fd5907df203aa23a56210a8 #v4.3.0 + if: always() + with: + name: workflow-args + path: | + comments
diff --git a/.github/workflows/regression-benchmark.yml b/.github/workflows/regression-benchmark.yml new file mode 100644 index 0000000..c7a0970 --- /dev/null +++ b/.github/workflows/regression-benchmark.yml
@@ -0,0 +1,59 @@ +name: Regression benchmark + +on: + workflow_run: + workflows: [Build and test] + types: [completed] + # branches: [main] + +permissions: + contents: write + deployments: write + +jobs: + benchmark: + name: Performance regression check + runs-on: ubuntu-latest + env: + CACHE_DIR: ${{ github.workspace }}/.ccache + steps: + - uses: actions/checkout@v4 + with: + submodules: recursive + token: ${{ secrets.GITHUB_TOKEN }} + + - name: Setup Python Version + uses: actions/setup-python@v5 + with: + python-version: 3.11 # Install the python version needed + + - name: Set PYTHONPATH + run: export PYTHONPATH=build/tools/mpact/python_packages/mpact + shell: bash + + - name: Build + run: cmake --build build --target build-benchmark-mpact + + - name: Install requirements + run: | + export CCACHE_DIR=${{ env.CACHE_DIR }} + python -m pip install --upgrade pip + python -m pip install pytest pytest-benchmark + + - name: Run benchmark + run: pytest benchmark/python/benchmarks/regression_benchmark.py --benchmark-json output.json + + - name: Store benchmark result + uses: benchmark-action/github-action-benchmark@v1 + with: + tool: 'pytest' + output-file-path: output.json + fail-on-alert: true + # GitHub API token to make a commit comment + github-token: ${{ secrets.GITHUB_TOKEN }} + # Enable alert commit comment + comment-on-alert: true + # Mention @rhysd in the commit comment + alert-comment-cc-users: '@yinying-lisa-li' + # Push and deploy GitHub pages branch automatically + auto-push: true \ No newline at end of file
diff --git a/.gitignore b/.gitignore new file mode 100644 index 0000000..a6d9b9f --- /dev/null +++ b/.gitignore
@@ -0,0 +1,7 @@ +*_venv/ +__pycache__ +/build*/ + +# lsp files +.cache/ +compile_commands.json
diff --git a/.gitmodules b/.gitmodules new file mode 100644 index 0000000..28810d3 --- /dev/null +++ b/.gitmodules
@@ -0,0 +1,6 @@ +[submodule "torch-mlir"] + path = externals/torch-mlir + url = https://github.com/llvm/torch-mlir.git +[submodule "externals/Enzyme"] + path = externals/Enzyme + url = https://github.com/EnzymeAD/Enzyme.git
diff --git a/CMakeLists.txt b/CMakeLists.txt new file mode 100644 index 0000000..e09570e --- /dev/null +++ b/CMakeLists.txt
@@ -0,0 +1,73 @@ +#------------------------------------------------------------------------------- +# The MPACT Compiler +#------------------------------------------------------------------------------- + +cmake_minimum_required(VERSION 3.12) + +project(mpact VERSION 1.0 LANGUAGES CXX C) + +set(CMAKE_C_STANDARD 11) +set(CMAKE_CXX_STANDARD 17) + +set(MPACT_SOURCE_DIR "${CMAKE_CURRENT_SOURCE_DIR}") +set(MPACT_BINARY_DIR "${CMAKE_CURRENT_BINARY_DIR}") +message(STATUS "Building the MPACT compiler at ${MPACT_SOURCE_DIR} (into ${MPACT_BINARY_DIR})") + +set(MPACT_PYTHON_PACKAGES_DIR "${MPACT_BINARY_DIR}/python_packages") + +#------------------------------------------------------------------------------- +# Configure out-of-tree vs in-tree build +#------------------------------------------------------------------------------- + +if(CMAKE_SOURCE_DIR STREQUAL CMAKE_CURRENT_SOURCE_DIR) + message(STATUS "MPACT out-of-tree build.") + message(FATAL_ERROR "TODO") +else() + message(STATUS "MPACT in-tree build.") + # In-tree build with LLVM_EXTERNAL_PROJECTS=mpact + option(MLIR_ENABLE_BINDINGS_PYTHON "Enables MLIR Python Bindings" OFF) + set(MLIR_MAIN_SRC_DIR ${LLVM_MAIN_SRC_DIR}/../mlir) + set(MLIR_INCLUDE_DIR ${LLVM_MAIN_SRC_DIR}/../mlir/include) + set(MLIR_GENERATED_INCLUDE_DIR ${LLVM_BINARY_DIR}/tools/mlir/include) + set(MLIR_INCLUDE_DIRS "${MLIR_INCLUDE_DIR};${MLIR_GENERATED_INCLUDE_DIR}") +endif() + +include_directories(${LLVM_INCLUDE_DIRS}) +include_directories(${MLIR_INCLUDE_DIRS}) +include_directories(${CMAKE_CURRENT_SOURCE_DIR}/include) +include_directories(${CMAKE_CURRENT_BINARY_DIR}/include) + +# Needed to build TorchMLIRExtensions. +include_directories(${CMAKE_CURRENT_SOURCE_DIR}/externals/torch-mlir/include) + +function(mpact_target_includes target) + set(_dirs + $<BUILD_INTERFACE:${MLIR_INCLUDE_DIRS}> + $<BUILD_INTERFACE:${MPACT_SOURCE_DIR}/include> + $<BUILD_INTERFACE:${MPACT_BINARY_DIR}/include> + ) + target_include_directories(${target} PUBLIC ${_dirs}) + if(TARGET obj.${target}) + target_include_directories(obj.${target} PRIVATE ${_dirs}) + endif() +endfunction() + +list(APPEND CMAKE_MODULE_PATH ${MLIR_MAIN_SRC_DIR}/cmake/modules) +list(APPEND CMAKE_MODULE_PATH ${LLVM_MAIN_SRC_DIR}/cmake) +list(APPEND CMAKE_MODULE_PATH ${CMAKE_CURRENT_SOURCE_DIR}/build_tools/cmake) + +include(TableGen) +include(AddLLVM) +include(AddMLIR) +include(AddMLIRPython) + +include(MLIRDetectPythonEnv) +mlir_configure_python_dev_packages() + +add_subdirectory(include) +add_subdirectory(lib) +add_subdirectory(tools) + +add_subdirectory(benchmark) +add_subdirectory(python) +add_subdirectory(test)
diff --git a/LICENSE b/LICENSE new file mode 100644 index 0000000..261eeb9 --- /dev/null +++ b/LICENSE
@@ -0,0 +1,201 @@ + Apache License + Version 2.0, January 2004 + http://www.apache.org/licenses/ + + TERMS AND CONDITIONS FOR USE, REPRODUCTION, AND DISTRIBUTION + + 1. Definitions. + + "License" shall mean the terms and conditions for use, reproduction, + and distribution as defined by Sections 1 through 9 of this document. + + "Licensor" shall mean the copyright owner or entity authorized by + the copyright owner that is granting the License. + + "Legal Entity" shall mean the union of the acting entity and all + other entities that control, are controlled by, or are under common + control with that entity. For the purposes of this definition, + "control" means (i) the power, direct or indirect, to cause the + direction or management of such entity, whether by contract or + otherwise, or (ii) ownership of fifty percent (50%) or more of the + outstanding shares, or (iii) beneficial ownership of such entity. + + "You" (or "Your") shall mean an individual or Legal Entity + exercising permissions granted by this License. + + "Source" form shall mean the preferred form for making modifications, + including but not limited to software source code, documentation + source, and configuration files. + + "Object" form shall mean any form resulting from mechanical + transformation or translation of a Source form, including but + not limited to compiled object code, generated documentation, + and conversions to other media types. + + "Work" shall mean the work of authorship, whether in Source or + Object form, made available under the License, as indicated by a + copyright notice that is included in or attached to the work + (an example is provided in the Appendix below). + + "Derivative Works" shall mean any work, whether in Source or Object + form, that is based on (or derived from) the Work and for which the + editorial revisions, annotations, elaborations, or other modifications + represent, as a whole, an original work of authorship. For the purposes + of this License, Derivative Works shall not include works that remain + separable from, or merely link (or bind by name) to the interfaces of, + the Work and Derivative Works thereof. + + "Contribution" shall mean any work of authorship, including + the original version of the Work and any modifications or additions + to that Work or Derivative Works thereof, that is intentionally + submitted to Licensor for inclusion in the Work by the copyright owner + or by an individual or Legal Entity authorized to submit on behalf of + the copyright owner. For the purposes of this definition, "submitted" + means any form of electronic, verbal, or written communication sent + to the Licensor or its representatives, including but not limited to + communication on electronic mailing lists, source code control systems, + and issue tracking systems that are managed by, or on behalf of, the + Licensor for the purpose of discussing and improving the Work, but + excluding communication that is conspicuously marked or otherwise + designated in writing by the copyright owner as "Not a Contribution." + + "Contributor" shall mean Licensor and any individual or Legal Entity + on behalf of whom a Contribution has been received by Licensor and + subsequently incorporated within the Work. + + 2. Grant of Copyright License. Subject to the terms and conditions of + this License, each Contributor hereby grants to You a perpetual, + worldwide, non-exclusive, no-charge, royalty-free, irrevocable + copyright license to reproduce, prepare Derivative Works of, + publicly display, publicly perform, sublicense, and distribute the + Work and such Derivative Works in Source or Object form. + + 3. Grant of Patent License. Subject to the terms and conditions of + this License, each Contributor hereby grants to You a perpetual, + worldwide, non-exclusive, no-charge, royalty-free, irrevocable + (except as stated in this section) patent license to make, have made, + use, offer to sell, sell, import, and otherwise transfer the Work, + where such license applies only to those patent claims licensable + by such Contributor that are necessarily infringed by their + Contribution(s) alone or by combination of their Contribution(s) + with the Work to which such Contribution(s) was submitted. If You + institute patent litigation against any entity (including a + cross-claim or counterclaim in a lawsuit) alleging that the Work + or a Contribution incorporated within the Work constitutes direct + or contributory patent infringement, then any patent licenses + granted to You under this License for that Work shall terminate + as of the date such litigation is filed. + + 4. Redistribution. You may reproduce and distribute copies of the + Work or Derivative Works thereof in any medium, with or without + modifications, and in Source or Object form, provided that You + meet the following conditions: + + (a) You must give any other recipients of the Work or + Derivative Works a copy of this License; and + + (b) You must cause any modified files to carry prominent notices + stating that You changed the files; and + + (c) You must retain, in the Source form of any Derivative Works + that You distribute, all copyright, patent, trademark, and + attribution notices from the Source form of the Work, + excluding those notices that do not pertain to any part of + the Derivative Works; and + + (d) If the Work includes a "NOTICE" text file as part of its + distribution, then any Derivative Works that You distribute must + include a readable copy of the attribution notices contained + within such NOTICE file, excluding those notices that do not + pertain to any part of the Derivative Works, in at least one + of the following places: within a NOTICE text file distributed + as part of the Derivative Works; within the Source form or + documentation, if provided along with the Derivative Works; or, + within a display generated by the Derivative Works, if and + wherever such third-party notices normally appear. The contents + of the NOTICE file are for informational purposes only and + do not modify the License. You may add Your own attribution + notices within Derivative Works that You distribute, alongside + or as an addendum to the NOTICE text from the Work, provided + that such additional attribution notices cannot be construed + as modifying the License. + + You may add Your own copyright statement to Your modifications and + may provide additional or different license terms and conditions + for use, reproduction, or distribution of Your modifications, or + for any such Derivative Works as a whole, provided Your use, + reproduction, and distribution of the Work otherwise complies with + the conditions stated in this License. + + 5. Submission of Contributions. Unless You explicitly state otherwise, + any Contribution intentionally submitted for inclusion in the Work + by You to the Licensor shall be under the terms and conditions of + this License, without any additional terms or conditions. + Notwithstanding the above, nothing herein shall supersede or modify + the terms of any separate license agreement you may have executed + with Licensor regarding such Contributions. + + 6. Trademarks. This License does not grant permission to use the trade + names, trademarks, service marks, or product names of the Licensor, + except as required for reasonable and customary use in describing the + origin of the Work and reproducing the content of the NOTICE file. + + 7. Disclaimer of Warranty. Unless required by applicable law or + agreed to in writing, Licensor provides the Work (and each + Contributor provides its Contributions) on an "AS IS" BASIS, + WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or + implied, including, without limitation, any warranties or conditions + of TITLE, NON-INFRINGEMENT, MERCHANTABILITY, or FITNESS FOR A + PARTICULAR PURPOSE. You are solely responsible for determining the + appropriateness of using or redistributing the Work and assume any + risks associated with Your exercise of permissions under this License. + + 8. Limitation of Liability. In no event and under no legal theory, + whether in tort (including negligence), contract, or otherwise, + unless required by applicable law (such as deliberate and grossly + negligent acts) or agreed to in writing, shall any Contributor be + liable to You for damages, including any direct, indirect, special, + incidental, or consequential damages of any character arising as a + result of this License or out of the use or inability to use the + Work (including but not limited to damages for loss of goodwill, + work stoppage, computer failure or malfunction, or any and all + other commercial damages or losses), even if such Contributor + has been advised of the possibility of such damages. + + 9. Accepting Warranty or Additional Liability. While redistributing + the Work or Derivative Works thereof, You may choose to offer, + and charge a fee for, acceptance of support, warranty, indemnity, + or other liability obligations and/or rights consistent with this + License. However, in accepting such obligations, You may act only + on Your own behalf and on Your sole responsibility, not on behalf + of any other Contributor, and only if You agree to indemnify, + defend, and hold each Contributor harmless for any liability + incurred by, or claims asserted against, such Contributor by reason + of your accepting any such warranty or additional liability. + + END OF TERMS AND CONDITIONS + + APPENDIX: How to apply the Apache License to your work. + + To apply the Apache License to your work, attach the following + boilerplate notice, with the fields enclosed by brackets "[]" + replaced with your own identifying information. (Don't include + the brackets!) The text should be enclosed in the appropriate + comment syntax for the file format. We also recommend that a + file or class name and description of purpose be included on the + same "printed page" as the copyright notice for easier + identification within third-party archives. + + Copyright [yyyy] [name of copyright owner] + + Licensed under the Apache License, Version 2.0 (the "License"); + you may not use this file except in compliance with the License. + You may obtain a copy of the License at + + http://www.apache.org/licenses/LICENSE-2.0 + + Unless required by applicable law or agreed to in writing, software + distributed under the License is distributed on an "AS IS" BASIS, + WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + See the License for the specific language governing permissions and + limitations under the License.
diff --git a/README.md b/README.md new file mode 100644 index 0000000..9c764b7 --- /dev/null +++ b/README.md
@@ -0,0 +1,118 @@ +# The MPACT Project + +## Introduction + +The MPACT project's main objective is to dramatically reduce the effort +required to create highly optimizing HPC and ML compilers for a large +class of architectures using LLVM and MLIR. We do this by providing +a declarative language-based mechanism for collecting and expressing +critical aspects of a target architecture in a way that can be reasoned +about and leveraged by all passes in both MLIR and LLVM. + +## Building the MPACT compiler + +To build and run the MPACT compiler from source (for developers), +please follow the steps below. + +### Check out code and sync submodules + +Use the following commands to clone the MPACT compiler repository. + +```shell +git clone https://github.com/MPACT-ORG/mpact-compiler.git +cd mpact-compiler +git submodule update --init --recursive --progress +``` + +To always get updated submodules through `git pull`, set the following flag: + +```shell +git config --global submodule.recurse true +``` + +NOTE: All following commands assume you remain in the `mpact-compiler` directory. + +### Setup Python virtual environment + +The following commands initialize a virtual environment under bash/sh/etc. For other shells, see Note 1, [below](README.md#notes). + +```shell +python3.11 -m venv mpact_venv # one time set up +source mpact_venv/bin/activate # MUST BE REPEATED FOR EVERY SESSION +``` + +Next, set the Python paths as follows; for shells not in the bash/sh family, see Note 2, [below](README.md#notes). +```shell +export PYTHONPATH=`pwd`/build/tools/mpact/python_packages/mpact +``` + +### Install build requirements + +Note that currently we rely on `torch-mlir` requirements defined in that +submodule to ensure all the build requirements are consistent. + +```shell +python -m pip install --upgrade pip +python -m pip install -r externals/torch-mlir/requirements.txt +python -m pip install -r externals/torch-mlir/torchvision-requirements.txt +``` +For shells not in the bash/sh family, see Note 3, [below](README.md#notes). + +### Building the MPACT compiler in-tree + +The following command generates configuration files to build the MPACT compiler +project completely *in-tree*, which means that both LLVM as well as torch-mlir +are built from source. + +```shell +cmake -GNinja -Bbuild \ + -DCMAKE_BUILD_TYPE=Release \ + -DPython3_FIND_VIRTUALENV=ONLY \ + -DLLVM_ENABLE_PROJECTS=mlir \ + -DLLVM_EXTERNAL_PROJECTS="torch-mlir;mpact" \ + -DLLVM_EXTERNAL_TORCH_MLIR_SOURCE_DIR="${PWD}/externals/torch-mlir" \ + -DLLVM_EXTERNAL_MPACT_SOURCE_DIR="${PWD}" \ + -DLLVM_TARGETS_TO_BUILD=host \ + -DMLIR_ENABLE_BINDINGS_PYTHON=ON \ + externals/torch-mlir/externals/llvm-project/llvm +``` + +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: + +```shell +-DCMAKE_C_COMPILER_LAUNCHER=ccache -DCMAKE_CXX_COMPILER_LAUNCHER=ccache +``` + +Run the following to ensure the MPACT compiler builds and runs correctly. + +```shell +cmake --build build --target check-mpact +``` + +And the following to run all benchmarks +(see [Benchmarks](benchmark/README.md) for more details). + +```shell +cmake --build build --target benchmark-mpact +``` + +## Notes + +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. + - For csh/tcsh/etc.: + ```shell + source `pwd`/mpact_venv/bin/activate.csh + ``` + - For fish/etc.: + ```shell + source <path_to_mpact_compiler>/mpact_venv/bin/activate.fish + ``` +2. Shells other than bash/sh/etc. set their environment variables differently: + - For csh/tcsh/etc.: + ```shell + setenv PYTHONPATH `pwd`/build/tools/mpact/python_packages/mpact + ``` +3. If using csh/tcsh/etc., run the following command before trying to build the compiler: +```shell +rehash +```
diff --git a/benchmark/CMakeLists.txt b/benchmark/CMakeLists.txt new file mode 100644 index 0000000..f52d4bb --- /dev/null +++ b/benchmark/CMakeLists.txt
@@ -0,0 +1,47 @@ +#------------------------------------------------------------------------------- +# The MPACT Compiler Python Benchmarks +#------------------------------------------------------------------------------- + +declare_mlir_python_sources(MPACTBenchmarkPythonSources) + +declare_mlir_python_sources(MPACTBenchmarkPythonSources.BenchmarkSuite + ROOT_DIR "${CMAKE_CURRENT_SOURCE_DIR}/python" + ADD_TO_PARENT MPACTBenchmarkPythonSources + SOURCES_GLOB + benchmarks/*.py +) + +declare_mlir_python_sources(MPACTBenchmarkPythonSources.BenchmarkUtils + ROOT_DIR "${CMAKE_CURRENT_SOURCE_DIR}/python" + ADD_TO_PARENT MPACTBenchmarkPythonSources + SOURCES_GLOB + utils/*.py +) + +#------------------------------------------------------------------------------- +# Python Modules +#------------------------------------------------------------------------------- + +add_mlir_python_modules(MPACTBenchmarkPythonModules + ROOT_PREFIX "${MPACT_PYTHON_PACKAGES_DIR}/mpact/mpact_benchmark" + INSTALL_PREFIX "python_packages/mpact/mpact_benchmark" + DECLARED_SOURCES MPACTBenchmarkPythonSources +) + +add_custom_target(build-benchmark-mpact) +add_dependencies(build-benchmark-mpact MPACTPythonModules MPACTBenchmarkPythonModules) + +add_custom_target(benchmark-mpact) +add_dependencies(benchmark-mpact build-benchmark-mpact) +file(GLOB PYTHON_FILES "${CMAKE_CURRENT_SOURCE_DIR}/python/benchmarks/*.py") + +# Loop over each matched .py file and create a custom command to run it. +foreach(PY_FILE IN LISTS PYTHON_FILES) + add_custom_command( + TARGET benchmark-mpact + COMMAND cmake -E echo "Running ${PY_FILE}" + COMMAND python ${PY_FILE} + DEPENDS ${PY_FILE} + USES_TERMINAL + ) +endforeach()
diff --git a/benchmark/README.md b/benchmark/README.md new file mode 100644 index 0000000..4749547 --- /dev/null +++ b/benchmark/README.md
@@ -0,0 +1,26 @@ +### Run benchmarks + +To run all benchmarks: + +```shell +cmake --build build --target benchmark-mpact +``` + +To run selected benchmarks, build the benchmark modules first: + +```shell +cmake --build build --target build-benchmark-mpact +``` + +And then run the benchmark file: + +```shell +python path/to/the/_benchmark.py +``` + +If you would like to run selected kernels in kernels_benchmark.py, +you can use `--benchmark-filter` flag like the following example: + +```shell +python path/to/the/kernels_benchmark.py --benchmark-filter=add +```
diff --git a/benchmark/python/benchmarks/gcn_benchmark.py b/benchmark/python/benchmarks/gcn_benchmark.py new file mode 100755 index 0000000..ef7a3bd --- /dev/null +++ b/benchmark/python/benchmarks/gcn_benchmark.py
@@ -0,0 +1,33 @@ +import torch +import numpy as np +from mpact.models.gcn import GraphConv +from mpact_benchmark.utils.benchmark_utils import benchmark, Backends + + +@benchmark( + [ + { + "name": f"{fmt}_{shape}_{dtype.__name__}", + "shape": shape, + "formats": fmt, + "dtype": dtype, + "drange": (1, 100), + "sparsity": [0, 0.5, 0.9, 0.99], + "backends": [b for b in Backends], + } + for shape in [ + [[128, 128], [128, 128]], + [[512, 512], [512, 512]], + [[1024, 1024], [1024, 1024]], + ] + for fmt in [["dense", "csr"]] + for dtype in [np.float32] + ] +) +def GCN() -> torch.nn.Module: + """Graph Convolution Network.""" + return GraphConv + + +if __name__ == "__main__": + GCN()
diff --git a/benchmark/python/benchmarks/kernels_benchmark.py b/benchmark/python/benchmarks/kernels_benchmark.py new file mode 100644 index 0000000..4dee945 --- /dev/null +++ b/benchmark/python/benchmarks/kernels_benchmark.py
@@ -0,0 +1,225 @@ +import torch +import argparse +import numpy as np +from mpact.models.kernels import * +from mpact_benchmark.utils.benchmark_utils import benchmark, Backends + + +@benchmark( + [ + { + "name": f"{lhs_fmt}_{rhs_fmt}_{shape}_{dtype.__name__}", + "shape": shape, + "formats": (lhs_fmt, rhs_fmt), + "dtype": dtype, + "backends": [b for b in Backends], + "drange": (1, 100), + "sparsity": [0, 0.5, 0.9, 0.99], + } + for shape in [([2**i, 2**i], [2**i, 2**i]) for i in range(5, 8)] + for lhs_fmt in ["dense", "csr"] + for rhs_fmt in ["dense", "csr"] + for dtype in [np.float64] + ] +) +def matmul() -> torch.nn.Module: + """Matrix multiplication.""" + return MMNet() + + +@benchmark( + [ + { + "name": f"{lhs_fmt}_{rhs_fmt}_{shape}_{dtype.__name__}", + "shape": shape, + "formats": (lhs_fmt, rhs_fmt), + "dtype": dtype, + "backends": [b for b in Backends], + "drange": (1, 100), + "sparsity": [0, 0.5, 0.9, 0.99], + } + for shape in [([2**i, 2**i], [2**i]) for i in range(5, 8)] + for lhs_fmt in ["dense", "csr"] + for rhs_fmt in ["dense"] # torch.mv only supports dense vector for now. + for dtype in [np.float64] + ] +) +def matvec() -> torch.nn.Module: + """Matrix-vector multiplication.""" + return MVNet() + + +@benchmark( + [ + { + "name": f"{lhs_fmt}_{rhs_fmt}_{shape}_{dtype.__name__}", + "shape": shape, + "formats": (lhs_fmt, rhs_fmt), + "dtype": dtype, + "backends": [b for b in Backends], + "drange": (1, 100), + "sparsity": [0, 0.5, 0.9, 0.99], + } + for shape in [ + ([2**i, 2**i], [2**i, 2**i]) for i in range(5, 8) + ] # 512x512 crashes runtime. + for lhs_fmt in ["dense", "csr"] + for rhs_fmt in ["dense", "csr"] + for dtype in [np.float64] + ] +) +def add() -> torch.nn.Module: + """Element-wise addition.""" + return AddNet() + + +@benchmark( + [ + { + "name": f"{lhs_fmt}_{rhs_fmt}_{shape}_{dtype.__name__}", + "shape": shape, + "formats": (lhs_fmt, rhs_fmt), + "dtype": dtype, + "backends": [b for b in Backends], + "drange": (1, 100), + "sparsity": [0, 0.5, 0.9, 0.99], + } + for shape in [([2**i, 2**i], [2**i, 2**i]) for i in range(5, 8)] + for lhs_fmt in ["dense", "csr"] + for rhs_fmt in ["dense", "csr"] + for dtype in [np.float64] + ] +) +def elt_mul() -> torch.nn.Module: + """Element-wise addition.""" + return MulNet() + + +@benchmark( + [ + { + "name": f"{fmt}_{shape}_{dtype.__name__}", + "shape": shape, + "formats": (fmt,), + "dtype": dtype, + "backends": [b for b in Backends], + "drange": (1, 100), + "sparsity": [0, 0.5, 0.9, 0.99], + } + for shape in [([2**i, 2**i],) for i in range(2, 3)] + for fmt in ["dense", "csr"] + for dtype in [np.float64] + ] +) +def nop() -> torch.nn.Module: + """Returns matrix unmodified (speed of light).""" + return SelfNet() + + +@benchmark( + [ + { + "name": f"{sample_fmt}_sample_{shape}_{dtype.__name__}", + "shape": shape, + "formats": (sample_fmt, "dense", "dense"), + "dtype": dtype, + "backends": [b for b in Backends], + "drange": (1, 100), + "sparsity": [0, 0.5, 0.9, 0.99], + } + for shape in [ + ([2**i, 2**i], [2**i, 2**i], [2**i, 2**i]) for i in range(5, 8) + ] + for sample_fmt in ["dense", "csr"] + for dtype in [np.float64] + ] +) +def sddmm() -> torch.nn.Module: + """SDDMM: C = S â—‹ (A X B) Sampled dense-dense matrix-matrix multiplication.""" + return SDDMMNet() + + +@benchmark( + [ + { + "name": f"{fmt}_{shape}_{dtype.__name__}", + "shape": shape, + "formats": (fmt,), + "dtype": dtype, + # TODO: add mpact and torch inductor once they work. + "backends": [ + b + for b in Backends + if b.value + in ( + Backends.TORCH_SPARSE_EAGER.value, + Backends.TORCH_DENSE_EAGER.value, + ) + ], + "drange": (1, 100), + "sparsity": [0, 0.5, 0.9, 0.99], + } + for shape in [([2**i, 2**i],) for i in range(5, 8)] + for fmt in ["dense"] + for dtype in [np.float64] + ] +) +def feature_scale() -> torch.nn.Module: + """Scales feature matrix in GNN.""" + return FeatureScale() + + +@benchmark( + [ + { + "name": f"{fmt}_{shape}_{dtype.__name__}", + "shape": shape, + "formats": (fmt,), + "dtype": dtype, + # TODO: add mpact and torch inductor once they work. + "backends": [ + b + for b in Backends + if b.value + in ( + Backends.TORCH_SPARSE_EAGER.value, + Backends.TORCH_DENSE_EAGER.value, + ) + ], + "drange": (1, 100), + "sparsity": [0, 0.5, 0.9, 0.99], + } + for shape in [([2**i, 2**i],) for i in range(5, 8)] + for fmt in ["dense"] + for dtype in [np.float64] + ] +) +def normalization() -> torch.nn.Module: + """Normalizes adjacency matrix in GNN.""" + return Normalization() + + +if __name__ == "__main__": + parser = argparse.ArgumentParser( + prog="pytorch_kernel_benchmarks", + description="Run a set of given PyTorch kernel benchmarks", + ) + parser.add_argument("--benchmark-filter", type=str, default="", required=False) + arguments = parser.parse_args() + + benchmark_list = [ + "nop", + "add", + "matmul", + "matvec", + "elt_mul", + "sddmm", + "feature_scale", + "normalization", + ] + if arguments.benchmark_filter: + benchmark_list = arguments.benchmark_filter.split(",") + + # Run selected benchmarks. + for benchmark_name in benchmark_list: + globals()[benchmark_name]()
diff --git a/benchmark/python/benchmarks/lif_benchmark.py b/benchmark/python/benchmarks/lif_benchmark.py new file mode 100644 index 0000000..3fe9784 --- /dev/null +++ b/benchmark/python/benchmarks/lif_benchmark.py
@@ -0,0 +1,44 @@ +import torch +import numpy as np +from mpact.models.lif import LIFSumOfSq +from mpact_benchmark.utils.benchmark_utils import benchmark, Backends + + +@benchmark( + [ + { + "name": f"{fmt}_{shape}_{dtype.__name__}", + "shape": shape, + "formats": fmt, + "dtype": dtype, + # Simulate batch normalization. + "drange": (-1, 1), + "sparsity": [0, 0.5, 0.9, 0.99], + # to_dense() in LIF prop hack is not supported in torch inductor. + # TODO: add torch inductor once prop hack is no longer needed. + "backends": [ + b + for b in Backends + if b.value + not in ( + Backends.TORCH_SPARSE_INDUCTOR.value, + Backends.TORCH_DENSE_INDUCTOR.value, + ) + ], + } + for shape in [ + [[64, 3, 32, 32, 1]], + [[32, 3, 64, 64, 1]], + [[16, 3, 224, 224, 1]], + ] + for fmt in [["dense"]] + for dtype in [np.float64] + ] +) +def LifSumOfSq() -> torch.nn.Module: + """LIF feeding into sum of squares.""" + return LIFSumOfSq() + + +if __name__ == "__main__": + LifSumOfSq()
diff --git a/benchmark/python/benchmarks/regression_benchmark.py b/benchmark/python/benchmarks/regression_benchmark.py new file mode 100644 index 0000000..6f5196e --- /dev/null +++ b/benchmark/python/benchmarks/regression_benchmark.py
@@ -0,0 +1,51 @@ +import pytest +from mpact.models.kernels import * +from mpact_benchmark.utils.tensor_generator import generate_tensor + +SHAPE = (1024, 1024) +SPARSITY = 0.8 + +dense_tensor1 = generate_tensor(0, SHAPE, SPARSITY) +dense_tensor2 = generate_tensor(1, SHAPE, SPARSITY) +dense_tensor3 = generate_tensor(2, SHAPE, SPARSITY) +dense_vector = generate_tensor(1, (SHAPE[0],), SPARSITY) + +sparse_tensor1 = dense_tensor1.to_sparse_csr() +sparse_tensor2 = dense_tensor2.to_sparse_csr() +sparse_tensor3 = dense_tensor3.to_sparse_csr() + +def test_mv_dense(benchmark): + benchmark(MVNet(), dense_tensor1, dense_vector) + +def test_mm_dense(benchmark): + benchmark(MMNet(), dense_tensor1, dense_tensor2) + +def test_add_dense(benchmark): + benchmark(AddNet(), dense_tensor1, dense_tensor2) + +def test_mul_dense(benchmark): + benchmark(MulNet(), dense_tensor1, dense_tensor2) + +def test_nop_dense(benchmark): + benchmark(SelfNet(), dense_tensor1) + +def test_sddmm_dense(benchmark): + benchmark(SDDMMNet(), dense_tensor1, dense_tensor2, dense_tensor3) + +def test_mv_sparse(benchmark): + benchmark(MVNet(), sparse_tensor1, dense_vector) + +def test_mm_sparse(benchmark): + benchmark(MMNet(), sparse_tensor1, sparse_tensor2) + +def test_add_sparse(benchmark): + benchmark(AddNet(), sparse_tensor1, sparse_tensor2) + +def test_mul_sparse(benchmark): + benchmark(MulNet(), sparse_tensor1, sparse_tensor2) + +def test_nop_sparse(benchmark): + benchmark(SelfNet(), sparse_tensor1) + +def test_sddmm_sparse(benchmark): + benchmark(SDDMMNet(), sparse_tensor1, dense_tensor2, dense_tensor3) \ No newline at end of file
diff --git a/benchmark/python/benchmarks/resnet_benchmark.py b/benchmark/python/benchmarks/resnet_benchmark.py new file mode 100644 index 0000000..89b317e --- /dev/null +++ b/benchmark/python/benchmarks/resnet_benchmark.py
@@ -0,0 +1,36 @@ +import torch +import numpy as np +from mpact.models.resnet import resnet_20 +from mpact_benchmark.utils.benchmark_utils import benchmark, Backends + + +@benchmark( + [ + { + "name": f"{fmt}_{shape}_{dtype.__name__}", + "shape": shape, + "formats": fmt, + "dtype": dtype, + "drange": (1, 100), + "sparsity": [0.5, 0.9], + # TODO: Torch inductor requires lower precision with larger input size, + # such as [8, 3, 32, 32]. + "precision": 1e-3, + "backends": [b for b in Backends], + } + for shape in [ + [[1, 3, 16, 16]], + ] + for fmt in [["dense"]] + for dtype in [np.float32] + ] +) +def resnet() -> torch.nn.Module: + """Restnet20 model.""" + resnet_model = resnet_20() + resnet_model.train(False) + return resnet_model + + +if __name__ == "__main__": + resnet()
diff --git a/benchmark/python/utils/benchmark_utils.py b/benchmark/python/utils/benchmark_utils.py new file mode 100644 index 0000000..ec29309 --- /dev/null +++ b/benchmark/python/utils/benchmark_utils.py
@@ -0,0 +1,231 @@ +import functools +import torch +from enum import Enum +from typing import Any, Callable +from torch.utils import benchmark as torch_benchmark +from mpact.mpactbackend import mpact_jit, mpact_jit_compile, mpact_jit_run +from mpact_benchmark.utils.tensor_generator import generate_inputs + + +class Backends(Enum): + TORCH_SPARSE_EAGER = 1 + TORCH_DENSE_EAGER = 2 + TORCH_SPARSE_INDUCTOR = 3 + TORCH_DENSE_INDUCTOR = 4 + MPACT_SPARSE = 5 + MPACT_DENSE = 6 + + +def timer(stmt: str, description: str, setup: str = "", **kwargs: Any) -> Any: + """Timer for benchmark.""" + return torch_benchmark.Timer( + stmt=stmt, + globals=kwargs["variables"], + setup=setup, + num_threads=1, + label=kwargs["variables"]["label"], + sub_label=kwargs["variables"]["sub_label"], + description=description, + ).adaptive_autorange() + + +def get_dynamo_compile_time(sub_label: str, label: str, description: str) -> Any: + """Get compile time from dynamo and create a benchmark measurement object.""" + try: + compile_time = torch_benchmark.Measurement( + 1, + [ + float( + torch._dynamo.utils.compile_times(repr="csv")[1][0] + .split(",")[-1] + .strip() + ) + ], + torch_benchmark.TaskSpec( + sub_label, + None, + description=description, + label=label, + ), + ) + return compile_time + except ValueError: + print(f"No compilation happened for {description}: {sub_label}.") + return None + + +def run_benchmark( + sparse_inputs: tuple[torch.Tensor, ...], + dense_inputs: tuple[torch.Tensor, ...], + torch_net: torch.nn.Module, + variables: dict[str, Any], + backends: tuple[Backends, ...], + runtime_results: list[torch_benchmark.Measurement], + compile_time_results: list[torch_benchmark.Measurement], +): + """Run benchmark with specified backends.""" + output = [] + + with torch.no_grad(): + for backend in backends: + match backend: + case Backends.TORCH_SPARSE_EAGER: + output.append(torch_net(*sparse_inputs)) + runtime_results.append( + timer( + "torch_net(*sparse_inputs)", + "torch-sparse-eager", + variables=variables, + ) + ) + case Backends.TORCH_DENSE_EAGER: + output.append(torch_net(*dense_inputs)) + runtime_results.append( + timer( + "torch_net(*dense_inputs)", + "torch-dense-eager", + variables=variables, + ) + ) + case Backends.TORCH_SPARSE_INDUCTOR: + torch_inductor = torch.compile(torch_net) + torch_out = torch_inductor(*sparse_inputs) + output.append(torch_out) + compile_time = get_dynamo_compile_time( + variables["sub_label"], + variables["label"], + "torch-sparse-inductor-compile", + ) + if compile_time: + compile_time_results.append(compile_time) + runtime_results.append( + timer( + "torch_inductor(*sparse_inputs)", + "torch-sparse-inductor-runtime", + variables=dict(variables, **locals()), + ) + ) + case Backends.TORCH_DENSE_INDUCTOR: + torch_inductor = torch.compile(torch_net) + output.append(torch_inductor(*dense_inputs)) + compile_time = get_dynamo_compile_time( + variables["sub_label"], + variables["label"], + "torch-dense-inductor-compile", + ) + if compile_time: + compile_time_results.append(compile_time) + runtime_results.append( + timer( + "torch_inductor(*dense_inputs)", + "torch-dense-inductor-runtime", + variables=dict(variables, **locals()), + ) + ) + case Backends.MPACT_SPARSE: + sp_out = mpact_jit(torch_net, *sparse_inputs) + # Construct sparse csr tensor if the output type is csr. + # TODO: return sparse tensor directly instead of a tuple of arrays. + if type(sp_out) is tuple: + # torch.sparse_csr_tensor could deduce the size incorrectly, + # so pass the dense_out's shape explicitly. + dense_out = mpact_jit(torch_net, *dense_inputs) + output.append( + torch.sparse_csr_tensor(*sp_out, size=dense_out.shape) + ) + else: + output.append(torch.from_numpy(sp_out)) + invoker, f = mpact_jit_compile(torch_net, *sparse_inputs) + compile_time_results.append( + timer( + "mpact_jit_compile(torch_net, *sparse_inputs)", + "mpact-sparse-compile", + "from mpact.mpactbackend import mpact_jit_compile", + variables=dict(variables, **locals()), + ) + ) + runtime_results.append( + timer( + "mpact_jit_run(invoker, f, *sparse_inputs)", + "mpact-sparse-runtime", + "from mpact.mpactbackend import mpact_jit_run", + variables=dict(variables, **locals()), + ) + ) + case Backends.MPACT_DENSE: + output.append(torch.from_numpy(mpact_jit(torch_net, *dense_inputs))) + invoker, f = mpact_jit_compile(torch_net, *dense_inputs) + compile_time_results.append( + timer( + "mpact_jit_compile(torch_net, *dense_inputs)", + "mpact-dense-compile", + "from mpact.mpactbackend import mpact_jit_compile", + variables=dict(variables, **locals()), + ) + ) + runtime_results.append( + timer( + "mpact_jit_run(invoker, f, *dense_inputs)", + "mpact-dense-runtime", + "from mpact.mpactbackend import mpact_jit_run", + variables=dict(variables, **locals()), + ) + ) + case _: + print(f"{backend} is not supported yet.") + + # Sanity check. + if output: + rtol = variables["precision"] if "precision" in variables else 1e-5 + assert all( + torch.allclose(output[0].to_dense(), out.to_dense(), rtol=rtol) + for out in output + ) + + +def benchmark(*args: Any) -> Callable: + """Wrapper for benchmark.""" + + def decorator(func): + @functools.wraps(func) + def wrapper(test_cases=args[0]): + runtime_results = [] + compile_time_results = [] + torch_net = net = func() + for test_case in test_cases: + label = func.__name__ + for sparsity in test_case["sparsity"]: + sub_label = f"{test_case['name']}_{sparsity}" + dense_inputs, sparse_inputs = generate_inputs( + test_case["shape"], + sparsity, + test_case["formats"], + test_case["dtype"], + test_case["drange"], + ) + + if "GCN" in label: + torch_net = net(*test_case["shape"][0]) + if "precision" in test_case: + precision = test_case["precision"] + + run_benchmark( + sparse_inputs, + dense_inputs, + torch_net, + locals(), + test_case["backends"], + runtime_results, + compile_time_results, + ) + + compare1 = torch_benchmark.Compare(runtime_results) + compare1.print() + compare2 = torch_benchmark.Compare(compile_time_results) + compare2.print() + + return func + + return wrapper + + return decorator
diff --git a/benchmark/python/utils/tensor_generator.py b/benchmark/python/utils/tensor_generator.py new file mode 100644 index 0000000..5f9d304 --- /dev/null +++ b/benchmark/python/utils/tensor_generator.py
@@ -0,0 +1,74 @@ +import torch +import math +import numpy as np +from typing import Any + + +def generate_inputs( + shapes: tuple[Any, ...], + sparsity: float, + formats: tuple[str, ...], + dtype: Any = np.float64, + drange: tuple[Any, ...] = (1, 100), +) -> tuple[tuple[torch.Tensor, ...], tuple[torch.Tensor, ...]]: + """Generates dense and sparse tensor inputs. + + Args: + shapes: Shape for each input. + sparsity: Sparsity level for the inputs. + formats: Sparsity format for each input. + dtype: Data type of the generated inputs. Default is np.float64. + drange: Data range of the non-zero values. Default is (1, 100). + + Returns: + dense_inputs: all dense tensors. + sparse_inputs: inputs are of the specified sparsity format, such as CSR. + """ + dense_inputs = [] + sparse_inputs = [] + # Each input has a different seed. + for seed, shape in enumerate(shapes): + dense_inputs.append(generate_tensor(seed, shape, sparsity, dtype, drange)) + for idx, dense_input in enumerate(dense_inputs): + if formats[idx] == "dense": + sparse_inputs.append(dense_input) + else: + # TODO: support more sparsity formats. + sparse_inputs.append(dense_input.to_sparse_csr()) + return dense_inputs, sparse_inputs + + +def generate_tensor( + seed: int, + shape: tuple[Any, ...], + sparsity: float, + dtype: Any = np.float64, + drange: tuple[Any, ...] = (1, 100), +) -> torch.Tensor: + """Generates a tensor given sparsity level, shape and data type. + + Args: + seed: Seed value for np.random. + shape: A tuple for the shape of tensor. + sparsity: Sparsity level in the range of [0, 1]. + dtype: Data type of the generated tensor. Default is np.float64. + drange: Data range of the non-zero values. Default is (1, 100). + + Returns: + A dense torch tensor with the specified shape, sparsity level and type. + + Note: the tensor generated doesn't guarantee each batch will have the same + number of specified elements. Therefore, for batched CSR, torch.cat can be + used to concatenate generated tensors in the specified dimension. + """ + np.random.seed(seed) + size = math.prod(shape) + nse = size - int(math.ceil(sparsity * size)) + + flat_output = np.zeros(size) + indices = np.random.choice(size, nse, replace=False) + values = np.random.uniform(drange[0], drange[1], nse) + flat_output[indices] = values + + result = np.reshape(flat_output, shape).astype(dtype) + return torch.from_numpy(result)
diff --git a/externals/Enzyme b/externals/Enzyme new file mode 160000 index 0000000..cf89592 --- /dev/null +++ b/externals/Enzyme
@@ -0,0 +1 @@ +Subproject commit cf89592eb10c2e94352954a127fc8697aef40953
diff --git a/externals/torch-mlir b/externals/torch-mlir new file mode 160000 index 0000000..c7d52f6 --- /dev/null +++ b/externals/torch-mlir
@@ -0,0 +1 @@ +Subproject commit c7d52f63b482b2c30f4efb435ce0cc2efeab25d9
diff --git a/include/CMakeLists.txt b/include/CMakeLists.txt new file mode 100644 index 0000000..711b39d --- /dev/null +++ b/include/CMakeLists.txt
@@ -0,0 +1 @@ +add_subdirectory(mpact)
diff --git a/include/mpact-c/Registration.h b/include/mpact-c/Registration.h new file mode 100644 index 0000000..877b06c --- /dev/null +++ b/include/mpact-c/Registration.h
@@ -0,0 +1,27 @@ +/*===-- mpact-c/Registration.h - Registration functions -----*- C -*-===*\ +|* *| +|* Part of the MPACT Project, under the Apache License v2.0 with LLVM *| +|* Exceptions. *| +|* See https://llvm.org/LICENSE.txt for license information. *| +|* SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception *| +|* *| +\*===----------------------------------------------------------------------===*/ + +#ifndef MPACT_C_REGISTRATION_H +#define MPACT_C_REGISTRATION_H + +#include "mlir-c/IR.h" +#include "mlir-c/Support.h" + +#ifdef __cplusplus +extern "C" { +#endif + +/** Registers all passes for symbolic access with the global registry. */ +MLIR_CAPI_EXPORTED void mpactRegisterAllPasses(void); + +#ifdef __cplusplus +} +#endif + +#endif // MPACT_C_REGISTRATION_H
diff --git a/include/mpact/CMakeLists.txt b/include/mpact/CMakeLists.txt new file mode 100644 index 0000000..e31af32 --- /dev/null +++ b/include/mpact/CMakeLists.txt
@@ -0,0 +1 @@ +add_subdirectory(Transforms)
diff --git a/include/mpact/Transforms/CMakeLists.txt b/include/mpact/Transforms/CMakeLists.txt new file mode 100644 index 0000000..17587b1 --- /dev/null +++ b/include/mpact/Transforms/CMakeLists.txt
@@ -0,0 +1,5 @@ +set(LLVM_TARGET_DEFINITIONS Passes.td) +mlir_tablegen(Passes.h.inc -gen-pass-decls) +add_public_tablegen_target(MPACTTransformsPassIncGen) + +add_mlir_doc(Passes MPACTTransformsPass ./ -gen-pass-doc)
diff --git a/include/mpact/Transforms/Passes.h b/include/mpact/Transforms/Passes.h new file mode 100644 index 0000000..389184b --- /dev/null +++ b/include/mpact/Transforms/Passes.h
@@ -0,0 +1,22 @@ +//===------------------------------------------------------------*- C++ -*-===// +// +// Part of the MPACT Project, under the Apache License v2.0 with LLVM +// Exceptions. See https://llvm.org/LICENSE.txt for license information. +// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception +// Also available under a BSD-style license. See LICENSE. +// +//===----------------------------------------------------------------------===// + +#ifndef MPACT_TRANSFORMS_PASSES_H +#define MPACT_TRANSFORMS_PASSES_H + +namespace mlir { +namespace mpact { + +/// Registers all mpact transform passes. +void registerTransformPasses(); + +} // namespace mpact +} // namespace mlir + +#endif // MPACT_TRANSFORMS_PASSES_H
diff --git a/include/mpact/Transforms/Passes.td b/include/mpact/Transforms/Passes.td new file mode 100644 index 0000000..c83f5d4 --- /dev/null +++ b/include/mpact/Transforms/Passes.td
@@ -0,0 +1,51 @@ +//===-- Passes.td - Transforms pass definition file --------*- tablegen -*-===// +// +// Part of the MPACT Project, under the Apache License v2.0 with LLVM Exceptions. +// See https://llvm.org/LICENSE.txt for license information. +// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception +// +//===----------------------------------------------------------------------===// +// +// This file contains definitions for passes within the Transforms/ directory. +// +//===----------------------------------------------------------------------===// + +#ifndef MPACT_TRANSFORMS_PASSES +#define MPACT_TRANSFORMS_PASSES + +include "mlir/Pass/PassBase.td" + +def SparseEncodingPropagation : Pass<"sparse-encoding-propagation", "func::FuncOp"> { + let summary = "Propagate sparse tensor encodings"; + let description = [{ + A pass that propagates sparse tensor encodings. + + Background: To avoid introducing repetitive operations, sparse tensors + in MLIR try to reuse tensor operations whenever available. However, most + tensor operations are canonicalized/transformed without the knowledge + of sparsity. The pass tries to propagate missing sparse encodings. + + For example: + ```mlir + %s = tensor.extract_slice %input[0, 0,] [2, 1] [1, 1] + : tensor<2x3xf32, #sparse> to tensor<2x1xf32, #sparse> + + // After rank reducing (by tensor dialect transformation) + %t = tensor.extract_slice %input[0, 0,] [2, 1] [1, 1] + : tensor<2x3xf32, #sparse> to tensor<2xf32> + %s = tensor.expand_shape [[0, 1]] %t + : tensor<2xf32> to tensor<2x1xf32, #sparse> + + // After sparsity propagation + %t = tensor.extract_slice %input[0, 0,] [2, 1] [1, 1] + : tensor<2x3xf32, #sparse> to tensor<2xf32, #sparse1> + %s = tensor.expand_shape [[0, 1]] %t + : tensor<2xf32, #sparse1> to tensor<2x1xf32, #sparse> + ``` + }]; + + let constructor = "mlir::mpact::createSparseEncodingPropagationPass()"; + let dependentDialects = []; +} + +#endif // MPACT_TRANSFORMS_PASSES
diff --git a/include/mpact/Transforms/Sparsity/SparseEncodingPropagate.h b/include/mpact/Transforms/Sparsity/SparseEncodingPropagate.h new file mode 100644 index 0000000..a829400 --- /dev/null +++ b/include/mpact/Transforms/Sparsity/SparseEncodingPropagate.h
@@ -0,0 +1,24 @@ +//===------------------------------------------------------------*- C++ -*-===// +// +// Part of the MPACT Project, under the Apache License v2.0 with LLVM +// Exceptions. See https://llvm.org/LICENSE.txt for license information. +// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception +// Also available under a BSD-style license. See LICENSE. +// +//===----------------------------------------------------------------------===// + +#ifndef MPACT_TRANSFORMS_SPARSITY_SPARSEENCODINGPROPAGATE_H +#define MPACT_TRANSFORMS_SPARSITY_SPARSEENCODINGPROPAGATE_H + +#include "mlir/Dialect/Func/IR/FuncOps.h" +#include "mlir/IR/BuiltinOps.h" +#include "mlir/Pass/Pass.h" + +namespace mlir { +namespace mpact { +std::unique_ptr<OperationPass<func::FuncOp>> +createSparseEncodingPropagationPass(); +} +} // namespace mlir + +#endif // MPACT_TRANSFORMS_SPARSITY_SPARSEENCODINGPROPAGATE_H
diff --git a/lib/CAPI/CMakeLists.txt b/lib/CAPI/CMakeLists.txt new file mode 100644 index 0000000..f45c469 --- /dev/null +++ b/lib/CAPI/CMakeLists.txt
@@ -0,0 +1,12 @@ +add_mlir_public_c_api_library(MPACTCAPI + Registration.cpp + + ENABLE_AGGREGATION + + LINK_LIBS PUBLIC + MLIRIR + MLIRSupport + MPACTTransformPasses +) + +mpact_target_includes(MPACTCAPI)
diff --git a/lib/CAPI/Registration.cpp b/lib/CAPI/Registration.cpp new file mode 100644 index 0000000..51234fb --- /dev/null +++ b/lib/CAPI/Registration.cpp
@@ -0,0 +1,20 @@ +//===- Registration.cpp - C Interface for MLIR Registration ---------------===// +// +// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions. +// See https://llvm.org/LICENSE.txt for license information. +// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception +// Also available under a BSD-style license. See LICENSE. +// +//===----------------------------------------------------------------------===// + +#include "mpact-c/Registration.h" + +#include "mlir/CAPI/IR.h" +#include "mlir/Conversion/Passes.h" +#include "mlir/Dialect/Linalg/Passes.h" +#include "mlir/Transforms/Passes.h" +#include "mpact/Transforms/Passes.h" + +MLIR_CAPI_EXPORTED void mpactRegisterAllPasses() { + mlir::mpact::registerTransformPasses(); +}
diff --git a/lib/CMakeLists.txt b/lib/CMakeLists.txt new file mode 100644 index 0000000..9e7dcd1 --- /dev/null +++ b/lib/CMakeLists.txt
@@ -0,0 +1,2 @@ +add_subdirectory(CAPI) +add_subdirectory(Transforms)
diff --git a/lib/Transforms/CMakeLists.txt b/lib/Transforms/CMakeLists.txt new file mode 100644 index 0000000..edd737f --- /dev/null +++ b/lib/Transforms/CMakeLists.txt
@@ -0,0 +1,13 @@ +add_subdirectory(Sparsity) + +set(linked_libs MPACTSparsityPropagation) + +add_mlir_library(MPACTTransformPasses + Passes.cpp + + DEPENDS + MPACTTransformsPassIncGen + + LINK_LIBS PUBLIC + ${linked_libs} +)
diff --git a/lib/Transforms/Passes.cpp b/lib/Transforms/Passes.cpp new file mode 100644 index 0000000..2363ac1 --- /dev/null +++ b/lib/Transforms/Passes.cpp
@@ -0,0 +1,22 @@ +//===----------------------------------------------------------------------===// +// +// Part of the MPACT Project, under the Apache License v2.0 with LLVM +// Exceptions. See https://llvm.org/LICENSE.txt for license information. +// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception +// Also available under a BSD-style license. See LICENSE. +// +//===----------------------------------------------------------------------===// + +#include "mpact/Transforms/Passes.h" +#include "mpact/Transforms/Sparsity/SparseEncodingPropagate.h" + +//===----------------------------------------------------------------------===// +// Pass registration +//===----------------------------------------------------------------------===// + +namespace { +#define GEN_PASS_REGISTRATION +#include "mpact/Transforms/Passes.h.inc" +} // end namespace + +void mlir::mpact::registerTransformPasses() { ::registerPasses(); }
diff --git a/lib/Transforms/Sparsity/CMakeLists.txt b/lib/Transforms/Sparsity/CMakeLists.txt new file mode 100644 index 0000000..9323021 --- /dev/null +++ b/lib/Transforms/Sparsity/CMakeLists.txt
@@ -0,0 +1,15 @@ +add_mlir_conversion_library(MPACTSparsityPropagation + SparseEncodingPropagate.cpp + + ADDITIONAL_HEADER_DIRS + ${PROJECT_SOURCE_DIR}/include/mpact/Transforms/Sparsity + + DEPENDS + MPACTTransformsPassIncGen + + LINK_LIBS PUBLIC + MLIRIR + MLIRPass +) + +mpact_target_includes(MPACTSparsityPropagation)
diff --git a/lib/Transforms/Sparsity/SparseEncodingPropagate.cpp b/lib/Transforms/Sparsity/SparseEncodingPropagate.cpp new file mode 100644 index 0000000..f42db3b --- /dev/null +++ b/lib/Transforms/Sparsity/SparseEncodingPropagate.cpp
@@ -0,0 +1,35 @@ +//===- SparseEncodingPropagate.cpp ---------------------------------------===// +// +// Part of the MPACT Project, under the Apache License v2.0 with LLVM +// Exceptions. See https://llvm.org/LICENSE.txt for license information. +// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception +// +//===----------------------------------------------------------------------===// + +#include "mpact/Transforms/Sparsity/SparseEncodingPropagate.h" + +namespace mlir { +#define GEN_PASS_DEF_SPARSEENCODINGPROPAGATION +#include "mpact/Transforms/Passes.h.inc" +} // namespace mlir + +using namespace mlir; + +// ----------------------------------------------------------------------------- +// The pass +// ----------------------------------------------------------------------------- + +namespace { +struct SparseEncodingPropagation + : public impl::SparseEncodingPropagationBase<SparseEncodingPropagation> { + SparseEncodingPropagation() = default; + SparseEncodingPropagation(const SparseEncodingPropagation &pass) = default; + + void runOnOperation() override {} +}; +} // namespace + +std::unique_ptr<OperationPass<func::FuncOp>> +mlir::mpact::createSparseEncodingPropagationPass() { + return std::make_unique<SparseEncodingPropagation>(); +}
diff --git a/python/CMakeLists.txt b/python/CMakeLists.txt new file mode 100644 index 0000000..aab4f9a --- /dev/null +++ b/python/CMakeLists.txt
@@ -0,0 +1,78 @@ +#------------------------------------------------------------------------------- +# The MPACT Compiler Python Modules +#------------------------------------------------------------------------------- + +# Disables generation of "version soname" (i.e. libFoo.so.<version>). +set(CMAKE_PLATFORM_NO_VERSIONED_SONAME ON) + +# The directory at which the Python import tree begins. +set(MPACT_PYTHON_ROOT_DIR "${CMAKE_CURRENT_SOURCE_DIR}/mpact") + +# We vendor our own MLIR instance in the `mpact` namespace. +add_compile_definitions("MLIR_PYTHON_PACKAGE_PREFIX=mpact.") + +declare_mlir_python_sources(MPACTPythonSources) +declare_mlir_python_sources(MPACTPythonExtensions) + +declare_mlir_python_sources(MPACTPythonSources.PublicAPI + ROOT_DIR "${MPACT_PYTHON_ROOT_DIR}" + ADD_TO_PARENT MPACTPythonSources + SOURCES + mpactbackend.py +) + +declare_mlir_python_sources(MPACTPythonSources.SampleModels + ROOT_DIR "${MPACT_PYTHON_ROOT_DIR}" + ADD_TO_PARENT MPACTPythonSources + SOURCES_GLOB + models/*.py +) + +#------------------------------------------------------------------------------- +# Extensions +#------------------------------------------------------------------------------- + +declare_mlir_python_extension(MPACTPythonExtensions.Main + MODULE_NAME _mpact + ADD_TO_PARENT MPACTPythonExtensions + SOURCES + MPACTModule.cpp + EMBED_CAPI_LINK_LIBS + MPACTCAPI + PRIVATE_LINK_LIBS + LLVMSupport +) + +#------------------------------------------------------------------------------- +# Python Modules +#------------------------------------------------------------------------------- + +set(_source_components + MLIRPythonSources + MLIRPythonExtension.Core + MLIRPythonExtension.RegisterEverything + + # We need the FxImporter from torch-mlir + TorchMLIRPythonSources.Importers + TorchMLIRPythonSources.Dialects + TorchMLIRPythonExtensions + + MPACTPythonSources + MPACTPythonExtensions +) + +add_mlir_python_common_capi_library(MPACTAggregateCAPI + INSTALL_COMPONENT MPACTPythonModules + INSTALL_DESTINATION python_packages/mpact/mpact/_mlir_libs + OUTPUT_DIRECTORY "${MPACT_PYTHON_PACKAGES_DIR}/mpact/mpact/_mlir_libs" + RELATIVE_INSTALL_ROOT ".." + DECLARED_SOURCES ${_source_components} +) + +add_mlir_python_modules(MPACTPythonModules + ROOT_PREFIX "${MPACT_PYTHON_PACKAGES_DIR}/mpact/mpact" + INSTALL_PREFIX "python_packages/mpact/mpact" + DECLARED_SOURCES ${_source_components} + COMMON_CAPI_LINK_LIBS + MPACTAggregateCAPI +)
diff --git a/python/MPACTModule.cpp b/python/MPACTModule.cpp new file mode 100644 index 0000000..5751287 --- /dev/null +++ b/python/MPACTModule.cpp
@@ -0,0 +1,17 @@ +//===-- MPACTModule.cpp ------------------------------------*- cpp -*-===// +// +// Part of the MPACT Project, under the Apache License v2.0 with LLVM +// Exceptions. See https://llvm.org/LICENSE.txt for license information. +// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception +// Also available under a BSD-style license. See LICENSE. +// +//===----------------------------------------------------------------------===// + +#include "mlir/Bindings/Python/PybindAdaptors.h" +#include "mpact-c/Registration.h" + +PYBIND11_MODULE(_mpact, m) { + mpactRegisterAllPasses(); + + m.doc() = "mpact main python extension"; +}
diff --git a/python/mpact/models/gat.py b/python/mpact/models/gat.py new file mode 100644 index 0000000..b8ab229 --- /dev/null +++ b/python/mpact/models/gat.py
@@ -0,0 +1,87 @@ +import torch +import torch.nn.functional as F + + +class GraphAttentionLayer(torch.nn.Module): + def __init__( + self, + in_features: int, + out_features: int, + n_heads: int, + dropout: float = 0.4, + leaky_relu_slope: float = 0.2, + ): + super(GraphAttentionLayer, self).__init__() + self.n_heads = n_heads + self.dropout = dropout + self.n_hidden = out_features + self.W = torch.nn.Parameter( + torch.empty(size=(in_features, self.n_hidden * n_heads)) + ) + self.a = torch.nn.Parameter(torch.empty(size=(n_heads, 2 * self.n_hidden, 1))) + self.leakyrelu = torch.nn.LeakyReLU(leaky_relu_slope) + self.softmax = torch.nn.Softmax(dim=1) + torch.nn.init.ones_(self.W) + torch.nn.init.ones_(self.a) + + def forward(self, h: torch.Tensor, adj_mat: torch.Tensor): + n_nodes = h.shape[0] + h_transformed = torch.mm(h, self.W) + h_transformed = F.dropout(h_transformed, self.dropout, training=self.training) + h_transformed = h_transformed.view( + n_nodes, self.n_heads, self.n_hidden + ).permute(1, 0, 2) + e = self._get_attention_scores(h_transformed) + connectivity_mask = -9e16 * torch.ones_like(e) + e = torch.where(adj_mat > 0, e, connectivity_mask) + attention = F.softmax(e, dim=-1) + attention = F.dropout(attention, self.dropout, training=self.training) + h_prime = torch.matmul(attention, h_transformed) + return h_prime.mean(dim=0) + + def _get_attention_scores(self, h_transformed: torch.Tensor): + source_scores = torch.matmul(h_transformed, self.a[:, : self.n_hidden, :]) + target_scores = torch.matmul(h_transformed, self.a[:, self.n_hidden :, :]) + e = source_scores + target_scores.mT + return self.leakyrelu(e) + + +class GAT(torch.nn.Module): + """ + Graph Attention Network (GAT) inspired by <https://arxiv.org/pdf/1710.10903.pdf>. + """ + + def __init__( + self, + in_features, + n_hidden, + n_heads, + num_classes, + dropout=0.4, + leaky_relu_slope=0.2, + ): + super(GAT, self).__init__() + self.gat1 = GraphAttentionLayer( + in_features=in_features, + out_features=n_hidden, + n_heads=n_heads, + dropout=dropout, + leaky_relu_slope=leaky_relu_slope, + ) + self.gat2 = GraphAttentionLayer( + in_features=n_hidden, + out_features=num_classes, + n_heads=1, + dropout=dropout, + leaky_relu_slope=leaky_relu_slope, + ) + + def forward(self, input_tensor: torch.Tensor, adj_mat: torch.Tensor): + x = self.gat1(input_tensor, adj_mat) + x = F.elu(x) + x = self.gat2(x, adj_mat) + return F.log_softmax(x, dim=1) + + +def gat_4_64_8_3(): + return GAT(in_features=4, n_hidden=64, n_heads=8, num_classes=3)
diff --git a/python/mpact/models/gcn.py b/python/mpact/models/gcn.py new file mode 100644 index 0000000..c41e6d9 --- /dev/null +++ b/python/mpact/models/gcn.py
@@ -0,0 +1,47 @@ +import torch +import torch.nn.functional as F + + +class GraphConv(torch.nn.Module): + def __init__(self, input_dim, output_dim): + super(GraphConv, self).__init__() + self.kernel = torch.nn.Parameter(torch.Tensor(input_dim, output_dim)) + torch.nn.init.ones_(self.kernel) + self.bias = torch.nn.Parameter(torch.Tensor(output_dim)) + torch.nn.init.ones_(self.bias) + + def forward(self, inp, adj_mat): + # Input matrix times weight matrix. + support = torch.mm(inp, self.kernel) + # Sparse adjacency matrix times support matrix. + output = torch.spmm(adj_mat, support) + # Add bias. + output = output + self.bias + return output + + +class GCN(torch.nn.Module): + """ + Graph Convolutional Network (GCN) inspired by <https://arxiv.org/pdf/1609.02907.pdf>. + """ + + def __init__(self, input_dim, hidden_dim, output_dim, dropout_p=0.1): + super(GCN, self).__init__() + self.gc1 = GraphConv(input_dim, hidden_dim) + self.gc2 = GraphConv(hidden_dim, output_dim) + self.dropout = torch.nn.Dropout(dropout_p) + + def forward(self, input_tensor, adj_mat): + x = self.gc1(input_tensor, adj_mat) + x = F.relu(x) + x = self.dropout(x) + x = self.gc2(x, adj_mat) + return F.log_softmax(x, dim=1) + + +def graphconv_4_4(): + return GraphConv(input_dim=4, output_dim=4) + + +def gcn_4_16_4(): + return GCN(input_dim=4, hidden_dim=16, output_dim=4)
diff --git a/python/mpact/models/kernels.py b/python/mpact/models/kernels.py new file mode 100644 index 0000000..d18d88a --- /dev/null +++ b/python/mpact/models/kernels.py
@@ -0,0 +1,54 @@ +import torch + + +class MVNet(torch.nn.Module): + def forward(self, x, v): + return torch.mv(x, v) + + +class MMNet(torch.nn.Module): + def forward(self, x, v): + return torch.mm(x, v) + + +class AddNet(torch.nn.Module): + def forward(self, x, v): + return torch.add(x, v) + + +class MulNet(torch.nn.Module): + def forward(self, x, v): + return torch.mul(x, v) + + +class SelfNet(torch.nn.Module): + def forward(self, x): + return x + + +class SDDMMNet(torch.nn.Module): + def forward(self, x, y, z): + return torch.mul(x, torch.mm(y, z)) + + +class SqSum(torch.nn.Module): + def forward(self, x): + return (x * x).sum() + + +class FeatureScale(torch.nn.Module): + def forward(self, F): + sum_vector = torch.sum(F, dim=1) + reciprocal_vector = 1 / sum_vector + reciprocal_vector[reciprocal_vector == float("inf")] = 0 + scaling_diagonal = torch.diag(reciprocal_vector).to_sparse() + return scaling_diagonal @ F + + +class Normalization(torch.nn.Module): + def forward(self, A): + sum_vector = torch.sum(A, dim=1) + reciprocal_vector = 1 / sum_vector + reciprocal_vector[reciprocal_vector == float("inf")] = 0 + scaling_diagonal = torch.diag(reciprocal_vector).to_sparse() + return scaling_diagonal @ A @ scaling_diagonal
diff --git a/python/mpact/models/lif.py b/python/mpact/models/lif.py new file mode 100644 index 0000000..fcb5a55 --- /dev/null +++ b/python/mpact/models/lif.py
@@ -0,0 +1,58 @@ +import torch + + +def spike(input): + return (input >= 0).float() + + +def sqSum(input): + return (input * input).sum() + + +class LIF(torch.nn.Module): + def __init__(self): + super(LIF, self).__init__() + self.thresh = 1.0 + self.decay = 0.5 + self.act = spike + + def forward(self, X): + """A filter that yields a binary-valued sparse tensor.""" + mem = 0 + spike_pot = [] + T = X.size(-1) + for t in range(T): + mem = mem * self.decay + X[..., t] + spike = self.act(mem - self.thresh) + spike = spike.to_sparse().to_dense() # prop hack + mem = mem * (1.0 - spike) + spike_pot.append(spike) + spike_pot = torch.stack(spike_pot, dim=-1) + return spike_pot + + +class tdLayer(torch.nn.Module): + def __init__(self, layer): + super(tdLayer, self).__init__() + self.layer = layer + + def forward(self, X): + T = X.size(-1) + out = [] + for t in range(T): + m = self.layer(X[..., t]) + out.append(m) + out = torch.stack(out, dim=-1) + return out + + +class LIFSumOfSq(torch.nn.Module): + def __init__(self): + super(LIFSumOfSq, self).__init__() + self.spike = LIF() + self.layer = tdLayer(sqSum) + + def forward(self, X): + out = self.spike(X) + out = self.layer(out) + return out
diff --git a/python/mpact/models/resnet.py b/python/mpact/models/resnet.py new file mode 100644 index 0000000..2556597 --- /dev/null +++ b/python/mpact/models/resnet.py
@@ -0,0 +1,255 @@ +import torch +import numpy as np + + +def spike(input): + return (input >= 0).float() + + +class Straight(torch.nn.Module): + def forward(self, input): + return input + + +class tdLayer(torch.nn.Module): + def __init__(self, layer, bn=None): + super(tdLayer, self).__init__() + self.layer = layer + self.bn = bn if bn is not None else Straight() + + def forward(self, X): + T = X.size(-1) + out = [] + for t in range(T): + m = self.layer(X[..., t]) + out.append(m) + out = torch.stack(out, dim=-1) + out = self.bn(out) + return out + + +class LIF(torch.nn.Module): + def __init__(self): + super(LIF, self).__init__() + self.thresh = 1.0 + self.decay = 0.5 + self.act = spike + self.gama = 1.0 + + def forward(self, X, gama=1): + mem = 0 + spike_pot = [] + T = X.size(-1) + for t in range(T): + mem = mem * self.decay + X[..., t] + spike = self.act(mem - self.thresh) + mem = mem * (1.0 - spike) + spike_pot.append(spike) + spike_pot = torch.stack(spike_pot, dim=-1) + return spike_pot + + +class tdBatchNorm(torch.nn.BatchNorm2d): + def __init__( + self, + num_features, + eps=1e-05, + momentum=0.1, + alpha=1, + affine=True, + track_running_stats=True, + ): + super(tdBatchNorm, self).__init__( + num_features, eps, momentum, affine, track_running_stats + ) + self.alpha = alpha + + def forward(self, input): + exponential_average_factor = 0.0 + mean = self.running_mean + var = self.running_var + input = ( + self.alpha + * (input - mean[None, :, None, None, None]) + / (torch.sqrt(var[None, :, None, None, None] + self.eps)) + ) + if self.affine: + input = ( + input * self.weight[None, :, None, None, None] + + self.bias[None, :, None, None, None] + ) + return input + + +def conv3x3(in_planes, out_planes, stride=1, groups=1, dilation=1): + return torch.nn.Conv2d( + in_planes, + out_planes, + kernel_size=3, + stride=stride, + padding=dilation, + groups=groups, + bias=False, + dilation=dilation, + ) + + +def conv1x1(in_planes, out_planes, stride=1): + return torch.nn.Conv2d( + in_planes, out_planes, kernel_size=1, stride=stride, bias=False + ) + + +class BasicBlock(torch.nn.Module): + expansion = 1 + + def __init__( + self, + inplanes, + planes, + stride=1, + downsample=None, + groups=1, + base_width=64, + dilation=1, + norm_layer=None, + ): + super(BasicBlock, self).__init__() + if norm_layer is None: + norm_layer = tdBatchNorm + # norm_layer = nn.BatchNorm2d + if groups != 1 or base_width != 64: + raise ValueError("BasicBlock only supports groups=1 and base_width=64") + if dilation > 1: + raise NotImplementedError("Dilation > 1 not supported in BasicBlock") + # Both self.conv1 and self.downsample layers downsample the input when stride != 1 + self.conv1 = conv3x3(inplanes, planes, stride) + self.bn1 = norm_layer(planes) + self.conv2 = conv3x3(planes, planes) + self.bn2 = norm_layer(planes) + self.downsample = downsample + self.stride = stride + self.conv1_s = tdLayer(self.conv1, self.bn1) + self.conv2_s = tdLayer(self.conv2, self.bn2) + self.spike1 = LIF() + self.spike2 = LIF() + + def forward(self, x): + identity = x + + out = self.conv1_s(x) + out = self.spike1(out) + out = self.conv2_s(out) + + if self.downsample is not None: + identity = self.downsample(x) + + out += identity + out = self.spike2(out) + + return out + + +class ResNety(torch.nn.Module): + def __init__( + self, + block, + layers, + num_classes=10, + zero_init_residual=False, + groups=1, + width_per_group=64, + replace_stride_with_dilation=None, + norm_layer=None, + ): + super(ResNety, self).__init__() + if norm_layer is None: + norm_layer = tdBatchNorm + # norm_layer = nn.BatchNorm2d + self._norm_layer = norm_layer + self.inplanes = 64 + self.dilation = 1 + self.groups = groups + self.base_width = width_per_group + self.pre = torch.nn.Sequential( + tdLayer( + layer=torch.nn.Conv2d( + 3, self.inplanes, kernel_size=(3, 3), stride=(1, 1) + ), + bn=self._norm_layer(self.inplanes), + ), + LIF(), + ) + self.layer1 = self._make_layer(block, 64, layers[0], stride=2) + self.layer2 = self._make_layer(block, 128, layers[1], stride=2) + self.layer3 = self._make_layer(block, 256, layers[2], stride=2) + self.avgpool = tdLayer(torch.nn.AdaptiveAvgPool2d((1, 1))) + self.fc = tdLayer(torch.nn.Linear(256, num_classes)) + self.T = 6 + for m in self.modules(): + if isinstance(m, torch.nn.Conv2d): + torch.nn.init.kaiming_normal_( + m.weight, mode="fan_out", nonlinearity="relu" + ) + + def _make_layer(self, block, planes, blocks, stride=1, dilate=False): + norm_layer = self._norm_layer + downsample = None + previous_dilation = self.dilation + if dilate: + self.dilation *= stride + stride = 1 + if stride != 1 or self.inplanes != planes * block.expansion: + downsample = tdLayer( + conv1x1(self.inplanes, planes * block.expansion, stride), + norm_layer(planes * block.expansion), + ) + + layers = [] + layers.append( + block( + self.inplanes, + planes, + stride, + downsample, + self.groups, + self.base_width, + previous_dilation, + norm_layer, + ) + ) + self.inplanes = planes * block.expansion + for _ in range(1, blocks): + layers.append( + block( + self.inplanes, + planes, + groups=self.groups, + base_width=self.base_width, + dilation=self.dilation, + norm_layer=norm_layer, + ) + ) + + return torch.nn.Sequential(*layers) + + def _forward_impl(self, input): + out = [] + input = input.unsqueeze(-1).repeat(1, 1, 1, 1, self.T) + x = self.pre(input) + x = self.layer1(x) + x = self.layer2(x) + x = self.layer3(x) + x = self.avgpool(x) + x = x.view(x.size(0), -1, x.size(-1)) + x = self.fc(x) + for t in range(self.T): + out.append(x[..., t]) + return torch.stack(out, dim=1) + + def forward(self, x): + return self._forward_impl(x) + + +def resnet_20(): + return ResNety(block=BasicBlock, layers=[2, 2, 2], num_classes=10)
diff --git a/python/mpact/mpactbackend.py b/python/mpact/mpactbackend.py new file mode 100644 index 0000000..3e8ae4b --- /dev/null +++ b/python/mpact/mpactbackend.py
@@ -0,0 +1,530 @@ +# Initialize mpact python extension. +import mpact._mlir_libs._mpact + +import ctypes +from io import StringIO +import numpy as np +import os +import sys +import tempfile +import torch + +from typing import Any, Callable, Optional, Tuple, Dict, TypeVar, Union + +from mpact import ir +from mpact.ir import Module +from mpact.dialects import torch as torch_d +from mpact.execution_engine import * +from mpact.extras.fx_importer import FxImporter, SparsityMeta +from mpact.ir import * +from mpact.passmanager import * +from mpact.runtime import * + +# One time set up of support library. +SUPPORT_LIB = os.getenv("SUPPORT_LIB", default=None) +SHARED_LIBS = [] if SUPPORT_LIB is None else [SUPPORT_LIB] + +# The result of MPACT compile() and input to load(). +MpactCompiledArtifact = TypeVar("MpactCompiledArtifact") + + +def get_module_name_for_debug_dump(module): + """Gets a name suitable for a debug dump. + + The name is not guaranteed to be unique. + """ + if not "torch.debug_module_name" in module.operation.attributes: + return "UnnammedModule" + return StringAttr(module.operation.attributes["torch.debug_module_name"]).value + + +class MPACTCompilerError(Exception): + pass + + +def run_pipeline_with_repro_report( + module, pipeline: str, description: str, enable_ir_printing: bool = False +): + """Runs `pipeline` on `module`, with a nice repro report if it fails.""" + module_name = get_module_name_for_debug_dump(module) + original_stderr = sys.stderr + try: + sys.stderr = StringIO() + asm_for_error_report = module.operation.get_asm( + large_elements_limit=10, enable_debug_info=True + ) + # Lower module in place to make it ready for compiler backends. + with module.context as ctx: + pm = PassManager.parse(pipeline) + if enable_ir_printing: + ctx.enable_multithreading(False) + pm.enable_ir_printing() + pm.run(module.operation) + except Exception as e: + # TODO: More robust. + # - don't arbitrarily clutter up /tmp. When a test suite has many + # tests, this can be a big disk cost (also, /tmp/ is frequently a + # RAM fs, which increases worries about capacity). + # - don't have colliding filenames (hard to do without cluttering + # up /tmp) + # - if we do have have colliding filenames, writes should at least + # avoid being racy. + filename = os.path.join(tempfile.gettempdir(), module_name + ".mlir") + with open(filename, "w") as f: + f.write(asm_for_error_report) + debug_options = "-mlir-print-ir-after-all -mlir-disable-threading" + # Put something descriptive here even if description is empty. + description = description or f"{module_name} compile" + + message = f"""\ + {description} failed with the following diagnostics: + {sys.stderr.getvalue()} + + python exception: {e} + + The error can be reproduced with: + $ mpact-opt -pass-pipeline='{pipeline}' {filename} + Add '{debug_options}' to get the IR dump for debugging purpose. + """ + trimmed_message = "\n".join([m.lstrip() for m in message.split("\n")]) + raise MPACTCompilerError(trimmed_message) from None + finally: + sys.stderr = original_stderr + + +def assert_arg_type_is_supported(ty): + SUPPORTED = [ + np.float16, + np.float32, + np.float64, + np.uint8, + np.int8, + np.int32, + np.int64, + np.bool_, + np.complex64, + np.complex128, + ] + assert ( + ty in SUPPORTED + ), f"Only numpy arrays with dtypes in {SUPPORTED} are supported, but got {ty}" + + +memref_type_to_np_dtype = { + "mrf16": np.float16, + "mrf32": np.float32, + "mrf64": np.float64, + "mri1": np.bool_, + "mri8": np.int8, + "mri32": np.int32, + "mri64": np.int64, + "mrc32": np.complex64, + "mrc64": np.complex128, +} +elemental_type_to_ctype = { + "i1": ctypes.c_bool, + "i8": ctypes.c_byte, + "i64": ctypes.c_int, + "f32": ctypes.c_float, + "f64": ctypes.c_double, +} + +CONSUME_RETURN_FUNC_PREFIX = "refbackend_consume_func_return_" + +SPARSE_LAYOUTS = [ + torch.sparse_coo, + torch.sparse_csr, + torch.sparse_csc, + torch.sparse_bsr, + torch.sparse_bsc, +] + + +def get_return_funcs(module): + return_prefix_len = len(CONSUME_RETURN_FUNC_PREFIX) + return_funcs = [] + with module.context: + for func in module.body: + # Returns strings of the form `"refbackend.."` so `"` is deleted. + func_name = str(func.attributes["sym_name"]).replace('"', "") + if func_name[:return_prefix_len] == CONSUME_RETURN_FUNC_PREFIX: + return_funcs.append(func_name) + + return return_funcs + + +def get_ctype_func(func_name): + return_prefix_len = len(CONSUME_RETURN_FUNC_PREFIX) + ret_types = func_name[return_prefix_len:].split("_") + ctypes_arg = [None] + for type in ret_types: + if type in elemental_type_to_ctype: + ctypes_arg.append(elemental_type_to_ctype[type]) + elif type in memref_type_to_np_dtype: + ctypes_arg.append(ctypes.POINTER(UnrankedMemRefDescriptor)) + else: + assert False, f"Not supported type: {type}" + + return ctypes.CFUNCTYPE(*ctypes_arg), ret_types + + +class MpactBackendInvoker: + def __init__(self, module, opt_level): + self.ee = ExecutionEngine(module, opt_level=opt_level, shared_libs=SHARED_LIBS) + self.result = None + + return_funcs = get_return_funcs(module) + + for ret_func in return_funcs: + ctype_wrapper, ret_types = get_ctype_func(ret_func) + + def consume_return_funcs(*args): + self.result = tuple( + [ + ( + arg + if type in elemental_type_to_ctype + else unranked_memref_to_numpy( + arg, memref_type_to_np_dtype[type] + ) + ) + for arg, type in zip(args, ret_types) + ] + ) + if len(self.result) == 1: + self.result = self.result[0] + + self.ee.register_runtime(ret_func, ctype_wrapper(consume_return_funcs)) + + def __getattr__(self, function_name: str): + def invoke(*args): + ffi_args = [] + for arg in args: + assert_arg_type_is_supported(arg.dtype) + ffi_args.append( + ctypes.pointer(ctypes.pointer(get_unranked_memref_descriptor(arg))) + ) + + self.ee.invoke(function_name, *ffi_args) + result = self.result + assert result is not None, "Invocation didn't produce a result" + self.result = None + return result + + return invoke + + +LOWERING_PIPELINE_TEMPLATE = ( + "builtin.module(" + + ",".join( + [ + "func.func(linalg-generalize-named-ops)", + "func.func(linalg-fuse-elementwise-ops)", + "convert-shape-to-std", + # Propagate sparse encodings before sparsifier mini-pipeline. + # TODO: the following pass currently contains no pattern. Will be + # added as needed. + "func.func(sparse-encoding-propagation)", + # MLIR Sparsifier mini-pipeline: + # use the PyTorch assembler conventions + # enable vectorization with VL=16 (more or less assumes AVX512 for float) + # allow 32-bit index optimizations (unsafe for very large dimensions) + "sparse-assembler{{direct-out}}", + "sparsification-and-bufferization{{{sp_options}}}", + "sparse-storage-specifier-to-llvm", + # Buffer deallocation pass does not know how to handle realloc. + "func.func(expand-realloc)", + # Generalize pad and concat after sparse compiler, as they are handled + # differently when the operations involve sparse operands. + "func.func(refback-generalize-tensor-pad)", + "func.func(refback-generalize-tensor-concat)", + # Bufferize. + "func.func(tm-tensor-bufferize)", + "one-shot-bufferize{{copy-before-write bufferize-function-boundaries function-boundary-type-conversion=identity-layout-map}}", + "refback-mlprogram-bufferize", + "func.func(finalizing-bufferize)", + "func.func(buffer-deallocation)", + # Inline sparse helper methods where useful (but after dealloc). + "inline", + "refback-munge-calling-conventions", + "func.func(tm-tensor-to-loops)", + "func.func(refback-munge-memref-copy)", + "func.func(convert-linalg-to-loops)", + "func.func(lower-affine)", + "convert-scf-to-cf", + "func.func(refback-expand-ops-for-llvm)", + "func.func(arith-expand)", + "func.func(convert-math-to-llvm)", + "convert-math-to-libm", + "expand-strided-metadata", + "finalize-memref-to-llvm", + "lower-affine", + "convert-bufferization-to-memref", + "finalize-memref-to-llvm", + "func.func(convert-arith-to-llvm)", + # Vector code (SIMD): + # allow fp reductions to reassociate + # allow 32-bit index optimizations (unsafe for very large dimensions) + # assume we are running on a good ol' Intel X86 (disable for ARM/other) + "convert-vector-to-llvm{{reassociate-fp-reductions force-32bit-vector-indices enable-x86vector}}", + "convert-func-to-llvm", + "convert-cf-to-llvm", + "convert-complex-to-llvm", + "reconcile-unrealized-casts", + ] + ) + + ")" +) + + +class MpactBackendCompiler: + """Main entry-point for the MPACT backend compiler.""" + + def __init__(self, opt_level, use_sp_it): + self.opt_level = opt_level + self.use_sp_it = use_sp_it + + def compile(self, imported_module: Module) -> MpactCompiledArtifact: + sp_options = ( + "sparse-emit-strategy=sparse-iterator" + if self.use_sp_it + else "vl=16 enable-simd-index32" + ) + LOWERING_PIPELINE = LOWERING_PIPELINE_TEMPLATE.format(sp_options=sp_options) + """Compiles an imported module, with a flat list of functions. + The module is expected to be in linalg-on-tensors + scalar code form. + + Args: + imported_module: The MLIR module in the torch dialect. + Returns: + An opaque artifact that can be passed to `load`. + """ + run_pipeline_with_repro_report( + imported_module, + LOWERING_PIPELINE, + "Lowering Linalg-on-Tensors IR to LLVM with MpactBackendCompiler", + enable_ir_printing=False, + ) + return imported_module + + def load(self, module: MpactCompiledArtifact) -> MpactBackendInvoker: + """Loads a compiled artifact into the runtime. + + Args: + module: The result of a previous call to `compile`. + Returns: + MPactInvoker to call a compiled method (viz `invoker.foo(...)`). + """ + return MpactBackendInvoker(module, self.opt_level) + + +def sparse_metadata(a: torch.Tensor) -> SparsityMeta: + """ + Returns a meta data tuple for the given sparse tensor. + + NOTE: this will be fully replaced by fx graph SparseTensorMetadata + """ + sparse_dim = a.sparse_dim() + dense_dim = a.dense_dim() + batch_dim = a.ndim - dense_dim - sparse_dim + blocksize = None + if a.layout is torch.sparse_coo: + return SparsityMeta( + a.layout, + batch_dim, + sparse_dim, + dense_dim, + blocksize, + a._indices().dtype, + a._indices().dtype, + ) + elif a.layout is torch.sparse_csr or a.layout is torch.sparse_bsr: + if a.layout is torch.sparse_bsr: + blocksize = a.values().shape[batch_dim + 1 : batch_dim + 3] + return SparsityMeta( + a.layout, + batch_dim, + sparse_dim, + dense_dim, + blocksize, + a.crow_indices().dtype, + a.col_indices().dtype, + ) + elif a.layout is torch.sparse_csc or a.layout is torch.sparse_bsc: + if a.layout is torch.sparse_bsc: + blocksize = a.values().shape[batch_dim + 1 : batch_dim + 3] + return SparsityMeta( + a.layout, + batch_dim, + sparse_dim, + dense_dim, + blocksize, + a.ccol_indices().dtype, + a.row_indices().dtype, + ) + else: + raise RuntimeError(f"Unsupported sparse layout for {a}") + + +def sparse_arg(args, i): + if isinstance(args[i], torch.fx.node.Node): + return args[i].meta.get("sparsity", None) + return None + + +def sparse_export( + f: Callable, args: Tuple[Any, ...], kwargs: Optional[Dict[str, Any]] = None +) -> torch.export.ExportedProgram: + """ + This is a ***temporary*** wrapper around `torch.export.export` + that eventually should be removed and simply replaced by the + standard API for exporting traced graphs. + + But until issue + + https://github.com/pytorch/pytorch/pull/117907 + + is addressed, this wrapper provides support for the sparse + tensor types by first converting all operands to dense tensors, + building the traced graph as for the dense case, then annotating + sparse parameters with their actual sparse layout attributes, + followed by some simple propagation rules. This temporary solution + accelerates testing torch-mlir with PyTorch sparse tensors until + the issue is resolved upstream. + """ + # Convert all arguments to dense. + dargs = tuple(a.to_dense() if a.layout in SPARSE_LAYOUTS else a for a in args) + mask = [a.layout in SPARSE_LAYOUTS for a in args] + # Build the regular FX traced graph with only dense arguments + # (the current version would crash otherwise, see issue above). + prog = torch.export.export(f, dargs, kwargs) + # Annotate sparse arguments in the graph and apply some very + # basic propagation rules for sparsity. + specs = prog.graph_signature.input_specs + alen = len(specs) + k = 0 + for i, node in enumerate(prog.graph.nodes): + if node.op == "placeholder": + # Argument. + spec = specs[i] + if spec.kind is torch.export.graph_signature.InputKind.USER_INPUT: + if mask[k]: + node.meta["sparsity"] = sparse_metadata(args[k]) + k = k + 1 + elif node.op == "call_function": + # TODO: use upstream _opname implementation when available + opname = node.target._schema.name.split("::")[1] + # Zero preserving elt-wise unary op. + if opname in {"abs", "neg", "relu", "sin"}: + node.meta["sparsity"] = sparse_arg(node.args, 0) + # Some simplistic rules for preserving sparsity. Soon + # to be replaced by proper FX graph propagation. + elif opname in {"mul"}: + m0 = sparse_arg(node.args, 0) + m1 = sparse_arg(node.args, 1) + if m0 is not None: + node.meta["sparsity"] = m0 + elif m1 is not None: + node.meta["sparsity"] = m1 + elif opname in {"add", "mm"}: + m0 = sparse_arg(node.args, 0) + m1 = sparse_arg(node.args, 1) + if m0 is not None and m1 is not None: + node.meta["sparsity"] = m0 + elif opname == "_to_sparse" or opname == "to_sparse": + dim = len(node.meta.get("val").shape) + node.meta["sparsity"] = SparsityMeta( + torch.sparse_coo, 0, dim, 0, None, torch.int64, torch.int64 + ) + # TODO: Uncomment this to hack sparsity into the network. + # elif opname == "_to_dense" or opname == "to_dense": + # # hack (assumes we never really want the to_dense for now) + # node.meta["sparsity"] = sparse_arg(node.args, 0) + elif opname == "select" and sparse_arg(node.args, 0): + dim = len(node.meta.get("val").shape) + node.meta["sparsity"] = SparsityMeta( + torch.sparse_coo, 0, dim, 0, None, torch.int64, torch.int64 + ) + elif opname == "stack" and sparse_arg(node.args[0], 0): + dim = len(node.meta.get("val").shape) + node.meta["sparsity"] = SparsityMeta( + torch.sparse_coo, 0, dim - 1, 1, None, torch.int64, torch.int64 + ) + return prog + + +def export_and_import(f, *args, **kwargs): + """This method implements Stella's importer, stripped down to essentials.""" + context = ir.Context() + torch_d.register_dialect(context) + fx_importer = FxImporter(context=context) + prog = sparse_export(f, args, kwargs) + fx_importer.import_frozen_program(prog) + return fx_importer.module + + +def mpact_jit_compile(f, *args, opt_level=2, use_sp_it=False, **kwargs): + """This method compiles the given callable using the MPACT backend.""" + # Import module and lower into Linalg IR. + module = export_and_import(f, *args, **kwargs) + run_pipeline_with_repro_report( + module, + ( + "builtin.module(" + "func.func(torch-decompose-complex-ops)," + "torch-backend-to-linalg-on-tensors-backend-pipeline)" + ), + "Lowering TorchFX IR -> Linalg IR", + enable_ir_printing=False, + ) + # Compile with MPACT backend compiler. + backend = MpactBackendCompiler(opt_level=opt_level, use_sp_it=use_sp_it) + compiled = backend.compile(module) + invoker = backend.load(compiled) + return invoker, f + + +def mpact_jit_run(invoker, f, *args, **kwargs): + """This method runs the given callable using the given MPACT invoker.""" + xargs = [] + # Prepare all the named buffer parameters (assume all dense). + # All scalar arguments are filtered out since they appear inline. + params = dict(f.named_buffers(remove_duplicate=True)) + params_flat, params_spec = torch.utils._pytree.tree_flatten(params) + for p in params_flat: + if len(p.shape) > 0: + xargs.append(p.numpy()) + # Prepare input parameters. Sparse input tensors are split into + # their composite tensors. All PyTorch tensors are converted + # to their backing numpy arrays. Note that the output consists + # of numpy arrays as well, which can trivially be reconstructed + # into PyTorch tensors (dense and sparse). + for a in args: + if a.layout is torch.sparse_coo: + # Construct the additional position array required by MLIR with data + # array([0, nnz]). The COO format always uses int64 indices. + xargs.append(np.array([0, a._nnz()], dtype=np.int64)) + # Transform a tensor<ndim x nnz> into ndim x tensor<nnz> to conform + # to the MLIR SoA COO representation. + for idx in a._indices(): + xargs.append(idx.numpy()) + xargs.append(a._values().numpy()) + elif a.layout is torch.sparse_csr or a.layout is torch.sparse_bsr: + xargs.append(a.crow_indices().numpy()) + xargs.append(a.col_indices().numpy()) + xargs.append(a.values().numpy()) + elif a.layout is torch.sparse_csc or a.layout is torch.sparse_bsc: + xargs.append(a.ccol_indices().numpy()) + xargs.append(a.row_indices().numpy()) + xargs.append(a.values().numpy()) + else: + xargs.append(a.numpy()) + # Invoke. + return invoker.main(*xargs) + + +# Convenience wrapper. +def mpact_jit(f, *args, **kwargs): + """This method compiles and runs the given callable using the MPACT backend.""" + invoker, fn = mpact_jit_compile(f, *args, **kwargs) + return mpact_jit_run(invoker, fn, *args, **kwargs)
diff --git a/setup.py b/setup.py new file mode 100644 index 0000000..17bdc0b --- /dev/null +++ b/setup.py
@@ -0,0 +1,187 @@ +# Script for generating the mpact wheel. +# ``` +# $ python setup.py bdist_wheel +# ``` +# Environment variables you are probably interested in: +# +# CMAKE_BUILD_TYPE: +# specify the build type: DEBUG/RelWithDebInfo/Release +# +# MPACT_CMAKE_ALREADY_BUILT: +# the `MPACT_CMAKE_BUILD_DIR` directory has already been compiled, +# and the CMake compilation process will not be executed again. +# On CIs, it is often advantageous to re-use/control the CMake build directory. +# +# It is recommended to build with Ninja and ccache. To do so, set environment +# variables by prefixing to above invocations: +# ``` +# CMAKE_GENERATOR=Ninja CMAKE_C_COMPILER_LAUNCHER=ccache CMAKE_CXX_COMPILER_LAUNCHER=ccache +# ``` +# +# Implementation notes: +# The contents of the wheel is just the contents of the `python_packages` +# directory that our CMake build produces. We go through quite a bit of effort +# on the CMake side to organize that directory already, so we avoid duplicating +# that here, and just package up its contents. + +import os +import pathlib +import shutil +import subprocess +import sys + +from datetime import date +from distutils.command.build import build as _build +from setuptools import setup, Extension +from setuptools.command.build_ext import build_ext +from setuptools.command.build_py import build_py + + +def _check_env_flag(name: str, default=None) -> bool: + return str(os.getenv(name, default)).upper() in ["ON", "1", "YES", "TRUE", "Y"] + + +PACKAGE_VERSION = "".join(str(date.today()).split("-")) +SRC_DIR = pathlib.Path(__file__).parent.absolute() +CMAKE_BUILD_TYPE = os.getenv("CMAKE_BUILD_TYPE", "Release") +MPACT_CMAKE_ALREADY_BUILT = _check_env_flag("MPACT_CMAKE_ALREADY_BUILT", False) +MPACT_CMAKE_BUILD_DIR = os.path.join(SRC_DIR, "build") + + +# Build phase discovery is unreliable. Just tell it what phases to run. +class CustomBuild(_build): + def initialize_options(self): + _build.initialize_options(self) + # Make setuptools not steal the build directory name, + # because the mlir c++ developers are quite + # used to having build/ be for cmake + self.build_base = "setup_build" + + def run(self): + self.run_command("build_py") + self.run_command("build_ext") + self.run_command("build_scripts") + + +class CMakeBuild(build_py): + def cmake_build(self, cmake_build_dir): + llvm_dir = str( + SRC_DIR / "externals" / "torch-mlir" / "externals" / "llvm-project" / "llvm" + ) + cmake_config_args = [ + f"cmake", + f"-GNinja", + f"-DCMAKE_BUILD_TYPE=Release", + f"-DPython3_FIND_VIRTUALENV=ONLY", + f"-DLLVM_ENABLE_PROJECTS=mlir", + f"-DLLVM_EXTERNAL_PROJECTS='torch-mlir;mpact'", + f"-DLLVM_EXTERNAL_TORCH_MLIR_SOURCE_DIR='{SRC_DIR}/externals/torch-mlir'", + f"-DLLVM_EXTERNAL_MPACT_SOURCE_DIR='{SRC_DIR}'", + f"-DLLVM_TARGETS_TO_BUILD=host", + f"-DMLIR_ENABLE_BINDINGS_PYTHON=ON", + # Optimization options for building wheels. + f"-DCMAKE_VISIBILITY_INLINES_HIDDEN=ON", + f"-DCMAKE_C_VISIBILITY_PRESET=hidden", + f"-DCMAKE_CXX_VISIBILITY_PRESET=hidden", + f"{llvm_dir}", + ] + + cmake_build_args = [ + f"cmake", + f"--build", + f".", + f"--target", + f"MPACTPythonModules", + f"MPACTBenchmarkPythonModules", + ] + + try: + subprocess.check_call(cmake_config_args, cwd=cmake_build_dir) + subprocess.check_call(cmake_build_args, cwd=cmake_build_dir) + except subprocess.CalledProcessError as e: + print("cmake build failed with\n", e) + print("debug by follow cmake command:") + sys.exit(e.returncode) + finally: + print(f"cmake config: {' '.join(cmake_config_args)}") + print(f"cmake build: {' '.join(cmake_build_args)}") + print(f"cmake workspace: {cmake_build_dir}") + print(SRC_DIR) + + def run(self): + target_dir = self.build_lib + cmake_build_dir = MPACT_CMAKE_BUILD_DIR + if not cmake_build_dir: + cmake_build_dir = os.path.abspath(os.path.join(target_dir, "..", "build")) + + python_package_dir = os.path.join( + cmake_build_dir, "tools", "mpact", "python_packages", "mpact" + ) + if not MPACT_CMAKE_ALREADY_BUILT: + os.makedirs(cmake_build_dir, exist_ok=True) + cmake_cache_file = os.path.join(cmake_build_dir, "CMakeCache.txt") + if os.path.exists(cmake_cache_file): + os.remove(cmake_cache_file) + # NOTE: With repeated builds for different Python versions, the + # prior version binaries will continue to accumulate. Here we just + # delete the directory where we build native extensions to keep + # this from happening but still take advantage of most of the + # build cache. + mlir_libs_dir = os.path.join(python_package_dir, "mpact", "_mlir_libs") + if os.path.exists(mlir_libs_dir): + print(f"Removing _mlir_mlibs dir to force rebuild: {mlir_libs_dir}") + shutil.rmtree(mlir_libs_dir) + else: + print(f"Not removing _mlir_libs dir (does not exist): {mlir_libs_dir}") + self.cmake_build(cmake_build_dir) + + if os.path.exists(target_dir): + shutil.rmtree(target_dir, ignore_errors=False, onerror=None) + + shutil.copytree(python_package_dir, target_dir, symlinks=False) + + +class CMakeExtension(Extension): + def __init__(self, name, sourcedir=""): + Extension.__init__(self, name, sources=[]) + self.sourcedir = os.path.abspath(sourcedir) + + +class NoopBuildExtension(build_ext): + def build_extension(self, ext): + pass + + +with open("README.md", "r", encoding="utf-8") as fh: + long_description = fh.read() + + +# Requires and extension modules depend on whether building PyTorch +# extensions. +INSTALL_REQUIRES = [ + "numpy", + "packaging", +] +EXT_MODULES = [ + CMakeExtension("mpact._mlir_libs._mpact"), +] + +setup( + name="mpact", + version=f"{PACKAGE_VERSION}", + author="Reid Tatge", + author_email="tatge@google.com", + description="MPACT retargetable ML compiler", + long_description=long_description, + long_description_content_type="text/markdown", + include_package_data=True, + cmdclass={ + "build": CustomBuild, + "built_ext": NoopBuildExtension, + "build_py": CMakeBuild, + }, + ext_modules=EXT_MODULES, + python_requires=">=3.8", + install_requires=INSTALL_REQUIRES, + zip_safe=False, +)
diff --git a/test/CMakeLists.txt b/test/CMakeLists.txt new file mode 100644 index 0000000..019d820 --- /dev/null +++ b/test/CMakeLists.txt
@@ -0,0 +1,25 @@ +#------------------------------------------------------------------------------- +# The MPACT Compiler Tests +#------------------------------------------------------------------------------- + +configure_lit_site_cfg( + ${CMAKE_CURRENT_SOURCE_DIR}/lit.site.cfg.py.in + ${CMAKE_CURRENT_BINARY_DIR}/lit.site.cfg.py + MAIN_CONFIG + ${CMAKE_CURRENT_SOURCE_DIR}/lit.cfg.py +) + +set(MPACT_TEST_DEPENDS + FileCheck count not + MPACTPythonModules + TorchMLIRPythonModules + torch-mlir-opt + ) + +add_lit_testsuite(check-mpact "Running the MPACT regression tests" + ${CMAKE_CURRENT_BINARY_DIR} + DEPENDS ${MPACT_TEST_DEPENDS} + ) +set_target_properties(check-mpact PROPERTIES FOLDER "Tests") + +add_lit_testsuites(MPACT ${CMAKE_CURRENT_SOURCE_DIR} DEPENDS ${TORCH_MLIR_TEST_DEPENDS})
diff --git a/test/lit.cfg.py b/test/lit.cfg.py new file mode 100644 index 0000000..910acca --- /dev/null +++ b/test/lit.cfg.py
@@ -0,0 +1,79 @@ +#------------------------------------------------------------------------------- +# The MPACT Compiler LIT Configuration +#------------------------------------------------------------------------------- + +import os +import platform +import re +import subprocess +import tempfile + +import lit.formats +import lit.util + +from lit.llvm import llvm_config +from lit.llvm.subst import ToolSubst +from lit.llvm.subst import FindTool + +# The name of this test suite. +config.name = "MPACT" + +# The test format. +config.test_format = lit.formats.ShTest(not llvm_config.use_lit_shell) + +# A list of file extensions to treat as test files. +config.suffixes = [".py"] + +# A list of files to exclude from the test suite. +config.excludes = [ + "CMakeLists.txt", + "README.txt", + "LICENSE.txt", + "lit.cfg.py", + "lit.site.cfg.py", +] + +# The root path where tests are located. +config.test_source_root = os.path.dirname(__file__) + +# The root path where tests should be run. +config.test_exec_root = os.path.join(config.mpact_obj_root, "test") +config.standalone_tools_dir = os.path.join(config.mpact_obj_root, "bin") + +# Substitutions. +config.substitutions.append(("%PATH%", config.environment["PATH"])) +config.substitutions.append(("%shlibext", config.llvm_shlib_ext)) + +# Tweak the PATH to include the tools dir. +llvm_config.with_environment("PATH", config.llvm_tools_dir, append_path=True) +llvm_config.with_environment( + "PATH", os.path.join(config.llvm_build_dir, "bin"), append_path=True +) +llvm_config.with_system_environment(["HOME", "INCLUDE", "LIB", "TMP", "TEMP"]) + +# On Windows the path to python could contain spaces in which case it needs to +# be provided in quotes. This is the equivalent of how %python is setup in +# llvm/utils/lit/lit/llvm/config.py. +if "Windows" in config.host_os: + config.python_executable = '"%s"' % (config.python_executable) + +# Tools. +tool_dirs = [ + config.standalone_tools_dir, + config.llvm_tools_dir, + config.mpact_obj_root, +] +tools = [ + "mpact-opt", + ToolSubst("%PYTHON", config.python_executable, unresolved="ignore"), +] + +llvm_config.add_tool_substitutions(tools, tool_dirs) + +llvm_config.with_environment( + "PYTHONPATH", + [ + os.path.join(config.mpact_obj_root, "python_packages/mpact"), + ], + append_path=True, +)
diff --git a/test/lit.site.cfg.py.in b/test/lit.site.cfg.py.in new file mode 100644 index 0000000..01f3d01 --- /dev/null +++ b/test/lit.site.cfg.py.in
@@ -0,0 +1,23 @@ +@LIT_SITE_CFG_IN_HEADER@ + +import sys + +config.host_os = "@HOST_OS@" +config.mpact_src_root = "@MPACT_SOURCE_DIR@" +config.mpact_obj_root = "@MPACT_BINARY_DIR@" +config.torch_mlir_obj_root = "@LLVM_BINARY_DIR@/tools/torch-mlir" +config.llvm_src_root = "@LLVM_SOURCE_DIR@" +config.llvm_obj_root = "@LLVM_BINARY_DIR@" +config.llvm_tools_dir = "@LLVM_TOOLS_DIR@" +config.llvm_build_dir = "@CMAKE_BINARY_DIR@" +config.llvm_lib_dir = "@LLVM_LIBS_DIR@" +config.llvm_shlib_dir = "@SHLIBDIR@" +config.llvm_shlib_ext = "@SHLIBEXT@" +config.llvm_exe_ext = "@EXEEXT@" +config.python_executable = "@Python3_EXECUTABLE@" + +import lit.llvm +lit.llvm.initialize(lit_config, config) + +# Let the main config do the real work. +lit_config.load_config(config, "@MPACT_SOURCE_DIR@/test/lit.cfg.py")
diff --git a/test/python/add.py b/test/python/add.py new file mode 100644 index 0000000..00d4d62 --- /dev/null +++ b/test/python/add.py
@@ -0,0 +1,89 @@ +# RUN: %PYTHON %s | FileCheck %s + +import torch +import numpy as np + +from mpact.mpactbackend import mpact_jit, mpact_jit_compile, mpact_jit_run + +from mpact.models.kernels import AddNet + + +def print_sparse(res): + print(res[0]) + print(res[1]) + print(res[2]) + + +net = AddNet() + +# Construct dense and sparse matrices. +X = torch.arange(0, 16, dtype=torch.float32).view(4, 4) +Y = torch.arange(16, 32, dtype=torch.float32).view(4, 4) +A = torch.tensor( + [ + [0.0, 1.0, 0.0, 0.0], + [0.0, 0.0, 0.0, 2.0], + [0.0, 0.0, 0.0, 0.0], + [3.0, 0.0, 0.0, 0.0], + ], + dtype=torch.float32, +) +S = A.to_sparse_csr() + +# +# CHECK: pytorch +# CHECK: tensor({{\[}}[16., 18., 20., 22.], +# CHECK: [24., 26., 28., 30.], +# CHECK: [32., 34., 36., 38.], +# CHECK: [40., 42., 44., 46.]{{\]}}) +# CHECK: tensor({{\[}}[16., 18., 18., 19.], +# CHECK: [20., 21., 22., 25.], +# CHECK: [24., 25., 26., 27.], +# CHECK: [31., 29., 30., 31.]{{\]}}) +# CHECK: tensor({{\[}}[ 0., 2., 2., 3.], +# CHECK: [ 4., 5., 6., 9.], +# CHECK: [ 8., 9., 10., 11.], +# CHECK: [15., 13., 14., 15.]{{\]}}) +# CHECK: tensor(crow_indices=tensor([0, 1, 2, 2, 3]), +# CHECK: col_indices=tensor([1, 3, 0]), +# CHECK: values=tensor([2., 4., 6.]), size=(4, 4), nnz=3, +# CHECK: layout=torch.sparse_csr) +# CHECK: mpact +# CHECK: {{\[}}[16. 18. 20. 22.] +# CHECK: [24. 26. 28. 30.] +# CHECK: [32. 34. 36. 38.] +# CHECK: [40. 42. 44. 46.]{{\]}} +# CHECK: {{\[}}[16. 18. 18. 19.] +# CHECK: [20. 21. 22. 25.] +# CHECK: [24. 25. 26. 27.] +# CHECK: [31. 29. 30. 31.]{{\]}} +# CHECK: {{\[}}[ 0. 2. 2. 3.] +# CHECK: [ 4. 5. 6. 9.] +# CHECK: [ 8. 9. 10. 11.] +# CHECK: [15. 13. 14. 15.]{{\]}} +# CHECK: [0 1 2 2 3] +# CHECK: [1 3 0] +# CHECK: [2. 4. 6.] +# + +# Run it with PyTorch. +print("pytorch") +res = net(X, Y) +print(res) +res = net(S, Y) +print(res) +res = net(X, S) +print(res) +res = net(S, S) +print(res) + +# Run it with MPACT. +print("mpact") +res = mpact_jit(net, X, Y) +print(res) +res = mpact_jit(net, S, Y) +print(res) +res = mpact_jit(net, X, S) +print(res) +res = mpact_jit(net, S, S) +print_sparse(res)
diff --git a/test/python/gat.py b/test/python/gat.py new file mode 100644 index 0000000..283c36f --- /dev/null +++ b/test/python/gat.py
@@ -0,0 +1,49 @@ +# RUN: %PYTHON %s | FileCheck %s + +import torch +import numpy as np + +from mpact.mpactbackend import mpact_jit, mpact_jit_compile, mpact_jit_run + +from mpact.models.gat import gat_4_64_8_3 + +net = gat_4_64_8_3() +net.eval() # Switch to inference. + +# Sparse input. +idx = torch.tensor([[0, 0, 1, 2], [0, 2, 3, 1]], dtype=torch.int64) +val = torch.tensor([14.0, 3.0, -8.0, 11.0], dtype=torch.float32) +S = torch.sparse_coo_tensor(idx, val, size=[4, 4]) + +# Construct adjacency matrix. +V = 4 +edges = np.array([[0, 1], [0, 2], [1, 2], [1, 3], [2, 3]], dtype=np.int32) +E = edges.shape[0] +adj_mat = torch.sparse_coo_tensor(edges.T, torch.ones(E), (V, V), dtype=torch.int64) +adj_mat = ( + torch.eye(V) + adj_mat +) # Add self-loops to the adjacency matrix (becomes dense) + + +# +# CHECK: pytorch gat +# CHECK: tensor({{\[}}[-1.0986, -1.0986, -1.0986], +# CHECK: [-1.0986, -1.0986, -1.0986], +# CHECK: [-1.0986, -1.0986, -1.0986], +# CHECK: [-1.0986, -1.0986, -1.0986]{{\]}} +# CHECK: mpact gat +# CHECK: {{\[}}[-1.0986123 -1.0986123 -1.0986123] +# CHECK: [-1.0986123 -1.0986123 -1.0986123] +# CHECK: [-1.0986123 -1.0986123 -1.0986123] +# CHECK: [-1.0986123 -1.0986123 -1.0986123]{{\]}} +# +with torch.no_grad(): + # Run it with PyTorch. + print("pytorch gat") + res = net(S, adj_mat) + print(res) + + # Run it with MPACT. + print("mpact gat") + res = mpact_jit(net, S, adj_mat) + print(res)
diff --git a/test/python/gcn.py b/test/python/gcn.py new file mode 100644 index 0000000..0453f83 --- /dev/null +++ b/test/python/gcn.py
@@ -0,0 +1,95 @@ +# RUN: %PYTHON %s | FileCheck %s + +import torch + +from mpact.mpactbackend import mpact_jit, mpact_jit_compile, mpact_jit_run + +from mpact.models.gcn import graphconv_4_4, gcn_4_16_4 + +net = graphconv_4_4() +net.eval() # Switch to inference. + +# Get random (but reproducible) matrices. +torch.manual_seed(0) +inp = torch.rand(4, 4) +adj_mat = torch.rand(4, 4).to_sparse() + +# +# CHECK: pytorch +# CHECK: tensor({{\[}}[4.4778, 4.4778, 4.4778, 4.4778], +# CHECK: [5.7502, 5.7502, 5.7502, 5.7502], +# CHECK: [4.6980, 4.6980, 4.6980, 4.6980], +# CHECK: [3.6407, 3.6407, 3.6407, 3.6407]{{\]}}) +# CHECK: mpact compile and run +# CHECK: {{\[}}[4.477828 4.477828 4.477828 4.477828 ] +# CHECK: [5.7501717 5.7501717 5.7501717 5.7501717] +# CHECK: [4.697952 4.697952 4.697952 4.697952 ] +# CHECK: [3.640687 3.640687 3.640687 3.640687 ]{{\]}} +# CHECK: mpact compile +# CHECK: mpact run +# CHECK: {{\[}}[4.477828 4.477828 4.477828 4.477828 ] +# CHECK: [5.7501717 5.7501717 5.7501717 5.7501717] +# CHECK: [4.697952 4.697952 4.697952 4.697952 ] +# CHECK: [3.640687 3.640687 3.640687 3.640687 ]{{\]}} +# CHECK: mpact compile opt=3 +# CHECK: mpact run +# CHECK: {{\[}}[4.477828 4.477828 4.477828 4.477828 ] +# CHECK: [5.7501717 5.7501717 5.7501717 5.7501717] +# CHECK: [4.697952 4.697952 4.697952 4.697952 ] +# +with torch.no_grad(): + # Run it with PyTorch. + print("pytorch") + res = net(inp, adj_mat) + print(res) + + # Run it with MPACT (compile and run at once). + print("mpact compile and run") + res = mpact_jit(net, inp, adj_mat) + print(res) + + # Run it with MPACT (with separate compile and run steps). + print("mpact compile") + invoker, fn = mpact_jit_compile(net, inp, adj_mat) + print("mpact run") + res = mpact_jit_run(invoker, fn, inp, adj_mat) + print(res) + + # Run it with MPACT (with separate compile and run steps, given opt_level). + print("mpact compile opt=3") + invoker, fn = mpact_jit_compile(net, inp, adj_mat, opt_level=3) + print("mpact run") + res = mpact_jit_run(invoker, fn, inp, adj_mat) + print(res) + +net = gcn_4_16_4() +net.eval() # Switch to inference. + + +# Sparse input. +idx = torch.tensor([[0, 0, 1, 2], [0, 2, 3, 1]], dtype=torch.int64) +val = torch.tensor([14.0, 3.0, -8.0, 11.0], dtype=torch.float32) +S = torch.sparse_coo_tensor(idx, val, size=[4, 4]) + +# +# CHECK: pytorch gcn +# CHECK: tensor({{\[}}[-1.3863, -1.3863, -1.3863, -1.3863], +# CHECK: [-1.3863, -1.3863, -1.3863, -1.3863], +# CHECK: [-1.3863, -1.3863, -1.3863, -1.3863], +# CHECK: [-1.3863, -1.3863, -1.3863, -1.3863]]) +# CHECK: mpact gcn +# CHECK: {{\[}}[-1.3862944 -1.3862944 -1.3862944 -1.3862944] +# CHECK: [-1.3862944 -1.3862944 -1.3862944 -1.3862944] +# CHECK: [-1.3862944 -1.3862944 -1.3862944 -1.3862944] +# CHECK: [-1.3862944 -1.3862944 -1.3862944 -1.3862944]{{\]}} +# +with torch.no_grad(): + # Run it with PyTorch. + print("pytorch gcn") + res = net(S, adj_mat) + print(res) + + # Run it with MPACT. + print("mpact gcn") + res = mpact_jit(net, S, adj_mat) + print(res)
diff --git a/test/python/lif.py b/test/python/lif.py new file mode 100644 index 0000000..7dc797a --- /dev/null +++ b/test/python/lif.py
@@ -0,0 +1,31 @@ +# RUN: %PYTHON %s | FileCheck %s + +import torch + +from mpact.mpactbackend import mpact_jit, mpact_jit_compile, mpact_jit_run + +from mpact.models.lif import LIFSumOfSq + +net = LIFSumOfSq() + +# Get a random (but reproducible) input, so that a +# general sparse tensor appears after LIF. +torch.manual_seed(0) +x = torch.rand(2, 3, 8, 8) + +# +# CHECK: pytorch +# CHECK: tensor([ 0., 11., 9., 11., 13., 11., 10., 12.]) +# CHECK: mpact +# CHECK: [ 0. 11. 9. 11. 13. 11. 10. 12.] +# + +# Run it with PyTorch. +print("pytorch") +res = net(x) +print(res) + +# Run it with MPACT. +print("mpact") +res = mpact_jit(net, x) +print(res)
diff --git a/test/python/mm.py b/test/python/mm.py new file mode 100644 index 0000000..3c51c37 --- /dev/null +++ b/test/python/mm.py
@@ -0,0 +1,89 @@ +# RUN: %PYTHON %s | FileCheck %s + +import torch +import numpy as np + +from mpact.mpactbackend import mpact_jit, mpact_jit_compile, mpact_jit_run + +from mpact.models.kernels import MMNet + + +def print_sparse(res): + print(res[0]) + print(res[1]) + print(res[2]) + + +net = MMNet() + +# Construct dense and sparse matrices. +X = torch.arange(0, 16, dtype=torch.float32).view(4, 4) +Y = torch.arange(16, 32, dtype=torch.float32).view(4, 4) +A = torch.tensor( + [ + [0.0, 1.0, 0.0, 0.0], + [0.0, 0.0, 0.0, 2.0], + [0.0, 0.0, 0.0, 0.0], + [3.0, 0.0, 0.0, 0.0], + ], + dtype=torch.float32, +) +S = A.to_sparse_csr() + +# +# CHECK: pytorch +# CHECK: tensor({{\[}}[ 152., 158., 164., 170.], +# CHECK: [ 504., 526., 548., 570.], +# CHECK: [ 856., 894., 932., 970.], +# CHECK: [1208., 1262., 1316., 1370.]{{\]}}) +# CHECK: tensor({{\[}}[20., 21., 22., 23.], +# CHECK: [56., 58., 60., 62.], +# CHECK: [ 0., 0., 0., 0.], +# CHECK: [48., 51., 54., 57.]{{\]}}) +# CHECK: tensor({{\[}}[ 9., 0., 0., 2.], +# CHECK: [21., 4., 0., 10.], +# CHECK: [33., 8., 0., 18.], +# CHECK: [45., 12., 0., 26.]{{\]}}) +# CHECK: tensor(crow_indices=tensor([0, 1, 2, 2, 3]), +# CHECK: col_indices=tensor([3, 0, 1]), +# CHECK: values=tensor([2., 6., 3.]), size=(4, 4), nnz=3, +# CHECK: layout=torch.sparse_csr) +# CHECK: mpact +# CHECK: {{\[}}[ 152. 158. 164. 170.] +# CHECK: [ 504. 526. 548. 570.] +# CHECK: [ 856. 894. 932. 970.] +# CHECK: [1208. 1262. 1316. 1370.]{{\]}} +# CHECK: {{\[}}[20. 21. 22. 23.] +# CHECK: [56. 58. 60. 62.] +# CHECK: [ 0. 0. 0. 0.] +# CHECK: [48. 51. 54. 57.]{{\]}} +# CHECK: {{\[}}[ 9. 0. 0. 2.] +# CHECK: [21. 4. 0. 10.] +# CHECK: [33. 8. 0. 18.] +# CHECK: [45. 12. 0. 26.]{{\]}} +# CHECK: [0 1 2 2 3] +# CHECK: [3 0 1] +# CHECK: [2. 6. 3.] +# + +# Run it with PyTorch. +print("pytorch") +res = net(X, Y) +print(res) +res = net(S, Y) +print(res) +res = net(X, S) +print(res) +res = net(S, S) +print(res) + +# Run it with MPACT. +print("mpact") +res = mpact_jit(net, X, Y) +print(res) +res = mpact_jit(net, S, Y) +print(res) +res = mpact_jit(net, X, S) +print(res) +res = mpact_jit(net, S, S) +print_sparse(res)
diff --git a/test/python/mul.py b/test/python/mul.py new file mode 100644 index 0000000..fd8692f --- /dev/null +++ b/test/python/mul.py
@@ -0,0 +1,87 @@ +# RUN: %PYTHON %s | FileCheck %s + +import torch +import numpy as np + +from mpact.mpactbackend import mpact_jit, mpact_jit_compile, mpact_jit_run + +from mpact.models.kernels import MulNet + + +def print_sparse(res): + print(res[0]) + print(res[1]) + print(res[2]) + + +net = MulNet() + +# Construct dense and sparse matrices. +X = torch.arange(0, 16, dtype=torch.float32).view(4, 4) +Y = torch.arange(16, 32, dtype=torch.float32).view(4, 4) +A = torch.tensor( + [ + [0.0, 1.0, 0.0, 0.0], + [0.0, 0.0, 0.0, 2.0], + [0.0, 0.0, 0.0, 0.0], + [3.0, 0.0, 0.0, 0.0], + ], + dtype=torch.float32, +) +S = A.to_sparse_csr() + +# +# CHECK: pytorch +# CHECK: tensor({{\[}}[ 0., 17., 36., 57.], +# CHECK: [ 80., 105., 132., 161.], +# CHECK: [192., 225., 260., 297.], +# CHECK: [336., 377., 420., 465.]{{\]}}) +# CHECK: tensor(crow_indices=tensor([0, 1, 2, 2, 3]), +# CHECK: col_indices=tensor([1, 3, 0]), +# CHECK: values=tensor([17., 46., 84.]), size=(4, 4), nnz=3, +# CHECK: layout=torch.sparse_csr) +# CHECK: tensor(crow_indices=tensor([0, 1, 2, 2, 3]), +# CHECK: col_indices=tensor([1, 3, 0]), +# CHECK: values=tensor([ 1., 14., 36.]), size=(4, 4), nnz=3, +# CHECK: layout=torch.sparse_csr) +# CHECK: tensor(crow_indices=tensor([0, 1, 2, 2, 3]), +# CHECK: col_indices=tensor([1, 3, 0]), +# CHECK: values=tensor([1., 4., 9.]), size=(4, 4), nnz=3, +# CHECK: layout=torch.sparse_csr) +# CHECK: mpact +# CHECK: {{\[}}[ 0. 17. 36. 57.] +# CHECK: [ 80. 105. 132. 161.] +# CHECK: [192. 225. 260. 297.] +# CHECK: [336. 377. 420. 465.]{{\]}} +# CHECK: [0 1 2 2 3] +# CHECK: [1 3 0] +# CHECK: [17. 46. 84.] +# CHECK: [0 1 2 2 3] +# CHECK: [1 3 0] +# CHECK: [ 1. 14. 36.] +# CHECK: [0 1 2 2 3] +# CHECK: [1 3 0] +# CHECK: [1. 4. 9.] +# + +# Run it with PyTorch. +print("pytorch") +res = net(X, Y) +print(res) +res = net(S, Y) +print(res) +res = net(X, S) +print(res) +res = net(S, S) +print(res) + +# Run it with MPACT. +print("mpact") +res = mpact_jit(net, X, Y) +print(res) +res = mpact_jit(net, S, Y) +print_sparse(res) +res = mpact_jit(net, X, S) +print_sparse(res) +res = mpact_jit(net, S, S) +print_sparse(res)
diff --git a/test/python/norm.py b/test/python/norm.py new file mode 100644 index 0000000..f589e46 --- /dev/null +++ b/test/python/norm.py
@@ -0,0 +1,45 @@ +# RUN: %PYTHON %s | FileCheck %s + +import torch +import numpy as np + +from mpact.mpactbackend import mpact_jit, mpact_jit_compile, mpact_jit_run + +from mpact.models.kernels import Normalization + +net = Normalization() + +# Construct adjacency matrix. +V = 8 +edges = np.array([[0, 1], [0, 4], [1, 4], [3, 4], [4, 3]], dtype=np.int32) +E = edges.shape[0] +adj_mat = torch.sparse_coo_tensor(edges.T, torch.ones(E), (V, V), dtype=torch.int64) +adj_mat = ( + torch.eye(V) + adj_mat +) # Add self-loops to the adjacency matrix (becomes dense) + +# +# CHECK: pytorch +# CHECK: tensor({{\[}}[0.1111, 0.1667, 0.0000, 0.0000, 0.1667, 0.0000, 0.0000, 0.0000], +# CHECK: [0.0000, 0.2500, 0.0000, 0.0000, 0.2500, 0.0000, 0.0000, 0.0000], +# CHECK: [0.0000, 0.0000, 1.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000], +# CHECK: [0.0000, 0.0000, 0.0000, 0.2500, 0.2500, 0.0000, 0.0000, 0.0000], +# CHECK: [0.0000, 0.0000, 0.0000, 0.2500, 0.2500, 0.0000, 0.0000, 0.0000], +# CHECK: [0.0000, 0.0000, 0.0000, 0.0000, 0.0000, 1.0000, 0.0000, 0.0000], +# CHECK: [0.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000, 1.0000, 0.0000], +# CHECK: [0.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000, 1.0000]{{\]}}) +# CHECK: mpact +# + +# Run it with PyTorch. +print("pytorch") +res = net(adj_mat) +print(res) + +# Run it with MPACT. +# +# TODO: make this work, crashes in TORCH-MLIR +# +print("mpact") +# res = mpact_jit(net, adj_mat) +# print(res)
diff --git a/test/python/resnet.py b/test/python/resnet.py new file mode 100644 index 0000000..7ac317b --- /dev/null +++ b/test/python/resnet.py
@@ -0,0 +1,37 @@ +# RUN: %PYTHON %s | FileCheck %s + +import torch +import numpy as np + +from mpact.mpactbackend import mpact_jit, mpact_jit_compile, mpact_jit_run + +from mpact.models.resnet import resnet_20 + +resnet = resnet_20() +resnet.eval() # Switch to inference. + +# Get a random input. +# B x RGB x H x W +x = torch.rand(1, 3, 16, 16) + +# +# CHECK: pytorch +# CHECK: mpact +# CHECK: passed +# + +with torch.no_grad(): + # Run it with PyTorch. + print("pytorch") + res1 = resnet(x) + print(res1) + + # Run it with MPACT. + print("mpact") + res2 = mpact_jit(resnet, x) + print(res2) + +# Completely different inputs and weights for each run, +# so we simply verify the two results are the same. +np.testing.assert_allclose(res1.numpy(), res2, rtol=1e-5, atol=0) +print("passed")
diff --git a/test/python/scale.py b/test/python/scale.py new file mode 100644 index 0000000..880fdf2 --- /dev/null +++ b/test/python/scale.py
@@ -0,0 +1,39 @@ +# RUN: %PYTHON %s | FileCheck %s + +import torch +import numpy as np + +from mpact.mpactbackend import mpact_jit, mpact_jit_compile, mpact_jit_run + +from mpact.models.kernels import FeatureScale + +net = FeatureScale() + +# Get random (but reproducible) matrices. +torch.manual_seed(0) +features = torch.rand(7, 7) + +# +# CHECK: pytorch +# CHECK: tensor({{\[}}[0.1702, 0.2634, 0.0303, 0.0453, 0.1054, 0.2174, 0.1680], +# CHECK: [0.3064, 0.1557, 0.2161, 0.1192, 0.1373, 0.0076, 0.0577], +# CHECK: [0.0856, 0.1510, 0.2031, 0.2329, 0.0469, 0.0822, 0.1984], +# CHECK: [0.2207, 0.0957, 0.2108, 0.1011, 0.1333, 0.2297, 0.0087], +# CHECK: [0.0774, 0.1561, 0.1275, 0.3896, 0.0735, 0.1128, 0.0630], +# CHECK: [0.0093, 0.0611, 0.2731, 0.2124, 0.2180, 0.1546, 0.0716], +# CHECK: [0.2026, 0.0115, 0.0481, 0.0839, 0.2826, 0.2749, 0.0964]{{\]}}) +# CHECK: mpact +# + +# Run it with PyTorch. +print("pytorch") +res = net(features) +print(res) + +# Run it with MPACT. +# +# TODO: make this work, crashes in TORCH-MLIR +# +print("mpact") +# res = mpact_jit(net, features) +# print(res)
diff --git a/test/python/sddmm.py b/test/python/sddmm.py new file mode 100644 index 0000000..a0890c7 --- /dev/null +++ b/test/python/sddmm.py
@@ -0,0 +1,65 @@ +# RUN: %PYTHON %s | FileCheck %s + +import torch +import numpy as np + +from mpact.mpactbackend import mpact_jit, mpact_jit_compile, mpact_jit_run + +from mpact.models.kernels import MMNet, SDDMMNet + + +def print_sparse(res): + print(res[0]) + print(res[1]) + print(res[2]) + print(res[3]) + + +mmnet = MMNet() +sddmmnet = SDDMMNet() + +# Construct very sparse matrix. +idx = torch.tensor([[0, 4], [0, 4]], dtype=torch.int64) +val = torch.tensor([2.0, 3.0], dtype=torch.float64) +S = torch.sparse_coo_tensor(idx, val, size=[5, 5]) + +# Trivial dense inputs. +A = torch.arange(0, 25, dtype=torch.float32).view(5, 5) +B = torch.arange(25, 50, dtype=torch.float32).view(5, 5) + +# +# CHECK: pytorch +# CHECK: tensor({{\[}}[ 400., 410., 420., 430., 440.], +# CHECK: [1275., 1310., 1345., 1380., 1415.], +# CHECK: [2150., 2210., 2270., 2330., 2390.], +# CHECK: [3025., 3110., 3195., 3280., 3365.], +# CHECK: [3900., 4010., 4120., 4230., 4340.]{{\]}}) +# CHECK: tensor(indices=tensor({{\[}}[0, 4], +# CHECK: [0, 4]{{\]}}), +# CHECK: values=tensor([ 800., 13020.]), +# CHECK: size=(5, 5), nnz=2, dtype=torch.float64, layout=torch.sparse_coo) +# CHECK: mpact +# CHECK: {{\[}}[ 400. 410. 420. 430. 440.] +# CHECK: [1275. 1310. 1345. 1380. 1415.] +# CHECK: [2150. 2210. 2270. 2330. 2390.] +# CHECK: [3025. 3110. 3195. 3280. 3365.] +# CHECK: [3900. 4010. 4120. 4230. 4340.]{{\]}} +# CHECK: [0 2] +# CHECK: [0 4] +# CHECK: [0 4] +# CHECK: [ 800. 13020.] +# + +# Run it with PyTorch. +print("pytorch") +dense = mmnet(A, B) +print(dense) +res = sddmmnet(S, A, B) +print(res) + +# Run it with MPACT. +print("mpact") +dense = mpact_jit(mmnet, A, B) +print(dense) +res = mpact_jit(sddmmnet, S, A, B) +print_sparse(res)
diff --git a/test/python/spmv.py b/test/python/spmv.py new file mode 100644 index 0000000..4f52ea0 --- /dev/null +++ b/test/python/spmv.py
@@ -0,0 +1,31 @@ +# RUN: %PYTHON %s | FileCheck %s + +import torch + +from mpact.mpactbackend import mpact_jit, mpact_jit_compile, mpact_jit_run + +from mpact.models.kernels import MVNet + +net = MVNet() + +# Get a fixed vector and matrix (which we make 2x2 block "sparse"). +dense_vector = torch.arange(1, 11, dtype=torch.float32) +dense_input = torch.arange(1, 101, dtype=torch.float32).view(10, 10) +sparse_matrix = dense_input.to_sparse_bsr(blocksize=(2, 2)) + +# +# CHECK: pytorch +# CHECK: tensor([ 385., 935., 1485., 2035., 2585., 3135., 3685., 4235., 4785., 5335.]) +# CHECK: mpact +# CHECK: [ 385. 935. 1485. 2035. 2585. 3135. 3685. 4235. 4785. 5335.] +# + +# Run it with PyTorch. +print("pytorch") +res = net(sparse_matrix, dense_vector) +print(res) + +# Run it with MPACT. +print("mpact") +res = mpact_jit(net, sparse_matrix, dense_vector) +print(res)
diff --git a/test/python/sqsum.py b/test/python/sqsum.py new file mode 100644 index 0000000..2d21204 --- /dev/null +++ b/test/python/sqsum.py
@@ -0,0 +1,35 @@ +# RUN: %PYTHON %s | FileCheck %s + +import torch +import numpy as np + +from mpact.mpactbackend import mpact_jit, mpact_jit_compile, mpact_jit_run + +from mpact.models.kernels import SqSum + +net = SqSum() + +# Construct adjacency matrix. +V = 8 +edges = np.array([[0, 1], [0, 4], [1, 4], [3, 4], [4, 3]], dtype=np.int32) +E = edges.shape[0] +adj_mat = torch.sparse_coo_tensor(edges.T, torch.ones(E), (V, V), dtype=torch.int64) + +# +# CHECK: pytorch +# CHECK: tensor(5) +# CHECK: mpact +# CHECK: 5 + +# Run it with PyTorch. +print("pytorch") +res = net(adj_mat) +print(res) + +# Run it with MPACT. +print("mpact") +# TODO: make this work, expose `sparse-emit-strategy=sparse-iterator` to +# mini-pipeline. +# res = mpact_jit(net, adj_mat, use_sp_it=True) +res = mpact_jit(net, adj_mat) +print(res)
diff --git a/tools/CMakeLists.txt b/tools/CMakeLists.txt new file mode 100644 index 0000000..b7c3ca5 --- /dev/null +++ b/tools/CMakeLists.txt
@@ -0,0 +1 @@ +add_subdirectory(mpact-opt)
diff --git a/tools/mpact-opt/CMakeLists.txt b/tools/mpact-opt/CMakeLists.txt new file mode 100644 index 0000000..a26b6ca --- /dev/null +++ b/tools/mpact-opt/CMakeLists.txt
@@ -0,0 +1,17 @@ +add_llvm_executable(mpact-opt mpact_opt.cpp) + +set(dependency_libraries) +if(TORCH_MLIR_ENABLE_STABLEHLO) + list(APPEND dependency_libraries StablehloRegister) +endif() + +target_link_libraries(mpact-opt PRIVATE + MLIROptLib + MLIRTransforms + TorchMLIRInitAll + TorchMLIRTorchDialect + TorchMLIRTorchPasses + + MPACTTransformPasses + ${dependency_libraries} +)
diff --git a/tools/mpact-opt/mpact_opt.cpp b/tools/mpact-opt/mpact_opt.cpp new file mode 100644 index 0000000..87d88c3 --- /dev/null +++ b/tools/mpact-opt/mpact_opt.cpp
@@ -0,0 +1,49 @@ +//===- mpact-opt.cpp - MLIR Optimizer Driver -------------------------===// +// +// Part of the MPACT Project, under the Apache License v2.0 with LLVM +// Exceptions. +// See https://llvm.org/LICENSE.txt for license information. +// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception +// Also available under a BSD-style license. See LICENSE. +// +//===----------------------------------------------------------------------===// + +#include "mlir/Dialect/SCF/IR/SCF.h" +#include "mlir/Dialect/Tensor/IR/Tensor.h" +#include "mlir/Tools/mlir-opt/MlirOptMain.h" +#include "mlir/Transforms/Passes.h" +#include "mpact/Transforms/Passes.h" +#include "torch-mlir/InitAll.h" + +#ifdef TORCH_MLIR_ENABLE_STABLEHLO +#include "stablehlo/dialect/Register.h" +#endif + +using namespace mlir; + +int main(int argc, char **argv) { + mlir::mpact::registerTransformPasses(); + + mlir::torch::registerAllPasses(); + + // Core Transforms + registerCanonicalizerPass(); + registerCSEPass(); + registerInlinerPass(); + registerLocationSnapshotPass(); + registerLoopInvariantCodeMotionPass(); + registerPrintOpStatsPass(); + registerViewOpGraphPass(); + registerStripDebugInfoPass(); + registerSymbolDCEPass(); + + DialectRegistry registry; + mlir::torch::registerAllDialects(registry); + mlir::torch::registerOptionalInputDialects(registry); + +#ifdef TORCH_MLIR_ENABLE_STABLEHLO + mlir::stablehlo::registerAllDialects(registry); +#endif + return mlir::asMainReturnCode(mlir::MlirOptMain( + argc, argv, "MLIR modular optimizer driver\n", registry)); +}