[mlir][python] Segment MLIR Python test dialect to avoid testonly dependency.

[stellaraccident]

Summary

With https://reviews.llvm.org/rG14c9207063bb00823a5126131e50c93f6e288bd3, the build is broken with -DMLIR_INCLUDE_TESTS=OFF. This patch fixes the build and we may want to do a better fix to the layering in a followup.

[hanchung]

Thanks for the fix!

mlir/python/CMakeLists.txt

 include(AddMLIRPython)
 
 ################################################################################
 # Structural groupings.
 ################################################################################
 
 declare_mlir_python_sources(MLIRPythonSources)
 declare_mlir_python_sources(MLIRPythonSources.Dialects
   ADD_TO_PARENT MLIRPythonSources)
 
-declare_mlir_python_sources(MLIRPythonTestSources)
-declare_mlir_python_sources(MLIRPythonTestSources.Dialects
-  ADD_TO_PARENT MLIRPythonTestSources)
-
 ################################################################################
 # Pure python sources and generated code
 ################################################################################
 
 declare_mlir_python_sources(MLIRPythonSources.Core
   ROOT_DIR "${CMAKE_CURRENT_SOURCE_DIR}/mlir"
   ADD_TO_PARENT MLIRPythonSources
   SOURCES
declare_mlir_dialect_python_bindings(
   ADD_TO_PARENT MLIRPythonSources.Dialects
   ROOT_DIR "${CMAKE_CURRENT_SOURCE_DIR}/mlir"
   TD_FILE dialects/MemRefOps.td
   SOURCES
     dialects/memref.py
     dialects/_memref_ops_ext.py
   DIALECT_NAME memref)
 
-# TODO: this uses a tablegen file from the test directory and should be
-# decoupled from here.
-declare_mlir_python_sources(
-  MLIRPythonSources.Dialects.PythonTest
-  ROOT_DIR "${CMAKE_CURRENT_SOURCE_DIR}/mlir"
-  ADD_TO_PARENT MLIRPythonSources.Dialects
-  SOURCES dialects/python_test.py)
-set(LLVM_TARGET_DEFINITIONS
-  "${MLIR_MAIN_SRC_DIR}/test/python/python_test_ops.td")
-mlir_tablegen(
-  "dialects/_python_test_ops_gen.py"
-  -gen-python-op-bindings
-  -bind-dialect=python_test)
-add_public_tablegen_target(PythonTestDialectPyIncGen)
-declare_mlir_python_sources(
-  MLIRPythonSources.Dialects.PythonTest.ops_gen
-  ROOT_DIR "${CMAKE_CURRENT_BINARY_DIR}"
-  ADD_TO_PARENT MLIRPythonSources.Dialects.PythonTest
-  SOURCES "dialects/_python_test_ops_gen.py")
-
 declare_mlir_dialect_python_bindings(
   ADD_TO_PARENT MLIRPythonSources.Dialects
   ROOT_DIR "${CMAKE_CURRENT_SOURCE_DIR}/mlir"
   TD_FILE dialects/SCFOps.td
   SOURCES
     dialects/scf.py
     dialects/_scf_ops_ext.py
   DIALECT_NAME scf)
declare_mlir_python_extension(MLIRPythonExtension.Transforms
   SOURCES
     ${PYTHON_SOURCE_DIR}/Transforms/Transforms.cpp
   PRIVATE_LINK_LIBS
     LLVMSupport
   EMBED_CAPI_LINK_LIBS
     MLIRCAPITransforms
 )
 
-# TODO: This should not be included in the main Python extension. However,
+# TODO: Figure out how to put this in the test tree.
+# This should not be included in the main Python extension. However,
 # putting it into MLIRPythonTestSources along with the dialect declaration
 # above confuses Python module loader when running under lit.
-declare_mlir_python_extension(MLIRPythonExtension.PythonTest
+set(_ADDL_TEST_SOURCES)
+if(MLIR_INCLUDE_TESTS)
+  set(_ADDL_TEST_SOURCES MLIRPythonTestSources)
+  declare_mlir_python_sources(MLIRPythonTestSources)
+  declare_mlir_python_sources(MLIRPythonTestSources.Dialects
+    ADD_TO_PARENT MLIRPythonTestSources)
+
+  # TODO: this uses a tablegen file from the test directory and should be
+  # decoupled from here.
+  declare_mlir_python_sources(
+    MLIRPythonTestSources.Dialects.PythonTest
+    ROOT_DIR "${CMAKE_CURRENT_SOURCE_DIR}/mlir"
+    ADD_TO_PARENT MLIRPythonTestSources.Dialects
+    SOURCES dialects/python_test.py)
+  set(LLVM_TARGET_DEFINITIONS
+    "${MLIR_MAIN_SRC_DIR}/test/python/python_test_ops.td")
+  mlir_tablegen(
+    "dialects/_python_test_ops_gen.py"
+    -gen-python-op-bindings
+    -bind-dialect=python_test)
+  add_public_tablegen_target(PythonTestDialectPyIncGen)
+  declare_mlir_python_sources(
+    MLIRPythonTestSources.Dialects.PythonTest.ops_gen
+    ROOT_DIR "${CMAKE_CURRENT_BINARY_DIR}"
+    ADD_TO_PARENT MLIRPythonTestSources.Dialects.PythonTest
+    SOURCES "dialects/_python_test_ops_gen.py")
+
+  declare_mlir_python_extension(MLIRPythonTestSources.PythonTestExtension
     MODULE_NAME _mlirPythonTest
-  ADD_TO_PARENT MLIRPythonSources.Dialects
+    ADD_TO_PARENT MLIRPythonTestSources.Dialects
     SOURCES
       ${MLIR_SOURCE_DIR}/test/python/lib/PythonTestModule.cpp
     PRIVATE_LINK_LIBS
       LLVMSupport
     EMBED_CAPI_LINK_LIBS
       MLIRCAPIPythonTestDialect
   )
+endif()
 
 ################################################################################
 # Common CAPI dependency DSO.
 # All python extensions must link through one DSO which exports the CAPI, and
 # this must have a globally unique name amongst all embeddors of the python
 # library since it will effectively have global scope.
 #
 # The presence of this aggregate library is part of the long term plan, but its
 # use needs to be made more flexible.
 #
 # TODO: Upgrade to the aggregate utility in https://reviews.llvm.org/D106419
 # once ready.
 ################################################################################
 
 add_mlir_python_common_capi_library(MLIRPythonCAPI
   INSTALL_COMPONENT MLIRPythonModules
   INSTALL_DESTINATION python_packages/mlir_core/mlir/_mlir_libs
   OUTPUT_DIRECTORY "${MLIR_BINARY_DIR}/python_packages/mlir_core/mlir/_mlir_libs"
   RELATIVE_INSTALL_ROOT "../../../.."
   DECLARED_SOURCES
     MLIRPythonSources
     MLIRPythonExtension.AllPassesRegistration
+    ${_ADDL_TEST_SOURCES}
 )
 
 ################################################################################
 # The fully assembled package of modules.
 # This must come last.
 ################################################################################
 
 add_mlir_python_modules(MLIRPythonModules
   ROOT_PREFIX "${MLIR_BINARY_DIR}/python_packages/mlir_core/mlir"
   INSTALL_PREFIX "python_packages/mlir_core/mlir"
   DECLARED_SOURCES
     MLIRPythonSources
     MLIRPythonExtension.AllPassesRegistration
     MLIRPythonCAPIHeaderSources
+    ${_ADDL_TEST_SOURCES}
   COMMON_CAPI_LINK_LIBS
     MLIRPythonCAPI
   )
-
-add_mlir_python_modules(MLIRPythonTestModules
-  ROOT_PREFIX "${MLIR_BINARY_DIR}/python_packages/mlir_test/mlir"
-  INSTALL_PREFIX "python_packages/mlir_test/mlir"
-  DECLARED_SOURCES
-    MLIRPythonTestSources
-  )

mlir/test/CMakeLists.txt

if(MLIR_ENABLE_VULKAN_RUNNER)
   list(APPEND MLIR_TEST_DEPENDS
     mlir-vulkan-runner
   )
 endif()
 
 if(MLIR_ENABLE_BINDINGS_PYTHON)
   list(APPEND MLIR_TEST_DEPENDS
     MLIRPythonModules
-    MLIRPythonTestModules
   )
 endif()
 
 # This target can be used to just build the dependencies
 # for the check-mlir target without executing the tests.
 # This is useful for bots when splitting the build step
 # from the testing step.
 add_custom_target(check-mlir-build-only