[llvm] [mlir][bazel] Build MLIR bindings extensions with `nanobind_bazel` (PR #169306)
Nicholas Junge via llvm-commits
llvm-commits at lists.llvm.org
Mon Nov 24 01:56:21 PST 2025
https://github.com/nicholasjng created https://github.com/llvm/llvm-project/pull/169306
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`.
-------------
Follow-up of #164891, just wanted to draw this up to visualize how building bindings with `nanobind_bazel` from the BCR might look like.
This change is not entirely equivalent to the current state, since I'm not setting macros in nanobind to enable shared library exports - I'm tracking that in https://github.com/nicholasjng/nanobind-bazel/issues/81. It does come with a "free" robin-map bump to v1.4.0, though, hence the removed TODO.
>From 889dc92ed9ecdedec8e90456f6fe860f76cc9a58 Mon Sep 17 00:00:00 2001
From: Nicholas Junge <nicholas.junge at web.de>
Date: Mon, 24 Nov 2025 10:43:23 +0100
Subject: [PATCH] [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`.
---
utils/bazel/MODULE.bazel | 3 +-
utils/bazel/extensions.bzl | 17 ---
.../llvm-project-overlay/mlir/BUILD.bazel | 100 ++++++------------
utils/bazel/third_party_build/nanobind.BUILD | 25 -----
utils/bazel/third_party_build/robin_map.BUILD | 12 ---
5 files changed, 33 insertions(+), 124 deletions(-)
delete mode 100644 utils/bazel/third_party_build/nanobind.BUILD
delete mode 100644 utils/bazel/third_party_build/robin_map.BUILD
diff --git a/utils/bazel/MODULE.bazel b/utils/bazel/MODULE.bazel
index 1a8327c33d246..a3f5872599acf 100644
--- a/utils/bazel/MODULE.bazel
+++ b/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",
)
diff --git a/utils/bazel/extensions.bzl b/utils/bazel/extensions.bzl
index bb5ce1955f916..9db08b02bf66f 100644
--- a/utils/bazel/extensions.bzl
+++ b/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,
)
diff --git a/utils/bazel/llvm-project-overlay/mlir/BUILD.bazel b/utils/bazel/llvm-project-overlay/mlir/BUILD.bazel
index 1421ec553f251..ac8d2b10581ac 100644
--- a/utils/bazel/llvm-project-overlay/mlir/BUILD.bazel
+++ b/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,7 +1131,7 @@ cc_library(
],
)
-cc_library(
+nanobind_library(
name = "MLIRBindingsPythonNanobindHeaders",
includes = [
"include",
@@ -1138,12 +1139,11 @@ cc_library(
textual_hdrs = [":MLIRBindingsPythonHeaderFiles"],
deps = [
":CAPIIRHeaders",
- "@nanobind",
"@rules_python//python/cc:current_py_cc_headers",
],
)
-cc_library(
+nanobind_library(
name = "MLIRBindingsPythonNanobindHeadersAndDeps",
includes = [
"include",
@@ -1151,18 +1151,11 @@ cc_library(
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",
],
)
diff --git a/utils/bazel/third_party_build/nanobind.BUILD b/utils/bazel/third_party_build/nanobind.BUILD
deleted file mode 100644
index 262d14a040b87..0000000000000
--- a/utils/bazel/third_party_build/nanobind.BUILD
+++ /dev/null
@@ -1,25 +0,0 @@
-cc_library(
- name = "nanobind",
- srcs = glob(
- [
- "src/*.cpp",
- ],
- exclude = ["src/nb_combined.cpp"],
- ),
- defines = [
- "NB_BUILD=1",
- "NB_SHARED=1",
- ],
- includes = ["include"],
- textual_hdrs = glob(
- [
- "include/**/*.h",
- "src/*.h",
- ],
- ),
- visibility = ["//visibility:public"],
- deps = [
- "@robin_map",
- "@rules_python//python/cc:current_py_cc_headers",
- ],
-)
diff --git a/utils/bazel/third_party_build/robin_map.BUILD b/utils/bazel/third_party_build/robin_map.BUILD
deleted file mode 100644
index b8d04beaed81f..0000000000000
--- a/utils/bazel/third_party_build/robin_map.BUILD
+++ /dev/null
@@ -1,12 +0,0 @@
-cc_library(
- name = "robin_map",
- hdrs = [
- "include/tsl/robin_growth_policy.h",
- "include/tsl/robin_hash.h",
- "include/tsl/robin_map.h",
- "include/tsl/robin_set.h",
- ],
- includes = ["."],
- strip_include_prefix = "include",
- visibility = ["//visibility:public"],
-)
More information about the llvm-commits
mailing list