[mlir][sparse][taco] Support the use of index values in tensor expressions.

[bixia]

Summary

PyTACO DSL doesn't support the use of index values as in A[i] = B[i]+ i.
We extend the DSL to support such a use in MLIR-PyTACO.

Remove an obsolete unit test. Add unit tests and PyTACO tests.

[aartbik]

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

46

how about one test with a "conjunction" as well, so it is only done for the nonzeros?

H[i] = F[i] * i

[bixia]

Add a test for conjunction.

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

46

Good idea. Added

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

 C.insert([0, 1], 5)
 C.insert([1, 2], 7)
 D[i, j] = A[i, j] + B[i, j] - C[i, j]
 
 indices, values = D.get_coordinates_and_values()
 passed = np.array_equal(indices, [[0, 0], [0, 1], [1, 2]])
 passed += np.allclose(values, [20.0, 5.0, 63.0])
 
-# CHECK: Number of passed: 2
+# PyTACO doesn't allow the use of index values, but MLIR-PyTACO removes this
+# restriction.
+E = pt.tensor([3])
+E[i] = i
+indices, values = E.get_coordinates_and_values()
+passed += np.array_equal(indices, [[0], [1], [2]])
+passed += np.allclose(values, [0.0, 1.0, 2.0])
+
+F = pt.tensor([3])
+G = pt.tensor([3])
+F.insert([0], 10)
+F.insert([2], 40)
+G[i] = F[i] + i
+indices, values = G.get_coordinates_and_values()
+passed += np.array_equal(indices, [[0], [1], [2]])
+passed += np.allclose(values, [10.0, 1.0, 42.0])
+
+H = pt.tensor([3])
+I = pt.tensor([3])
+H.insert([0], 10)
+H.insert([2], 40)
+I[i] = H[i] * i
+indices, values = I.get_coordinates_and_values()
+passed += np.array_equal(indices, [[0], [2]])
+passed += np.allclose(values, [0.0, 80.0])
+
+# CHECK: Number of passed: 8
 print("Number of passed:", passed)

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

if not structop_roots or structop_roots[-1] != self:
       # The top level expression is not a reduction. Add the top level
       # expression as a structured op root.
       structop_roots.append(self)
 
     # Use user specified information for the destination tensor to build an
     # _StructOpInfo for the top level expression.
     expr_to_info[self].structop_info = _StructOpInfo(dst_indices,
                                                      tuple(dst.shape),
-                                                     self.dtype(), dst.name,
+                                                     dst.dtype, dst.name,
                                                      dst.format)
 
     return structop_roots
 
   def _validate_and_collect_expr_info(
       self,
       dst: "Tensor",
       dst_indices: Tuple["IndexVar", ...],
) -> _ExprInfoDict:
     # Validate the destination dimension information.
     info = expr_to_info[self]
     index_to_dim_info = {i: d for i, d in zip(info.src_indices, info.dim_infos)}
     for i, d, in zip(dst_indices, dst.shape):
       if i not in index_to_dim_info:
         raise ValueError("Destination IndexVar not used in the "
                          f"source expression: {i}")
       else:
