mlir/test/Integration/Dialect/SparseTensor/taco/tools/mlir_pytaco_utils.py

   1 #  Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
   2 #  See https://llvm.org/LICENSE.txt for license information.
   3 #  SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
   4
   5 # This file contains the utilities to process sparse tensor outputs.
   6
   7 from typing import Tuple
   8 import ctypes
   9 import functools
  10 import numpy as np
  11
  12
  13 @functools.lru_cache()
  14 def _get_c_shared_lib(lib_name: str) -> ctypes.CDLL:
  15   """Loads and returns the requested C shared library.
  16
  17   Args:
  18     lib_name: A string representing the C shared library.
  19
  20   Returns:
  21     The C shared library.
  22
  23   Raises:
  24     OSError: If there is any problem in loading the shared library.
  25     ValueError: If the shared library doesn't contain the needed routines.
  26   """
  27   # This raises OSError exception if there is any problem in loading the shared
  28   # library.
  29   c_lib = ctypes.CDLL(lib_name)
  30
  31   try:
  32     c_lib.convertToMLIRSparseTensor.restype = ctypes.c_void_p
  33   except Exception as e:
  34     raise ValueError("Missing function convertToMLIRSparseTensor from "
  35                      f"the supporting C shared library: {e} ") from e
  36
  37   try:
  38     c_lib.convertFromMLIRSparseTensor.restype = ctypes.c_void_p
  39   except Exception as e:
  40     raise ValueError("Missing function convertFromMLIRSparseTensor from "
  41                      f"the C shared library: {e} ") from e
  42
  43   return c_lib
  44
  45
  46 def sparse_tensor_to_coo_tensor(
  47     lib_name: str,
  48     sparse_tensor: ctypes.c_void_p,
  49     dtype: np.dtype,
  50 ) -> Tuple[int, int, np.ndarray, np.ndarray, np.ndarray]:
  51   """Converts an MLIR sparse tensor to a COO-flavored format tensor.
  52
  53   Args:
  54      lib_name: A string for the supporting C shared library.
  55      sparse_tensor: A ctypes.c_void_p to the MLIR sparse tensor descriptor.
  56      dtype: The numpy data type for the tensor elements.
  57
  58   Returns:
  59     A tuple that contains the following values for the COO-flavored format
  60     tensor:
  61     rank: An integer for the rank of the tensor.
  62     nse: An interger for the number of non-zero values in the tensor.
  63     shape: A 1D numpy array of integers, for the shape of the tensor.
  64     values: A 1D numpy array, for the non-zero values in the tensor.
  65     indices: A 2D numpy array of integers, representing the indices for the
  66       non-zero values in the tensor.
  67
  68   Raises:
  69     OSError: If there is any problem in loading the shared library.
  70     ValueError: If the shared library doesn't contain the needed routines.
  71   """
  72   c_lib = _get_c_shared_lib(lib_name)
  73
  74   rank = ctypes.c_ulonglong(0)
  75   nse = ctypes.c_ulonglong(0)
  76   shape = ctypes.POINTER(ctypes.c_ulonglong)()
  77   values = ctypes.POINTER(np.ctypeslib.as_ctypes_type(dtype))()
  78   indices = ctypes.POINTER(ctypes.c_ulonglong)()
  79   c_lib.convertFromMLIRSparseTensor(sparse_tensor, ctypes.byref(rank),
  80                                     ctypes.byref(nse), ctypes.byref(shape),
  81                                     ctypes.byref(values), ctypes.byref(indices))
  82
  83   # Convert the returned values to the corresponding numpy types.
  84   shape = np.ctypeslib.as_array(shape, shape=[rank.value])
  85   values = np.ctypeslib.as_array(values, shape=[nse.value])
  86   indices = np.ctypeslib.as_array(indices, shape=[nse.value, rank.value])
  87   return rank, nse, shape, values, indices
  88
  89
  90 def coo_tensor_to_sparse_tensor(lib_name: str, np_shape: np.ndarray,
  91                                 np_values: np.ndarray,
  92                                 np_indices: np.ndarray) -> int:
  93   """Converts a COO-flavored format sparse tensor to an MLIR sparse tensor.
  94
  95   Args:
  96      lib_name: A string for the supporting C shared library.
  97      np_shape: A 1D numpy array of integers, for the shape of the tensor.
  98      np_values: A 1D numpy array, for the non-zero values in the tensor.
  99      np_indices: A 2D numpy array of integers, representing the indices for the
 100        non-zero values in the tensor.
 101
 102   Returns:
 103      An integer for the non-null ctypes.c_void_p to the MLIR sparse tensor
 104      descriptor.
 105
 106   Raises:
 107     OSError: If there is any problem in loading the shared library.
 108     ValueError: If the shared library doesn't contain the needed routines.
 109   """
 110
 111   rank = ctypes.c_ulonglong(len(np_shape))
 112   nse = ctypes.c_ulonglong(len(np_values))
 113   shape = np_shape.ctypes.data_as(ctypes.POINTER(ctypes.c_ulonglong))
 114   values = np_values.ctypes.data_as(
 115       ctypes.POINTER(np.ctypeslib.as_ctypes_type(np_values.dtype)))
 116   indices = np_indices.ctypes.data_as(ctypes.POINTER(ctypes.c_ulonglong))
 117
 118   c_lib = _get_c_shared_lib(lib_name)
 119   ptr = c_lib.convertToMLIRSparseTensor(rank, nse, shape, values, indices)
 120   assert ptr is not None, "Problem with calling convertToMLIRSparseTensor"
 121   return ptr