Update docs and .env file to reflect PYTHONPATH variability

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" \

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

build_tools/write_env_file.sh

@@ -13,7 +13,21 @@ portable_realpath() {
 
 td="$(portable_realpath "$(dirname "$0")"/..)"
 build_dir="$(portable_realpath "${TORCH_MLIR_BUILD_DIR:-$td/build}")"
-python_packages_dir="$build_dir/python_packages"
+
+in_tree_pkg_dir="${build_dir}/tools/torch-mlir/python_packages"
+out_of_tree_pkg_dir="${build_dir}/python_packages"
+
+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."
+  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
 
 write_env_file() {
   echo "Updating $build_dir/.env file"

docs/development.md

@@ -194,30 +194,48 @@ TIP: add multiple target options to stack build phases
 
 ### Setup Python Environment to export the built Python packages
 
+When CMake is configured with `-DMLIR_ENABLE_BINDINGS_PYTHON=ON`, the python packages will typically be located in either:
+
+1. `./build/tools/torch-mlir/python_packages/`  if doing an in-tree build.
+2. `./build/python_packages/` if doing an out-of-tree build.
+
+For the following sections, let `python_pkg_dir` represent whichever of the above is relevant for your build setup. On Linux and macOS, you can run `./build_tools/write_env_file.sh` to generate a file `./.env` in your root source directory with the correct `PYTHONPATH`.
+
 #### Linux and macOS
 
+To get the base `PYTHONPATH`, run:
+
 ```shell
-export PYTHONPATH=`pwd`/build/python_packages/torch_mlir:`pwd`/test/python/fx_importer
+./build_tools/write_env_file.sh
+source ./.env && export PYTHONPATH
+```
+
+To run fx_importer tests, you can append the following:
+
+```
+export PYTHONPATH="${PYTHONPATH}":/test/python/fx_importer"
 ```
 
 #### Windows PowerShell
 
+To get the base `PYTHONPATH`, identify your `python_pkg_dir` from above and set this variable in your environment:
+
+```shell
+$env:PYTHONPATH = "<python_pkg_dir>/torch-mlir"
+```
+
+To run fx_importer tests, you can append the following:
+
 ```shell
-$env:PYTHONPATH = "$PWD/build/tools/torch-mlir/python_packages/torch_mlir;$PWD/test/python/fx_importer"
+$env:PYTHONPATH += ";$PWD/test/python/fx_importer"
 ```
 
 ### Testing MLIR output in various dialects
 
 To test the MLIR output to torch dialect, you can use `test/python/fx_importer/basic_test.py`.
 
 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
-```
+(if running on Windows, modify the environment variable as shown above).
 
 This will display the basic example in TORCH dialect.
 
@@ -226,10 +244,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:
+The base `PYTHONPATH` should be set as above, then the example can be run with the following command (similar on Windows):
 
 ```shell
-export PYTHONPATH=`pwd`/build/tools/torch-mlir/python_packages/torch_mlir:`pwd`/projects/pt1/examples
+export PYTHONPATH="${PYTHONPATH}:$PWD/projects/pt1/examples"
 python projects/pt1/examples/torchscript_resnet18_all_output_types.py
 ```