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)
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.])
# 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)
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]
passed = (B._assignment is not None)
passed += (B._engine is None)
try:
B.compute()
except ValueError as e:
passed += (str(e).startswith("Need to invoke compile"))
B.compile()
passed += (B._engine is not None)
B.compute()
passed += (B._assignment is None)
passed += (B._engine is None)
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: 9
print("Number of passed:", passed)
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.array_equal(indices, [[0, 1], [1, 2]])
passed += np.allclose(values, [5.0, 6.0])
# CHECK: Number of passed: 2
print("Number of passed:", passed)
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
print("Number of passed:", passed)
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)
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)
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)
def test_write_packed_tns():
a = mlir_pytaco.Tensor([2, 3])
a.insert([0, 1], 10)
a.insert([1, 2], 40)
a.insert([0, 0], 20)
b = mlir_pytaco.Tensor([2, 3])
i, j = mlir_pytaco.get_index_vars(2)
b[i, j] = a[i, j] + a[i, j]
with tempfile.TemporaryDirectory() as test_dir:
file_name = os.path.join(test_dir, "data.tns")
mlir_pytaco_io.write(file_name, b)
with open(file_name, "r") as file:
lines = file.readlines()
passed = 0
# Skip the comment line in the output.
if lines[1:] == ["2 3\n", "2 3\n", "1 1 40\n", "1 2 20\n", "2 3 80\n"]:
passed = 1
# CHECK: 1
print(passed)
def test_write_unpacked_tns():
a = mlir_pytaco.Tensor([2, 3])
a.insert([0, 1], 10)
a.insert([1, 2], 40)
a.insert([0, 0], 20)
with tempfile.TemporaryDirectory() as test_dir:
file_name = os.path.join(test_dir, "data.tns")
try:
mlir_pytaco_io.write(file_name, a)
except ValueError as e:
# CHECK: Writing unpacked sparse tensors to file is not supported
print(e)
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)
