[mlir][Python] Fix conversion of non-zero offset memrefs to np.arrays

[ubfx]

Summary

Memref descriptors contain an offset field that denotes the start of
the content of the memref relative to the alignedPtr. This offset is
not considered when converting a memref descriptor to a np.array in the
Python runtime library, essentially treating all memrefs as if they had
an offset of zero. This patch introduces the necessary pointer arithmetic
to find the actual beginning of the memref contents to the memref->numpy
conversion functions.

There is an ongoing discussion about whether the offset field is needed
at all in the memref descriptor.
Until that is decided, the Python runtime and CRunnerUtils should
still correctly implement the offset handling.

Related: https://reviews.llvm.org/D157008

[ubfx]

Related forum thread: https://discourse.llvm.org/t/setting-memref-elements-in-python-callback/

[ubfx]

Thank you for the review @ftynse, would you mind landing this revision for me?

mlir/python/mlir/runtime/np_to_memref.py

 def get_unranked_memref_descriptor(nparray):
     """Returns a generic/unranked memref descriptor for the given numpy array."""
     d = UnrankedMemRefDescriptor()
     d.rank = nparray.ndim
     x = get_ranked_memref_descriptor(nparray)
     d.descriptor = ctypes.cast(ctypes.pointer(x), ctypes.c_void_p)
     return d
 
+def move_aligned_ptr_by_offset(aligned_ptr, offset):
+    """Moves the supplied ctypes pointer ahead by `offset` elements."""
+    aligned_addr = ctypes.addressof(aligned_ptr.contents)
+    elem_size = ctypes.sizeof(aligned_ptr.contents)
+    shift = offset * elem_size
+    content_ptr = ctypes.cast(aligned_addr + shift, type(aligned_ptr))
+    return content_ptr
 
 def unranked_memref_to_numpy(unranked_memref, np_dtype):
     """Converts unranked memrefs to numpy arrays."""
     ctp = as_ctype(np_dtype)
     descriptor = make_nd_memref_descriptor(unranked_memref[0].rank, ctp)
     val = ctypes.cast(unranked_memref[0].descriptor, ctypes.POINTER(descriptor))
-    np_arr = np.ctypeslib.as_array(val[0].aligned, shape=val[0].shape)
+    content_ptr = move_aligned_ptr_by_offset(val[0].aligned, val[0].offset)
+    np_arr = np.ctypeslib.as_array(content_ptr, shape=val[0].shape)
     strided_arr = np.lib.stride_tricks.as_strided(
         np_arr,
         np.ctypeslib.as_array(val[0].shape),
         np.ctypeslib.as_array(val[0].strides) * np_arr.itemsize,
     )
     return to_numpy(strided_arr)
 
 
 def ranked_memref_to_numpy(ranked_memref):
     """Converts ranked memrefs to numpy arrays."""
+    content_ptr = move_aligned_ptr_by_offset(ranked_memref[0].aligned, ranked_memref[0].offset)
     np_arr = np.ctypeslib.as_array(
-        ranked_memref[0].aligned, shape=ranked_memref[0].shape
+        content_ptr, shape=ranked_memref[0].shape
     )
     strided_arr = np.lib.stride_tricks.as_strided(
         np_arr,
         np.ctypeslib.as_array(ranked_memref[0].shape),
         np.ctypeslib.as_array(ranked_memref[0].strides) * np_arr.itemsize,
     )
     return to_numpy(strided_arr)

mlir/test/python/execution_engine.py

