[mlir][bazel] Build MLIR bindings extensions with nanobind_bazel

Instead of maintaining custom nanobind and robin-map build files, and bootstrapping
these archives with a custom module extension, use `nanobind_bazel` from the BCR.

The advantage of this is that less maintenance is needed inside the LLVM Bazel tree,
and usage of nanobind_bazel's wrapper macros save a few lines of code while directly
expressing the extensions as `nanobind_extension`s, which is more indicative of the
target's purpose than just a `cc_binary`.

Author: nicholasjng

utils/bazel/MODULE.bazel

@@ -8,6 +8,7 @@ module(name = "llvm-project-overlay")
 
 bazel_dep(name = "apple_support", version = "1.24.1", repo_name = "build_bazel_apple_support")
 bazel_dep(name = "bazel_skylib", version = "1.8.2")
+bazel_dep(name = "nanobind_bazel", version = "2.9.2")
 bazel_dep(name = "platforms", version = "1.0.0")
 bazel_dep(name = "protobuf", version = "31.1", repo_name = "com_google_protobuf")
 bazel_dep(name = "rules_android", version = "0.6.6")
@@ -25,11 +26,9 @@ use_repo(
     "llvm_zstd",
     "mpc",
     "mpfr",
-    "nanobind",
     "pfm",
     "pybind11",
     "pyyaml",
-    "robin_map",
     "vulkan_headers",
     "vulkan_sdk",
 )

utils/bazel/extensions.bzl

@@ -105,23 +105,6 @@ def _llvm_repos_extension_impl(module_ctx):
         build_file = "@llvm-raw//utils/bazel/third_party_build:pyyaml.BUILD",
     )
 
-    # TODO: bump to robin-map-1.4.0
-    http_archive(
-        name = "robin_map",
-        build_file = "@llvm-raw//utils/bazel/third_party_build:robin_map.BUILD",
-        sha256 = "a8424ad3b0affd4c57ed26f0f3d8a29604f0e1f2ef2089f497f614b1c94c7236",
-        strip_prefix = "robin-map-1.3.0",
-        url = "https://github.com/Tessil/robin-map/archive/refs/tags/v1.3.0.tar.gz",
-    )
-
-    http_archive(
-        name = "nanobind",
-        build_file = "@llvm-raw//utils/bazel/third_party_build:nanobind.BUILD",
-        sha256 = "8ce3667dce3e64fc06bfb9b778b6f48731482362fb89a43da156632266cd5a90",
-        strip_prefix = "nanobind-2.9.2",
-        url = "https://github.com/wjakob/nanobind/archive/refs/tags/v2.9.2.tar.gz",
-    )
-
 llvm_repos_extension = module_extension(
     implementation = _llvm_repos_extension_impl,
 )

utils/bazel/llvm-project-overlay/mlir/BUILD.bazel

@@ -7,6 +7,7 @@
 
 load("@bazel_skylib//rules:common_settings.bzl", "bool_flag")
 load("@bazel_skylib//rules:expand_template.bzl", "expand_template")
+load("@nanobind_bazel//:build_defs.bzl", "nanobind_extension", "nanobind_library")
 load("@rules_cc//cc:defs.bzl", "cc_binary", "cc_library")
 load("@rules_python//python:defs.bzl", "py_binary")
 load("//llvm:targets.bzl", "llvm_targets")
@@ -1130,39 +1131,31 @@ cc_library(
     ],
 )
 
