[mlir:python] Prevent crash in DenseElementsAttr.

[ingomueller-net]

This PR fixes a crash in the bf_getbuffer implementation of PyDenseElementsAttribute that occurred when an element type was not supported, such as bf16. I believe that supportion bf16 is not possible with that protocol but that's out of the scope of this PR. Previsouly, the code raised an std::exception out of bf_getbuffer that nanobind does not catch (see also pybind/pybind11#3336). The PR makes the function catch all std::exceptions and manually raises a Python exception instead.

[llvmbot]

@llvm/pr-subscribers-mlir

Author: Ingo Müller (ingomueller-net)

Changes

This PR fixes a crash in the bf_getbuffer implementation of PyDenseElementsAttribute that occurred when an element type was not supported, such as bf16. I believe that supportion bf16 is not possible with that protocol but that's out of the scope of this PR. Previsouly, the code raised an std::exception out of bf_getbuffer that nanobind does not catch (see also pybind/pybind11#3336). The PR makes the function catch all std::exceptions and manually raises a Python exception instead.

---

Full diff: https://github.com/llvm/llvm-project/pull/163564.diff

1 Files Affected:

• (modified) mlir/lib/Bindings/Python/IRAttributes.cpp (+4)

diff --git a/mlir/lib/Bindings/Python/IRAttributes.cpp b/mlir/lib/Bindings/Python/IRAttributes.cpp
index 045c0fbf4630f..c0a945e3f4f3b 100644
--- a/mlir/lib/Bindings/Python/IRAttributes.cpp
+++ b/mlir/lib/Bindings/Python/IRAttributes.cpp
@@ -1306,6 +1306,10 @@ PyType_Slot PyDenseElementsAttribute::slots[] = {
     e.restore();
     nb::chain_error(PyExc_BufferError, "Error converting attribute to buffer");
     return -1;
+  } catch (std::exception &e) {
+    nb::chain_error(PyExc_BufferError,
+                    "Error converting attribute to buffer: %s", e.what());
+    return -1;
   }
   view->obj = obj;
   view->ndim = 1;

[makslevental]

LGTM but I'm curious - this doesn't work even when constructing the original denseattr using ml_dtypes?

[ingomueller-net]

TBH, I have not tried but I am pretty confident that it doesn't matter how the attribute has been constructed. The bf_getbuffer protocol entails setting the format indicating the element type. Unfortunately, there is such format for BF16, so we can't use that protocol for that datatype.

It could be that the ml_dtypes packages provides some other way, though; maybe that's a good starting point for expanding the type support here.

[ingomueller-net]

I have just stumbled upon DLPack. That's probably a good way to replace the bf_getbuffer protocol with support for more data types. It's what nanobind uses for nb::ndarray.

[makslevental]

I have just stumbled upon DLPack. That's probably a good way to replace the bf_getbuffer protocol with support for more data types. It's what nanobind uses for nb::ndarray.

nanobind's dlpack isn't "real" - they don't actually use the dlpack header and have a reimplementation. So it'll solve your bf16 problem but won't work for the other dtypes that dlpack should support (I know this because I had your same exact problem last 2 weeks ago but for f8).

[ingomueller-net]

Right, but I think that that has a few workable solutions. For example, we can include dlpack.h. It should even be enough to use just enum values; DLPack is an ABI and it doesn't really matter whether we take the values from the DLPack enum or the nanobind enum. Second, we can even just copy the enum values; DLPack can only ever add values to that enum. Finally, we can also add the enum values to nanobind; however, we'd have to increase the required version again, which may cause friction for some users/developers. WDYT?

[makslevental]

It should even be enough to use just enum values

Tried that, didn't work. Nanobind doesn't handle it correctly (I didn't investigate deeply).

Finally, we can also add the enum values to nanobind; however, we'd have to increase the required version again,

This is the thing to do but I can't do it (I can't do it - I don't have permission to contribute to nanobind).

which may cause friction for some users/developers

🤷‍♂️