Add support for numpy arrays to memref conversions.

[pashu123]

Summary

This offers the ability to pass numpy arrays to the corresponding
memref argument.

[pashu123]

Minor Changes.

[pashu123]

Python formating.

[mehdi_amini]

mlir/lib/Bindings/Python/mlir/runtime/memref/np_to_memref.py

27 ↗

(On Diff #335948)

mlir/test/Bindings/Python/execution_engine.py

144

This will not work with a memref that has strides I believe. You'll need to use something like np.lib.stride_tricks.as_strided.
I think this should be a utility in the runtime as well memref_to_numpy_view or something like that.

(If you can add a test with an array and a view with strides that could show this)

183

Debug left-over? (same below)

[pashu123]

Clean left over.

[bondhugula]

I'd really like to see this support! Thanks for implementing this.

mlir/lib/Bindings/Python/mlir/runtime/memref/np_to_memref.py

4 ↗

(On Diff #335952)

Please add a line on what this file is about.

41 ↗

(On Diff #335952)

Doc comment.

mlir/test/Bindings/Python/execution_engine.py

136

Terminate all comments with a full stop.

144

Is there a way to restrict this patch to identity layout map? This would enable and unlock *many* things even if it just handles the default/row major identity layout. Can we check and bail out / assert on non-default strides?

191

Unrelated to this PR: If the wrong number of arguments are provided, will ExecutionEngine catch it? I think from a developer standpoint, this is quite useful since folks might sometimes lower the MLIR module to a form where there is a mismatch in the number of function arguments.

[bondhugula]

mlir/lib/Bindings/Python/mlir/runtime/memref/np_to_memref.py

47 ↗

(On Diff #335952)

get_ranked_memref_descriptor?

Here and elsewhere, can we name these things descriptively?

[mehdi_amini]

mlir/test/Bindings/Python/execution_engine.py

144

The stride is really easy to add though, maybe easier than bailing out (and adding the tests for all the bail out).

But if you prefer to add these tests, I can implement the strides in a follow up.

191

This is a layer of consistency that is important, but we'll have to build above this, these are all quite low-level APIs: the execution engine does not have any knowledge of the APIs present in the Module.

[stellaraccident]

The only comment I have (not covered elsewhere) is a weak preference to unnest by one directory level (i.e. remove the memref directory and put this in the parent). All of the exported APIs already have memref in the name and fewer imports/less repeating for these kind of things is better, imo. Also fine to see how things evolve and rework the namespace once/if there is more there.

[pashu123]

Added support for strided memrefs.

[nicolasvasilache]

mlir/lib/Bindings/Python/mlir/runtime/memref/np_to_memref.py

68 ↗

(On Diff #336556)

Note that numpy uses byte quantities to express strides.

MLIR OTOH uses the torch abstraction which specifies strides in terms of elements; the ConversionToLLVM takes care of generating the right addresses (which will also requirer furrther hooking the data layout better).

Bottom line, I experimented with memref<...xf32> by hacking the following on top of this:

+    # x.strides = nparray.ctypes.strides
+   strides_ctype_t = ctypes.c_longlong * nparray.ndim
+    x.strides = strides_ctype_t(*[t // 4 for t in nparray.strides])

It seemed to help a bit but I still saw issues that I have not yet debugged.
I am not sure whether the ctype lifetime is reasonable the code I wrote.

mlir/test/Bindings/Python/execution_engine.py

144

Also note the discrepancy between np.strides and memref / torch strides I signal above.

[pashu123]

Remove the memref directory and put files in the parent.

[nicolasvasilache]

mlir/lib/Bindings/Python/mlir/runtime/memref/np_to_memref.py

110 ↗

(On Diff #336556)

nice, you also want to pointwise-multiply by the element size in bytes.

[pashu123]

Cleaning up.

[nicolasvasilache]

mlir/lib/Bindings/Python/mlir/runtime/np_to_memref.py

75

The 4 is only valid for 4-byte type like f32, you need to get the size in bytes of the type (here and other places you updated).

[pashu123]

mlir/lib/Bindings/Python/mlir/runtime/np_to_memref.py

75

I see, thanks for pointing this out.

[pashu123]

Fix itemsize.

[pashu123]

Formatting.

[bondhugula]

mlir/lib/Bindings/Python/mlir/runtime/np_to_memref.py

75

This also means that test cases were missing - you don't seem to have f64 test cases.

mlir/test/Bindings/Python/execution_engine.py

159

You are missing test cases on float64 and so this isn't exercising all of the code well.

[nicolasvasilache]

mlir/lib/Bindings/Python/mlir/runtime/np_to_memref.py

60

This is not the element size but the displacement from the "base pointer address" at which the relevant data lives.
You probably picked the value here: https://numpy.org/doc/stable/reference/generated/numpy.ndarray.html ?
This shows that similary to strides, numpy represents this "offset in bytes" whereas MLIR is in number of elements.

I verified that setting this to 0 in my local experiment (here and below), worked in simple cases.

In the general case you need to translate between MLIR and NP offsets similarly as you do for strides.

[pashu123]

mlir/lib/Bindings/Python/mlir/runtime/np_to_memref.py

60

I tried setting this to multiple different values and it worked. If offset is represented as a number of elements in MLIR, then the offset should be 1. Similarly, it should be 1 * nparray.itemsize while converting back to NumPy from memrefs world.

[nicolasvasilache]

mlir/lib/Bindings/Python/mlir/runtime/np_to_memref.py

60

I am unclear what "it worked" means in the absence of more context on what you tried?

Please note that lowering will do the right thing in the static case here: https://github.com/llvm/llvm-project/blob/main/mlir/lib/Conversion/StandardToLLVM/StandardToLLVM.cpp#L3399.
The problem I mention is only exhibited by dynamic offsets which I do not see your tests implementing.

For instance, this commit will fail with an offset different than 0 (at runtime, in this case I hardcoded 0 for expediency while awaiting for this PR to land) : https://github.com/google/iree-llvm-sandbox/commit/87015b29c5f7cbc445bd85e1ce4a5d7597e80361.

The underlying reason is that after tiling, the offset becomes dependence on the loop IV (i.e. the memref has a static ? in its type) and then the dynamic branch of the StandardToLLVM.cpp will kick in.
Address computation will compute an offset that is 1 + actual value and will segfault at runtime.

[pashu123]

mlir/lib/Bindings/Python/mlir/runtime/np_to_memref.py

60

Thanks for pointing this out. I didn't try on Dynamic memrefs.

[pashu123]

Redirecting prints to stderr.

[mehdi_amini]

LGTM, seems like a good basis to iterate on!

[mehdi_amini]

@bondhugula @nicolasvasilache : can you confirm the recent changes address your previous comments?

[nicolasvasilache]

@mehdi_amini not yet, the author should at the very least set the 2 offsets to 0 to as described in the commit message of https://github.com/google/iree-llvm-sandbox/commit/87015b29c5f7cbc445bd85e1ce4a5d7597e80361, so that it passes without crashing.

I'm fine with iterating in-tree once that is fixed.

[mehdi_amini]

mlir/lib/Bindings/Python/mlir/runtime/np_to_memref.py

60

I don't quite indeed this x.offset = ctypes.c_longlong(4)? Shouldn't it always be zero? I suspect Numpy is baking the offset into the pointer we get with nparray.ctypes.data_as.

[nicolasvasilache]

@mehdi_amini indeed, in https://numpy.org/doc/stable/reference/generated/numpy.ndarray.html we can see:

>>> np.ndarray((2,), buffer=np.array([1,2,3]),
...            offset=np.int_().itemsize,
...            dtype=int) # offset = 1*itemsize, i.e. skip first element
array([2, 3])

as well as the default = 0 for offset:

class numpy.ndarray(shape, dtype=float, buffer=None, offset=0, strides=None, order=None)[source]¶

0 is the right value to use for now.
When we do morer fancy numpy / linalg level subviews etc it will be a different story; but for now anything else than 0 will crash non-trivial cases.

[mehdi_amini]

In D100077#2687200, @nicolasvasilache wrote:

@mehdi_amini indeed, in https://numpy.org/doc/stable/reference/generated/numpy.ndarray.html we can see:

>>> np.ndarray((2,), buffer=np.array([1,2,3]),
...            offset=np.int_().itemsize,
...            dtype=int) # offset = 1*itemsize, i.e. skip first element
array([2, 3])

Right, but the ctypes interface does not expose the offset, so I wonder if in such cases the offset that you provide on construction will not just be computed in the "data" field internally (I could check the source code for ndarray...)

[pashu123]

Changing the default value of offset to zero.

[pashu123]

In D100077#2687134, @nicolasvasilache wrote:

@mehdi_amini not yet, the author should at the very least set the 2 offsets to 0 to as described in the commit message of https://github.com/google/iree-llvm-sandbox/commit/87015b29c5f7cbc445bd85e1ce4a5d7597e80361, so that it passes without crashing.

I'm fine with iterating in-tree once that is fixed.

Done. I have changed it to zero. I had pushed the wrong changes.

[bondhugula]

In D100077#2687129, @mehdi_amini wrote:

@bondhugula @nicolasvasilache : can you confirm the recent changes address your previous comments?

Is this revision also meant to handle dynamic memrefs? If yes, a test case is missing. If not, we should check and assert (and add a TODO)?

[nicolasvasilache]

In D100077#2687846, @bondhugula wrote:

In D100077#2687129, @mehdi_amini wrote:

@bondhugula @nicolasvasilache : can you confirm the recent changes address your previous comments?

Is this revision also meant to handle dynamic memrefs? If yes, a test case is missing. If not, we should check and assert (and add a TODO)?

It supports a subset of cases (e.g. anything that takes subviews and crosses the NP -> MLIR or MLIR -> NP boundary should be considered broken atm).
+1 on adding a simple mixed static / dynamic 2-D add and compare against NP's.

[pashu123]

Adding test case of element wise addition of dynamic and static memrefs.

[nicolasvasilache]

Great, thanks for your contribution, let's land this ! :)

mlir/lib/Bindings/Python/mlir/runtime/__init__.py

+from .np_to_memref import *

mlir/lib/Bindings/Python/mlir/runtime/np_to_memref.py

+#  Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
+#  See https://llvm.org/LICENSE.txt for license information.
+#  SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+
+#  This file contains functions to convert between Memrefs and NumPy arrays and vice-versa.
+
+import numpy as np
+import ctypes
+
+def make_nd_memref_descriptor(rank, dtype):
+    class MemRefDescriptor(ctypes.Structure):
+        """
+        Build an empty descriptor for the given rank/dtype, where rank>0.
+        """
+
+        _fields_ = [
+            ("allocated", ctypes.c_longlong),
+            ("aligned", ctypes.POINTER(dtype)),
+            ("offset", ctypes.c_longlong),
+            ("shape", ctypes.c_longlong * rank),
+            ("strides", ctypes.c_longlong * rank),
+        ]
+
+    return MemRefDescriptor
+
+
+def make_zero_d_memref_descriptor(dtype):
+    class MemRefDescriptor(ctypes.Structure):
+        """
+        Build an empty descriptor for the given dtype, where rank=0.
+        """
+
+        _fields_ = [
+            ("allocated", ctypes.c_longlong),
+            ("aligned", ctypes.POINTER(dtype)),
+            ("offset", ctypes.c_longlong),
+        ]
+
+    return MemRefDescriptor
+
+
+class UnrankedMemRefDescriptor(ctypes.Structure):
+    """ Creates a ctype struct for memref descriptor"""
+
+    _fields_ = [("rank", ctypes.c_longlong), ("descriptor", ctypes.c_void_p)]
+
+
+def get_ranked_memref_descriptor(nparray):
+    """
+    Return a ranked memref descriptor for the given numpy array.
+    """
+    if nparray.ndim == 0:
+        x = make_zero_d_memref_descriptor(np.ctypeslib.as_ctypes_type(nparray.dtype))()
+        x.allocated = nparray.ctypes.data
+        x.aligned = nparray.ctypes.data_as(
+            ctypes.POINTER(np.ctypeslib.as_ctypes_type(nparray.dtype))
+        )
+        x.offset = ctypes.c_longlong(nparray.dtype.itemsize)
+        return x
+
+    x = make_nd_memref_descriptor(
+        nparray.ndim, np.ctypeslib.as_ctypes_type(nparray.dtype)
+    )()
+    x.allocated = nparray.ctypes.data
+    x.aligned = nparray.ctypes.data_as(
+        ctypes.POINTER(np.ctypeslib.as_ctypes_type(nparray.dtype))
+    )
+    x.offset = ctypes.c_longlong(nparray.dtype.itemsize)
+    x.shape = nparray.ctypes.shape
+
+    # Numpy uses byte quantities to express strides, MLIR OTOH uses the 
+    # torch abstraction which specifies strides in terms of elements.
+    strides_ctype_t = ctypes.c_longlong * nparray.ndim
+    x.strides = strides_ctype_t(*[x//nparray.itemsize for x in nparray.strides])
+    return x
+
+
+def get_unranked_memref_descriptor(nparray):
+    """
+    Return 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 unranked_memref_to_numpy(unranked_memref, np_dtype):
+    """
+    Converts unranked memrefs to numpy arrays.
+    """
+    descriptor = make_nd_memref_descriptor(
+        unranked_memref[0].rank, np.ctypeslib.as_ctypes_type(np_dtype)
+    )
+    val = ctypes.cast(unranked_memref[0].descriptor, ctypes.POINTER(descriptor))
+    np_arr = np.ctypeslib.as_array(val[0].aligned, 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 strided_arr
+
+
+def ranked_memref_to_numpy(ranked_memref):
+    """
+    Converts ranked memrefs to numpy arrays.
+    """
+    np_arr = np.ctypeslib.as_array(
+        ranked_memref[0].aligned, 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 strided_arr

mlir/test/Bindings/Python/execution_engine.py

 # RUN: %PYTHON %s 2>&1 | FileCheck %s
 
 import gc, sys
 from mlir.ir import *
 from mlir.passmanager import *
 from mlir.execution_engine import *
+from mlir.runtime import *
 
 # Log everything to stderr and flush so that we have a unified stream to match
 # errors/info emitted by MLIR to stderr.
 def log(*args):
   print(*args, file=sys.stderr)
   sys.stderr.flush()
 
 def run(f):
func private @some_callback_into_python(f32, i32) -> f32 attributes { llvm.emit_c_interface }
     arg0 = c_float_p(42.)
     arg1 = c_int_p(2)
     res = c_float_p(-1.)
     execution_engine.invoke("add", arg0, arg1, res)
     # CHECK: 42.0 + 2 = 44.0
     log("{0} + {1} = {2}".format(arg0[0], arg1[0], res[0]*2))
 
 run(testBasicCallback)
+
+# Test callback with an unranked memref
+# CHECK-LABEL: TEST: testUnrankedMemRefCallback
+def testUnrankedMemRefCallback():
+    # 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)
+        print("Inside callback: ")
+        print(arr)
+
+    with Context():
+        # The module just forwards to a runtime function known as "some_callback_into_python".
+        module = Module.parse(
+            r"""
+func @callback_memref(%arg0: memref<*xf32>) attributes { llvm.emit_c_interface } {
+  call @some_callback_into_python(%arg0) : (memref<*xf32>) -> ()
+  return
+}
+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.0, 2.0], [3.0, 4.0]], np.float32)
+        # CHECK: Inside callback: 
+        # CHECK{LITERAL}: [[1. 2.]
+        # CHECK{LITERAL}:  [3. 4.]]
+        execution_engine.invoke(
+            "callback_memref",
+            ctypes.pointer(ctypes.pointer(get_unranked_memref_descriptor(inp_arr))),
+        )
+        inp_arr_1 = np.array([5, 6, 7], dtype=np.float32)
+        strided_arr = np.lib.stride_tricks.as_strided(
+            inp_arr_1, strides=(4, 0), shape=(3, 4)
+        )
+        # CHECK: Inside callback: 
+        # CHECK{LITERAL}: [[5. 5. 5. 5.]
+        # CHECK{LITERAL}:  [6. 6. 6. 6.]
+        # CHECK{LITERAL}:  [7. 7. 7. 7.]]
+        execution_engine.invoke(
+            "callback_memref",
+            ctypes.pointer(
+                ctypes.pointer(get_unranked_memref_descriptor(strided_arr))
+            ),
+        )
+
+run(testUnrankedMemRefCallback)
+
+# Test callback with a ranked memref.
+# CHECK-LABEL: TEST: testRankedMemRefCallback
+def testRankedMemRefCallback():
+    # 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(2, np.ctypeslib.as_ctypes_type(np.float32))
+        ),
+    )
+    def callback(a):
+        arr = ranked_memref_to_numpy(a)
+        print("Inside Callback: ")
+        print(arr)
+
+    with Context():
+        # The module just forwards to a runtime function known as "some_callback_into_python".
+        module = Module.parse(
+            r"""
+func @callback_memref(%arg0: memref<2x2xf32>) attributes { llvm.emit_c_interface } {
+  call @some_callback_into_python(%arg0) : (memref<2x2xf32>) -> ()
+  return
+}
+func private @some_callback_into_python(memref<2x2xf32>) -> () attributes { llvm.emit_c_interface }
+"""
+        )
+        execution_engine = ExecutionEngine(lowerToLLVM(module))
+        execution_engine.register_runtime("some_callback_into_python", callback)
+        inp_arr = np.array([[1.0, 5.0], [6.0, 7.0]], np.float32)
+        # CHECK: Inside Callback:
+        # CHECK{LITERAL}: [[1. 5.]
+        # CHECK{LITERAL}:  [6. 7.]]
+        execution_engine.invoke(
+            "callback_memref", ctypes.pointer(ctypes.pointer(get_ranked_memref_descriptor(inp_arr)))
+        )
+
+run(testRankedMemRefCallback)
+
+#  Test addition of two memref
+# CHECK-LABEL: TEST: testMemrefAdd
+def testMemrefAdd():
+    with Context():
+        module = Module.parse(
+            """
+      module  {
+      func @main(%arg0: memref<1xf32>, %arg1: memref<f32>, %arg2: memref<1xf32>) attributes { llvm.emit_c_interface } {
+        %0 = constant 0 : index
+        %1 = memref.load %arg0[%0] : memref<1xf32>
+        %2 = memref.load %arg1[] : memref<f32>
+        %3 = addf %1, %2 : f32
+        memref.store %3, %arg2[%0] : memref<1xf32>
+        return
+      }
+     } """
+        )
+        arg1 = np.array([32.5]).astype(np.float32)
+        arg2 = np.array(6).astype(np.float32)
+        res = np.array([0]).astype(np.float32)
+
+        arg1_memref_ptr = ctypes.pointer(ctypes.pointer(get_ranked_memref_descriptor(arg1)))
+        arg2_memref_ptr = ctypes.pointer(ctypes.pointer(get_ranked_memref_descriptor(arg2)))
+        res_memref_ptr = ctypes.pointer(ctypes.pointer(get_ranked_memref_descriptor(res)))
+
+        execution_engine = ExecutionEngine(lowerToLLVM(module))
+        execution_engine.invoke(
+            "main", arg1_memref_ptr, arg2_memref_ptr, res_memref_ptr
+        )
+        # CHECK: [32.5] + 6.0 = [38.5]
+        log("{0} + {1} = {2}".format(arg1, arg2, res))
+
+run(testMemrefAdd)