[mlir][sparse][taco] Handle tensor copy and trivial reduction expression.

[bixia]

Summary

Handle tensor copy, such as A[i, j] = B[i, j]. Also, handle trivial
reduction expression, such as A[i] = B[i, j].

Add unit tests.

[aartbik]

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

54

Here and below. It seems a bit odd to test indices with an "all close" since surely they are not subject to rounding errors. Can we make the test on indices just exact?

[bixia]

Replace np.allClose with np.array_equal to compare indices.

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

) -> None:
   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.
   """
+  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
+
   assert (isinstance(expr, _BinaryExpr))
   a_info = expr_to_info[expr.a]
   b_info = expr_to_info[expr.b]
-  expr_info = expr_to_info[expr]
 
   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 reduce_index in a_info.src_indices:
     _mark_structured_op_root(expr.a, reduce_index, expr_to_info)
   elif reduce_index in b_info.src_indices:
) -> None:
   if isinstance(expr, _BinaryExpr):
     a_info = expr_to_info[expr.a]
     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)
   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(
     expr: IndexExpr,
     expr_to_info: _ExprInfoDict,
     structop_roots: List[IndexExpr],
 ) -> None:
   """Adds structured op root expression information to structop_roots.
) -> None:
 
   Args:
     expr: The IndexExpr being visited.
     root: The root of the current structured op.
     expr_to_info: The dictionary to look up _ExprInfo for IndexExpr.
     structop_inputs: The resulting list of IndexExpr that provide input to the
       current structured op.
   """
-  if (expr != root and expr not in structop_inputs) and (
-      isinstance(expr, Access) or
-      (expr in expr_to_info and expr_to_info[expr].structop_info)):
+  if ((expr != root or isinstance(expr, Access)) and
+      expr not in structop_inputs) and (isinstance(expr, Access) or
+                                        (expr in expr_to_info and
+                                         expr_to_info[expr].structop_info)):
     structop_inputs.append(expr)
 
 
 def _emit_structured_op_input(
     expr: IndexExpr,
     expr_to_info: _ExprInfoDict,
     op_def: lang.LinalgOpDef,
 ) -> lang.OperandDef:
   """Emits OperandDef in the linalg dialect for the input IndexExpr.
 
   Args:
     expr: The input IndexExpr for the current structured op.
     expr_to_info: The dictionary to look up _ExprInfo for IndexExpr.
     op_def: The linalg operation for the current structured op.
 
   Returns:
     An OperandDef in the linalg dialect for the input IndexExpr.
   """
   op_info = expr_to_info[expr].structop_info
-  if op_info:
+  if op_info and not isinstance(expr, Access):
     # The input is a temporary tensor produced by another structured op.
     indices = op_info.dst_indices
     name = op_info.dst_name
   else:
     # The input is a user provided tensor.
     assert isinstance(expr, Access)
     indices = expr.indices
     name = expr.tensor.name
 
   dim_sym = _mlir_symbols_from_index_vars(indices)
   opnd = lang.OperandDef(lang.OperandKind.InputTensor, lang.T, dim_sym)
   op_def.add_operand(name, opnd)
   return opnd

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

 @testing_utils.run_test
 def test_tensor_true_dense():
   a = mlir_pytaco.Tensor.from_array(np.random.uniform(size=5))
   passed = a.is_dense()
   passed += (a.order == 1)
   passed += (a.shape[0] == 5)
   # CHECK: Number of passed: 3
   print("Number of passed:", passed)
+
+
+# CHECK-LABEL: test_tensor_copy
+@testing_utils.run_test
+def test_tensor_copy():
+  i, j = mlir_pytaco.get_index_vars(2)
+  I = 2
+  J = 3
+  A = mlir_pytaco.Tensor([I, J])
+  A.insert([0, 1], 5.0)
+  A.insert([1, 2], 6.0)
+  B = mlir_pytaco.Tensor([I, J])
+  B[i, j] = A[i, j]
+  indices, values = B.get_coordinates_and_values()
+  passed = np.allclose(indices, [[0, 1], [1, 2]])
+  passed += np.allclose(values, [5.0, 6.0])
+
+  # CHECK: Number of passed: 2
+  print("Number of passed:", passed)
+
+
+# CHECK-LABEL: test_tensor_trivial_reduction
+@testing_utils.run_test
+def test_tensor_trivial_reduction():
+  i, j = mlir_pytaco.get_index_vars(2)
+  I = 2
+  J = 3
+  A = mlir_pytaco.Tensor([I, J])
+  A.insert([0, 1], 5.0)
+  A.insert([0, 2], 3.0)
+  A.insert([1, 2], 6.0)
+  B = mlir_pytaco.Tensor([I])
+  B[i] = A[i, j]
+  indices, values = B.get_coordinates_and_values()
+  passed = np.allclose(indices, [[0], [1]])
+  passed += np.allclose(values, [8.0, 6.0])
+
+  # CHECK: Number of passed: 2
+  print("Number of passed:", passed)