[nix] add Python binding support in nix

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
-)
+)

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

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;
     };
 }

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
+)

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.

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"
+  '';
 }

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 { };
 }

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<InstructionMatcher &> 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)