[Mlir-commits] [mlir] 4b7745c - [mlir][sparse][taco] Add more unit tests.
Bixia Zheng
llvmlistbot at llvm.org
Mon Mar 7 10:10:08 PST 2022
Author: Bixia Zheng
Date: 2022-03-07T10:10:01-08:00
New Revision: 4b7745c176a6db91deaa9c701f4031c5425c046e
URL: https://github.com/llvm/llvm-project/commit/4b7745c176a6db91deaa9c701f4031c5425c046e
DIFF: https://github.com/llvm/llvm-project/commit/4b7745c176a6db91deaa9c701f4031c5425c046e.diff
LOG: [mlir][sparse][taco] Add more unit tests.
These unit tests resides in an internal repository. Porting the tests to the
public repository.
Reviewed By: aartbik
Differential Revision: https://reviews.llvm.org/D121021
Added:
Modified:
mlir/test/Integration/Dialect/SparseTensor/taco/unit_test_tensor_core.py
Removed:
################################################################################
diff --git a/mlir/test/Integration/Dialect/SparseTensor/taco/unit_test_tensor_core.py b/mlir/test/Integration/Dialect/SparseTensor/taco/unit_test_tensor_core.py
index 8703ef9126c8d..6b539ad06d1d4 100644
--- a/mlir/test/Integration/Dialect/SparseTensor/taco/unit_test_tensor_core.py
+++ b/mlir/test/Integration/Dialect/SparseTensor/taco/unit_test_tensor_core.py
@@ -17,6 +17,235 @@
_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
+ at 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
+ at 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
+ at 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
+ at 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():
@@ -53,8 +282,10 @@ def test_tensor_copy():
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)
@@ -73,6 +304,212 @@ def test_tensor_trivial_reduction():
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
+ at 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
+ at 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
+ at 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
+ at 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
+ at 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
+ at 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
+ at 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)
More information about the Mlir-commits
mailing list