[Mlir-commits] [mlir] c25f3df - [mlir][sparse][taco] Support tensor dimension storage ordering and more general

[Bixia Zheng]

Author: Bixia Zheng
Date: 2022-03-01T15:36:38-08:00
New Revision: c25f3dfff3978db8c7728e094e2dd61ab8db5ec5

URL: https://github.com/llvm/llvm-project/commit/c25f3dfff3978db8c7728e094e2dd61ab8db5ec5
DIFF: https://github.com/llvm/llvm-project/commit/c25f3dfff3978db8c7728e094e2dd61ab8db5ec5.diff

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

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.

Reviewed By: aartbik

Differential Revision: https://reviews.llvm.org/D120777

Added: 
    

Modified: 
    mlir/test/Integration/Dialect/SparseTensor/taco/test_Tensor.py
    mlir/test/Integration/Dialect/SparseTensor/taco/tools/mlir_pytaco.py
    mlir/test/Integration/Dialect/SparseTensor/taco/tools/mlir_pytaco_utils.py

Removed: 
    


################################################################################
diff  --git a/mlir/test/Integration/Dialect/SparseTensor/taco/test_Tensor.py b/mlir/test/Integration/Dialect/SparseTensor/taco/test_Tensor.py
index 5b8b687562de7..59483fa16919e 100644
--- a/mlir/test/Integration/Dialect/SparseTensor/taco/test_Tensor.py
+++ b/mlir/test/Integration/Dialect/SparseTensor/taco/test_Tensor.py
@@ -17,9 +17,10 @@
 # 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)

diff  --git a/mlir/test/Integration/Dialect/SparseTensor/taco/tools/mlir_pytaco.py b/mlir/test/Integration/Dialect/SparseTensor/taco/tools/mlir_pytaco.py
index e93c9343dd497..da521bd33504e 100644
--- a/mlir/test/Integration/Dialect/SparseTensor/taco/tools/mlir_pytaco.py
+++ b/mlir/test/Integration/Dialect/SparseTensor/taco/tools/mlir_pytaco.py
@@ -325,6 +325,13 @@ 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 = (
@@ -1017,7 +1024,9 @@ def ctype_pointer(self) -> ctypes.pointer:
       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
 

diff  --git a/mlir/test/Integration/Dialect/SparseTensor/taco/tools/mlir_pytaco_utils.py b/mlir/test/Integration/Dialect/SparseTensor/taco/tools/mlir_pytaco_utils.py
index 90a12cc329dd6..4e2e8ba43930a 100644
--- a/mlir/test/Integration/Dialect/SparseTensor/taco/tools/mlir_pytaco_utils.py
+++ b/mlir/test/Integration/Dialect/SparseTensor/taco/tools/mlir_pytaco_utils.py
@@ -133,7 +133,8 @@ def sparse_tensor_to_coo_tensor(
 
 
 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:
@@ -141,6 +142,10 @@ def coo_tensor_to_sparse_tensor(np_shape: np.ndarray, np_values: np.ndarray,
      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
@@ -159,8 +164,6 @@ def coo_tensor_to_sparse_tensor(np_shape: np.ndarray, np_values: np.ndarray,
       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))