[Mlir-commits] [mlir] 6f07191 - [mlir][sparse][taco] Support reduction to scalar tensors.

[Bixia Zheng]

Author: Bixia Zheng
Date: 2022-02-25T14:17:45-08:00
New Revision: 6f071911010323f23f45eadff0d037b1d04a8e7d

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

LOG: [mlir][sparse][taco] Support reduction to scalar tensors.

The PyTACO DSL doesn't support reduction to scalars. This change
enhances the MLIR-PyTACO implementation to support reduction to scalars.

Extend an existing test to show the syntax of reduction to scalars and
two methods to retrieve the scalar values.

Reviewed By: aartbik

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

Added: 
    

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

Removed: 
    


################################################################################
diff  --git a/mlir/test/Integration/Dialect/SparseTensor/taco/test_scalar_tensor_algebra.py b/mlir/test/Integration/Dialect/SparseTensor/taco/test_scalar_tensor_algebra.py
index d559943c2c3ec..aa77d888b28d6 100644
--- a/mlir/test/Integration/Dialect/SparseTensor/taco/test_scalar_tensor_algebra.py
+++ b/mlir/test/Integration/Dialect/SparseTensor/taco/test_scalar_tensor_algebra.py
@@ -24,5 +24,13 @@
 passed = np.array_equal(indices, [[0, 1], [1, 2]])
 passed += np.array_equal(values, [30.0, 120.0])
 
-# CHECK: Number of passed: 2
+# Sum all the values in A.
+S[0] = A[i, j]
+passed += (S.get_scalar_value() == 50.0)
+
+indices, values = S.get_coordinates_and_values()
+passed += (len(indices)==0)
+passed += (values == 50.0)
+
+# CHECK: Number of passed: 5
 print("Number of passed:", passed)

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 2fdb9f451012f..3a936e6cda50a 100644
--- a/mlir/test/Integration/Dialect/SparseTensor/taco/tools/mlir_pytaco.py
+++ b/mlir/test/Integration/Dialect/SparseTensor/taco/tools/mlir_pytaco.py
@@ -477,7 +477,7 @@ def __post_init__(self) -> None:
 
   def emit_tensor_init(self) -> ir.RankedTensorType:
     """Returns an initialization for the destination tensor."""
-    if self.dst_format is None:
+    if self.dst_format is None or self.dst_format.rank() == 0:
       # Initialize the dense tensor.
       ir_type = _mlir_type_from_taco_type(self.dst_dtype)
       tensor = linalg.InitTensorOp(self.dst_dims, ir_type).result
@@ -1023,11 +1023,26 @@ def ctype_pointer(self) -> ctypes.pointer:
 
     return ctypes.pointer(ctypes.cast(ptr, ctypes.c_void_p))
 
+  def get_scalar_value(self) -> _AnyRuntimeType:
+    """Returns the value for the scalar tensor.
+
+    This method also evaluates the assignment to the tensor.
+
+    Raises:
+      ValueError: If the tensor is not a scalar.
+    """
+    if self.order != 0:
+      raise ValueError(f"Expected a scalar tensor, got: rank={self.order}")
+
+    self._sync_value()
+    return self._dense_storage
+
+
   def get_coordinates_and_values(
       self) -> Tuple[List[Tuple[int, ...]], List[_AnyRuntimeType]]:
     """Returns the coordinates and values for the non-zero elements.
 
-    This method also evaluate the assignment to the tensor and unpack the
+    This method also evaluates the assignment to the tensor and unpack the
     sparse tensor.
     """
     self._sync_value()
@@ -1036,6 +1051,9 @@ def get_coordinates_and_values(
       self.unpack()
       return (self._coords, self._values)
 
+    if self.order == 0:
+      return ([], self._dense_storage)
+
     # Coordinates for non-zero elements, grouped by dimensions.
     coords_by_dims = self._dense_storage.nonzero()
     # Coordinates for non-zero elements, grouped by elements.