[mlir][sparse][taco] Add more unit tests.

[bixia]

Summary

These unit tests resides in an internal repository. Porting the tests to the
public repository.

[aartbik]

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

39

then verifies? or something else?

[bixia]

Fixed a comment.

[bixia]

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

39

Good catch. Should be "then verified the".

[bixia]

Removed debug print.

[aartbik]

Ship it!

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

 from tools import mlir_pytaco
 from tools import testing_utils as testing_utils
 
 # Define the aliases to shorten the code.
 _COMPRESSED = mlir_pytaco.ModeFormat.COMPRESSED
 _DENSE = mlir_pytaco.ModeFormat.DENSE
 
 
+def _init_3d(T, I, J, K):
+  for i in range(I):
+    for j in range(J):
+      for k in range(K):
+        T.insert([i, j, k], i + j + k + 1)
+
+
+def _init_2d(T, I, J):
+  for i in range(I):
+    for j in range(J):
+      T.insert([i, j], i + j + 1)
+
+
+def _init_1d_with_value(T, I, v):
+  for i in range(I):
+    T.insert([i], v)
+
+
+def test_expect_error(name, code, error):
+  """Executes the code then verifies the expected error message."""
+  try:
+    exec(code)
+  except ValueError as e:
+    passed = "passed" if (str(e).startswith(error)) else "failed"
+    print(f"test_{name}: {passed}")
+
+
+# CHECK-LABEL: test_tensor_dtype
+@testing_utils.run_test
+def test_tensor_dtype():
+  passed = mlir_pytaco.DType(mlir_pytaco.Type.INT16).is_int()
+  passed += mlir_pytaco.DType(mlir_pytaco.Type.INT32).is_int()
+  passed += mlir_pytaco.DType(mlir_pytaco.Type.INT64).is_int()
+  passed += mlir_pytaco.DType(mlir_pytaco.Type.FLOAT32).is_float()
+  passed += mlir_pytaco.DType(mlir_pytaco.Type.FLOAT64).is_float()
+  # CHECK: Number of passed: 5
+  print("Number of passed:", passed)
+
+
+# CHECK: test_mode_ordering_not_int: passed
+test_expect_error("mode_ordering_not_int",
+                  "m = mlir_pytaco.ModeOrdering(['x'])",
+                  "Ordering must be a list of integers")
+
+# CHECK: test_mode_ordering_not_permutation: passed
+test_expect_error("mode_ordering_not_permutation",
+                  "m = mlir_pytaco.ModeOrdering([2, 1])", "Invalid ordering")
+
+# CHECK: test_mode_format_invalid: passed
+test_expect_error("mode_format_invalid",
+                  "m = mlir_pytaco.ModeFormatPack(['y'])",
+                  "Formats must be a list of ModeFormat")
+
+# CHECK: test_expect_mode_format_pack: passed
+test_expect_error("expect_mode_format_pack", ("""
+mode_ordering = mlir_pytaco.ModeOrdering([0, 1, 2])
+f = mlir_pytaco.Format(["x"], mode_ordering)
+    """), "Expected a list of ModeFormat")
+
+# CHECK: test_expect_mode_ordering: passed
+test_expect_error("expect_mode_ordering", ("""
+mode_format_pack = mlir_pytaco.ModeFormatPack([_COMPRESSED, _COMPRESSED])
+f = mlir_pytaco.Format(mode_format_pack, "x")
+    """), "Expected ModeOrdering")
+
+# CHECK: test_inconsistent_mode_format_pack_and_mode_ordering: passed
+test_expect_error("inconsistent_mode_format_pack_and_mode_ordering", ("""
+mode_format_pack = mlir_pytaco.ModeFormatPack([_COMPRESSED, _COMPRESSED])
+mode_ordering = mlir_pytaco.ModeOrdering([0, 1, 2])
+f = mlir_pytaco.Format(mode_format_pack, mode_ordering)
+    """), "Inconsistent ModeFormatPack and ModeOrdering")
+
+
+# CHECK-LABEL: test_format_default_ordering
+@testing_utils.run_test
+def test_format_default_ordering():
+  f = mlir_pytaco.Format([_COMPRESSED, _COMPRESSED])
+  passed = 0
+  passed += np.array_equal(f.ordering.ordering, [0, 1])
+  # CHECK: Number of passed: 1
+  print("Number of passed:", passed)
+
+
+# CHECK-LABEL: test_format_explicit_ordering
+@testing_utils.run_test
+def test_format_explicit_ordering():
+  f = mlir_pytaco.Format([_COMPRESSED, _DENSE], [1, 0])
+  passed = 0
+  passed += np.array_equal(f.ordering.ordering, [1, 0])
+  # CHECK: Number of passed: 1
+  print("Number of passed:", passed)
+
+
+# CHECK-LABEL: test_index_var
+@testing_utils.run_test
+def test_index_var():
+  i = mlir_pytaco.IndexVar()
+  j = mlir_pytaco.IndexVar()
+  passed = (i.name != j.name)
+
+  vars = mlir_pytaco.get_index_vars(10)
+  passed += (len(vars) == 10)
+  passed += (all([isinstance(e, mlir_pytaco.IndexVar) for e in vars]))
+
+  # CHECK: Number of passed: 3
+  print("Number of passed:", passed)
+
+
+# CHECK: test_tensor_invalid_first_argument: passed
+test_expect_error("tensor_invalid_first_argument",
+                  "t = mlir_pytaco.Tensor('f')", "Invalid first argument")
+
+# CHECK: test_tensor_inconsistent_shape_and_format: passed
+test_expect_error("tensor_inconsistent_shape_and_format", ("""
+mode_format_pack = mlir_pytaco.ModeFormatPack([_COMPRESSED, _COMPRESSED])
+mode_ordering = mlir_pytaco.ModeOrdering([0, 1])
+f = mlir_pytaco.Format(mode_format_pack, mode_ordering)
+t = mlir_pytaco.Tensor([3], f)
+    """), "Inconsistent shape and format")
+
+# CHECK: test_tensor_invalid_format: passed
+test_expect_error("tensor_invalid_format", "t = mlir_pytaco.Tensor([3], 'f')",
+                  "Invalid format argument")
+
+# CHECK: test_tensor_insert_nonlist_coordinate: passed
+test_expect_error("tensor_insert_nonlist_coordinate", ("""
+t = mlir_pytaco.Tensor([3])
+t.insert(1, 0)
+    """), "Non list coordinate detected")
+
+# CHECK: test_tensor_insert_too_much_coordinate: passed
+test_expect_error("tensor_insert_too_much_coordinate", ("""
+t = mlir_pytaco.Tensor([3])
+t.insert([0, 0], 0)
+    """), "Invalid coordinate")
+
+# CHECK: test_tensor_insert_coordinate_outof_range: passed
+test_expect_error("tensor_insert_coordinate_outof_range", ("""
+t = mlir_pytaco.Tensor([1, 1])
+t.insert([1, 0], 0)
+    """), "Invalid coordinate")
+
+# CHECK: test_tensor_insert_coordinate_nonint: passed
+test_expect_error("tensor_insert_coordinate_nonint", ("""
+t = mlir_pytaco.Tensor([1, 1])
+t.insert([0, "xy"], 0)
+    """), "Non integer coordinate detected")
+
+# CHECK: test_tensor_insert_invalid_value: passed
+test_expect_error("tensor_insert_invalid_value", ("""
+t = mlir_pytaco.Tensor([1, 1])
+t.insert([0, 0], "x")
+    """), "Value is neither int nor float")
+
+# CHECK: test_access_non_index_var_index: passed
+test_expect_error("access_non_index_var_index", ("""
+t = mlir_pytaco.Tensor([5, 6])
+i = mlir_pytaco.IndexVar()
+a = mlir_pytaco.Access(t, (i, "j"))
+    """), "Indices contain non IndexVar")
+
+# CHECK: test_access_inconsistent_rank_indices: passed
+test_expect_error("access_inconsistent_rank_indices", ("""
+t = mlir_pytaco.Tensor([5, 6])
+i = mlir_pytaco.IndexVar()
+a = mlir_pytaco.Access(t, (i,))
+    """), "Invalid indices for rank")
+
+# CHECK: test_access_invalid_indices_for_rank: passed
+test_expect_error("access_invalid_indices_for_rank", ("""
+t = mlir_pytaco.Tensor([5, 6])
+i, j, k = mlir_pytaco.get_index_vars(3)
+a = mlir_pytaco.Access(t, (i,j, k))
+    """), "Invalid indices for rank")
+
+# CHECK: test_invalid_indices: passed
+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")
+
+# CHECK: test_destination_index_not_used_in_source: passed
+test_expect_error("destination_index_not_used_in_source", ("""
+i, j = mlir_pytaco.get_index_vars(2)
+A = mlir_pytaco.Tensor([3])
+C = mlir_pytaco.Tensor([3], _DENSE)
+C[j] = A[i]
+C.evaluate()
+    """), "Destination IndexVar not used in the source expression")
+
+# CHECK: test_destination_dim_not_consistent_with_source: passed
+test_expect_error("destination_dim_not_consistent_with_source", ("""
+i = mlir_pytaco.IndexVar()
+A = mlir_pytaco.Tensor([3])
+C = mlir_pytaco.Tensor([5], _DENSE)
+C[i] = A[i]
+C.evaluate()
+    """), "Inconsistent destination dimension for IndexVar")
+
+# CHECK: test_inconsistent_source_dim: passed
+test_expect_error("inconsistent_source_dim", ("""
+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-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)
   # CHECK: Number of passed: 3
def test_tensor_copy():
   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.array_equal(indices, [[0, 1], [1, 2]])
   passed += np.allclose(values, [5.0, 6.0])
+  # No temporary tensor is used.
+  passed += (B._stats.get_total() == 0)
 
-  # CHECK: Number of passed: 2
+  # CHECK: Number of passed: 3
   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.array_equal(indices, [[0], [1]])
   passed += np.allclose(values, [8.0, 6.0])
+  # No temporary tensor is used.
+  passed += (B._stats.get_total() == 0)
+
+  # CHECK: Number of passed: 3
+  print("Number of passed:", passed)
+
+
+# CHECK-LABEL: test_binary_add
+@testing_utils.run_test
+def test_binary_add():
+  i = mlir_pytaco.IndexVar()
+  A = mlir_pytaco.Tensor([4])
+  B = mlir_pytaco.Tensor([4])
+  C = mlir_pytaco.Tensor([4])
+  A.insert([1], 10)
+  A.insert([2], 1)
+  B.insert([3], 20)
+  B.insert([2], 2)
+  C[i] = A[i] + B[i]
+  indices, values = C.get_coordinates_and_values()
+  passed = np.array_equal(indices, [[1], [2], [3]])
+  passed += np.array_equal(values, [10., 3., 20.])
+  # No temporary tensor is used.
+  passed += (C._stats.get_total() == 0)
+  # CHECK: Number of passed: 3
+  print("Number of passed:", passed)
+
+
+# CHECK-LABEL: test_binary_add_sub
+@testing_utils.run_test
+def test_binary_add_sub():
+  i = mlir_pytaco.IndexVar()
+  j = mlir_pytaco.IndexVar()
+  A = mlir_pytaco.Tensor([2, 3])
+  B = mlir_pytaco.Tensor([2, 3])
+  C = mlir_pytaco.Tensor([2, 3])
+  D = mlir_pytaco.Tensor([2, 3])
+  A.insert([0, 1], 10)
+  A.insert([1, 2], 40)
+  B.insert([0, 0], 20)
+  B.insert([1, 2], 30)
+  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.array_equal(values, [20., 5., 63.])
+  # No temporary tensor is used.
+  passed += (D._stats.get_total() == 0)
+  # CHECK: Number of passed: 3
+  print("Number of passed:", passed)
+
+
+# CHECK-LABEL: test_binary_mul_add
+@testing_utils.run_test
+def test_binary_mul_add():
+  i = mlir_pytaco.IndexVar()
+  j = mlir_pytaco.IndexVar()
+  A = mlir_pytaco.Tensor([2, 3])
+  B = mlir_pytaco.Tensor([2, 3])
+  C = mlir_pytaco.Tensor([2, 3])
+  D = mlir_pytaco.Tensor([2, 3])
+  A.insert([0, 1], 10)
+  A.insert([1, 2], 40)
+  B.insert([0, 0], 20)
+  B.insert([1, 2], 30)
+  C.insert([0, 1], 5)
+  C.insert([1, 2], 7)
+  D[i, j] = A[i, j] * C[i, j] + B[i, j]
+  indices, values = D.get_coordinates_and_values()
+  passed = np.array_equal(indices, [[0, 0], [0, 1], [1, 2]])
+  passed += np.array_equal(values, [20., 50., 310.])
+  # No temporary tensor is used.
+  passed += (D._stats.get_total() == 0)
+  # CHECK: Number of passed: 3
+  print("Number of passed:", passed)
+
+
+# CHECK-LABEL: test_binary_add_reduce_at_root
+@testing_utils.run_test
+def test_binary_add_reduce_at_root():
+  i = mlir_pytaco.IndexVar()
+  j = mlir_pytaco.IndexVar()
+  A = mlir_pytaco.Tensor([2, 3])
+  B = mlir_pytaco.Tensor([2, 3])
+  C = mlir_pytaco.Tensor([2], _DENSE)
+  A.insert([0, 1], 10)
+  A.insert([1, 2], 40)
+  B.insert([0, 0], 20)
+  B.insert([1, 2], 30)
+  C[i] = A[i, j] + B[i, j]
+  indices, values = C.get_coordinates_and_values()
+  passed = np.array_equal(indices, [[0], [1]])
+  passed += np.array_equal(values, [30., 70.])
+  # No temporary tensor is used.
+  passed += (C._stats.get_total() == 0)
+  # CHECK: Number of passed: 3
+  print("Number of passed:", passed)
+
+
+# CHECK-LABEL: test_binary_add_reduce_at_child
+@testing_utils.run_test
+def test_binary_add_reduce_at_child():
+  i = mlir_pytaco.IndexVar()
+  j = mlir_pytaco.IndexVar()
+  I = 2
+  J = 3
+  A = mlir_pytaco.Tensor([I, J])
+  B = mlir_pytaco.Tensor([J])
+  C = mlir_pytaco.Tensor([I])
+  D = mlir_pytaco.Tensor([I], _DENSE)
+
+  _init_2d(A, I, J)
+  _init_1d_with_value(C, I, 2)
+  _init_1d_with_value(B, J, 1)
+
+  D[i] = A[i, j] * B[j] + C[i]
+  indices, values = D.get_coordinates_and_values()
+  passed = np.array_equal(indices, [[0], [1]])
+  passed += np.array_equal(values, [8., 11.])
+
+  # The expression is implemented as:
+  #    temp0[i] = A[i, j] * B[i]
+  #    D[i] = temp0[i] + C[i]
+  # Check the temporary tensor introduced by the implementation.
+  stats = D._stats
+  passed += (stats.get_total() == 1)
+  passed += (stats.get_formats(0) == (_COMPRESSED,))
+  passed += (stats.get_dimensions(0) == (I,))
+  # CHECK: Number of passed: 5
+  print("Number of passed:", passed)
+
+
+# CHECK-LABEL: test_binary_add_reduce_3d_1
+@testing_utils.run_test
+def test_binary_add_reduce_3d_1():
+  i, j, k, l = mlir_pytaco.get_index_vars(4)
+  I = 2
+  J = 3
+  K = 4
+  L = 5
+  A = mlir_pytaco.Tensor([I, J, K])
+  B = mlir_pytaco.Tensor([I, J, L])
+  C = mlir_pytaco.Tensor([K])
+  D = mlir_pytaco.Tensor([L])
+  E = mlir_pytaco.Tensor([I], _DENSE)
+
+  _init_3d(A, I, J, K)
+  _init_3d(B, I, J, L)
+  _init_1d_with_value(C, K, 1)
+  _init_1d_with_value(D, L, 2)
+
+  E[i] = A[i, j, k] * C[k] + B[i, j, l] * D[l]
+  indices, values = E.get_coordinates_and_values()
+  passed = np.array_equal(indices, [[0], [1]])
+  passed += np.array_equal(values, [162., 204.])
+
+  # The expression is implemented as:
+  #    temp0[i, j] = A[i, j, k] * C[k]
+  #    temp1[i, j] = B[i, j, l] * D[l]
+  #    E[i] = temp0[i, j] + temp1[i, j]
+  # Check the two temporary tensors introduced by the implementation.
+  stats = E._stats
+  passed += (stats.get_total() == 2)
+  passed += (stats.get_formats(0) == (_COMPRESSED, _COMPRESSED))
+  passed += (stats.get_dimensions(0) == (I, J))
+  passed += (stats.get_formats(1) == (_COMPRESSED, _COMPRESSED))
+  passed += (stats.get_dimensions(1) == (I, J))
+  # CHECK: Number of passed: 7
+  print("Number of passed:", passed)
+
+
+# CHECK-LABEL: test_binary_add_reduce_3d_2
+@testing_utils.run_test
+def test_binary_add_reduce_3d_2():
+  i, j, k, l = mlir_pytaco.get_index_vars(4)
+  I = 2
+  J = 3
+  K = 4
+  L = 5
+  A = mlir_pytaco.Tensor([I, J, K], [_COMPRESSED, _COMPRESSED, _DENSE])
+  B = mlir_pytaco.Tensor([I, L, K], [_DENSE, _COMPRESSED, _COMPRESSED])
+  C = mlir_pytaco.Tensor([J, K], [_COMPRESSED, _COMPRESSED])
+  D = mlir_pytaco.Tensor([L])
+  E = mlir_pytaco.Tensor([I], _DENSE)
+
+  _init_3d(A, I, J, K)
+  _init_3d(B, I, L, K)
+  _init_2d(C, J, K)
+  _init_1d_with_value(D, L, 2)
+
+  E[i] = A[i, j, k] + C[j, k] + B[i, l, k] * D[l]
+  indices, values = E.get_coordinates_and_values()
+  passed = np.array_equal(indices, [[0], [1]])
+  passed += np.array_equal(values, [264., 316.])
 
-  # CHECK: Number of passed: 2
+  # The expression is implemented as:
+  #    temp0[i, k] = A[i, j, k] + C[j, k]
+  #    temp1[i, k] = B[i, l, k] * D[l]
+  #    E[i] = temp0[i, k] + temp1[i, k]
+  # Check the two temporary tensors introduced by the implementation.
+  stats = E._stats
+  passed += (stats.get_total() == 2)
+  passed += (stats.get_formats(0) == (_COMPRESSED, _DENSE))
+  passed += (stats.get_dimensions(0) == (I, K))
+  passed += (stats.get_formats(1) == (_DENSE, _COMPRESSED))
+  passed += (stats.get_dimensions(1) == (I, K))
+  # CHECK: Number of passed: 7
   print("Number of passed:", passed)