[Mlir-commits] [mlir] [mlir][python] Handle rank-zero MemRefs in NumPy conversion (PR #231403)

[Anutosh Bhat]

https://github.com/anutosh491 created https://github.com/llvm/llvm-project/pull/231403

I found this while testing MLIR's NumPy MemRef helpers in wasm32 and wasm64 xeus-python kernels.(Basically while working on my RFC : https://discourse.llvm.org/t/rfc-bringing-interactive-mlir-compilation-and-execution-to-webassembly/91566/3)

Rank-zero descriptors have no `shape` or `strides`, but `ranked_memref_to_numpy()` accesses those fields and raises `AttributeError`. Handle them as scalar NumPy views, preserving the offset and shared storage.

The regression covers float32, float16 and complex64, including nonzero offsets. Added a new file. Can club with the execution_engine.py tests too if required.

>From 4c3470bbb12d0dfeb2760a3f6b234eab712d11b5 Mon Sep 17 00:00:00 2001
From: anutosh491 <andersonbhat491 at gmail.com>
Date: Tue, 13 Oct 2026 13:28:20 +0530
Subject: [PATCH] [mlir][python] Handle rank-zero MemRefs in NumPy conversion

---
 mlir/python/mlir/runtime/np_to_memref.py |  2 ++
 mlir/test/python/np_to_memref.py         | 34 ++++++++++++++++++++++++
 2 files changed, 36 insertions(+)
 create mode 100644 mlir/test/python/np_to_memref.py

diff --git a/mlir/python/mlir/runtime/np_to_memref.py b/mlir/python/mlir/runtime/np_to_memref.py
index d65ba51afdb90..2987c3107bf4a 100644
--- a/mlir/python/mlir/runtime/np_to_memref.py
+++ b/mlir/python/mlir/runtime/np_to_memref.py
@@ -194,6 +194,8 @@ def ranked_memref_to_numpy(ranked_memref):
     content_ptr = move_aligned_ptr_by_offset(
         ranked_memref[0].aligned, ranked_memref[0].offset
     )
+    if not hasattr(ranked_memref[0], "shape"):
+        return to_numpy(np.ctypeslib.as_array(content_ptr, shape=()))
     np_arr = np.ctypeslib.as_array(content_ptr, shape=ranked_memref[0].shape)
     strided_arr = np.lib.stride_tricks.as_strided(
         np_arr,
diff --git a/mlir/test/python/np_to_memref.py b/mlir/test/python/np_to_memref.py
new file mode 100644
index 0000000000000..c243ff32f5226
--- /dev/null
+++ b/mlir/test/python/np_to_memref.py
@@ -0,0 +1,34 @@
+# RUN: %PYTHON %s | FileCheck %s
+
+import ctypes
+
+import numpy as np
+
+from mlir.runtime import get_ranked_memref_descriptor, ranked_memref_to_numpy
+
+
+# CHECK: rank-zero round trip and offset: float32
+# CHECK: rank-zero round trip and offset: float16
+# CHECK: rank-zero round trip and offset: complex64
+for dtype in (np.float32, np.float16, np.complex64):
+    storage = np.arange(4, dtype=dtype)
+    scalar = storage[2:3].reshape(())
+    descriptor = get_ranked_memref_descriptor(scalar)
+    result = ranked_memref_to_numpy(ctypes.pointer(descriptor))
+    assert result.shape == ()
+    assert result.dtype == scalar.dtype
+    np.testing.assert_array_equal(result, scalar)
+    assert np.shares_memory(result, storage)
+    result[...] = 9
+    assert storage[2] == 9
+
+    descriptor.aligned = storage.ctypes.data_as(type(descriptor.aligned))
+    descriptor.offset = 1
+    result = ranked_memref_to_numpy(ctypes.pointer(descriptor))
+    assert result.shape == ()
+    assert result.dtype == scalar.dtype
+    np.testing.assert_array_equal(result, storage[1])
+    assert np.shares_memory(result, storage)
+    result[...] = 7
+    assert storage[1] == 7
+    print(f"rank-zero round trip and offset: {np.dtype(dtype)}")