[mlir][sparse][taco] Support tensor dimension storage ordering and more general sparsity values.

[bixia]

Summary

Previously, we can't properly handle input tensors with a dimension
ordering that is different from the natural ordering or with a mixed of
compressed and dense dimensions. This change fixes the problems by
passing the dimension ordering and sparsity values to the runtime
routine.

Modify an existing test to test the situation.

mlir/test/Integration/Dialect/SparseTensor/taco/test_Tensor.py

 from tools import mlir_pytaco_api as pt
 from tools import testing_utils as utils
 
 i, j, k = pt.get_index_vars(3)
 
 # Set up scalar and sparse tensors.
 alpha = pt.tensor(42.0)
 S = pt.tensor([8, 8, 8],
-              pt.format([pt.compressed, pt.compressed, pt.compressed]))
+              pt.format([pt.compressed, pt.dense, pt.compressed], [1, 0, 2]))
 X = pt.tensor([8, 8, 8],
-              pt.format([pt.compressed, pt.compressed, pt.compressed]))
+              pt.format([pt.compressed, pt.compressed, pt.compressed],
+                        [1, 0, 2]))
 S.insert([0, 0, 0], 2.0)
 S.insert([1, 1, 1], 3.0)
 S.insert([4, 4, 4], 4.0)
 S.insert([7, 7, 7], 5.0)
 
 X[i, j, k] = alpha[0] * S[i, j, k]
 
 expected = """; extended FROSTT format

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

if self.format_pack.rank() != self.ordering.rank():
       raise ValueError("Inconsistent ModeFormatPack and ModeOrdering: "
                        f"len({self.format_pack}) != "
                        f"len({self.ordering})")
 
   def rank(self) -> int:
     """Returns the number of dimensions represented by the format."""
     return self.format_pack.rank()
 
+  def get_permutation_and_sparsity(self) -> Tuple[np.ndarray, np.ndarray]:
+    """Constructs the numpy arrays for the permutation and sparsity."""
+    perm = np.array(self.ordering.ordering, dtype=np.ulonglong)
+    a = [0 if s == ModeFormat.DENSE else 1 for s in self.format_pack.formats]
+    sparse = np.array(a, dtype=np.uint8)
+    return (perm, sparse)
+
   def mlir_tensor_attr(self) -> Optional[sparse_tensor.EncodingAttr]:
     """Constructs the MLIR attributes for the tensor format."""
     order = (
         range(self.rank()) if
         (self.ordering is None) else self.ordering.ordering)
     mlir_storage_format = [f.value for f in self.format_pack.formats]
     return sparse_tensor.EncodingAttr.get(mlir_storage_format,
                                           ir.AffineMap.get_permutation(order),
if self.is_dense():
       if self._dense_storage is None:
         self._dense_storage = np.zeros(self._shape, self._dtype.value)
       return _ctype_pointer_from_array(self._dense_storage)
 
     if self.is_unpacked():
       shape = np.array(self._shape, np.int64)
       indices = np.array(self._coords, np.int64)
       values = np.array(self._values, self._dtype.value)
-      ptr = utils.coo_tensor_to_sparse_tensor(shape, values, indices)
+      perm, sparse = self.format.get_permutation_and_sparsity()
+      ptr = utils.coo_tensor_to_sparse_tensor(shape, values, indices, perm,
+                                              sparse)
     else:
       ptr = self._packed_sparse_value
 
     return ctypes.pointer(ctypes.cast(ptr, ctypes.c_void_p))
 
   def get_scalar_value(self) -> _AnyRuntimeType:
     """Returns the value for the scalar tensor.
 

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

) -> Tuple[int, int, np.ndarray, np.ndarray, np.ndarray]:
   # Convert the returned values to the corresponding numpy types.
   shape = np.ctypeslib.as_array(shape, shape=[rank.value])
   values = np.ctypeslib.as_array(values, shape=[nse.value])
   indices = np.ctypeslib.as_array(indices, shape=[nse.value, rank.value])
   return rank.value, nse.value, shape, values, indices
 
 
 def coo_tensor_to_sparse_tensor(np_shape: np.ndarray, np_values: np.ndarray,
-                                np_indices: np.ndarray) -> int:
+                                np_indices: np.ndarray, np_perm: np.ndarray,
+                                np_sparse: np.ndarray) -> int:
   """Converts a COO-flavored format sparse tensor to an MLIR sparse tensor.
 
   Args:
      np_shape: A 1D numpy array of integers, for the shape of the tensor.
      np_values: A 1D numpy array, for the non-zero values in the tensor.
      np_indices: A 2D numpy array of integers, representing the indices for the
        non-zero values in the tensor.
+     np_perm: A 1D numpy array of integers, representing the storage ordering
+       for the dimensions.
+     np_sparse: A 1D numpy array of uint8, representing the sparsity values
+       for the dimensions.
 
   Returns:
      An integer for the non-null ctypes.c_void_p to the MLIR sparse tensor
      descriptor.
 
   Raises:
     OSError: If there is any problem in loading the shared library.
     ValueError: If the shared library doesn't contain the needed routines.
   """
 
   r = len(np_shape)
   rank = ctypes.c_ulonglong(r)
   nse = ctypes.c_ulonglong(len(np_values))
   shape = np_shape.ctypes.data_as(ctypes.POINTER(ctypes.c_ulonglong))
   values = np_values.ctypes.data_as(
       ctypes.POINTER(np.ctypeslib.as_ctypes_type(np_values.dtype)))
   indices = np_indices.ctypes.data_as(ctypes.POINTER(ctypes.c_ulonglong))
 
-  np_perm = np.arange(r, dtype=np.ulonglong)
-  np_sparse = np.full(r, 1, dtype=np.uint8)
   perm = np_perm.ctypes.data_as(ctypes.POINTER(ctypes.c_ulonglong))
   sparse = np_sparse.ctypes.data_as(ctypes.POINTER(ctypes.c_uint8))
 
   convert_to = _get_support_func_locator()(np_values.dtype.type)[0]
   ptr = convert_to(rank, nse, shape, values, indices, perm, sparse)
   assert ptr is not None, "Problem with calling convertToMLIRSparseTensorF64"
   return ptr