-cc_library(
+nanobind_library(
     name = "MLIRBindingsPythonNanobindHeaders",
     includes = [
         "include",
     ],
     textual_hdrs = [":MLIRBindingsPythonHeaderFiles"],
     deps = [
         ":CAPIIRHeaders",
-        "@nanobind",
         "@rules_python//python/cc:current_py_cc_headers",
     ],
 )
 
-cc_library(
+nanobind_library(
     name = "MLIRBindingsPythonNanobindHeadersAndDeps",
     includes = [
         "include",
     ],
     textual_hdrs = [":MLIRBindingsPythonHeaderFiles"],
     deps = [
         ":CAPIIR",
-        "@nanobind",
         "@rules_python//python/cc:current_py_cc_headers",
     ],
 )
 
-# These flags are needed for pybind11 to work.
-PYBIND11_COPTS = [
-    "-fexceptions",
-    "-frtti",
-]
-
-PYBIND11_FEATURES = [
+NB_FEATURES = [
     # Cannot use header_modules (parse_headers feature fails).
     "-use_header_modules",
 ]
@@ -1188,11 +1181,10 @@ filegroup(
     ]),
 )
 
-cc_library(
+nanobind_library(
     name = "MLIRBindingsPythonCore",
     srcs = [":MLIRBindingsPythonSourceFiles"],
-    copts = PYBIND11_COPTS,
-    features = PYBIND11_FEATURES,
+    features = NB_FEATURES,
     includes = [
         "lib/Bindings/Python",
     ],
@@ -1207,16 +1199,14 @@ cc_library(
         ":Support",
         ":config",
         "//llvm:Support",
-        "@nanobind",
         "@rules_python//python/cc:current_py_cc_headers",
     ],
 )
 
-cc_library(
+nanobind_library(
     name = "MLIRBindingsPythonCoreNoCAPI",
     srcs = [":MLIRBindingsPythonSourceFiles"],
-    copts = PYBIND11_COPTS,
-    features = PYBIND11_FEATURES,
+    features = NB_FEATURES,
     includes = [
         "lib/Bindings/Python",
     ],
@@ -1230,7 +1220,6 @@ cc_library(
         ":Support",
         ":config",
         "//llvm:Support",
-        "@nanobind",
         "@rules_python//python/cc:current_py_cc_headers",
     ],
 )
@@ -1249,130 +1238,105 @@ cc_library(
 )
 
 # Dynamic library with the MLIR Python extension.
-cc_binary(
-    name = "_mlir.so",
+nanobind_extension(
+    name = "_mlir",
     srcs = ["lib/Bindings/Python/MainModule.cpp"],
-    # These flags are needed for pybind11 to work.
-    copts = PYBIND11_COPTS,
-    features = PYBIND11_FEATURES,
+    features = NB_FEATURES,
     includes = [
         "lib/Bindings/Python",
     ],
-    linkshared = 1,
     linkstatic = 0,
     deps = [
         ":MLIRBindingsPythonCore",
         ":MLIRBindingsPythonHeadersAndDeps",
-        "@nanobind",
     ],
 )
 
-cc_binary(
-    name = "_mlirDialectsIRDL.so",
+nanobind_extension(
+    name = "_mlirDialectsIRDL",
     srcs = ["lib/Bindings/Python/DialectIRDL.cpp"],
-    copts = PYBIND11_COPTS,
-    features = PYBIND11_FEATURES,
-    linkshared = 1,
+    features = NB_FEATURES,
     linkstatic = 0,
     deps = [
         ":CAPIIR",
         ":MLIRBindingsPythonNanobindHeadersAndDeps",
-        "@nanobind",
     ],
 )
 
-cc_binary(
-    name = "_mlirDialectsLinalg.so",
+nanobind_extension(
+    name = "_mlirDialectsLinalg",
     srcs = ["lib/Bindings/Python/DialectLinalg.cpp"],
-    copts = PYBIND11_COPTS,
-    features = PYBIND11_FEATURES,
+    features = NB_FEATURES,
     includes = [
         "lib/Bindings/Python",
     ],
-    linkshared = 1,
     linkstatic = 0,
     deps = [
         ":CAPIIR",
         ":CAPILinalg",
         ":MLIRBindingsPythonNanobindHeadersAndDeps",
-        "@nanobind",
     ],
 )
 
-cc_binary(
-    name = "_mlirDialectsLLVM.so",
+nanobind_extension(
+    name = "_mlirDialectsLLVM",
     srcs = ["lib/Bindings/Python/DialectLLVM.cpp"],
-    copts = PYBIND11_COPTS,
-    features = PYBIND11_FEATURES,
-    linkshared = 1,
+    features = NB_FEATURES,
     linkstatic = 0,
     deps = [
         ":CAPIIR",
         ":CAPILLVM",
         ":CAPITarget",
         ":MLIRBindingsPythonNanobindHeadersAndDeps",
-        "@nanobind",
     ],
 )
 
-cc_binary(
-    name = "_mlirDialectsQuant.so",
+nanobind_extension(
+    name = "_mlirDialectsQuant",
     srcs = ["lib/Bindings/Python/DialectQuant.cpp"],
-    copts = PYBIND11_COPTS,
-    features = PYBIND11_FEATURES,
-    linkshared = 1,
+    features = NB_FEATURES,
     linkstatic = 0,
     deps = [
         ":CAPIIR",
         ":CAPIQuant",
         ":MLIRBindingsPythonNanobindHeadersAndDeps",
-        "@nanobind",
     ],
 )
 
-cc_binary(
-    name = "_mlirDialectsSparseTensor.so",
+nanobind_extension(
+    name = "_mlirDialectsSparseTensor",
     srcs = ["lib/Bindings/Python/DialectSparseTensor.cpp"],
-    copts = PYBIND11_COPTS,
-    features = PYBIND11_FEATURES,
-    linkshared = 1,
+    features = NB_FEATURES,
     linkstatic = 0,
     deps = [
         ":CAPIIR",
         ":CAPISparseTensor",
         ":MLIRBindingsPythonNanobindHeadersAndDeps",
-        "@nanobind",
     ],
 )
 
 # Dynamic library with the MLIR Conversions Python extension.
-cc_binary(
-    name = "_mlirExecutionEngine.so",
+nanobind_extension(
+    name = "_mlirExecutionEngine",
     srcs = ["lib/Bindings/Python/ExecutionEngineModule.cpp"],
-    copts = PYBIND11_COPTS,
-    features = PYBIND11_FEATURES,
-    linkshared = 1,
+    features = NB_FEATURES,
     linkstatic = 0,
     deps = [
         ":CAPIExecutionEngine",
         ":MLIRBindingsPythonNanobindHeadersAndDeps",
-        "@nanobind",
         "@rules_python//python/cc:current_py_cc_headers",
     ],
 )
 
 # Dynamic library with the MLIR Linalg dialect+passes Python extension.
-cc_binary(
-    name = "_mlirLinalgPasses.so",
+nanobind_extension(
+    name = "_mlirLinalgPasses",
     srcs = ["lib/Bindings/Python/LinalgPasses.cpp"],
-    copts = PYBIND11_COPTS,
-    features = PYBIND11_FEATURES,
-    linkshared = 1,
+    features = NB_FEATURES,
     linkstatic = 0,
     deps = [
         ":CAPILinalg",
         ":MLIRBindingsPythonNanobindHeadersAndDeps",
-        "@nanobind",
         "@rules_python//python/cc:current_py_cc_headers",
     ],
 )