[Mlir-commits] [mlir] [mlir][python] Fix segfault in DenseResourceElementsAttr.get_from_buffer for 0-d tensors (PR #183070)

[llvmlistbot at llvm.org]

llvmbot wrote:


<!--LLVM PR SUMMARY COMMENT-->

@llvm/pr-subscribers-mlir

Author: Rahul Kayaith (rkayaith)

<details>
<summary>Changes</summary>

When `ndim == 0`, `view->strides[view->ndim - 1]` is an out-of-bounds access (unsigned underflow to `SIZE_MAX`). Use `view->itemsize` for alignment instead, since a scalar buffer is trivially aligned to its element size.

Fixes iree-org/iree-turbine#<!-- -->1312.

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


2 Files Affected:

- (modified) mlir/lib/Bindings/Python/IRAttributes.cpp (+2) 
- (modified) mlir/test/python/ir/array_attributes.py (+16) 


``````````diff
diff --git a/mlir/lib/Bindings/Python/IRAttributes.cpp b/mlir/lib/Bindings/Python/IRAttributes.cpp
index 8c4a2dcd5a7f7..d322933eccd1c 100644
--- a/mlir/lib/Bindings/Python/IRAttributes.cpp
+++ b/mlir/lib/Bindings/Python/IRAttributes.cpp
@@ -1086,6 +1086,8 @@ PyDenseResourceElementsAttribute::getFromBuffer(
   size_t inferredAlignment;
   if (alignment)
     inferredAlignment = *alignment;
+  else if (view->ndim == 0)
+    inferredAlignment = view->itemsize;
   else
     inferredAlignment = view->strides[view->ndim - 1];
 
diff --git a/mlir/test/python/ir/array_attributes.py b/mlir/test/python/ir/array_attributes.py
index c1e2dd5f5ae5e..9b4249b0c686b 100644
--- a/mlir/test/python/ir/array_attributes.py
+++ b/mlir/test/python/ir/array_attributes.py
@@ -621,6 +621,22 @@ def test_attribute(context, mview):
     print("EXIT FUNCTION")
 
 
+# CHECK-LABEL: TEST: testGetDenseResourceElementsAttrScalar
+ at run
+def testGetDenseResourceElementsAttrScalar():
+    with Context(), Location.unknown():
+        element_type = IntegerType.get_signless(64)
+        tensor_type = RankedTensorType.get((), element_type)
+        resource = DenseResourceElementsAttr.get_from_buffer(
+            memoryview(np.array(42, dtype=np.int64)), "scalar", tensor_type
+        )
+        module = Module.create()
+        module.operation.attributes["test.resource"] = resource
+        # CHECK: test.resource = dense_resource<scalar> : tensor<i64>
+        # CHECK: scalar: "0x080000002A00000000000000"
+        print(module)
+
+
 # CHECK-LABEL: TEST: testDanglingResource
 print("TEST: testDanglingResource")
 # see https://github.com/llvm/llvm-project/pull/149414, https://github.com/llvm/llvm-project/pull/150137, https://github.com/llvm/llvm-project/pull/150561

``````````

</details>


https://github.com/llvm/llvm-project/pull/183070