-        if d != index_to_dim_info[i].dim:
+        if d != index_to_dim_info[i].dim and index_to_dim_info[i].dim != -1:
           raise ValueError(f"Inconsistent destination dimension for {i}: "
                            f"{d} vs {index_to_dim_info[i].dim}")
 
     return expr_to_info
 
   def _emit_assignment(
       self,
       module: ir.Module,
class _AtomicCounter:
   def increment(self) -> int:
     """Increments the counter by one and returns the old value."""
     old_value = self._counter
     with self._counter_lock:
       self._counter = self._counter + 1
     return old_value
 
 
-class IndexVar:
+class IndexVar(IndexExpr):
   """The tensor index class.
 
   We support the TACO API index_var class with an alias of this class.
 
   An IndexVar object represents an index variable in tensor index notation.
 
   Attributes:
     name: A unique string name of the IndexVar.
   """
   _counter = _AtomicCounter()
 
   def __init__(self):
     id = self._counter.increment()
     self._name = f"{_TACO_INDEX_PREFIX}{id}"
 
   def __repr__(self) -> str:
     return f"IndexVar(name={repr(self._name)})"
 
   @property
   def name(self) -> str:
     """Returns the name of the IndexVar."""
     return self._name
 
+  def _visit(self,
+             func: _ExprVisitor,
+             args,
+             *,
+             leaf_checker: _SubtreeLeafChecker = None) -> None:
+    """A post-order visitor."""
+    if leaf_checker:
+      assert leaf_checker(self, *args)
+    func(self, *args)
+
+  def _emit_expression(
+      self,
+      expr_to_opnd: Dict[IndexExpr, lang.OperandDef],
+      expr_to_info: _ExprInfoDict,
+  ) -> lang.ScalarExpression:
+    """Emits a index value casted to the data type of the tensor expression."""
+    dim = getattr(lang.D, self.name)
+    index = lang.index(dim)
+    int_value = lang.TypeFn.cast_unsigned(lang.TV.I64, index)
+    return lang.TypeFn.cast_unsigned(lang.T, int_value)
+
+  def dtype(self) -> DType:
+    """Returns the data type for the index value.
+
+    This is unreachable for IndexVar.
+    """
+    assert 0
+
 
 def get_index_vars(n: int) -> List[IndexVar]:
   """Returns a list of n IndexVar.
 
   This routine is defined by the TACO API.
 
   Args:
     n: An integer representing the number of IndexVar to get.
class _DimInfo:
   Attributes:
     dim: An integer for the size of the dimension.
     mode_format: A ModeFormat for the dimension sparsity.
   """
   dim: int
   mode_format: ModeFormat
 
 
+def _get_dummy_dim_info() -> _DimInfo:
+  """Constructs the _DimInfo for an index used in tensor expressions."""
+  return _DimInfo(-1, ModeFormat.DENSE)
+
+
 @dataclasses.dataclass()
 class _ExprInfo:
   """Expression information for validation and code generation.
 
   Attributes:
     src_indices: A tuple of IndexVar for the indices used by the tensors in the
       expression tree.
     dim_infos: A tuple of _DimInfo, representing the dimension information
) -> None:
   Raises:
     ValueError if there is any problem in the IndexVars or dimensional values.
   """
   assert len(indices) == len(dim_infos)
   for i, d in zip(indices, dim_infos):
     if i not in index_to_dim_info:
       index_to_dim_info[i] = d
     else:
-      if d.dim != index_to_dim_info[i].dim:
+      dim = index_to_dim_info[i].dim
+      if dim == -1 or d.dim == -1:
+        dim = dim if dim != -1 else d.dim
+      elif dim != d.dim:
         raise ValueError(f"Inconsistent source dimension for {i}: "
-                         f"{d.dim} vs {index_to_dim_info[i].dim}")
+                         f"{d.dim} vs {dim}")
       mode_format = _mode_format_estimator(expr.op)(
           index_to_dim_info[i].mode_format, d.mode_format)
       index_to_dim_info[i] = _DimInfo(d.dim, mode_format)
 
 
 def _validate_and_collect_expr_info(
     expr: IndexExpr,
     expr_to_info: _ExprInfoDict,
Raises:
     ValueError if there is any problem in the IndexVars or dimensional values.
   """
   # Objects of class Access can be shared by different expressions. Avoid
   # processing Access objects multiple times by skipping the processing if expr
   # is already in the dictionary.
   if expr in expr_to_info:
     return
 
-  if isinstance(expr, Access):
+  if isinstance(expr, IndexVar):
+    src_indices = expr,  # A tuple with one element.
+    dim_infos = _get_dummy_dim_info(),  # A tuple with one element.
+  elif isinstance(expr, Access):
     src_indices = expr.indices
     src_dims = tuple(expr.tensor.shape)
     if expr.tensor.format is None:
       # Treat each dimension of a dense tensor as DENSE for the purpose of
       # calculating temporary tensor storage format.
       mode_formats = tuple([ModeFormat.DENSE] * len(src_dims))
     else:
       mode_formats = tuple(expr.tensor.format.format_pack.formats)
) -> None:
   the given IndexVar. If the reduction of the given IndexVar should be performed
   on expression X, then the IndexVar is added to expr_to_info[X].reduce_indices
 
   Args:
     expr: The root IndexExpr for the tensor algebra expression.
     reduce_index: The IndexVar which we want to find out the proper expression
       to perform a reduction.
     expr_to_info: The dictionary to look up _ExprInfo for IndexExpr.
+
+  Raises:
+      ValueError: If the expression is not proper or not supported.
   """
   expr_info = expr_to_info[expr]
   if isinstance(expr, Access):
     # Handle simple reduction expression in the format of A[i] = B[i, j].
     if reduce_index in expr_info.src_indices:
       expr_info.reduce_indices.add(reduce_index)
     return
+  elif isinstance(expr, IndexVar):
+    # A[i] = B[i] + j is not allowed.
+    raise ValueError(f"IndexVar is not part of the iteration domain: {expr}.")
 
   assert (isinstance(expr, _BinaryExpr))
   a_info = expr_to_info[expr.a]
   b_info = expr_to_info[expr.b]
 
   if reduce_index in a_info.src_indices and reduce_index in b_info.src_indices:
     expr_info.reduce_indices.add(reduce_index)
     return
if isinstance(expr, _BinaryExpr):
     b_info = expr_to_info[expr.b]
     expr_info.acc_reduce_indices = (
         a_info.acc_reduce_indices | b_info.acc_reduce_indices
         | expr_info.reduce_indices)
   elif isinstance(expr, _UnaryExpr):
     a_info = expr_to_info[expr.a]
     expr_info.acc_reduce_indices = (
         a_info.acc_reduce_indices | expr_info.reduce_indices)
+  elif isinstance(expr, IndexVar):
+    # If an IndexVar is reducing itself, it means the IndexVar is outside the
+    # iteration domain. This usage is now allowed and we should emit an error
+    # before reaching here.
+    assert not expr_info.reduce_indices
   else:
     assert isinstance(expr, Access)
     # Handle simple reduction expression in the format of A[i] = B[i, j].
     expr_info.acc_reduce_indices = expr_info.reduce_indices
 
 
 
 def _gather_structured_op(
Args:
     root: The root of the current structured op.
     expr_to_info: The dictionary to look up _ExprInfo for IndexExpr.
 
   Returns:
     True if the current IndexExpr is a leaf for the current structured op.
   """
   return (expr != root and
           expr_to_info[expr].structop_info is not None) or isinstance(
-              expr, Access)
+              expr, Access) or isinstance(expr, IndexVar)
 
 
 def _gather_structured_op_input(
     expr: IndexExpr,
     root: IndexExpr,
     expr_to_info: _ExprInfoDict,
     structop_inputs: List[IndexExpr],
 ) -> None:

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

 test_expect_error("invalid_indices", ("""
 i, j = mlir_pytaco.get_index_vars(2)
 A = mlir_pytaco.Tensor([2, 3])
 B = mlir_pytaco.Tensor([2, 3])
 C = mlir_pytaco.Tensor([2, 3], _DENSE)
 C[i, j] = A[1, j] + B[i, j]
     """), "Expected IndexVars")
 
-# CHECK: test_invalid_operation: passed
-test_expect_error("invalid_operation", ("""
-i, j = mlir_pytaco.get_index_vars(2)
-A = mlir_pytaco.Tensor([2, 3])
-C = mlir_pytaco.Tensor([2, 3], _DENSE)
-C[i, j] = A[i, j] + i
-    """), "Expected IndexExpr")
-
 # CHECK: test_inconsistent_rank_indices: passed
 test_expect_error("inconsistent_rank_indices", ("""
 i, j = mlir_pytaco.get_index_vars(2)
 A = mlir_pytaco.Tensor([2, 3])
 C = mlir_pytaco.Tensor([2, 3], _DENSE)
 C[i, j] = A[i]
     """), "Invalid indices for rank")
 

 i = mlir_pytaco.IndexVar()
 A = mlir_pytaco.Tensor([3])
 B = mlir_pytaco.Tensor([5])
 C = mlir_pytaco.Tensor([3], _DENSE)
 C[i] = A[i] + B[i]
 C.evaluate()
     """), "Inconsistent source dimension for IndexVar")
 
+# CHECK: test_index_var_outside_domain: passed
+test_expect_error("index_var_outside_domain", ("""
+i, j = mlir_pytaco.get_index_vars(2)
+A = mlir_pytaco.Tensor([3])
+B = mlir_pytaco.Tensor([3])
+B[i] = A[i] + j
+B.evaluate()
+    """), "IndexVar is not part of the iteration domain")
+
 
 # CHECK-LABEL: test_tensor_all_dense_sparse
 @testing_utils.run_test
 def test_tensor_all_dense_sparse():
   a = mlir_pytaco.Tensor([4], [_DENSE])
   passed = (not a.is_dense())
   passed += (a.order == 1)
   passed += (a.shape[0] == 4)