diff --git a/CMakeLists.txt b/CMakeLists.txt index 4528ce1a9f..c848d92152 100644 --- a/CMakeLists.txt +++ b/CMakeLists.txt @@ -180,16 +180,6 @@ if(BUDDY_MLIR_USE_MIMALLOC) find_package(mimalloc REQUIRED) endif() -#------------------------------------------------------------------------------- -# Initialize Python packages -#------------------------------------------------------------------------------- -if(BUDDY_MLIR_ENABLE_PYTHON_PACKAGES) - file(MAKE_DIRECTORY ${BUDDY_MLIR_PYTHON_PACKAGES_DIR}/buddy) - file(MAKE_DIRECTORY ${BUDDY_MLIR_PYTHON_PACKAGES_DIR}/buddy/compiler) - file(WRITE ${BUDDY_MLIR_PYTHON_PACKAGES_DIR}/buddy/__init__.py "") - file(WRITE ${BUDDY_MLIR_PYTHON_PACKAGES_DIR}/buddy/compiler/__init__.py "") -endif() - #------------------------------------------------------------------------------- # Directory setup #------------------------------------------------------------------------------- @@ -208,4 +198,4 @@ add_subdirectory(tests) add_custom_target(check-buddy DEPENDS check-examples check-tests -) \ No newline at end of file +) diff --git a/README.md b/README.md index cb9a5f1c24..a86a129cbf 100644 --- a/README.md +++ b/README.md @@ -128,21 +128,32 @@ $ ninja check-buddy This repository have nix flake support. You can follow the [nix installation instruction](https://nixos.org/manual/nix/stable/installation/installation.html) and enable the [flake features](https://nixos.wiki/wiki/Flakes#Other_Distros.2C_without_Home-Manager) to have nix setup. -- If you want to contribute to this project: +- If you want to use the buddy-mlir bintools ```bash -nix develop . +nix shell .#buddy-mlir +buddy-opt --version ``` -This will setup a bash shell with `clang`, `clangd`, `cmake`, `ninja`, and other necessary dependencies to build buddy-mlir from source. +- If you want to play with E2E -- If you want to use the buddy-mlir bintools +```bash +nix develop +python tests/Python/test_addmm.py +``` + +- If you want to contribute to this project but don't know how to setup the project: ```bash -nix build .#buddy-mlir -./result/bin/buddy-opt --version +nix develop .#buddy-mlir + +cmake -G Ninja -Bbuild ...etc ``` +This will setup a bash shell with `clang`, `cmake`, `ninja`, and other necessary dependencies to build buddy-mlir from source. +By default it will add Python dependencies and add a Python 3.10.* to your environment to build the Python3 bindings. +If you doesn't want that, you can override the `enablePython3Bindings` argument. + ## Dialects ### Bud Dialect diff --git a/flake.nix b/flake.nix index 8f94e2aec0..de1182c8a7 100644 --- a/flake.nix +++ b/flake.nix @@ -8,43 +8,24 @@ outputs = { self, nixpkgs, flake-utils }@inputs: let - overlay = import ./nix/overlay.nix; - pkgsForSys = system: import nixpkgs { overlays = [ overlay ]; inherit system; }; + myOverlay = import ./nix/overlay.nix; in + # Iterate over all nix supported system flake-utils.lib.eachDefaultSystem (system: let - pkgs = pkgsForSys system; - mkLLVMShell = pkgs.mkShell.override { stdenv = pkgs.llvmPkgs.stdenv; }; + pkgs = import nixpkgs { overlays = [ myOverlay ]; inherit system; }; in { - # Help other use packages in this flake legacyPackages = pkgs; - - devShells.default = mkLLVMShell { - buildInputs = with pkgs; [ - # buddy-mlir build tools - cmake - ninja - python3 - llvmPkgs.bintools # For ld.lld - - # buddy-mlir libraries - libjpeg - libpng - zlib-ng - ]; - - postHook = '' - export PATH="${pkgs.clang-tools}/bin:$PATH" - ''; - }; - - formatter = pkgs.nixpkgs-fmt; + devShells.default = pkgs.mkShell + { + buildInputs = [ pkgs.buddy-mlir ]; + }; }) // { # Other system-independent attr inherit inputs; - overlays.default = overlay; + overlays.default = myOverlay; }; } diff --git a/frontend/Python/CMakeLists.txt b/frontend/Python/CMakeLists.txt index 64cfbd471e..f0aa8c9aee 100644 --- a/frontend/Python/CMakeLists.txt +++ b/frontend/Python/CMakeLists.txt @@ -1,11 +1,19 @@ -# Recursively retrieve all python files from the current directory. -file(GLOB_RECURSE ALL_PY_FILES RELATIVE "${CMAKE_CURRENT_SOURCE_DIR}" "*.py") +include(AddMLIRPython) -foreach(FILE ${ALL_PY_FILES}) - # Get the directory of the current file. - get_filename_component(DIR "${FILE}" DIRECTORY) - # Set the destination directory for the target file. - set(DEST "${BUDDY_MLIR_PYTHON_PACKAGES_DIR}/buddy/compiler/${DIR}") - # Copy the file into the destination directory. - file(COPY ${FILE} DESTINATION ${DEST}) -endforeach() +declare_mlir_python_sources(BuddyPythonSources) +declare_mlir_python_sources(BuddyPythonSources.Compiler + ADD_TO_PARENT BuddyPythonSources + ROOT_DIR "${CMAKE_CURRENT_SOURCE_DIR}/buddy" + SOURCES + compiler/__init__.py + compiler/frontend.py + compiler/ops/linalg.py + compiler/ops/math.py + compiler/ops/tosa.py +) + +add_mlir_python_modules(BuddyPythonModules + ROOT_PREFIX "${CMAKE_CURRENT_BINARY_DIR}/python_packages/buddy" + INSTALL_PREFIX "python_packages/buddy" + DECLARED_SOURCES BuddyPythonSources +) diff --git a/frontend/Python/buddy/compiler/__init__.py b/frontend/Python/buddy/compiler/__init__.py new file mode 100644 index 0000000000..0edeefa5eb --- /dev/null +++ b/frontend/Python/buddy/compiler/__init__.py @@ -0,0 +1,13 @@ +# ===- compiler ------------------------------------------------------------- +# +# 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/frontend/Python/frontend.py b/frontend/Python/buddy/compiler/frontend.py similarity index 100% rename from frontend/Python/frontend.py rename to frontend/Python/buddy/compiler/frontend.py diff --git a/frontend/Python/ops/linalg.py b/frontend/Python/buddy/compiler/ops/linalg.py similarity index 100% rename from frontend/Python/ops/linalg.py rename to frontend/Python/buddy/compiler/ops/linalg.py diff --git a/frontend/Python/ops/math.py b/frontend/Python/buddy/compiler/ops/math.py similarity index 100% rename from frontend/Python/ops/math.py rename to frontend/Python/buddy/compiler/ops/math.py diff --git a/frontend/Python/ops/tosa.py b/frontend/Python/buddy/compiler/ops/tosa.py similarity index 100% rename from frontend/Python/ops/tosa.py rename to frontend/Python/buddy/compiler/ops/tosa.py diff --git a/nix/buddy-mlir.nix b/nix/buddy-mlir.nix index b59d82275f..68b07924a5 100644 --- a/nix/buddy-mlir.nix +++ b/nix/buddy-mlir.nix @@ -1,4 +1,36 @@ -{ cmake, ninja, python3, llvmPackages_16, fetchFromGitHub, libjpeg, libpng, zlib-ng }: +{ lib +, fetchFromGitHub +, cmake +, ninja +, llvmPackages_16 + + # Required by graphic frontend +, libjpeg +, libpng +, zlib-ng + + # Compile Options, can be disabled using override. Eg: + # + # ```nix + # # overlay.nix + # final: prev: + # { + # yourBuddyMlir = prev.buddy-mlir.override { + # enablePython3Bindings = false; + # }; + # } + # +, enablePython3Bindings ? true + # + # We need to pin python to 3.10.* for two reasons: + # + # 1. PyTorch in current Nixpkgs doens't support 3.11 yet. + # 2. The current mainline of PyTorch do have 3.11 support, but they bypass a large amount of tests for Python 3.11, + # which is not reliable to be used. + # +, python310 +, python310Packages +}: let # Using git submodule to obtain the llvm source is really slow. # So here I use tarball to save time from git index. @@ -33,7 +65,32 @@ llvmPackages_16.stdenv.mkDerivation { patches = [ ./tblgen.patch ]; postPatch = "popd"; - nativeBuildInputs = [ cmake ninja python3 llvmPackages_16.bintools libjpeg libpng zlib-ng ]; + nativeBuildInputs = [ + cmake + ninja + python310 + llvmPackages_16.bintools # use lld instead of ld to link + ]; + + buildInputs = [ + libjpeg + libpng + zlib-ng + ] ++ lib.optionals enablePython3Bindings [ + # Required by Python binding + python310Packages.pybind11 + ]; + + # Required by E2E + propagatedBuildInputs = lib.optionals enablePython3Bindings (with python310Packages; [ + # The CPython extension will need same major version of Python to load the share library, + # so it is better to add python 310 here for downstream user, to avoid potential problems. + python310 + numpy + transformers + pyyaml + torch + ]); cmakeDir = "../llvm"; cmakeFlags = [ @@ -45,7 +102,25 @@ llvmPackages_16.stdenv.mkDerivation { "-DLLVM_EXTERNAL_PROJECTS=buddy-mlir" "-DLLVM_EXTERNAL_BUDDY_MLIR_SOURCE_DIR=../../buddy-mlir" + ] ++ lib.optionals enablePython3Bindings [ + "-DMLIR_ENABLE_BINDINGS_PYTHON=ON" + "-DBUDDY_MLIR_ENABLE_PYTHON_PACKAGES=ON" ]; checkTarget = "check-mlir check-buddy"; + + # Python exeutable in Nix have wrapper that can automatically append libraries into PYTHONPATH + # But our Python bindings are not packaged in pip spec. Here is the manual fix for automatically + # adding Python module `buddy` and `mlir` into PYTHONPATH. + postFixup = lib.optionalString enablePython3Bindings '' + local sitePkgPath="$out/python_packages/${python310.sitePackages}" + mkdir -p $sitePkgPath + + mv $out/python_packages/buddy $sitePkgPath + mv $out/python_packages/mlir_core/mlir $sitePkgPath + rm -r $out/python_packages/mlir_core + + mkdir -p "$out/nix-support" + echo "$out/python_packages" >> "$out/nix-support/propagated-build-inputs" + ''; } diff --git a/nix/overlay.nix b/nix/overlay.nix index 19c97fc33c..cccc505d5f 100644 --- a/nix/overlay.nix +++ b/nix/overlay.nix @@ -1,6 +1,4 @@ final: prev: { - # Add an alias here can help future migration - llvmPkgs = final.llvmPackages_16; buddy-mlir = final.callPackage ./buddy-mlir.nix { }; } diff --git a/nix/tblgen.patch b/nix/tblgen.patch new file mode 100644 index 0000000000..4d3a9c0b9c --- /dev/null +++ b/nix/tblgen.patch @@ -0,0 +1,14 @@ +diff --git a/llvm/utils/TableGen/GlobalISelEmitter.cpp b/llvm/utils/TableGen/GlobalISelEmitter.cpp +index 8d9ded1b2ac5..e8765e8bc3c9 100644 +--- a/llvm/utils/TableGen/GlobalISelEmitter.cpp ++++ b/llvm/utils/TableGen/GlobalISelEmitter.cpp +@@ -770,9 +770,6 @@ Expected GlobalISelEmitter::createAndImportSelDAGMatcher( + + if (Predicate.hasGISelPredicateCode()) { + if (Predicate.usesOperands()) { +- assert(WaitingForNamedOperands == 0 && +- "previous predicate didn't find all operands or " +- "nested predicate that uses operands"); + TreePattern *TP = Predicate.getOrigPatFragRecord(); + WaitingForNamedOperands = TP->getNumArgs(); + for (unsigned i = 0; i < WaitingForNamedOperands; ++i)