"""
             "callback_memref",
             ctypes.pointer(ctypes.pointer(get_ranked_memref_descriptor(inp_arr))),
         )
 
 
 run(testRankedMemRefCallback)
 
 
+# Test callback with a ranked memref with non-zero offset.
+# CHECK-LABEL: TEST: testRankedMemRefWithOffsetCallback
+def testRankedMemRefWithOffsetCallback():
+    # Define a callback function that takes a ranked memref, converts it to a numpy array and prints it.
+    @ctypes.CFUNCTYPE(
+        None,
+        ctypes.POINTER(
+            make_nd_memref_descriptor(1, np.ctypeslib.as_ctypes_type(np.float32))
+        ),
+    )
+    def callback(a):
+        arr = ranked_memref_to_numpy(a)
+        log("Inside Callback: ")
+        log(arr)
+
+    with Context():
+        # The module takes a subview of the argument memref and calls the callback with it
+        module = Module.parse(
+            r"""
+func.func @callback_memref(%arg0: memref<5xf32>) attributes {llvm.emit_c_interface} {
+  %base_buffer, %offset, %sizes, %strides = memref.extract_strided_metadata %arg0 : memref<5xf32> -> memref<f32>, index, index, index
+  %reinterpret_cast = memref.reinterpret_cast %base_buffer to offset: [3], sizes: [2], strides: [1] : memref<f32> to memref<2xf32, strided<[1], offset: 3>>
+  %cast = memref.cast %reinterpret_cast : memref<2xf32, strided<[1], offset: 3>> to memref<?xf32, strided<[?], offset: ?>>
+  call @some_callback_into_python(%cast) : (memref<?xf32, strided<[?], offset: ?>>) -> ()
+  return
+}
+func.func private @some_callback_into_python(memref<?xf32, strided<[?], offset: ?>>) attributes {llvm.emit_c_interface}
+"""
+        )
+        execution_engine = ExecutionEngine(lowerToLLVM(module))
+        execution_engine.register_runtime("some_callback_into_python", callback)
+        inp_arr = np.array([0, 0, 0, 1, 2], np.float32)
+        # CHECK: Inside Callback:
+        # CHECK{LITERAL}: [1. 2.]
+        execution_engine.invoke(
+            "callback_memref",
+            ctypes.pointer(ctypes.pointer(get_ranked_memref_descriptor(inp_arr))),
+        )
+
+
+run(testRankedMemRefWithOffsetCallback)
+
+
+# Test callback with an unranked memref with non-zero offset
+# CHECK-LABEL: TEST: testUnrankedMemRefWithOffsetCallback
+def testUnrankedMemRefWithOffsetCallback():
+    # Define a callback function that takes an unranked memref, converts it to a numpy array and prints it.
+    @ctypes.CFUNCTYPE(None, ctypes.POINTER(UnrankedMemRefDescriptor))
+    def callback(a):
+        arr = unranked_memref_to_numpy(a, np.float32)
+        log("Inside callback: ")
+        log(arr)
+
+    with Context():
+        # The module takes a subview of the argument memref, casts it to an unranked memref and 
+        # calls the callback with it.
+        module = Module.parse(
+            r"""
+func.func @callback_memref(%arg0: memref<5xf32>) attributes {llvm.emit_c_interface} {
+    %base_buffer, %offset, %sizes, %strides = memref.extract_strided_metadata %arg0 : memref<5xf32> -> memref<f32>, index, index, index
+    %reinterpret_cast = memref.reinterpret_cast %base_buffer to offset: [3], sizes: [2], strides: [1] : memref<f32> to memref<2xf32, strided<[1], offset: 3>>
+    %cast = memref.cast %reinterpret_cast : memref<2xf32, strided<[1], offset: 3>> to memref<*xf32>
+    call @some_callback_into_python(%cast) : (memref<*xf32>) -> ()
+    return
+}
+func.func private @some_callback_into_python(memref<*xf32>) attributes {llvm.emit_c_interface}
+"""
+        )
+        execution_engine = ExecutionEngine(lowerToLLVM(module))
+        execution_engine.register_runtime("some_callback_into_python", callback)
+        inp_arr = np.array([1, 2, 3, 4, 5], np.float32)
+        # CHECK: Inside callback:
+        # CHECK{LITERAL}: [4. 5.]
+        execution_engine.invoke(
+            "callback_memref",
+            ctypes.pointer(ctypes.pointer(get_ranked_memref_descriptor(inp_arr))),
+        )
+
+run(testUnrankedMemRefWithOffsetCallback)
+
+
 #  Test addition of two memrefs.
 # CHECK-LABEL: TEST: testMemrefAdd
 def testMemrefAdd():
     with Context():
         module = Module.parse(
             """
     module  {
       func.func @main(%arg0: memref<1xf32>, %arg1: memref<f32>, %arg2: memref<1xf32>) attributes { llvm.emit_c_interface } {