llvm · zjgarvey · Oct 3, 2025 · Oct 3, 2025 · Oct 3, 2025 · Oct 3, 2025
diff --git a/CMakeLists.txt b/CMakeLists.txt
@@ -206,6 +206,8 @@ if(MLIR_ENABLE_BINDINGS_PYTHON)
   if(NOT TORCH_MLIR_PYTHON_PACKAGES_DIR)
     set(TORCH_MLIR_PYTHON_PACKAGES_DIR "${CMAKE_CURRENT_BINARY_DIR}/python_packages")
   endif()
+  file(WRITE "${TORCH_MLIR_SOURCE_DIR}/.env" "PYTHONPATH=${TORCH_MLIR_PYTHON_PACKAGES_DIR}/torch_mlir")
+  message(STATUS "PYTHONPATH to torch-mlir python_packages dir written to '${TORCH_MLIR_SOURCE_DIR}/.env'.")
 endif()
 
 add_subdirectory(test)

diff --git a/build_tools/ci/build_posix.sh b/build_tools/ci/build_posix.sh
@@ -41,6 +41,7 @@ cmake -S "$repo_root/externals/llvm-project/llvm" -B "$build_dir" \
   -GNinja \
   -DCMAKE_BUILD_TYPE=Release \
   -DPython3_EXECUTABLE="$(which python3)" \
+  -DPython_EXECUTABLE="$(which python3)" \
   -DLLVM_ENABLE_ASSERTIONS=ON \
   -DTORCH_MLIR_ENABLE_WERROR_FLAG=ON \
   -DCMAKE_INSTALL_PREFIX="$install_dir" \

