| # 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(0) |
| 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(0) |
| 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 |
