def test_deep_matrix_factorization(three_ranges):
"""Test a basic matrix multiplication factorization problem.
Similar to the shallow factorization test, the final expression to optimize
is mathematically
.. math::
(X - 2 Y) * U * V
Different from the shallow test case, here we concentrate more on the
treatment of depth complexity in the input. The sum intermediate needs to
be factored again.
"""
#
# Basic context setting-up.
#
dr = three_ranges
p = dr.names
m = p.m
a, b, c, d = p.a, p.b, p.c, p.d
# The indexed bases.
x = IndexedBase('X')
y = IndexedBase('Y')
u = IndexedBase('U')
v = IndexedBase('V')
t = IndexedBase('T')
# The target.
target = dr.define_einst(
t[a, b], x[a, c] * u[c, d] * v[d, b] - 2 * y[a, c] * u[c, d] * v[d, b])
targets = [target]
# The actual optimization.
res = optimize(targets)
assert len(res) == 3
# Test the correctness.
assert verify_eval_seq(res, targets, simplify=True)
# Test the cost.
cost = get_flop_cost(res)
leading_cost = get_flop_cost(res, leading=True)
assert cost == 4 * m**3 + m**2
assert leading_cost == 4 * m**3
cost = get_flop_cost(res, ignore_consts=False)
assert cost == 4 * m**3 + 2 * m**2
# Test disabling summation optimization.
res = optimize(targets, opt_sum=False)
assert verify_eval_seq(res, targets, simplify=True)
new_cost = get_flop_cost(res, ignore_consts=False)
assert new_cost - cost != 0
def test_optimization_of_common_terms(three_ranges):
"""Test optimization of common terms in summations.
In this test, there are just two matrices involved, X, Y. The target reads
.. math::
T[a, b] = X[a, b] - X[b, a] + 2 Y[a, b] - 2 Y[b, a]
Ideally, it should be evaluated as,
.. math::
I[a, b] = X[a, b] + 2 Y[a, b]
T[a, b] = I[a, b] - I[b, a]
or,
.. math::
I[a, b] = X[a, b] - 2 Y[b, a]
T[a, b] = I[a, b] - I[b, a]
"""
#
# Basic context setting-up.
#
dr = three_ranges
p = dr.names
a, b, c, d = p.a, p.b, p.c, p.d
# The indexed bases.
x = IndexedBase('x')
y = IndexedBase('y')
t = dr.define_einst(
IndexedBase('t')[a, b], x[a, b] - x[b, a] + 2 * y[a, b] - 2 * y[b, a])
targets = [t]
eval_seq = optimize(targets)
assert len(eval_seq) == 2
verify_eval_seq(eval_seq, targets)
# Check the result when the common symmetrization optimization is disabled.
eval_seq = optimize(targets, opt_symm=False)
assert len(eval_seq) == 1
verify_eval_seq(eval_seq, targets)
def test_ccd_doubles_terms(parthole_drudge):
"""Test optimization of selected terms in CCD equations.
This purpose of this test is mostly on the treatment of term in the present
of symmetries.
"""
dr = parthole_drudge
p = dr.names
a, b, c, d = p.V_dumms[:4]
i, j, k, l = p.O_dumms[:4]
u = dr.two_body
t = IndexedBase('t')
dr.set_dbbar_base(t, 2)
r = IndexedBase('r')
tensor = dr.define_einst(
r[a, b, i, j],
+ t[a, b, l, j] * t[c, d, i, k] * u[k, l, c, d]
+ t[a, d, i, j] * t[b, c, k, l] * u[k, l, c, d]
- t[a, b, i, l] * t[c, d, k, j] * u[k, l, c, d]
- t[a, c, k, l] * t[b, d, i, j] * u[k, l, c, d]
)
targets = [tensor]
eval_seq = optimize(targets, substs={p.nv: p.no * 10})
assert verify_eval_seq(eval_seq, targets)
def test_ccsd_singles_terms(parthole_drudge):
"""Test selected terms in CCSD singles equation.
The purpose of this test is the capability of recognition of repeated
appearance of the same summation intermediates.
"""
dr = parthole_drudge
p = dr.names
a, b, c = p.V_dumms[:3]
i, j, k = p.O_dumms[:3]
u = dr.two_body
f = dr.fock
t = IndexedBase('t')
dr.set_dbbar_base(t, 2)
r = IndexedBase('r')
tensor = dr.define_einst(
r[a, i],
t[a, b, i, j] * u[j, k, b, c] * t[c, k] + t[a, b, i, j] * f[j, b]
- t[a, j] * t[b, i] * f[j, b]
- t[a, j] * t[b, i] * t[c, k] * u[j, k, b, c]
)
targets = [tensor]
eval_seq = optimize(targets, substs={p.nv: p.no * 10})
assert verify_eval_seq(eval_seq, targets)
assert len(eval_seq) == 4
def test_disconnected_outer_product_factorization(three_ranges):
"""Test optimization of expressions with disconnected outer products.
"""
dr = three_ranges
p = dr.names
m = p.m
a, b, c, d, e = p.a, p.b, p.c, p.d, p.e
# The indexed bases.
u = IndexedBase('U')
x = IndexedBase('X')
y = IndexedBase('Y')
z = IndexedBase('Z')
t = IndexedBase('T')
# The target.
target = dr.define_einst(
t[a, b], u[a, b] * z[c, e] * x[e, c] + u[a, b] * z[c, e] * y[e, c])
targets = [target]
# The actual optimization.
res = optimize(targets)
assert len(res) == 3
# Test the correctness.
assert verify_eval_seq(res, targets, simplify=False)
# Test the cost.
cost = get_flop_cost(res)
leading_cost = get_flop_cost(res, leading=True)
assert cost == 4 * m**2
assert leading_cost == 4 * m**2
def test_factorization_needing_canonicalization(three_ranges):
"""Test a simple factorization needing canonicalization.
The inability of gristmill to fully optimize this test is the ultimate
reason why the above test is flaky.
"""
dr = three_ranges
p = dr.names
m = p.m
a, b = p.a, p.b
x = IndexedBase('X')
y = IndexedBase('Y')
z = IndexedBase('Z')
t = Symbol('T')
# The target.
target = dr.define_einst(t, x[b, a] * z[a, b] + y[a, b] * z[b, a])
targets = [target]
# The actual optimization.
res = optimize(targets)
assert len(res) == 2
# Test the correctness.
assert verify_eval_seq(res, targets, simplify=False)
def test_optimization_handles_nonlinear_factors(spark_ctx):
"""Test optimization of with nonlinear factors.
Here a factor is the square of an indexed quantity.
"""
dr = Drudge(spark_ctx)
n = symbols('n')
r = Range('r', 0, n)
dumms = symbols('a b c d e f g h')
dr.set_dumms(r, dumms)
a, b, c, d = dumms[:4]
dr.add_default_resolver(r)
u = symbols('u')
s = IndexedBase('s')
targets = [
dr.define(
u,
dr.sum((a, r), (b, r), (c, r), (d, r),
32 * s[a, c]**2 * s[b, d]**2 +
32 * s[a, c] * s[a, d] * s[b, c] * s[b, d]))
]
eval_seq = optimize(targets)
assert verify_eval_seq(eval_seq, targets)
def test_optimization_handles_coeffcients(spark_ctx):
"""Test optimization of scalar intermediates scaled by coefficients.
This test comes from PoST theory. It tests the optimization of tensor
evaluations with scalar intermediates scaled by a factor.
"""
dr = Drudge(spark_ctx)
n = symbols('n')
r = Range('r', 0, n)
a, b = symbols('a b')
dr.set_dumms(r, [a, b])
dr.add_default_resolver(r)
r = IndexedBase('r')
eps = IndexedBase('epsilon')
t = IndexedBase('t')
targets = [
dr.define(r[a, b],
dr.sum(2 * eps[a] * t[a, b]) - 2 * eps[b] * t[a, b])
]
eval_seq = optimize(targets)
assert verify_eval_seq(eval_seq, targets)
def test_optimization_handles_scalar_intermediates(spark_ctx):
"""Test optimization of scalar intermediates scaling other tensors.
This is set as a special test primarily since it would entail the same
collectible giving residues with different ranges.
"""
dr = Drudge(spark_ctx)
n = symbols('n')
r = Range('r', 0, n)
dumms = symbols('a b c d e')
dr.set_dumms(r, dumms)
a, b, c = dumms[:3]
dr.add_default_resolver(r)
u = IndexedBase('u')
eps = IndexedBase('epsilon')
t = IndexedBase('t')
s = IndexedBase('s')
targets = [
dr.define(
u, (a, r), (b, r),
dr.sum((c, r), 8 * s[a, b] * eps[c] * t[a]) -
8 * s[a, b] * eps[a] * t[a])
]
eval_seq = optimize(targets)
assert verify_eval_seq(eval_seq, targets)
def test_shallow_matrix_factorization(three_ranges, rand_constr):
"""Test a shallow matrix multiplication factorization problem.
In this test, there are four matrices involved, X, Y, U, and V. The final
expression to optimize is mathematically
.. math::
(2 X - Y) * (2 U + V)
Here, the expression is to be given in its expanded form originally, and
we test if it can be factorized into something similar to what we have
above. Here we have the signs and coefficients to have better code
coverage for these cases. This test case more concentrates on the
horizontal complexity in the input.
"""
#
# Basic context setting-up.
#
dr = three_ranges
p = dr.names
m = p.m
a, b, c, d = p.a, p.b, p.c, p.d
# The indexed bases.
x = IndexedBase('X')
y = IndexedBase('Y')
u = IndexedBase('U')
v = IndexedBase('V')
t = IndexedBase('T')
# The target.
target = dr.define_einst(
t[a, b], 4 * x[a, c] * u[c, b] + 2 * x[a, c] * v[c, b] -
2 * y[a, c] * u[c, b] - y[a, c] * v[c, b])
targets = [target]
# The actual optimization.
res = optimize(targets, rand_constr=rand_constr)
assert len(res) == 3
# Test the correctness.
assert verify_eval_seq(res, targets, simplify=False)
# Test the cost.
cost = get_flop_cost(res)
leading_cost = get_flop_cost(res, leading=True)
assert cost == 2 * m**3 + 2 * m**2
assert leading_cost == 2 * m**3
cost = get_flop_cost(res, ignore_consts=False)
assert cost == 2 * m**3 + 4 * m**2
def test_eval_compression(three_ranges):
"""Test compression of optimized evaluations.
Here we have two targets,
.. math::
U X V + U Y V
and
.. math::
U X W + U Y W
and it has been deliberately made such that the multiplication with U
should be carried out first. Then after the factorization of U, we have
an intermediate U (X + Y), which is a sum of a single product
intermediate. This test succeeds when we have two intermediates only,
without the unnecessary addition of a single product.
"""
# Basic context setting-up.
dr = three_ranges
p = dr.names
a = p.a # Small range
i, j, k = p.i, p.j, p.k # Big range
# The indexed bases.
u = IndexedBase('U')
v = IndexedBase('V')
w = IndexedBase('W')
x = IndexedBase('X')
y = IndexedBase('Y')
s = IndexedBase('S')
t1 = IndexedBase('T1')
t2 = IndexedBase('T2')
# The target.
s_def = dr.define_einst(s[i, j], u[i, k] * x[k, j] + u[i, k] * y[k, j])
targets = [
dr.define_einst(t1[i, j], s_def[i, a] * v[a, j]),
dr.define_einst(t2[i, j], s_def[i, a] * w[a, j])
]
# The actual optimization.
res = optimize(targets, substs=dr.substs)
assert len(res) == 4
# Test the correctness.
assert verify_eval_seq(res, targets, simplify=False)
def test_ccsd_doubles(parthole_drudge):
"""Test discovery of effective T in CCSD doubles equation.
The purpose of this test is similar to the CCSD energy test. Just here the
more complexity about the external indices necessitates using ``ALL``
strategy for optimization.
"""
dr = parthole_drudge
p = dr.names
a, b, c, d = p.V_dumms[:4]
i, j = p.O_dumms[:2]
u = dr.two_body
t = IndexedBase('t')
dr.set_dbbar_base(t, 2)
tensor = dr.define_einst(
IndexedBase('r')[a, b, i, j],
t[c, d, i, j] * u[a, b, c, d] + u[a, b, c, d] * t[c, i] * t[d, j]
)
targets = [tensor]
all_eval_seq = optimize(
targets, substs={p.nv: p.no * 10},
strategy=Strategy.ALL | Strategy.SUM | Strategy.COMMON
)
assert verify_eval_seq(all_eval_seq, targets)
assert len(all_eval_seq) == 2
all_cost = get_flop_cost(all_eval_seq)
best_eval_seq = optimize(
targets, substs={p.nv: p.no * 10},
strategy=Strategy.BEST | Strategy.SUM | Strategy.COMMON
)
assert verify_eval_seq(best_eval_seq, targets)
assert len(best_eval_seq) == 2
best_cost = get_flop_cost(best_eval_seq)
assert (best_cost - all_cost).xreplace({p.no: 1, p.nv: 10}) > 0
def test_simple_scalar_optimization(spark_ctx):
"""Test optimization of a simple scalar.
There is not much optimization that can be done for simple scalars. But we
need to ensure that we get correct result here.
"""
dr = Drudge(spark_ctx)
a, b, r = symbols('a b r')
targets = [dr.define(r, a * b)]
eval_seq = optimize(targets)
assert verify_eval_seq(eval_seq, targets)
def test_ccsd_energy(parthole_drudge):
"""Test discovery of effective T in CCSD energy equation.
The purpose of this test is the capability of using locally non-optimal
contractions in the final summation optimization. The equation is not CCSD
energy equation exactly.
"""
dr = parthole_drudge
p = dr.names
a, b = p.V_dumms[:2]
i, j = p.O_dumms[:2]
u = dr.two_body
t = IndexedBase('t')
energy = dr.define_einst(
Symbol('e'),
u[i, j, a, b] * t[a, b, i, j] * Rational(1, 2)
+ u[i, j, a, b] * t[a, i] * t[b, j]
)
targets = [energy]
searched_eval_seq = optimize(targets, substs={p.nv: p.no * 10})
assert verify_eval_seq(searched_eval_seq, targets)
assert len(searched_eval_seq) == 2
searched_cost = get_flop_cost(searched_eval_seq)
best_eval_seq = optimize(
targets, substs={p.nv: p.no * 10},
strategy=Strategy.BEST | Strategy.SUM | Strategy.COMMON
)
assert verify_eval_seq(best_eval_seq, targets)
assert len(best_eval_seq) == 2
best_cost = get_flop_cost(best_eval_seq)
assert (best_cost - searched_cost).xreplace({p.no: 1, p.nv: 10}) > 0
def test_ccsd_doubles_a_terms(parthole_drudge):
r"""Test optimization of the a tau term in CCSD doubles equation.
This test is for the optimization of the
.. math::
a^{kl}_{ij} \tau^{ab}{kl}
in Equation (8) of GE Scuseria and HF Schaefer: J Chem Phys 90 (7) 1989.
"""
dr = parthole_drudge
p = dr.names
a, b, c, d = p.V_dumms[:4]
i, j, k, l = p.O_dumms[:4]
u = dr.two_body
t = IndexedBase('t')
dr.set_dbbar_base(t, 2)
tau = dr.define_einst(
IndexedBase('tau')[a, b, i, j],
Rational(1, 2) * t[a, b, i, j] + t[a, i] * t[b, j]
)
a_i = dr.define_einst(
IndexedBase('ai')[k, l, i, j], u[i, c, k, l] * t[c, j]
)
a_ = dr.define(
IndexedBase('a')[k, l, i, j],
u[k, l, i, j] +
a_i[k, l, i, j] - a_i[k, l, j, i]
+ u[k, l, c, d] * tau[c, d, i, j]
)
tensor = dr.define_einst(
IndexedBase('r')[a, b, i, j],
a_[k, l, i, j] * tau[a, b, k, l]
)
targets = [tensor]
eval_seq = optimize(
targets, substs={p.nv: p.no * 10}, strategy=Strategy.ALL | Strategy.SUM
)
assert verify_eval_seq(eval_seq, targets)
# Here we just assert that the final step is a simple product.
assert len(eval_seq[-1].rhs_terms) == 1
def test_interleaving_res_interm(three_ranges, res_at_end):
r"""Test the interleaving of results and intermediates.
Here we have intermediate,
.. math::
I = X Y
and result
.. math::
R1 = I * 2
and result
.. math::
R2 = I * tr(R1)
"""
dr = three_ranges
p = dr.names
a, b, c, d, e = p.a, p.b, p.c, p.d, p.e
x = IndexedBase('X')
y = IndexedBase('Y')
r1 = IndexedBase('R1')
r2 = IndexedBase('R2')
r1_def = dr.define_einst(r1[a, b], x[a, c] * y[c, b] * 2)
r2_def = dr.define_einst(r2[a, b], x[a, c] * y[c, b] * x[d, e] * y[e, d])
origs = [r1_def, r2_def]
eval_seq = optimize(origs, res_at_end=res_at_end)
assert verify_eval_seq(eval_seq, origs)
assert len(eval_seq) == 4
if res_at_end:
assert eval_seq[2].base == r1
else:
assert eval_seq[1].base == r1
assert eval_seq[3].base == r2
for i in eval_seq:
assert i.if_interm == (not (str(i.base)[0] == 'R'))
continue
def test_factorization_of_two_products(three_ranges):
"""Test a sum where we have two disjoint products.
The final expression to optimize is
.. math::
2 X (3 U + 5 V) - 7 Y (11 U + 13 V) + 17 T
In this test case, we concentrate on the handling of multiple disjoint
possible factorization inside a single sum.
"""
#
# Basic context setting-up.
#
dr = three_ranges
p = dr.names
m = p.m
a, b, c = p.a, p.b, p.c
# The indexed bases.
x = IndexedBase('X')
y = IndexedBase('Y')
u = IndexedBase('U')
v = IndexedBase('V')
t = IndexedBase('T')
# The target.
target = dr.define_einst(
IndexedBase('r')[a,
b], 6 * x[a, c] * u[c, b] + 10 * x[a, c] * v[c, b] -
77 * y[a, c] * u[c, b] - 91 * y[a, c] * v[c, b] + 17 * t[a, b])
targets = [target]
# The actual optimization.
res = optimize(targets)
assert len(res) == 3
assert res[-1].n_terms == 3
# Test the correctness.
assert verify_eval_seq(res, targets, simplify=True)
# Test the cost.
cost = get_flop_cost(res)
assert cost == 4 * m**3 + 4 * m**2
def test_general_matrix_problem(three_ranges):
"""Test optimization of a very general matrix computation.
This is a very general problem trying to test and illustrate many different
aspects of the optimization, parenthesization, recursion to newly-formed
factors, and sum of disjoint factorizations. The target to evaluate reads
.. math::
(A + 2B) (3C + 5D) (7E + 13F) + (17P + 19Q) (23X + 29Y)
where
- A, B, P, Q is over ranges M, L
- C, D is over M, N
- E, F is over N, L
- X, Y is over L, N
"""
dr = three_ranges
p = dr.names
m, n, l = p.m, p.n, p.l
a, b = p.a, p.b
i = p.i
p = p.p
f1 = IndexedBase('A')[a, i] + 2 * IndexedBase('B')[a, i]
f2 = 3 * IndexedBase('C')[i, p] + 5 * IndexedBase('D')[i, p]
f3 = 7 * IndexedBase('E')[p, b] + 13 * IndexedBase('F')[p, b]
f4 = 17 * IndexedBase('P')[a, i] + 19 * IndexedBase('Q')[a, i]
f5 = 23 * IndexedBase('X')[i, b] + 29 * IndexedBase('Y')[i, b]
target = dr.define_einst(
IndexedBase('R')[a, b], (f1 * f2 * f3 + f4 * f5).expand())
targets = [target]
assert target.n_terms == 12
assert get_flop_cost(targets).subs(dr.substs) == (144 * m**4 + 16 * m**3 +
11 * m**2)
eval_seq = optimize(targets, substs=dr.substs)
# Check the correctness.
assert verify_eval_seq(eval_seq, targets)
assert len(eval_seq) == 7
cost = get_flop_cost(eval_seq)
assert cost.subs(dr.substs) == 20 * m**3 + 16 * m**2
def test_matrix_chain(three_ranges):
"""Test a basic matrix chain multiplication problem.
Here a very simple matrix chain multiplication problem with three
matrices are used to test the factorization facilities. In this simple
test, we will have three matrices :math:`x`, :math:`y`, and :math:`z`,
which are of shapes :math:`m\\times n`, :math:`n \\times l`, and :math:`l
\\times m` respectively. In the factorization, we are going to set
:math:`n = 2 m` and :math:`l = 3 m`.
"""
dr = three_ranges
p = dr.names
m, n, l = p.m, p.n, p.l
# The indexed bases.
x = IndexedBase('x', shape=(m, n))
y = IndexedBase('y', shape=(n, l))
z = IndexedBase('z', shape=(l, m))
target_base = IndexedBase('t')
target = dr.define_einst(target_base[p.a, p.b],
x[p.a, p.i] * y[p.i, p.p] * z[p.p, p.b])
# Perform the factorization.
targets = [target]
stats = {}
eval_seq = optimize(targets, substs=dr.substs, stats=stats)
assert stats['Number of nodes'] < 2**3
assert len(eval_seq) == 2
# Check the correctness.
assert verify_eval_seq(eval_seq, targets)
# Check the cost.
cost = get_flop_cost(eval_seq)
leading_cost = get_flop_cost(eval_seq, leading=True)
expected_cost = 2 * l * m * n + 2 * m**2 * n
assert cost == expected_cost
assert leading_cost == expected_cost
def test_disconnected_outer_product_factorization(spark_ctx):
"""Test optimization of expressions with disconnected outer products.
"""
# Basic context setting-up.
dr = Drudge(spark_ctx)
n = Symbol('n')
r = Range('R', 0, n)
dumms = symbols('a b c d e f g')
a, b, c, d, e = dumms[:5]
dr.set_dumms(r, dumms)
dr.add_resolver_for_dumms()
# The indexed bases.
u = IndexedBase('U')
x = IndexedBase('X')
y = IndexedBase('Y')
z = IndexedBase('Z')
t = IndexedBase('T')
# The target.
target = dr.define_einst(
t[a, b],
u[a, b] * z[c, e] * x[e, c] + u[a, b] * z[c, e] * y[e, c]
)
targets = [target]
# The actual optimization.
res = optimize(targets)
assert len(res) == 3
# Test the correctness.
assert verify_eval_seq(res, targets, simplify=False)
# Test the cost.
cost = get_flop_cost(res)
leading_cost = get_flop_cost(res, leading=True)
assert cost == 4 * n ** 2
assert leading_cost == 4 * n ** 2
def test_conjugation_optimization(spark_ctx):
"""Test optimization of expressions containing complex conjugate.
"""
dr = Drudge(spark_ctx)
n = symbols('n')
r = Range('r', 0, n)
a, b, c, d = symbols('a b c d')
dr.set_dumms(r, [a, b, c, d])
dr.add_default_resolver(r)
p = IndexedBase('p')
x = IndexedBase('x')
y = IndexedBase('y')
z = IndexedBase('z')
targets = [
dr.define_einst(
p[a, b],
x[a, c] * conjugate(y[c, b]) + x[a, c] * conjugate(z[c, b]))
]
eval_seq = optimize(targets)
assert verify_eval_seq(eval_seq, targets)
def get_eval_seq():
"""Get the evaluation sequence for the benchmark."""
conf = SparkConf().setAppName('CCD-mosaic')
ctx = SparkContext(conf=conf)
dr = PartHoleDrudge(ctx)
dr.full_simplify = False
p = dr.names
a, b, c, d = p.V_dumms[:4]
i, j, k, l = p.O_dumms[:4]
u = dr.two_body
t = IndexedBase('t')
dr.set_dbbar_base(t, 2)
r = IndexedBase('r')
tensor = dr.define_einst(r[a, b, i, j],
t[a, b, l, j] * t[c, d, i, k] * u[k, l, c, d])
targets = [tensor]
eval_seq = optimize(targets,
substs={p.nv: p.no * NVNO_RATIO},
interm_fmt='tau{}')
return eval_seq, types.SimpleNamespace(no=p.no, nv=p.nv)
def test_matrix_factorization(spark_ctx):
"""Test a basic matrix multiplication factorization problem.
In this test, there are four matrices involved, X, Y, U, and V. And they
are used in two test cases for different scenarios.
"""
#
# Basic context setting-up.
#
dr = Drudge(spark_ctx)
n = Symbol('n')
r = Range('R', 0, n)
dumms = symbols('a b c d e f g')
a, b, c, d = dumms[:4]
dr.set_dumms(r, dumms)
dr.add_resolver_for_dumms()
# The indexed bases.
x = IndexedBase('X')
y = IndexedBase('Y')
u = IndexedBase('U')
v = IndexedBase('V')
t = IndexedBase('T')
#
# Test case 1.
#
# The final expression to optimize is mathematically
#
# .. math::
#
# (2 X - Y) * (2 U + V)
#
# Here, the expression is to be given in its extended form originally, and
# we test if it can be factorized into something similar to what we have
# above. Here we have the signs and coefficients to have better code
# coverage for these cases. This test case more concentrates on the
# horizontal complexity in the input.
#
# The target.
target = dr.define_einst(
t[a, b], 4 * x[a, c] * u[c, b] + 2 * x[a, c] * v[c, b] -
2 * y[a, c] * u[c, b] - y[a, c] * v[c, b])
targets = [target]
# The actual optimization.
res = optimize(targets)
assert len(res) == 3
# Test the correctness.
assert verify_eval_seq(res, targets, simplify=False)
# Test the cost.
cost = get_flop_cost(res)
leading_cost = get_flop_cost(res, leading=True)
assert cost == 2 * n**3 + 2 * n**2
assert leading_cost == 2 * n**3
cost = get_flop_cost(res, ignore_consts=False)
assert cost == 2 * n**3 + 4 * n**2
#
# Test case 2.
#
# The final expression to optimize is mathematically
#
# .. math::
#
# (X - 2 Y) * U * V
#
# Different from the first test case, here we concentrate more on the
# treatment of depth complexity in the input. The sum intermediate needs to
# be factored again.
#
# The target.
target = dr.define_einst(
t[a, b], x[a, c] * u[c, d] * v[d, b] - 2 * y[a, c] * u[c, d] * v[d, b])
targets = [target]
# The actual optimization.
res = optimize(targets)
assert len(res) == 3
# Test the correctness.
assert verify_eval_seq(res, targets, simplify=True)
# Test the cost.
cost = get_flop_cost(res)
leading_cost = get_flop_cost(res, leading=True)
assert cost == 4 * n**3 + n**2
assert leading_cost == 4 * n**3
cost = get_flop_cost(res, ignore_consts=False)
assert cost == 4 * n**3 + 2 * n**2
# Test disabling summation optimization.
res = optimize(targets, strategy=Strategy.BEST)
assert verify_eval_seq(res, targets, simplify=True)
new_cost = get_flop_cost(res, ignore_consts=False)
assert new_cost - cost != 0
def test_optimization_of_common_terms(spark_ctx):
"""Test optimization of common terms in summations.
In this test, there are just two matrices involved, X, Y. The target reads
.. math::
T[a, b] = X[a, b] - X[b, a] + 2 Y[a, b] - 2 Y[b, a]
Ideally, it should be evaluated as,
.. math::
I[a, b] = X[a, b] + 2 Y[a, b]
T[a, b] = I[a, b] - I[b, a]
or,
.. math::
I[a, b] = X[a, b] - 2 Y[b, a]
T[a, b] = I[a, b] - I[b, a]
Here, in order to emulate real cases where common term reference is in
interplay with factorization, the X and Y matrices are written as :math:`X =
S U` and :math:`Y = S V`.
"""
#
# Basic context setting-up.
#
dr = Drudge(spark_ctx)
n = Symbol('n')
r = Range('R', 0, n)
dumms = symbols('a b c d e f g')
a, b, c, d = dumms[:4]
dr.set_dumms(r, dumms)
dr.add_resolver_for_dumms()
# The indexed bases.
s = IndexedBase('S')
u = IndexedBase('U')
v = IndexedBase('V')
x = dr.define(IndexedBase('X')[a, b], s[a, c] * u[c, b])
y = dr.define(IndexedBase('Y')[a, b], s[a, c] * v[c, b])
t = dr.define_einst(
IndexedBase('t')[a, b], x[a, b] - x[b, a] + 2 * y[a, b] - 2 * y[b, a])
targets = [t]
eval_seq = optimize(targets)
assert len(eval_seq) == 3
# Check correctness.
verify_eval_seq(eval_seq, targets)
# Check cost.
cost = get_flop_cost(eval_seq)
assert cost == 2 * n**3 + 2 * n**2
cost = get_flop_cost(eval_seq, ignore_consts=False)
assert cost == 2 * n**3 + 3 * n**2
# Check the result when the common symmetrization optimization is disabled.
eval_seq = optimize(targets, strategy=Strategy.DEFAULT & ~Strategy.COMMON)
verify_eval_seq(eval_seq, targets)
new_cost = get_flop_cost(eval_seq, ignore_consts=True)
assert new_cost - cost != 0
def test_matrix_chain(spark_ctx):
"""Test a basic matrix chain multiplication problem.
Matrix chain multiplication problem is the classical problem that motivated
the algorithm for single-term optimization in this package. So here a very
simple matrix chain multiplication problem with three matrices are used to
test the factorization facilities. In this simple test, we will have three
matrices :math:`x`, :math:`y`, and :math:`z`, which are of shapes
:math:`m\\times n`, :math:`n \\times l`, and :math:`l \\times m`
respectively. In the factorization, we are going to set :math:`n = 2 m` and
:math:`l = 3 m`.
If we multiply the first two matrices first, the cost will be (two times)
.. math::
m n l + m^2 l
Or if we multiply the last two matrices first, the cost will be (two times)
.. math::
m n l + m^2 n
In addition to the classical matrix chain product, also tested is the
trace of their cyclic product.
.. math::
t = \\sum_i \\sum_j \\sum_k x_{i, j} y_{j, k} z_{k, i}
If we first take the product of :math:`Y Z`, the cost will be (two times)
:math:`n m l + n m`. For first multiplying :math:`X Y` and :math:`Z X`,
the costs will be (two times) :math:`n m l + m l` and :math:`n m l + n l`
respectively.
"""
#
# Basic context setting-up.
#
dr = Drudge(spark_ctx)
# The sizes.
m, n, l = symbols('m n l')
# The ranges.
m_range = Range('M', 0, m)
n_range = Range('N', 0, n)
l_range = Range('L', 0, l)
dr.set_dumms(m_range, symbols('a b c'))
dr.set_dumms(n_range, symbols('i j k'))
dr.set_dumms(l_range, symbols('p q r'))
dr.add_resolver_for_dumms()
# The indexed bases.
x = IndexedBase('x', shape=(m, n))
y = IndexedBase('y', shape=(n, l))
z = IndexedBase('z', shape=(l, m))
# The costs substitution.
substs = {n: m * 2, l: m * 3}
#
# Actual tests.
#
p = dr.names
target_base = IndexedBase('t')
target = dr.define_einst(target_base[p.a, p.b],
x[p.a, p.i] * y[p.i, p.p] * z[p.p, p.b])
# Perform the factorization.
targets = [target]
eval_seq = optimize(targets, substs=substs)
assert len(eval_seq) == 2
# Check the correctness.
assert verify_eval_seq(eval_seq, targets)
# Check the cost.
cost = get_flop_cost(eval_seq)
leading_cost = get_flop_cost(eval_seq, leading=True)
expected_cost = 2 * l * m * n + 2 * m**2 * n
assert cost == expected_cost
assert leading_cost == expected_cost
def test_removal_of_shallow_interms(spark_ctx):
"""Test removal of shallow intermediates.
Here we have two intermediates,
.. math::
U X V + U Y V
and
.. math::
U X W + U Y W
and it has been deliberately made such that the multiplication with U should
be carried out first. Then after the collection of U, we have a shallow
intermediate U (X + Y), which is a sum of a single product intermediate.
This test succeeds when we have two intermediates only without the shallow
ones.
"""
# Basic context setting-up.
dr = Drudge(spark_ctx)
n = Symbol('n')
r = Range('R', 0, n)
r_small = Range('R', 0, Rational(1 / 2) * n)
dumms = symbols('a b c d')
a, b, c = dumms[:3]
dumms_small = symbols('e f g h')
e = dumms_small[0]
dr.set_dumms(r, dumms)
dr.set_dumms(r_small, dumms_small)
dr.add_resolver_for_dumms()
# The indexed bases.
u = IndexedBase('U')
v = IndexedBase('V')
w = IndexedBase('W')
x = IndexedBase('X')
y = IndexedBase('Y')
s = IndexedBase('S')
t = IndexedBase('T')
# The target.
s_def = dr.define_einst(s[a, b], u[a, c] * x[c, b] + u[a, c] * y[c, b])
targets = [
dr.define_einst(t[a, b], s_def[a, e] * v[e, b]),
dr.define_einst(t[a, b], s_def[a, e] * w[e, b])
]
# The actual optimization.
res = optimize(targets)
assert len(res) == 4
# Test the correctness.
assert verify_eval_seq(res, targets, simplify=False)
