From a23d77100bf3029eec6b5a40abae3cd3499b9526 Mon Sep 17 00:00:00 2001
From: Stella Laurenzo <stellaraccident@gmail.com>
Date: Thu, 21 Oct 2021 21:18:08 -0700
Subject: [PATCH] Set some wheel building optimization options.

* Also adds a requirements.txt and updates docs to reference it versus stringy pip install.
* Adds doc with instructions on creating a wheel.

Fixes #370
---
 .gitignore            |  1 +
 README.md             | 15 +++++++++++++--
 python/CMakeLists.txt |  4 ++++
 requirements.txt      | 10 ++++++++++
 setup.py              |  4 ++++
 5 files changed, 32 insertions(+), 2 deletions(-)
 create mode 100644 requirements.txt

diff --git a/.gitignore b/.gitignore
index 327b932b5..8d533fa5b 100644
--- a/.gitignore
+++ b/.gitignore
@@ -4,6 +4,7 @@
 *.code-workspace
 .ipynb_checkpoints
 *.venv/
+mlir_venv/
 
 /build/
 __pycache__
diff --git a/README.md b/README.md
index 4b612adb4..59b44a7ba 100644
--- a/README.md
+++ b/README.md
@@ -46,8 +46,8 @@ python -m venv mlir_venv
 source mlir_venv/bin/activate
 # Some older pip installs may not be able to handle the recent PyTorch deps
 python -m pip install --upgrade pip
-# Install latest PyTorch nightlies
-python -m pip install --pre torch torchvision pybind11 -f https://download.pytorch.org/whl/nightly/cpu/torch_nightly.html
+# Install latest PyTorch nightlies and build requirements.
+python -m pip install -r requirements.txt
 ```
 
 ## Build
@@ -178,3 +178,14 @@ manually `source`'d in a shell.
 - `#torch-mlir` channel on the LLVM [Discord](https://discord.gg/xS7Z362)
 - Issues [here](https://github.com/llvm/torch-mlir/issues)
 - [`torch-mlir` section](https://llvm.discourse.group/c/projects-that-want-to-become-official-llvm-projects/torch-mlir/41) of LLVM Discourse
+
+## Build Python Packages
+
+We have preliminary support for building Python packages. This can be done
+with the following commands:
+
+```
+python -m pip install --upgrade pip
+python -m pip install -r requirements.txt
+CMAKE_GENERATOR=Ninja python setup.py bdist_wheel
+```
diff --git a/python/CMakeLists.txt b/python/CMakeLists.txt
index b48fdd28f..388aec994 100644
--- a/python/CMakeLists.txt
+++ b/python/CMakeLists.txt
@@ -1,5 +1,9 @@
 include(AddMLIRPython)
 
+# Disables generation of "version soname" (i.e. libFoo.so.<version>), which
+# causes pure duplication as part of Python wheels.
+set(CMAKE_PLATFORM_NO_VERSIONED_SONAME ON)
+
 # The directory at which the Python import tree begins.
 # See documentation for `declare_mlir_python_sources`'s ROOT_DIR
 # argument.
diff --git a/requirements.txt b/requirements.txt
new file mode 100644
index 000000000..45e653a35
--- /dev/null
+++ b/requirements.txt
@@ -0,0 +1,10 @@
+-f https://download.pytorch.org/whl/nightly/cpu/torch_nightly.html
+--pre
+
+numpy
+torch
+torchvision
+
+# Build requirements.
+pybind11
+wheel
diff --git a/setup.py b/setup.py
index c005b33fa..57b402b25 100644
--- a/setup.py
+++ b/setup.py
@@ -64,6 +64,10 @@ class CMakeBuild(build_py):
                 f"-DLLVM_ENABLE_PROJECTS=mlir",
                 f"-DLLVM_EXTERNAL_PROJECTS=torch-mlir",
                 f"-DLLVM_EXTERNAL_TORCH_MLIR_SOURCE_DIR={src_dir}",
+                # Optimization options for building wheels.
+                f"-DCMAKE_VISIBILITY_INLINES_HIDDEN=ON",
+                f"-DCMAKE_C_VISIBILITY_PRESET=hidden",
+                f"-DCMAKE_CXX_VISIBILITY_PRESET=hidden",
             ]
             os.makedirs(cmake_build_dir, exist_ok=True)
             cmake_cache_file = os.path.join(cmake_build_dir, "CMakeCache.txt")