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)