[nix] add Python module support for buddy-mlir

Signed-off-by: Avimitin <[email protected]>

nix/overlay.nix

@@ -19,7 +19,14 @@ rec {
   dramsim3 = final.callPackage ./pkgs/dramsim3.nix { };
   libspike = final.callPackage ./pkgs/libspike.nix { };
   libspike_interfaces = final.callPackage ../difftest/spike_interfaces { };
-  buddy-mlir = final.callPackage ./pkgs/buddy-mlir.nix { };
+
+  # DynamoCompiler doesn't support python 3.12+ yet
+  buddy-mlir = final.callPackage ./pkgs/buddy-mlir.nix { python3 = final.python311; };
+  buddy-mlir-pyenv = final.buddy-mlir.pythonModule.withPackages (ps: [
+    final.buddy-mlir
+    ps.torch
+  ]);
+
   fetchMillDeps = final.callPackage ./pkgs/mill-builder.nix { };
   circt-full = final.callPackage ./pkgs/circt-full.nix { };
   rvv-codegen = final.callPackage ./pkgs/rvv-codegen.nix { };

nix/pkgs/buddy-mlir.nix

@@ -3,13 +3,14 @@
 , llvmPackages_17
 , fetchFromGitHub
 , fetchpatch
+, python3
 , callPackage
 }:
 let
   stdenv = llvmPackages_17.stdenv;
   bintools = llvmPackages_17.bintools;
 
-  buddy-llvm = callPackage ./buddy-llvm.nix { inherit stdenv; };
+  buddy-llvm = callPackage ./buddy-llvm.nix { inherit stdenv python3; };
 in
 stdenv.mkDerivation {
   pname = "buddy-mlir";
@@ -37,4 +38,20 @@ stdenv.mkDerivation {
 
   # No need to do check, and it also takes too much time to finish.
   doCheck = false;
+
+  # Here we 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/
+  '';
+
+  passthru = {
+    llvm = buddy-llvm;
+
+    # Below three fields are black magic that allow site-packages automatically imported with nixpkgs hooks
+    pythonModule = python3;
+    pythonPath = [ ];
+    requiredPythonModules = [ ];
+  };
 }