From 2736e68ab05c56860f56d704c9cf11eb7ef7f814 Mon Sep 17 00:00:00 2001
From: Avimitin <dev@avimit.in>
Date: Fri, 9 Aug 2024 14:15:57 +0800
Subject: [PATCH] [nix] embedded pyenv to buddy-mlir sub attr

Signed-off-by: Avimitin <dev@avimit.in>
---
 nix/overlay.nix           |  4 --
 nix/pkgs/buddy-mlir.nix   | 91 +++++++++++++++++++++------------------
 tests/pytorch/README.md   |  2 +-
 tests/pytorch/default.nix |  3 +-
 4 files changed, 51 insertions(+), 49 deletions(-)

diff --git a/nix/overlay.nix b/nix/overlay.nix
index bd590d879..20b69d081 100644
--- a/nix/overlay.nix
+++ b/nix/overlay.nix
@@ -22,10 +22,6 @@ rec {
 
   # 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 { };
diff --git a/nix/pkgs/buddy-mlir.nix b/nix/pkgs/buddy-mlir.nix
index 4ec3a4ff9..17b426871 100644
--- a/nix/pkgs/buddy-mlir.nix
+++ b/nix/pkgs/buddy-mlir.nix
@@ -11,47 +11,54 @@ let
   bintools = llvmPackages_17.bintools;
 
   buddy-llvm = callPackage ./buddy-llvm.nix { inherit stdenv python3; };
-in
-stdenv.mkDerivation {
-  pname = "buddy-mlir";
-  version = "unstable-2024-07-18";
-
-  src = fetchFromGitHub {
-    owner = "buddy-compiler";
-    repo = "buddy-mlir";
-    rev = "d7d90a488ac0d6fc1e700e932f842c7b2bcad816";
-    hash = "sha256-MhykCa6Z7Z8PpAlNh+vMuWYEOZZDyWhtMzMnFlNbGIk=";
-  };
+  self = stdenv.mkDerivation {
+    pname = "buddy-mlir";
+    version = "unstable-2024-07-18";
+
+    src = fetchFromGitHub {
+      owner = "buddy-compiler";
+      repo = "buddy-mlir";
+      rev = "d7d90a488ac0d6fc1e700e932f842c7b2bcad816";
+      hash = "sha256-MhykCa6Z7Z8PpAlNh+vMuWYEOZZDyWhtMzMnFlNbGIk=";
+    };
+
+    nativeBuildInputs = [ cmake ninja bintools ];
+    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"
+      "-DCMAKE_BUILD_TYPE=Release"
+    ];
 
-  nativeBuildInputs = [ cmake ninja bintools ];
-  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"
-    "-DCMAKE_BUILD_TYPE=Release"
-  ];
-
-  # 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 = [ ];
+    # 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 = [ ];
+
+      # nix run buddy-mlir.pyenv to start a python with PyTorch/LLVM MLIR/Buddy Frontend support
+      pyenv = python3.withPackages (ps: [
+        self
+        ps.torch
+      ]);
+    };
   };
-}
+in
+self
diff --git a/tests/pytorch/README.md b/tests/pytorch/README.md
index cf638e951..5e53fb203 100644
--- a/tests/pytorch/README.md
+++ b/tests/pytorch/README.md
@@ -121,7 +121,7 @@ nix build '.#t1.blastoise.ip.cases.pytorch.demo.emu-result.with-trace' -L
 * How to debug the PyTorch code
 
 ```bash
-nix run '.#buddy-mlir-pyenv' -- demo.py
+nix run '.#buddy-mlir.pyenv' -- demo.py
 ```
 
 * How to run buddy compiler tools manually
diff --git a/tests/pytorch/default.nix b/tests/pytorch/default.nix
index 2da4a2609..2a4d60cc2 100644
--- a/tests/pytorch/default.nix
+++ b/tests/pytorch/default.nix
@@ -1,7 +1,6 @@
 { lib
 , linkerScript
 , buddy-mlir
-, buddy-mlir-pyenv
 , makeBuilder
 , findAndBuild
 , getTestRequiredFeatures
@@ -29,7 +28,7 @@ let
 
       featuresRequired = getTestRequiredFeatures sourcePath;
 
-      nativeBuildInputs = [ buddy-mlir-pyenv buddy-mlir ];
+      nativeBuildInputs = [ buddy-mlir.pyenv buddy-mlir ];
 
       pythonArgs = buddyBuildConfig.pythonArgs or [ ];
       buddyTranslateArgs = buddyBuildConfig.buddyTranslateArgs or defaultBuddyTranslateArgs;