[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

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 d559943..aa77d88 100644 (file)
--- 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 @@ indices, values = B.get_coordinates_and_values()
 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 2fdb9f4..3a936e6 100644 (file)
--- 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 @@ class _StructOpInfo:
 
   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 @@ class Tensor:
 
     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 @@ class Tensor:
       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.