From f94b4639a429835dcd7a1a37baaf44f81de9745b Mon Sep 17 00:00:00 2001
From: Joseph Melber <jgmelber@gmail.com>
Date: Thu, 20 Feb 2025 14:02:13 -0700
Subject: [PATCH] [Quick setup script] pip install milr_aie (#2055)

---
 utils/env_setup.sh   | 19 +++++++++++--------
 utils/quick_setup.sh | 24 +++++++-----------------
 2 files changed, 18 insertions(+), 25 deletions(-)

diff --git a/utils/env_setup.sh b/utils/env_setup.sh
index 6ea6ffec29..8d1d07f574 100755
--- a/utils/env_setup.sh
+++ b/utils/env_setup.sh
@@ -1,5 +1,5 @@
 #!/bin/bash
-##===- utils/env_setup.sh - Setup mlir-aie env post build to compile mlir-aie designs --*- Script -*-===##
+##===- utils/env_setup.sh - Setup mlir-aie env post build to compile IRON designs --*- Script -*-===##
 # 
 # This file licensed under the Apache License v2.0 with LLVM Exceptions.
 # See https://llvm.org/LICENSE.txt for license information.
@@ -7,8 +7,8 @@
 # 
 ##===----------------------------------------------------------------------===##
 #
-# This script sets up the environment to use mlir-aie tools.
-# The script will also download and set up llvm-aie (peano).
+# This script sets up the environment to compile IRON designs.
+# The script will download and set up mlir-aie and llvm-aie (peano).
 # 
 #
 # source env_setup.sh <mlir-aie install dir> 
@@ -19,16 +19,19 @@
 #
 ##===----------------------------------------------------------------------===##
 
-if [ "$#" -lt 1 ]; then
-    echo "ERROR: Needs an argument for <mlir-aie install dir>."
-    return 1
+if [ "$#" -ge 1 ]; then
+    export MLIR_AIE_INSTALL_DIR=`realpath $1`
 fi
 
-export MLIR_AIE_INSTALL_DIR=`realpath $1`
-if [ "$#" -eq 2 ]; then
+if [ "$#" -ge 2 ]; then
     export PEANO_INSTALL_DIR=`realpath $2`
 fi
 
+if [[ $MLIR_AIE_INSTALL_DIR == "" ]]; then
+  python3 -m pip install mlir_aie -f https://github.com/Xilinx/mlir-aie/releases/expanded_assets/latest-wheels/ 
+  export MLIR_AIE_INSTALL_DIR="$(pip show mlir_aie | grep ^Location: | awk '{print $2}')/mlir_aie"
+fi
+
 if [[ $PEANO_INSTALL_DIR == "" ]]; then
   python3 -m pip install llvm-aie -f https://github.com/Xilinx/llvm-aie/releases/expanded_assets/nightly
   export PEANO_INSTALL_DIR="$(pip show llvm-aie | grep ^Location: | awk '{print $2}')/llvm-aie"
diff --git a/utils/quick_setup.sh b/utils/quick_setup.sh
index d97a0a1aa4..4b7f065923 100755
--- a/utils/quick_setup.sh
+++ b/utils/quick_setup.sh
@@ -43,27 +43,16 @@ else
    echo "This script requires python3.10 or python3.12"
    return 1
 fi
-if ! hash unzip; then
-  echo "unzip is not installed"
-  return 1
-fi
+
 # if an install is already present, remove it to start from a clean slate
 rm -rf ironenv
 rm -rf my_install
 $my_python -m venv ironenv
 source ironenv/bin/activate
 python3 -m pip install --upgrade pip
-mkdir -p my_install
-pushd my_install
-pip download mlir_aie -f https://github.com/Xilinx/mlir-aie/releases/expanded_assets/latest-wheels/
-unzip -q mlir_aie-*_x86_64.whl
-pushd mlir_aie/python
-pip download aie_python_bindings -f https://github.com/Xilinx/mlir-aie/releases/expanded_assets/latest-wheels/
-unzip -q -o aie_python_bindings*.whl
-rm *.whl
-popd
-rm -rf mlir*.whl
-popd
+
+python3 -m pip install mlir_aie -f https://github.com/Xilinx/mlir-aie/releases/expanded_assets/latest-wheels/ 
+MLIR_AIE_DIR="$(pip show mlir_aie | grep ^Location: | awk '{print $2}')/mlir_aie"
 
 python3 -m pip install llvm-aie -f https://github.com/Xilinx/llvm-aie/releases/expanded_assets/nightly
 export PEANO_INSTALL_DIR="$(pip show llvm-aie | grep ^Location: | awk '{print $2}')/llvm-aie"
@@ -76,8 +65,6 @@ pre-commit install
 HOST_MLIR_PYTHON_PACKAGE_PREFIX=aie python3 -m pip install -r python/requirements_extras.txt
 python3 -m pip install -r python/requirements_ml.txt
 
-source utils/env_setup.sh my_install/mlir_aie
-
 # This creates an ipykernel (for use in notebooks) using the ironenv venv
 python3 -m ipykernel install --user --name ironenv
 
@@ -89,6 +76,9 @@ python3 -m ipykernel install --user --name ironenv
 venv_site_packages=`python3 -c 'import sysconfig; print(sysconfig.get_paths()["purelib"])'`
 echo ${MLIR_AIE_INSTALL_DIR}/python > $venv_site_packages/mlir-aie.pth
 
+# Setup environment and add tools to PATHs
+source utils/env_setup.sh
+
 pushd programming_examples
 echo "PATH              : $PATH"
 echo "LD_LIBRARY_PATH   : $LD_LIBRARY_PATH"