[nix] bump nixpkgs and fix buddy-mlir

.gitignore

@@ -14,4 +14,10 @@
 .cache
 
 # Clangd configurations
-.clangd
+.clangd
+
+# nix output
+result
+
+# ccls
+.ccls-cache/

README.md

@@ -155,16 +155,19 @@ This repository have nix flake support. You can follow the [nix installation ins
 - If you want to contribute to this project:
 
 ```bash
-nix develop .
+nix develop .#buddy-mlir
 ```
 
 This will setup a bash shell with `clang`, `ccls`, `cmake`, `ninja`, and other necessary dependencies to build buddy-mlir from source.
 
 - If you want to use the buddy-mlir bintools
 
 ```bash
-nix build .#buddy-mlir
-./result/bin/buddy-opt --version
+nix develop .
+buddy-opt --version
+llc --version
+mlir-tblgen --version
+python -c "import torch; print(torch.__version__)"
 ```
 
 ## Dialects

flake.nix

@@ -1,5 +1,5 @@
 {
-  description = "Generic devshell setup";
+  description = "Buddy MLIR nix flake setup";
 
   inputs = {
     nixpkgs.url = "github:NixOS/nixpkgs/nixos-unstable";
@@ -19,7 +19,31 @@
           # Help other use packages in this flake
           legacyPackages = pkgs;
 
-          devShells.default = pkgs.buddy-mlir.devShell;
+          # A shell with buddy-mlir tools and Python environment
+          # run: nix develop .
+          devShells.default = pkgs.mkShell {
+            nativeBuildInputs = with pkgs; [
+              buddy-mlir
+              buddy-mlir.llvm
+              buddy-mlir.pyenv
+            ];
+          };
+
+          # A shell for writing buddy-mlir code
+          # run: nix develop .#buddy-mlir
+          devShells.buddy-mlir = pkgs.mkShell {
+            nativeBuildInputs = with pkgs;
+              # Add following extra packages to buddy-mlir developement environment
+              pkgs.buddy-mlir.nativeBuildInputs ++ [
+                libjpeg
+                libpng
+                zlib-ng
+                ccls
+              ];
+          };
+
+          # run: nix build .
+          packages.default = pkgs.buddy-mlir;
 
           formatter = pkgs.nixpkgs-fmt;
         }) //

nix/buddy-mlir.nix

@@ -1,68 +1,78 @@
 { lib
 , stdenv
+, fetchFromGitHub
 , buddy-llvm
 , cmake
 , ninja
 , llvmPkgs
-, libjpeg
-, libpng
-, zlib-ng
-, ccls
+, python3
 }:
-let
-  self = stdenv.mkDerivation {
-    pname = "buddy-mlir";
-    version = "unstable-2024-07-18";
+stdenv.mkDerivation (finalAttrs: {
+  name = "buddy-mlir";
 
-    src = with lib.fileset; toSource {
-      root = ./..;
-      fileset = unions [
-        ./../backend
-        ./../cmake
-        ./../examples
-        ./../frontend
-        ./../midend
-        ./../tests
-        ./../tools
-        ./../thirdparty
-        ./../CMakeLists.txt
-        ./../flake.lock
-        ./../flake.nix
-      ];
-    };
-
-    nativeBuildInputs = [
-      cmake
-      ninja
-      llvmPkgs.bintools
+  # Manually specify buddy-mlir source here to avoic rebuild on miscellaneous files
+  src = with lib.fileset; toSource {
+    root = ./..;
+    fileset = unions [
+      ./../backend
+      ./../cmake
+      ./../examples
+      ./../frontend
+      ./../midend
+      ./../tests
+      ./../tools
+      ./../thirdparty
+      ./../CMakeLists.txt
     ];
+  };
 
-    buildInputs = [
-      buddy-llvm
-    ];
+  nativeBuildInputs = [
+    cmake
+    ninja
+    llvmPkgs.bintools
+  ];
 
-    cmakeFlags = [
-      "-DMLIR_DIR=${buddy-llvm.dev}/lib/cmake/mlir"
-      "-DLLVM_DIR=${buddy-llvm.dev}/lib/cmake/llvm"
-      "-DLLVM_MAIN_SRC_DIR=${buddy-llvm.src}/llvm"
-      "-DBUDDY_MLIR_ENABLE_PYTHON_PACKAGES=ON"
-      "-DCMAKE_BUILD_TYPE=Release"
-    ];
+  buildInputs = [
+    buddy-llvm
+  ];
+
+  cmakeFlags = [
+    "-DMLIR_DIR=${buddy-llvm.dev}/lib/cmake/mlir"
+    "-DLLVM_DIR=${buddy-llvm.dev}/lib/cmake/llvm"
+    "-DLLVM_MAIN_SRC_DIR=${buddy-llvm.src}/llvm"
+    "-DBUDDY_MLIR_ENABLE_PYTHON_PACKAGES=ON"
+  ];
 
-    passthru = {
-      llvm = buddy-llvm;
-      devShell = self.overrideAttrs (old: {
-        nativeBuildInputs = old.nativeBuildInputs ++ [
-          libjpeg
-          libpng
-          zlib-ng
-          ccls
-        ];
-      });
-    };
+  # This fixup script concatenate the LLVM and Buddy python module into one directory for easier import
+  postFixup = ''
+    mkdir -p $out/lib/python${python3.pythonVersion}/site-packages
+    cp -vr $out/python_packages/buddy $out/lib/python${python3.pythonVersion}/site-packages/
+    cp -vr ${buddy-llvm}/python_packages/mlir_core/mlir $out/lib/python${python3.pythonVersion}/site-packages/
+  '';
 
-    # No need to do check, and it also takes too much time to finish.
-    doCheck = false;
+  passthru = {
+    llvm = buddy-llvm;
+
+    # Below three fields are nixpkgs black magic that allow site-packages automatically imported with nixpkgs hooks
+    # Update here only when you know what will happen
+    pythonModule = python3;
+    pythonPath = [ ];
+    requiredPythonModules = [ ];
+
+    # nix run .#buddy-mlir.pyenv to start a python with PyTorch/LLVM MLIR/Buddy Frontend support
+    pyenv = python3.withPackages (ps: [
+      finalAttrs.finalPackage
+      ps.torch
+
+      # mobilenet
+      ps.torchvision
+
+      # tinyllama
+      ps.transformers
+      ps.accelerate
+    ]);
   };
-in
-self
+
+  # No need to do check, and it also takes too much time to finish.
+  doCheck = false;
+})