Enable Python bindings for MHLO (chlo and mhlo dialects).

PiperOrigin-RevId: 387406846

BUILD

@@ -6,10 +6,14 @@ package(
 )
 
 exports_files([
+    "include/mlir-hlo/Dialect/mhlo/IR/clo_ops.td",
     "include/mlir-hlo/Dialect/mhlo/IR/hlo_ops.td",
     "include/mlir-hlo/Dialect/mhlo/IR/lhlo_ops.td",
 ])
 
+# Python extension sources.
+exports_files(["python/MlirHloModule.cpp"])
+
 td_library(
     name = "hlo_ops_td_files",
     srcs = glob(["include/mlir-hlo/Dialect/mhlo/IR/*.td"]),
@@ -1744,6 +1748,20 @@ cc_library(
     ],
 )
 
+cc_library(
+    name = "CAPI",
+    srcs = [
+        "lib/CAPI/Dialects.cpp",
+    ],
+    hdrs = [
+        "include/mlir-hlo-c/Dialects.h",
+    ],
+    deps = [
+        ":hlo",
+        "@llvm-project//mlir:CAPIIR",
+    ],
+)
+
 cc_binary(
     name = "mlir-hlo-opt",
     srcs = [
@@ -1786,11 +1804,29 @@ gentbl_filegroup(
                 "-gen-python-op-bindings",
                 "-bind-dialect=mhlo",
             ],
-            "python/_mhlo_ops_gen.py",
+            "python/mlir/dialects/_mhlo_ops_gen.py",
+        ),
+    ],
+    tblgen = "@llvm-project//mlir:mlir-tblgen",
+    td_file = "python/mlir/dialects/MhloOps.td",
+    deps = [
+        ":MhloOpsPyTdFiles",
+    ],
+)
+
+gentbl_filegroup(
+    name = "ChloOpsPyGen",
+    tbl_outs = [
+        (
+            [
+                "-gen-python-op-bindings",
+                "-bind-dialect=chlo",
+            ],
+            "python/mlir/dialects/_chlo_ops_gen.py",
         ),
     ],
     tblgen = "@llvm-project//mlir:mlir-tblgen",
-    td_file = "python/MhloOps.td",
+    td_file = "python/mlir/dialects/ChloOps.td",
     deps = [
         ":MhloOpsPyTdFiles",
     ],
@@ -1799,7 +1835,9 @@ gentbl_filegroup(
 filegroup(
     name = "MhloOpsPyFiles",
     srcs = [
-        "python/mhlo.py",
+        "python/mlir/dialects/chlo.py",
+        "python/mlir/dialects/mhlo.py",
+        ":ChloOpsPyGen",
         ":MhloOpsPyGen",
     ],
 )

CMakeLists.txt

@@ -29,19 +29,30 @@ if(POLICY CMP0077)
 endif()
 
 #-------------------------------------------------------------------------------
-# Project setup and globals
+# Options and settings
 #-------------------------------------------------------------------------------
-
-project(mlir-hlo LANGUAGES CXX C)
-set(CMAKE_C_STANDARD 11)
-set(CMAKE_CXX_STANDARD 14)
-list(APPEND CMAKE_MODULE_PATH "${CMAKE_CURRENT_SOURCE_DIR}/cmake/modules")
+option(MHLO_BUILD_EMBEDDED "Build MHLO as part of another project" OFF)
+option(MHLO_ENABLE_BINDINGS_PYTHON "Enables MHLO python bindings" OFF)
 
 #-------------------------------------------------------------------------------
-# Options and settings
+# Project setup and globals
 #-------------------------------------------------------------------------------
+set(MHLO_EXTERNAL_PROJECT_BUILD OFF)
 
-option(MHLO_BUILD_EMBEDDED "Build MHLO as part of another project" OFF)
+if(PROJECT_NAME STREQUAL "LLVM")
+  # Building as part of LLVM via the external project mechanism.
+  set(MHLO_EXTERNAL_PROJECT_BUILD ON)
+else()
+  # Building standalone.
+  project(mlir-hlo LANGUAGES CXX C)
+  set(CMAKE_C_STANDARD 11)
+  set(CMAKE_CXX_STANDARD 14)
+  list(APPEND CMAKE_MODULE_PATH "${CMAKE_CURRENT_SOURCE_DIR}/cmake/modules")
+endif()
+
+if(MHLO_ENABLE_BINDINGS_PYTHON AND NOT MHLO_EXTERNAL_PROJECT_BUILD)
+  message(WARNING "MHLO python bindings are only supported in unified, external project builds")
+endif()
 
 #-------------------------------------------------------------------------------
 # MLIR/LLVM Configuration
@@ -50,12 +61,27 @@ option(MHLO_BUILD_EMBEDDED "Build MHLO as part of another project" OFF)
 # Find MLIR to install if we are building standalone. If building as part of
 # another project, let it handle the MLIR dependency. The dependent project
 # might use a bundled version of MLIR instead of installing, for instance.
-if(NOT MHLO_BUILD_EMBEDDED)
+if(MHLO_EXTERNAL_PROJECT_BUILD)
+  message(STATUS "Building MHLO as an external LLVM project")
+  set(MLIR_MAIN_SRC_DIR ${LLVM_MAIN_SRC_DIR}/../mlir ) # --src-root
+  set(MLIR_INCLUDE_DIR ${MLIR_MAIN_SRC_DIR}/include ) # --includedir
+  set(MLIR_GENERATED_INCLUDE_DIR ${CMAKE_BINARY_DIR}/tools/mlir/include)
+  set(MLIR_TABLEGEN_EXE $<TARGET_FILE:mlir-tblgen>)
+  include_directories(SYSTEM ${MLIR_INCLUDE_DIR})
+  include_directories(SYSTEM ${MLIR_GENERATED_INCLUDE_DIR})
+  include_directories(SYSTEM ${MLIR_TABLEGEN_OUTPUT_DIR})
+
+  set(BACKEND_PACKAGE_STRING "${PACKAGE_STRING}")
+  list(APPEND CMAKE_MODULE_PATH "${MLIR_MAIN_SRC_DIR}/cmake/modules")
+elseif(NOT MHLO_BUILD_EMBEDDED)
+  message(STATUS "Building MHLO with an installed MLIR")
   find_package(MLIR REQUIRED CONFIG)
   message(STATUS "Using MLIRConfig.cmake in: ${MLIR_DIR}")
   message(STATUS "Using LLVMConfig.cmake in: ${LLVM_DIR}")
   list(APPEND CMAKE_MODULE_PATH "${MLIR_CMAKE_DIR}")
   list(APPEND CMAKE_MODULE_PATH "${LLVM_CMAKE_DIR}")
+else()
+  message(STATUS "Building MHLO embedded in another project")
 endif()
 
 if(LLVM_ENABLE_ZLIB)
@@ -68,12 +94,24 @@ include(AddMLIR)
 include(HandleLLVMOptions)
 include_directories(${LLVM_INCLUDE_DIRS})
 include_directories(${MLIR_INCLUDE_DIRS})
-include_directories(${PROJECT_SOURCE_DIR}/include)
-include_directories(${PROJECT_BINARY_DIR}/include)
-include_directories(${PROJECT_BINARY_DIR}/)
+include_directories(${CMAKE_CURRENT_SOURCE_DIR}/include)
+include_directories(${CMAKE_CURRENT_BINARY_DIR}/include)
+include_directories(${CMAKE_CURRENT_BINARY_DIR}/)
 link_directories(${LLVM_BUILD_LIBRARY_DIR})
 add_definitions(${LLVM_DEFINITIONS})
 
+#-------------------------------------------------------------------------------
+# Python configuration
+#-------------------------------------------------------------------------------
+
+if(MHLO_ENABLE_BINDINGS_PYTHON)
+  include(MLIRDetectPythonEnv)
+  mlir_detect_pybind11_install()
+  find_package(Python3 ${LLVM_MINIMUM_PYTHON_VERSION}
+  COMPONENTS Interpreter Development NumPy REQUIRED)
+  find_package(pybind11 2.6 CONFIG REQUIRED)
+endif()
+
 #-------------------------------------------------------------------------------
 # Directory setup
 #-------------------------------------------------------------------------------
@@ -85,6 +123,9 @@ add_custom_target(check-mlir-hlo)
 
 add_subdirectory(include/mlir-hlo)
 add_subdirectory(lib)
-add_subdirectory(python)
 add_subdirectory(tools)
 add_subdirectory(tests)
+
+if(MHLO_ENABLE_BINDINGS_PYTHON)
+  add_subdirectory(python)
+endif()

README.md

@@ -235,3 +235,32 @@ Exit:
 ## End-to-End pipeline
 
 TODO
+
+## Alternative build setups
+
+### Building Python API
+
+Building the MHLO Python API requires building as an LLVM external project.
+The below instructions presume that you have this `mlir-hlo` repo and an
+`llvm-project` repo checked out side by side.
+
+Note that the python package produced by this procedure includes the `mlir`
+package and is not suitable for deployment as-is (but it can be included into
+a larger aggregate).
+
+```
+mkdir build && cd build
+cmake -GNinja -B. ${LLVM_SRC_DIR}/llvm \
+    -DCMAKE_BUILD_TYPE=Release \
+    -DLLVM_ENABLE_PROJECTS=mlir \
+    -DLLVM_EXTERNAL_PROJECTS=mlir_hlo \
+    -DLLVM_EXTERNAL_MLIR_HLO_SOURCE_DIR=${MLIR_HLO_SRC_DIR} \
+    -DLLVM_TARGETS_TO_BUILD=host \
+    -DPython3_EXECUTABLE=$(which python) \
+    -DMLIR_ENABLE_BINDINGS_PYTHON=ON \
+    -DMHLO_ENABLE_BINDINGS_PYTHON=ON
+
+ninja MLIRHLOPythonModules
+export PYTHONPATH=$PWD/tools/mlir_hlo/python_packages/mlir_hlo
+python -c "import mlir.dialects.mhlo"
+```

build_tools/smoketest_python_api.py

@@ -0,0 +1,43 @@
+# Copyright 2021 The TensorFlow Authors. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ==============================================================================
+
+"""Simple smoketest for the Python API."""
+
+# TODO(laurenzo): This should be a real test but we don't have enough build
+# support for it yet.
+
+# pylint: disable=wildcard-import,undefined-variable
+
+from mlir.dialects.chlo import *
+from mlir.dialects.mhlo import *
+from mlir.ir import *
+
+ASM = """
+func @dynamicBroadcast(%arg0: tensor<?xf32>, %arg1: tensor<?x?xf32>) -> tensor<?x?xf32> {
+  %0 = chlo.broadcast_add %arg0, %arg1 : (tensor<?xf32>, tensor<?x?xf32>) -> tensor<?x?xf32>
+  return %0 : tensor<?x?xf32>
+}
+"""
+
+with Context() as context:
+  register_chlo_dialect(context)
+  register_mhlo_dialect(context)
+
+  m = Module.parse(ASM)
+  print(m)
+  add_op = m.body.operations[0].regions[0].blocks[0].operations[0]
+  print(repr(add_op))
+  print(add_op)
+  print("Everything works")

include/mlir-hlo-c/Dialects.h

@@ -0,0 +1,29 @@
+/* Copyright 2021 The TensorFlow Authors. All Rights Reserved.
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+    http://www.apache.org/licenses/LICENSE-2.0
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+
+#ifndef TENSORFLOW_COMPILER_MLIR_HLO_INCLUDE_MLIR_HLO_C_DIALECTS_H_
+#define TENSORFLOW_COMPILER_MLIR_HLO_INCLUDE_MLIR_HLO_C_DIALECTS_H_
+
+#include "mlir-c/Registration.h"
+
+#ifdef __cplusplus
+extern "C" {
+#endif
+
+MLIR_DECLARE_CAPI_DIALECT_REGISTRATION(Chlo, chlo);
+MLIR_DECLARE_CAPI_DIALECT_REGISTRATION(Mhlo, mhlo);
+
+#ifdef __cplusplus
+}
+#endif
+
+#endif  // TENSORFLOW_COMPILER_MLIR_HLO_INCLUDE_MLIR_HLO_C_DIALECTS_H_

lib/CAPI/CMakeLists.txt

@@ -0,0 +1,6 @@
+add_mlir_public_c_api_library(MLIRHLOCAPIDialects
+  Dialects.cpp
+  LINK_LIBS PUBLIC
+  ChloDialect
+  MhloDialect
+)

lib/CAPI/Dialects.cpp

@@ -0,0 +1,20 @@
+/* Copyright 2021 The TensorFlow Authors. All Rights Reserved.
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+    http://www.apache.org/licenses/LICENSE-2.0
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+
+#include "mlir-hlo-c/Dialects.h"
+
+#include "mlir-hlo/Dialect/mhlo/IR/chlo_ops.h"
+#include "mlir-hlo/Dialect/mhlo/IR/hlo_ops.h"
+#include "mlir/CAPI/Registration.h"
+
+MLIR_DEFINE_CAPI_DIALECT_REGISTRATION(Chlo, chlo, mlir::chlo::HloClientDialect)
+MLIR_DEFINE_CAPI_DIALECT_REGISTRATION(Mhlo, mhlo, mlir::mhlo::MhloDialect)

lib/CMakeLists.txt

@@ -14,6 +14,7 @@
 # limitations under the License.
 #
 add_subdirectory(Analysis)
+add_subdirectory(CAPI)
 add_subdirectory(Dialect)
 add_subdirectory(Transforms)
 add_subdirectory(utils)

python/CMakeLists.txt

@@ -1,9 +1,67 @@
 include(AddMLIRPython)
 
+declare_mlir_python_sources(MLIRHLOPythonSources)
+declare_mlir_python_sources(MLIRHLOPythonExtensions)
+
 ################################################################################
 # Generate dialect-specific bindings.
 ################################################################################
 
-add_mlir_dialect_python_bindings(MLIRBindingsPythonMhloOps
-  TD_FILE MhloOps.td
-  DIALECT_NAME mhlo_dialect)
+declare_mlir_python_sources(MLIRHLOPythonSources.Dialects
+  ADD_TO_PARENT MLIRHLOPythonSources
+)
+
+declare_mlir_dialect_python_bindings(
+  ADD_TO_PARENT MLIRHLOPythonSources.Dialects
+  ROOT_DIR "${CMAKE_CURRENT_SOURCE_DIR}/mlir"
+  TD_FILE dialects/MhloOps.td
+  SOURCES dialects/mhlo.py
+  DIALECT_NAME mhlo)
+
+ declare_mlir_dialect_python_bindings(
+  ADD_TO_PARENT MLIRHLOPythonSources.Dialects
+  ROOT_DIR "${CMAKE_CURRENT_SOURCE_DIR}/mlir"
+  TD_FILE dialects/ChloOps.td
+  SOURCES dialects/chlo.py
+  DIALECT_NAME chlo)
+
+################################################################################
+# Extensions
+################################################################################
+
+declare_mlir_python_extension(MLIRHLOPythonExtensions.Main
+  MODULE_NAME _mlirHlo
+  ADD_TO_PARENT MLIRHLOPythonExtensions
+  SOURCES
+    ${CMAKE_CURRENT_SOURCE_DIR}/MlirHloModule.cpp
+  EMBED_CAPI_LINK_LIBS
+    MLIRHLOCAPIDialects
+)
+
+################################################################################
+# Generate packages and shared library
+################################################################################
+
+add_mlir_python_common_capi_library(MLIRHLOCAPI
+  INSTALL_COMPONENT MLIRHLOPythonModules
+  INSTALL_DESTINATION python_packages/mlir_hlo/mlir/_mlir_libs
+  OUTPUT_DIRECTORY "${MLIR_HLO_BINARY_DIR}/python_packages/mlir_hlo/mlir/_mlir_libs"
+  RELATIVE_INSTALL_ROOT "../../../.."
+  DECLARED_SOURCES
+    MLIRPythonSources
+    MLIRPythonExtension.AllPassesRegistration
+    MLIRHLOPythonSources
+    MLIRHLOPythonExtensions
+)
+
+add_mlir_python_modules(MLIRHLOPythonModules
+  ROOT_PREFIX "${MLIR_HLO_BINARY_DIR}/python_packages/mlir_hlo/mlir"
+  INSTALL_PREFIX "python_packages/mlir_hlo/mlir"
+  DECLARED_SOURCES
+    MLIRPythonSources
+    MLIRPythonExtension.AllPassesRegistration
+    MLIRHLOPythonSources
+    MLIRHLOPythonExtensions
+  COMMON_CAPI_LINK_LIBS
+    MLIRHLOCAPI
+  )