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

[bixia]

Summary

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.

[aartbik]

Nice!

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

 
 # Use [0] to index the scalar tensor.
 B[i, j] = A[i, j] * S[0]
 
 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)

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

class _StructOpInfo:
   dst_format: Optional[Format]
 
   def __post_init__(self) -> None:
     """Verifies the integrity of the attribute values."""
     assert len(self.dst_indices) == len(self.dst_dims)
 
   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
       zero = arith.ConstantOp(ir_type, 0.0)
       return linalg.FillOp(output=tensor, value=zero).results[0]
 
     # Initialize the sparse tensor.
     mlir_type = _mlir_tensor_type(self.dst_dtype, self.dst_dims,
if self.is_unpacked():
       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)
     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.
+
+    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()
 
     if not self.is_dense():
       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.
     coords = np.transpose(coords_by_dims)
     values = self._dense_storage[coords_by_dims]
     return (coords, values)
 
   def _record_stats(self, structop: "_StructOpInfo"):