[mlir][sparse][taco] Add support for scalar tensors.

[bixia]

Summary

This change allows the use of scalar tensors with index 0 in tensor index
expressions. In this case, the scalar value is broadcast to match the
dimensions of other tensors in the same expression.

Using scalar tensors as a destination in tensor index expressions is not
supported in the PyTACO DSL.

Add a PyTACO test to show the use of scalar tensors.

[aartbik]

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

25

no allclose for indices, they must be exact!

26

even these seem to be simple enough to check for equality

[bixia]

Use np.array_equal to compare indices.

[bixia]

Use np.array_equal to check the result.

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

+# RUN: SUPPORTLIB=%mlir_runner_utils_dir/libmlir_c_runner_utils%shlibext %PYTHON %s | FileCheck %s
+
+import numpy as np
+import os
+import sys
+
+_SCRIPT_PATH = os.path.dirname(os.path.abspath(__file__))
+sys.path.append(_SCRIPT_PATH)
+from tools import mlir_pytaco_api as pt
+
+compressed = pt.compressed
+
+i, j = pt.get_index_vars(2)
+A = pt.tensor([2, 3])
+S = pt.tensor(3) # S is a scalar tensor.
+B = pt.tensor([2, 3], compressed)
+A.insert([0, 1], 10)
+A.insert([1, 2], 40)
+
+# 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.allclose(values, [30.0, 120.0])
+
+# CHECK: Number of passed: 2
+print("Number of passed:", passed)

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

def _mlir_tensor_type(
 
   Returns:
     An MLIR ranked tensor type.
   """
   ir_type = _mlir_type_from_taco_type(dtype)
   return ir.RankedTensorType.get(shape, ir_type, attr)
 
 
-def _verify_and_normalize_indices(indices) -> Tuple[IndexVar, ...]:
-  """Verifies and normalizes the indices for a tensor access.
-
-  Args:
-    indices: The index expression used to access a tensor, which could be any
-      Python object from user inputs.
-
-  Returns:
-    A tuple of IndexVar.
-
-  Raises:
-    ValueError: If indices is not an IndexVar or a tuple of IndexVar.
-  """
-  if isinstance(indices, IndexVar):
-    return (indices,)
-  elif isinstance(indices, tuple) and _all_instance_of(indices, IndexVar):
-    return indices
-
-  raise ValueError(f"Expected IndexVars: {indices}")
-
-
 @dataclasses.dataclass(frozen=True)
 class _StructOpInfo:
   """Information for generating a structured op in the linalg dialect.
 
   This information is associated with an expression node that serves as the
   root for an expression subtree implemented with a structured op.
 
   Attributes:
def insert(self, coords: List[int], val: Union[float, int]) -> None:
     if not isinstance(val, int) and not isinstance(val, float):
       raise ValueError(f"Value is neither int nor float: {val}.")
 
     self._coords.append(tuple(coords))
     self._values.append(self._dtype.value(val))
 
   def is_dense(self) -> bool:
     """Returns true if the tensor doesn't have sparsity annotation."""
-    return self._format is None
+    return self.order == 0 or self._format is None
 
   def to_array(self) -> np.ndarray:
     """Returns the numpy array for the Tensor.
 
     This is currenly only implemented for dense Tensor.
     """
     if not self.is_dense():
       raise ValueError("Conversion from non-dense Tensor "
def order(self) -> int:
     """Returns the rank of the Tensor."""
     return len(self._shape)
 
   @property
   def shape(self) -> List[int]:
     """Returns the shape of the Tensor."""
     return self._shape
 
+  def _verify_and_normalize_indices(self, indices) -> Tuple[IndexVar, ...]:
+    """Verifies and normalizes the indices to access the tensor.
+
+    Args:
+      indices: The index expression used to access a tensor, which could be any
+        Python object from user inputs.
+
+    Returns:
+      A tuple of IndexVar.
+
+    Raises:
+      ValueError: If indices is not 0 for scalar tensors, or not an IndexVar or
+        a tuple of IndexVar for other tensors.
+    """
+    if self.order == 0:
+      if not isinstance(indices, int) or indices != 0:
+        raise ValueError(f"Expected 0 to index scalar tensors: {indices}")
+      return ()
+
+    if isinstance(indices, IndexVar):
+      return (indices,)
+    elif isinstance(indices, tuple) and _all_instance_of(indices, IndexVar):
+      return indices
+
+    raise ValueError(f"Expected IndexVars: {indices}")
+
   def __getitem__(self, key) -> "Access":
     """Verifies and processes a tensor access.
 
     In the tensor index notation, a tensor access T[i, j] is represented as
     retrieving a value with key (i, j) from the tensor object T in Python. This
     routine verifies the key for the tensor access and returns a tensor access
     object.
 
     Args:
       key: The key used to access the tensor, which could be any Python object
         from user inputs.
 
     Returns:
       The corresponding tensor access object.
 
     Raises:
       ValueError: If key is not an IndexVar or a tuple of IndexVar.
     """
-    indices = _verify_and_normalize_indices(key)
+    indices = self._verify_and_normalize_indices(key)
     return Access(self, indices)
 
   def __setitem__(self, key, value) -> None:
     """Verifies and processes a tensor assignment.
 
     In the tensor index notation, a tensor assignment "T[i, j] = ..." is
     represented as setting a value for a tensor object T via key (i, j) in
     Python. This routine verifies the key, evaluates the value, and assigns the
     value to the tensor.
 
     We only support assignment of dense tensor currently.
 
     Args:
       key: The key used to access the tensor, which could be any Python object
         from user inputs.
       value: The value assigned to the tensor, which could be any Python object
         from user inputs.
 
     Raises:
       ValueError: If tensor is not a dense tensor, or the key is not an IndexVar
         or a tuple of IndexVar, or the length of the indices is not the same as
         the rank of the tensor.
     """
-    indices = _verify_and_normalize_indices(key)
+    indices = self._verify_and_normalize_indices(key)
     if len(indices) != self.order:
       raise ValueError("Mismatch between indices and tensor rank: "
                        f"len({indices}) != {self.order}.")
 
     self._assignment = _Assignment(indices, value)
 
   def evaluate(self) -> None:
     """Evaluates the assignment to the tensor."""
     result = self._assignment.expression.evaluate(self,
                                                   self._assignment.indices)
     self._assignment = None
     if self.is_dense():
       assert isinstance(result, np.ndarray)
       self._dense_storage = result
     else:
       self._set_packed_sparse_tensor(result)
 
   def _sync_value(self) -> None:
     """Updates the tensor value by evaluating the pending assignment."""
     if self._assignment is not None:
       self.evaluate()
 
   def mlir_tensor_type(self) -> ir.RankedTensorType:
     """Returns the MLIR type for the tensor."""
-    mlir_attr = None if (
-        self._format is None) else self._format.mlir_tensor_attr()
+    mlir_attr = (None if (self._format is None or self.order == 0) else
+                 self._format.mlir_tensor_attr())
     return _mlir_tensor_type(self._dtype, tuple(self._shape), mlir_attr)
 
   def dense_dst_ctype_pointer(self) -> ctypes.pointer:
     """Returns the ctypes pointer for the pointer to an MemRefDescriptor.
 
     For a dense tensor output, the MLIR compiler allocates the storage for
     the tensor. This routine returns the pointer to an MLIR MemRefDescriptor for
     receiving the tensor.