Migrate StableHLO Python extension to nanobind.

I'm working towards moving the MLIR Python core code to use nanobind instead of pybind11:
* llvm/llvm-project#117922, which was merged recently, allows downstream Python dialect extensions to be defined using either pybind11 or nanobind.
* llvm/llvm-project#118583 is a PR in review that ports the Python core code to use nanobind instead of pybind11.

This PR migrates StableHLO and related dialects to use nanobind rather than pybind11, with the goal of migrating JAX away from pybind11.

PiperOrigin-RevId: 707537037

Author: hawkinsp

stablehlo/stablehlo/integrations/python/CheckModule.cpp

@@ -11,12 +11,13 @@ limitations under the License.
 ==============================================================================*/
 
 #include "mlir-c/IR.h"
-#include "mlir/Bindings/Python/PybindAdaptors.h"
+#include "mlir/Bindings/Python/NanobindAdaptors.h"
+#include "nanobind/nanobind.h"
 #include "stablehlo/integrations/c/CheckDialect.h"
 
-namespace py = pybind11;
+namespace nb = nanobind;
 
-PYBIND11_MODULE(_check, m) {
+NB_MODULE(_check, m) {
   m.doc() = "check main python extension";
 
   //
@@ -32,5 +33,5 @@ PYBIND11_MODULE(_check, m) {
           mlirDialectHandleLoadDialect(dialect, context);
         }
       },
-      py::arg("context"), py::arg("load") = true);
+      nb::arg("context"), nb::arg("load") = true);
 }

stablehlo/stablehlo/integrations/python/ChloModule.cpp

@@ -12,21 +12,23 @@ limitations under the License.
 ==============================================================================*/
 
 #include "mlir-c/IR.h"
-#include "mlir/Bindings/Python/PybindAdaptors.h"
+#include "mlir/Bindings/Python/NanobindAdaptors.h"
+#include "nanobind/nanobind.h"
+#include "nanobind/stl/string_view.h"
 #include "stablehlo/integrations/c/ChloAttributes.h"
 #include "stablehlo/integrations/c/ChloDialect.h"
 
-namespace py = pybind11;
+namespace nb = nanobind;
 
 namespace {
 
 auto toPyString(MlirStringRef mlirStringRef) {
-  return py::str(mlirStringRef.data, mlirStringRef.length);
+  return nb::str(mlirStringRef.data, mlirStringRef.length);
 }
 
 }  // namespace
 
-PYBIND11_MODULE(_chlo, m) {
+NB_MODULE(_chlo, m) {
   m.doc() = "chlo main python extension";
 
   //
@@ -42,35 +44,37 @@ PYBIND11_MODULE(_chlo, m) {
           mlirDialectHandleLoadDialect(dialect, context);
         }
       },
-      py::arg("context"), py::arg("load") = true);
+      nb::arg("context"), nb::arg("load") = true);
 
   //
   // Attributes.
   //
 
-  mlir::python::adaptors::mlir_attribute_subclass(
+  mlir::python::nanobind_adaptors::mlir_attribute_subclass(
       m, "ComparisonDirectionAttr", chloAttributeIsAComparisonDirectionAttr)
       .def_classmethod(
           "get",
-          [](py::object cls, const std::string &value, MlirContext ctx) {
+          [](nb::object cls, std::string_view value, MlirContext ctx) {
             return cls(chloComparisonDirectionAttrGet(
-                ctx, mlirStringRefCreate(value.c_str(), value.size())));
+                ctx, mlirStringRefCreate(value.data(), value.size())));
           },
-          py::arg("cls"), py::arg("value"), py::arg("context") = py::none(),
+          nb::arg("cls"), nb::arg("value"),
+          nb::arg("context").none() = nb::none(),
           "Creates a ComparisonDirection attribute with the given value.")
       .def_property_readonly("value", [](MlirAttribute self) {
         return toPyString(chloComparisonDirectionAttrGetValue(self));
       });
 