diff --git a/build_tools/python_deploy/build_linux_packages.sh b/build_tools/python_deploy/build_linux_packages.sh
@@ -244,6 +244,7 @@ function build_in_tree() {
       -DTORCH_MLIR_SRC_PYTORCH_BRANCH=${TORCH_MLIR_SRC_PYTORCH_BRANCH} \
       -DTM_PYTORCH_INSTALL_WITHOUT_REBUILD=${TM_PYTORCH_INSTALL_WITHOUT_REBUILD} \
       -DPython3_EXECUTABLE="$(which python3)" \
+      -DPython_EXECUTABLE="$(which python3)" \
       /main_checkout/torch-mlir/externals/llvm-project/llvm
   cmake --build /main_checkout/torch-mlir/build --target tools/torch-mlir/all
   ccache -s
@@ -387,6 +388,7 @@ function build_out_of_tree() {
         -DLLVM_TARGETS_TO_BUILD=host \
         -DMLIR_ENABLE_BINDINGS_PYTHON=ON \
         -DPython3_EXECUTABLE="$(which python3)" \
+        -DPython_EXECUTABLE="$(which python3)" \
         /main_checkout/torch-mlir/externals/llvm-project/llvm
     cmake --build /main_checkout/torch-mlir/llvm-build
   fi
@@ -409,6 +411,7 @@ function build_out_of_tree() {
       -DTORCH_MLIR_SRC_PYTORCH_BRANCH=${TORCH_MLIR_SRC_PYTORCH_BRANCH} \
       -DTM_PYTORCH_INSTALL_WITHOUT_REBUILD=${TM_PYTORCH_INSTALL_WITHOUT_REBUILD} \
       -DPython3_EXECUTABLE="$(which python3)" \
+      -DPython_EXECUTABLE="$(which python3)" \
       /main_checkout/torch-mlir
   cmake --build /main_checkout/torch-mlir/build_oot
   ccache -s

diff --git a/build_tools/update_abstract_interp_lib.sh b/build_tools/update_abstract_interp_lib.sh
@@ -13,38 +13,32 @@
 set -euo pipefail
 
 src_dir="$(realpath "$(dirname "$0")"/..)"
-build_dir="$(realpath "${TORCH_MLIR_BUILD_DIR:-$src_dir/build}")"
 torch_transforms_cpp_dir="${src_dir}/lib/Dialect/Torch/Transforms"
 
-in_tree_pkg_dir="${build_dir}/tools/torch-mlir/python_packages"
-out_of_tree_pkg_dir="${build_dir}/python_packages"
+env_file_path="${src_dir}/.env"
 
-if [[ ! -d "${in_tree_pkg_dir}" && ! -d "${out_of_tree_pkg_dir}" ]]; then
-  echo "Couldn't find in-tree or out-of-tree build, exiting."
+if [[ ! -f "${env_file_path}" ]]; then
+  echo "Couldn't find an env file at ${env_file_path}!"
   exit 1
 fi
 
-# The `-nt` check works even if one of the two directories is missing.
-if [[ "${in_tree_pkg_dir}" -nt "${out_of_tree_pkg_dir}" ]]; then
-  python_packages_dir="${in_tree_pkg_dir}"
-else
-  python_packages_dir="${out_of_tree_pkg_dir}"
-fi
-
+# Get PYTHONPATH from env file.
+source $env_file_path
+# Update PYTHONPATH with externals if specified.
 TORCH_MLIR_EXT_PYTHONPATH="${TORCH_MLIR_EXT_PYTHONPATH:-""}"
-pypath="${python_packages_dir}/torch_mlir"
 if [ ! -z ${TORCH_MLIR_EXT_PYTHONPATH} ]; then
-  pypath="${pypath}:${TORCH_MLIR_EXT_PYTHONPATH}"
+  PYTHONPATH="${PYTHONPATH}:${TORCH_MLIR_EXT_PYTHONPATH}"
 fi
 TORCH_MLIR_EXT_MODULES="${TORCH_MLIR_EXT_MODULES:-""}"
+ext_module="${ext_module:-""}"
 if [ ! -z ${TORCH_MLIR_EXT_MODULES} ]; then
-  ext_module="${TORCH_MLIR_EXT_MODULES} "
+  ext_module="${TORCH_MLIR_EXT_MODULES}"
 fi
 
 # To enable this python package, manually build torch_mlir with:
 #   -DTORCH_MLIR_ENABLE_JIT_IR_IMPORTER=ON
 # TODO: move this package out of JIT_IR_IMPORTER.
-PYTHONPATH="${pypath}" python3 \
+PYTHONPATH=${PYTHONPATH} python3 \
   -m torch_mlir.jit_ir_importer.build_tools.abstract_interp_lib_gen \
-  --pytorch_op_extensions=${ext_module:-""} \
+  --pytorch_op_extensions=${ext_module} \
   --torch_transforms_cpp_dir="${torch_transforms_cpp_dir}"
diff --git a/build_tools/update_torch_ods.sh b/build_tools/update_torch_ods.sh
@@ -12,40 +12,33 @@
 set -euo pipefail
 
 src_dir="$(realpath "$(dirname "$0")"/..)"
-build_dir="$(realpath "${TORCH_MLIR_BUILD_DIR:-$src_dir/build}")"
 torch_ir_include_dir="${src_dir}/include/torch-mlir/Dialect/Torch/IR"
 
-in_tree_pkg_dir="${build_dir}/tools/torch-mlir/python_packages"
-out_of_tree_pkg_dir="${build_dir}/python_packages"
+env_file_path="${src_dir}/.env"
 
-if [[ ! -d "${in_tree_pkg_dir}" && ! -d "${out_of_tree_pkg_dir}" ]]; then
-  echo "Couldn't find in-tree or out-of-tree build, exiting."
+if [[ ! -f "${env_file_path}" ]]; then
+  echo "Couldn't find an env file at ${env_file_path}!"
   exit 1
 fi
 
-# The `-nt` check works even if one of the two directories is missing.
-if [[ "${in_tree_pkg_dir}" -nt "${out_of_tree_pkg_dir}" ]]; then
-  python_packages_dir="${in_tree_pkg_dir}"
-else
-  python_packages_dir="${out_of_tree_pkg_dir}"
-fi
-
+# Get PYTHONPATH from env file.
+source $env_file_path
+# Update PYTHONPATH with externals if specified.
 TORCH_MLIR_EXT_PYTHONPATH="${TORCH_MLIR_EXT_PYTHONPATH:-""}"
-pypath="${python_packages_dir}/torch_mlir"
 if [ ! -z ${TORCH_MLIR_EXT_PYTHONPATH} ]; then
-  pypath="${pypath}:${TORCH_MLIR_EXT_PYTHONPATH}"
+  PYTHONPATH="${PYTHONPATH}:${TORCH_MLIR_EXT_PYTHONPATH}"
 fi
 TORCH_MLIR_EXT_MODULES="${TORCH_MLIR_EXT_MODULES:-""}"
 ext_module="${ext_module:-""}"
 if [ ! -z ${TORCH_MLIR_EXT_MODULES} ]; then
   ext_module="${TORCH_MLIR_EXT_MODULES}"
 fi
 
-set +u
 # To enable this python package, manually build torch_mlir with:
 #   -DTORCH_MLIR_ENABLE_PYTORCH_EXTENSIONS=ON
+#   -DTORCH_MLIR_ENABLE_JIT_IR_IMPORTER=ON
 # TODO: move this package out of JIT_IR_IMPORTER.
-PYTHONPATH="${PYTHONPATH}:${pypath}" python3 \
+PYTHONPATH=${PYTHONPATH} python3 \
   -m torch_mlir.jit_ir_importer.build_tools.torch_ods_gen \
   --torch_ir_include_dir="${torch_ir_include_dir}" \
   --pytorch_op_extensions="${ext_module}" \

diff --git a/build_tools/write_env_file.sh b/build_tools/write_env_file.sh
diff --git a/docs/development.md b/docs/development.md
@@ -197,13 +197,14 @@ TIP: add multiple target options to stack build phases
 #### Linux and macOS
 
 ```shell
-export PYTHONPATH=`pwd`/build/python_packages/torch_mlir:`pwd`/test/python/fx_importer
+source $PWD/.env && export PYTHONPATH="${PYTHONPATH}:${PWD}/test/python/fx_importer"
 ```
 
 #### Windows PowerShell
 
 ```shell
-$env:PYTHONPATH = "$PWD/build/tools/torch-mlir/python_packages/torch_mlir;$PWD/test/python/fx_importer"
+$env:PYTHONPATH = ... # Get from .env file in torch-mlir repo root.
+$env:PYTHONPATH += ";${PWD}/test/python/fx_importer"
 ```
 
 ### Testing MLIR output in various dialects
@@ -214,8 +215,6 @@ Make sure you have activated the virtualenv and set the `PYTHONPATH` above
 (if running on Windows, modify the environment variable as shown above):
 
 ```shell
-source mlir_venv/bin/activate
-export PYTHONPATH=`pwd`/build/tools/torch-mlir/python_packages/torch_mlir:`pwd`/test/python/fx_importer
 python test/python/fx_importer/basic_test.py
 ```
 
@@ -226,10 +225,10 @@ using torchscript with the example `projects/pt1/examples/torchscript_resnet18_a
 This path doesn't give access to the current generation work that is being driven via the fx_importer
 and may lead to errors.
 
-Same as above, but with different python path and example:
+Same as above, but you need to append to the PYTHONPATH:
 
 ```shell
-export PYTHONPATH=`pwd`/build/tools/torch-mlir/python_packages/torch_mlir:`pwd`/projects/pt1/examples
+export PYTHONPATH="${PYTHONPATH}:/projects/pt1/examples"
 python projects/pt1/examples/torchscript_resnet18_all_output_types.py
 ```
 
@@ -259,14 +258,6 @@ jupyter notebook
 [Example IR](https://gist.github.com/silvasean/e74780f8a8a449339aac05c51e8b0caa) for a simple 1 layer MLP to show the compilation steps from TorchScript.
 
 
-### Interactive Use
-
-The `build_tools/write_env_file.sh` script will output a `.env`
-file in the workspace folder with the correct PYTHONPATH set. This allows
-tools like VSCode to work by default for debugging. This file can also be
-manually `source`'d in a shell.
-
-
 ### Bazel Build
 
 > **NOTE** Our Bazel build follows LLVM's Bazel build policy: only the
@@ -478,10 +469,6 @@ Torch-MLIR has two types of tests:
 
 ### Running execution (end-to-end) tests:
 
-> **Note**
-> An `.env` file must be generated via `build_tools/write_env_file.sh` before these commands can be run.
-
-
 The following assumes you are in the `projects/pt1`  directory:
 
 ```shell