From 93863d3924f4d64ecb64b2792ed21e5c844bd1aa Mon Sep 17 00:00:00 2001
From: Jack Wolfard <31635014+JackWolfard@users.noreply.github.com>
Date: Tue, 1 Aug 2023 00:41:20 -0500
Subject: [PATCH] Fix MLIR Python bindings build in byo_llvm (#14411)

Depends on https://reviews.llvm.org/D155362. This commit points to my
LLVM fork for demonstration. Will need to pull in this commit once
merged into LLVM.
---
 build_tools/llvm/byo_llvm.sh       |  6 +-----
 build_tools/llvm/mlir_config.cmake | 10 ++++++----
 2 files changed, 7 insertions(+), 9 deletions(-)

diff --git a/build_tools/llvm/byo_llvm.sh b/build_tools/llvm/byo_llvm.sh
index 4bba7f242d7c..0f3d0fda3d4d 100755
--- a/build_tools/llvm/byo_llvm.sh
+++ b/build_tools/llvm/byo_llvm.sh
@@ -126,12 +126,8 @@ do_build_mlir() {
   echo "CMake Options: ${cmake_options}"
   cmake -GNinja -S "${LLVM_SOURCE_DIR}/mlir" -B "${mlir_build_dir}" \
     $cmake_options
-  # TODO: We should be able to just do install-mlirdevelopment-distribution to
-  # tightly control distribution, but this presently leaves out the Python
-  # sources needed to build downstream Python bindings. Once this is fixed,
-  # we should use the more fine grained install target.
   cmake --build "${mlir_build_dir}" \
-    --target install
+    --target install-mlirdevelopment-distribution
 }
 
 print_iree_config() {
diff --git a/build_tools/llvm/mlir_config.cmake b/build_tools/llvm/mlir_config.cmake
index 2214d9d33ed2..186f5e861c49 100644
--- a/build_tools/llvm/mlir_config.cmake
+++ b/build_tools/llvm/mlir_config.cmake
@@ -13,11 +13,7 @@
 #     -DCMAKE_C_COMPILER=clang \
 #     -DCMAKE_CXX_COMPILER=clang++ \
 #     -DCMAKE_LINKER=lld
-# TODO: For some reason Python sources aren't selectable as part of a
-# distribution and targets don't show up properly.
 # ninja -C $MLIR_BUILD_DIR install-mlirdevelopment-distribution
-# The full install target seems to get everything...
-# ninja -C $MLIR_BUILD_DIR install
 
 if(NOT EXISTS ${LLVM_DIR})
   message(FATAL_ERROR "LLVM_DIR (${LLVM_DIR}) does not exist")
@@ -49,10 +45,16 @@ set(LLVM_MLIR_TOOLS
       mlir-translate
     CACHE STRING "")
 
+set(LLVM_MLIR_Python_COMPONENTS
+      MLIRPythonModules
+      mlir-python-sources
+    CACHE STRING "")
+
 set(LLVM_MlirDevelopment_DISTRIBUTION_COMPONENTS
       mlir-cmake-exports
       mlir-headers
       mlir-libraries
       MLIRPythonModules
       ${LLVM_MLIR_TOOLS}
+      ${LLVM_MLIR_Python_COMPONENTS}
     CACHE STRING "")