-  mlir::python::adaptors::mlir_attribute_subclass(
+  mlir::python::nanobind_adaptors::mlir_attribute_subclass(
       m, "ComparisonTypeAttr", chloAttributeIsAComparisonTypeAttr)
       .def_classmethod(
           "get",
-          [](py::object cls, const std::string &value, MlirContext ctx) {
+          [](nb::object cls, std::string_view value, MlirContext ctx) {
             return cls(chloComparisonTypeAttrGet(
-                ctx, mlirStringRefCreate(value.c_str(), value.size())));
+                ctx, mlirStringRefCreate(value.data(), value.size())));
           },
-          py::arg("cls"), py::arg("value"), py::arg("context") = py::none(),
+          nb::arg("cls"), nb::arg("value"),
+          nb::arg("context").none() = nb::none(),
           "Creates a ComparisonType attribute with the given value.")
       .def_property_readonly("value", [](MlirAttribute self) {
         return toPyString(chloComparisonTypeAttrGetValue(self));

stablehlo/stablehlo/integrations/python/StablehloApi.cpp

@@ -15,17 +15,22 @@ limitations under the License.
 
 #include "stablehlo/integrations/python/StablehloApi.h"
 
+#include <stdexcept>
 #include <string>
 #include <string_view>
 
 #include "llvm/Support/raw_ostream.h"
 #include "mlir-c/BuiltinAttributes.h"
 #include "mlir-c/IR.h"
 #include "mlir-c/Support.h"
-#include "mlir/Bindings/Python/PybindAdaptors.h"
+#include "mlir/Bindings/Python/NanobindAdaptors.h"
+#include "nanobind/nanobind.h"
+#include "nanobind/stl/string.h"
+#include "nanobind/stl/string_view.h"
+#include "nanobind/stl/vector.h"
 #include "stablehlo/integrations/c/StablehloApi.h"
 
-namespace py = pybind11;
+namespace nb = nanobind;
 
 namespace mlir {
 namespace stablehlo {
@@ -63,14 +68,18 @@ static MlirStringRef toMlirStringRef(std::string_view s) {
   return mlirStringRefCreate(s.data(), s.size());
 }
 
-void AddStablehloApi(py::module &m) {
+static MlirStringRef toMlirStringRef(const nb::bytes &s) {
+  return mlirStringRefCreate(static_cast<const char *>(s.data()), s.size());
+}
+
+void AddStablehloApi(nb::module_ &m) {
   // Portable API is a subset of StableHLO API
   AddPortableApi(m);
 
   //
   // Utility APIs.
   //
-  py::enum_<MlirStablehloCompatibilityRequirement>(
+  nb::enum_<MlirStablehloCompatibilityRequirement>(
       m, "StablehloCompatibilityRequirement")
       .value("NONE", MlirStablehloCompatibilityRequirement::NONE)
       .value("WEEK_4", MlirStablehloCompatibilityRequirement::WEEK_4)
@@ -79,48 +88,57 @@ void AddStablehloApi(py::module &m) {
 
   m.def(
       "get_version_from_compatibility_requirement",
-      [](MlirStablehloCompatibilityRequirement requirement) -> py::str {
+      [](MlirStablehloCompatibilityRequirement requirement) -> std::string {
         StringWriterHelper accumulator;
         stablehloVersionFromCompatibilityRequirement(
             requirement, accumulator.getMlirStringCallback(),
             accumulator.getUserData());
         return accumulator.toString();
       },
-      py::arg("requirement"));
+      nb::arg("requirement"));
 
   //
   // Serialization APIs.
   //
   m.def(
       "serialize_portable_artifact",
-      [](MlirModule module, std::string_view target) -> py::bytes {
+      [](MlirModule module, std::string_view target) -> nb::bytes {
         StringWriterHelper accumulator;
         if (mlirLogicalResultIsFailure(
                 stablehloSerializePortableArtifactFromModule(
                     module, toMlirStringRef(target),
                     accumulator.getMlirStringCallback(),
                     accumulator.getUserData()))) {
-          PyErr_SetString(PyExc_ValueError, "failed to serialize module");
-          return "";
+          throw nb::value_error("failed to serialize module");
         }
 
-        return py::bytes(accumulator.toString());
+        std::string serialized = accumulator.toString();
+        return nb::bytes(serialized.data(), serialized.size());
       },
-      py::arg("module"), py::arg("target"));
+      nb::arg("module"), nb::arg("target"));
 
   m.def(
       "deserialize_portable_artifact",
       [](MlirContext context, std::string_view artifact) -> MlirModule {
         auto module = stablehloDeserializePortableArtifactNoError(
             toMlirStringRef(artifact), context);
         if (mlirModuleIsNull(module)) {
-          PyErr_SetString(PyExc_ValueError, "failed to deserialize module");
-          return {};
+          throw nb::value_error("failed to deserialize module");
         }
         return module;
       },
-      py::arg("context"), py::arg("artifact"));
-
+      nb::arg("context"), nb::arg("artifact"));
+  m.def(
+      "deserialize_portable_artifact",
+      [](MlirContext context, nb::bytes artifact) -> MlirModule {
+        auto module = stablehloDeserializePortableArtifactNoError(
+            toMlirStringRef(artifact), context);
+        if (mlirModuleIsNull(module)) {
+          throw nb::value_error("failed to deserialize module");
+        }
+        return module;
+      },
+      nb::arg("context"), nb::arg("artifact"));
   //
   // Reference APIs
   //
@@ -130,9 +148,7 @@ void AddStablehloApi(py::module &m) {
          std::vector<MlirAttribute> &args) -> std::vector<MlirAttribute> {
         for (auto arg : args) {
           if (!mlirAttributeIsADenseElements(arg)) {
-            PyErr_SetString(PyExc_ValueError,
-                            "input args must be DenseElementsAttr");
-            return {};
+            throw nb::value_error("input args must be DenseElementsAttr");
           }
         }
 
@@ -141,8 +157,7 @@ void AddStablehloApi(py::module &m) {
             stablehloEvalModule(module, args.size(), args.data(), &errorCode);
 
         if (errorCode != 0) {
-          PyErr_SetString(PyExc_ValueError, "interpreter failed");
-          return {};
+          throw nb::value_error("interpreter failed");
         }
 
         std::vector<MlirAttribute> pyResults;
@@ -151,39 +166,39 @@ void AddStablehloApi(py::module &m) {
         }
         return pyResults;
       },
-      py::arg("module"), py::arg("args"));
+      nb::arg("module"), nb::arg("args"));
 }
 
-void AddPortableApi(py::module &m) {
+void AddPortableApi(nb::module_ &m) {
   //
   // Utility APIs.
   //
   m.def("get_api_version", []() { return stablehloGetApiVersion(); });
 
   m.def(
       "get_smaller_version",
-      [](const std::string &version1, const std::string &version2) -> py::str {
+      [](const std::string &version1,
+         const std::string &version2) -> std::string {
         StringWriterHelper accumulator;
         if (mlirLogicalResultIsFailure(stablehloGetSmallerVersion(
                 toMlirStringRef(version1), toMlirStringRef(version2),
                 accumulator.getMlirStringCallback(),
                 accumulator.getUserData()))) {
-          PyErr_SetString(PyExc_ValueError,
-                          "failed to convert version to stablehlo version");
-          return "";
+          throw nb::value_error(
+              "failed to convert version to stablehlo version");
         }
         return accumulator.toString();
       },
-      py::arg("version1"), py::arg("version2"));
+      nb::arg("version1"), nb::arg("version2"));
 
-  m.def("get_current_version", []() -> py::str {
+  m.def("get_current_version", []() -> std::string {
     StringWriterHelper accumulator;
     stablehloGetCurrentVersion(accumulator.getMlirStringCallback(),
                                accumulator.getUserData());
     return accumulator.toString();
   });
 
-  m.def("get_minimum_version", []() -> py::str {
+  m.def("get_minimum_version", []() -> std::string {
     StringWriterHelper accumulator;
     stablehloGetMinimumVersion(accumulator.getMlirStringCallback(),
                                accumulator.getUserData());
@@ -196,34 +211,64 @@ void AddPortableApi(py::module &m) {
   m.def(
       "serialize_portable_artifact_str",
       [](std::string_view moduleStrOrBytecode,
-         std::string_view targetVersion) -> py::bytes {
+         std::string_view targetVersion) -> nb::bytes {
+        StringWriterHelper accumulator;
+        if (mlirLogicalResultIsFailure(
+                stablehloSerializePortableArtifactFromStringRef(
+                    toMlirStringRef(moduleStrOrBytecode),
+                    toMlirStringRef(targetVersion),
+                    accumulator.getMlirStringCallback(),
+                    accumulator.getUserData()))) {
+          throw nb::value_error("failed to serialize module");
+        }
+        std::string serialized = accumulator.toString();
+        return nb::bytes(serialized.data(), serialized.size());
+      },
+      nb::arg("module_str"), nb::arg("target_version"));
+  m.def(
+      "serialize_portable_artifact_str",
+      [](nb::bytes moduleStrOrBytecode,
+         std::string_view targetVersion) -> nb::bytes {
         StringWriterHelper accumulator;
         if (mlirLogicalResultIsFailure(
                 stablehloSerializePortableArtifactFromStringRef(
                     toMlirStringRef(moduleStrOrBytecode),
                     toMlirStringRef(targetVersion),
                     accumulator.getMlirStringCallback(),
                     accumulator.getUserData()))) {
-          PyErr_SetString(PyExc_ValueError, "failed to serialize module");
-          return "";
+          throw nb::value_error("failed to serialize module");
         }
-        return py::bytes(accumulator.toString());
+        std::string serialized = accumulator.toString();
+        return nb::bytes(serialized.data(), serialized.size());
       },
-      py::arg("module_str"), py::arg("target_version"));
+      nb::arg("module_str"), nb::arg("target_version"));
 
   m.def(
       "deserialize_portable_artifact_str",
-      [](std::string_view artifact) -> py::bytes {
+      [](std::string_view artifact) -> nb::bytes {
+        StringWriterHelper accumulator;
+        if (mlirLogicalResultIsFailure(stablehloDeserializePortableArtifact(
+                toMlirStringRef(artifact), accumulator.getMlirStringCallback(),
+                accumulator.getUserData()))) {
+          throw nb::value_error("failed to deserialize module");
+        }
+        std::string serialized = accumulator.toString();
+        return nb::bytes(serialized.data(), serialized.size());
+      },
+      nb::arg("artifact_str"));
+  m.def(
+      "deserialize_portable_artifact_str",
+      [](const nb::bytes& artifact) -> nb::bytes {
         StringWriterHelper accumulator;
         if (mlirLogicalResultIsFailure(stablehloDeserializePortableArtifact(
                 toMlirStringRef(artifact), accumulator.getMlirStringCallback(),
                 accumulator.getUserData()))) {
-          PyErr_SetString(PyExc_ValueError, "failed to deserialize module");
-          return "";
+          throw nb::value_error("failed to deserialize module");
         }
-        return py::bytes(accumulator.toString());
+        std::string serialized = accumulator.toString();
+        return nb::bytes(serialized.data(), serialized.size());
       },
-      py::arg("artifact_str"));
+      nb::arg("artifact_str"));
 }
 
 }  // namespace stablehlo

stablehlo/stablehlo/integrations/python/StablehloApi.h

@@ -16,20 +16,20 @@ limitations under the License.
 #ifndef STABLEHLO_INTEGRATIONS_PYTHON_API_STABLEHLOAPI_H
 #define STABLEHLO_INTEGRATIONS_PYTHON_API_STABLEHLOAPI_H
 
-#include "pybind11/pybind11.h"
+#include "nanobind/nanobind.h"
 
 namespace mlir {
 namespace stablehlo {
 
-// Add StableHLO APIs to the pybind11 module.
+// Add StableHLO APIs to the nanobind module.
 // Signatures of these APIs have no dependency on C++ MLIR types and all must
 // use C API passthrough.
-void AddStablehloApi(pybind11::module& m);
+void AddStablehloApi(nanobind::module_& m);
 
 // Adds a subset of the StableHLO API that doesn't use MLIR in any definitions,
 // and is methods only, introducing no new objects / enums to avoid potential
 // redefinition issues in complex build environments.
-void AddPortableApi(pybind11::module& m);
+void AddPortableApi(nanobind::module_& m);
 
 }  // namespace stablehlo
 }  // namespace mlir