Skip to content

Commit d38b314

Browse files
evaleevevaleev
authored
Merge pull request #564 from ValeevGroup/evaleev/feature/mixed-t-tot-trees
expressions: mixed T x ToT products in arbitrary expression trees (Phase F)
2 parents cf8e9f0 + 3f00dae commit d38b314

6 files changed

Lines changed: 489 additions & 144 deletions

File tree

‎src/TiledArray/expressions/binary_engine.h‎

Lines changed: 46 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -258,6 +258,52 @@ class BinaryEngine : public ExprEngine<Derived> {
258258
return demand;
259259
}
260260

261+
/// Deduce each child's index SET top-down from this node's target and
262+
/// initialize the children with the deduced demands (the product-engine
263+
/// down-pass). The index set an inner node must produce depends on what
264+
/// its ancestors need: a shared index of the two children is fused iff
265+
/// this node's target carries it, contracted here otherwise; an index of
266+
/// one child that neither the sibling nor the target carries is consumed
267+
/// entirely within that child's subtree and is not demanded of it (a
268+
/// genuinely orphaned index -- an implicit trace, unsupported -- surfaces
269+
/// as the leaf-level target-is-not-a-permutation error). Each child's
270+
/// demand is ordered target-kept-first followed by the contracted-here
271+
/// indices in leaf-availability order, then reordered by the child's own
272+
/// preferred_layout() (a product child reorders to its canonical
273+
/// (fused, left-free, right-free) layout -- a general product cannot host
274+
/// a result permutation, and contraction consumers absorb any child
275+
/// layout via the GEMM transpose forms).
276+
void init_children_indices(const BipartiteIndexList& target_indices) {
277+
auto const avail_l = left_.available_indices();
278+
auto const avail_r = right_.available_indices();
279+
auto demand = [](auto const& avail, auto const& sibling, auto const& tgt) {
280+
container::svector<std::string> r;
281+
for (auto&& idx : tgt)
282+
if (avail.count(idx)) r.push_back(idx);
283+
for (auto&& idx : avail) {
284+
if (tgt.count(idx)) continue;
285+
if (sibling.count(idx)) r.push_back(idx);
286+
}
287+
return r;
288+
};
289+
auto bipartite_demand = [&demand](auto const& avail, auto const& sib,
290+
auto const& tgt) {
291+
auto const out = demand(outer(avail), outer(sib), outer(tgt));
292+
auto const in = demand(inner(avail), inner(sib), inner(tgt));
293+
return BipartiteIndexList(IndexList(out.begin(), out.end()),
294+
IndexList(in.begin(), in.end()));
295+
};
296+
// Materialize the demands as named lvalues: some preferred_layout()
297+
// overloads (leaf/binary/unary) return a reference to their argument, so
298+
// binding that to a temporary demand would be needlessly fragile.
299+
const BipartiteIndexList left_demand =
300+
bipartite_demand(avail_l, avail_r, target_indices);
301+
const BipartiteIndexList right_demand =
302+
bipartite_demand(avail_r, avail_l, target_indices);
303+
left_.init_indices(left_.preferred_layout(left_demand));
304+
right_.init_indices(right_.preferred_layout(right_demand));
305+
}
306+
261307
/// Initialize result tensor structure
262308

263309
/// This function will initialize the permutation, tiled range, and shape

‎src/TiledArray/expressions/cont_engine.h‎

Lines changed: 85 additions & 51 deletions
Original file line numberDiff line numberDiff line change
@@ -674,27 +674,46 @@ class ContEngine : public BinaryEngine<Derived> {
674674
this->init_perm(target_indices);
675675
general_repermute_ = (outer(target_indices) != outer(indices_));
676676

677-
// the tile op operates on the folded (fused-mode-free) shapes
678-
const auto left_op = to_cblas_op(left_outer_permtype_);
677+
// A product with NO external (free) outer indices (every outer index
678+
// fused or contracted, e.g. C("i,j;a,b") = A("x,i,j;a") * B("x,i,j;b"))
679+
// folds to a GEMM with no free modes, i.e. rank-0 tensors, which the
680+
// tile kernels do not support. Evaluate it with a SYNTHETIC UNIT
681+
// left-external mode instead: the folded product becomes
682+
// (1,K) x (K) -> (1), the exact shape of the (supported) one-sided
683+
// neB == 0 case. The unit mode lives only in the tile op's GemmHelper;
684+
// tranges, shapes and tiles carry the true (external-free) ranks, and
685+
// BatchedContractReduce / SparseShape::gemm_batched detect the
686+
// synthetic mode from the one-rank mismatch and pad their folded views
687+
// with a unit extent.
688+
const unsigned int u = (outer_size(indices_) == nh) ? 1u : 0u;
689+
690+
// the tile op operates on the folded (fused-mode-free) shapes; the
691+
// synthetic unit mode leads the folded left operand, so it is NoTrans
692+
const auto left_op =
693+
u ? math::blas::NoTranspose : to_cblas_op(left_outer_permtype_);
679694
const auto right_op = to_cblas_op(right_outer_permtype_);
680695
if constexpr (!TiledArray::detail::is_tensor_of_tensor_v<value_type>) {
681-
op_ = op_type(left_op, right_op, factor_, outer_size(indices_) - nh,
682-
outer_size(left_indices_) - nh,
696+
op_ = op_type(left_op, right_op, factor_, outer_size(indices_) - nh + u,
697+
outer_size(left_indices_) - nh + u,
683698
outer_size(right_indices_) - nh);
684699
} else {
685700
// the batched tile op must be perm-free (BatchedContractReduce cannot
686-
// host the folded-rank result permutation); the outer perm is empty by
687-
// the interleaved-target gate above, so only an explicit inner result
688-
// permutation can require one
689-
if (!implicit_permute_inner_ && bool(inner(perm_)))
701+
// host the folded-rank result permutation); the outer perm is handled
702+
// by the streaming re-permute (general_repermute_), so only a genuine
703+
// (non-identity) explicit inner result permutation requires one. N.B.
704+
// perm_ may carry a non-null identity inner component when only the
705+
// outer modes are permuted (the bipartite perm is constructed whole).
706+
if (!implicit_permute_inner_ && bool(inner(perm_)) &&
707+
!inner(perm_).is_identity())
690708
TA_EXCEPTION(
691709
"general products of tensors-of-tensors: a non-identity inner "
692710
"result permutation is not yet supported; reorder the inner "
693711
"annotation of the result");
694712

695713
// factor_ is absorbed into element_nonreturn_op_
696714
op_ = op_type(left_op, right_op, scalar_type(1),
697-
outer_size(indices_) - nh, outer_size(left_indices_) - nh,
715+
outer_size(indices_) - nh + u,
716+
outer_size(left_indices_) - nh + u,
698717
outer_size(right_indices_) - nh, BipartitePermutation{},
699718
this->element_nonreturn_op_, std::move(this->arena_plan_));
700719
// ce+e, ce+ce_right and ce+ce_left are mutually exclusive; at most one
@@ -734,7 +753,11 @@ class ContEngine : public BinaryEngine<Derived> {
734753
trange_type make_trange_general() const {
735754
const unsigned int nh = n_fused_modes_;
736755
const unsigned int nc = op_.gemm_helper().num_contract_ranks();
737-
const unsigned int neA = op_.gemm_helper().left_rank() - nc;
756+
// the no-external case carries a synthetic unit left-external mode in
757+
// the GemmHelper only (see init_struct_general); the actual tranges do
758+
// not have it
759+
const unsigned int u = (outer_size(indices_) == n_fused_modes_) ? 1u : 0u;
760+
const unsigned int neA = op_.gemm_helper().left_rank() - nc - u;
738761
const unsigned int neB = op_.gemm_helper().right_rank() - nc;
739762

740763
typename trange_type::Ranges ranges(nh + neA + neB);
@@ -786,7 +809,11 @@ class ContEngine : public BinaryEngine<Derived> {
786809
std::shared_ptr<const pmap_interface> pmap) {
787810
const unsigned int nh = n_fused_modes_;
788811
const unsigned int nc = op_.gemm_helper().num_contract_ranks();
789-
const unsigned int neA = op_.gemm_helper().left_rank() - nc;
812+
// the no-external case carries a synthetic unit left-external mode in
813+
// the GemmHelper only (see init_struct_general); the actual tranges do
814+
// not have it
815+
const unsigned int u = (outer_size(indices_) == nh) ? 1u : 0u;
816+
const unsigned int neA = op_.gemm_helper().left_rank() - nc - u;
790817
const unsigned int neB = op_.gemm_helper().right_rank() - nc;
791818

792819
// Get pointers to the argument sizes
@@ -1707,7 +1734,11 @@ class ContEngine : public BinaryEngine<Derived> {
17071734
TiledArray::detail::is_contraction_arena_tot_v<
17081735
result_tile_type, left_tile_type, right_tile_type>;
17091736
if constexpr (arena_eligible_scale) {
1710-
if (this->outer_product_uses_summa()) {
1737+
// the fused arena scale ops are factor-free; a non-unit
1738+
// expression-level prefactor (ScalMult) takes the fallback op,
1739+
// which absorbs it
1740+
if (this->outer_product_uses_summa() &&
1741+
this->factor_ == scalar_type(1)) {
17111742
// The inner perm handed to the plan must match how the inner
17121743
// *result* permutation is applied for this result cell type --
17131744
// and the two cell types apply it in different places:
@@ -1741,45 +1772,48 @@ class ContEngine : public BinaryEngine<Derived> {
17411772
// cells. The Hadamard outer product is an assignment
17421773
// `result = (perm ^ tot) * scalar`, which needs value-returning
17431774
// `scale`; only owning inner cells support it.
1744-
auto fallback_op = [perm = !this->implicit_permute_inner_
1745-
? inner(this->perm_)
1746-
: Permutation{},
1747-
outer_uses_summa =
1748-
this->outer_product_uses_summa()](
1749-
result_tile_element_type& result,
1750-
const left_tile_element_type& left,
1751-
const right_tile_element_type& right) {
1752-
if (outer_uses_summa) {
1753-
using TiledArray::axpy_to;
1754-
if constexpr (tot_x_t) {
1755-
if (left.empty()) return; // absent cell: no contribution
1756-
if (perm)
1757-
axpy_to(result, left, right, perm);
1758-
else
1759-
axpy_to(result, left, right);
1760-
} else {
1761-
if (right.empty()) return; // absent cell: no contribution
1762-
if (perm)
1763-
axpy_to(result, right, left, perm);
1764-
else
1765-
axpy_to(result, right, left);
1766-
}
1767-
} else {
1768-
if constexpr (!TiledArray::is_tensor_view_v<
1769-
result_tile_element_type>) {
1770-
using TiledArray::scale;
1771-
if constexpr (tot_x_t)
1772-
result = perm ? scale(left, right, perm) : scale(left, right);
1773-
else
1774-
result = perm ? scale(right, left, perm) : scale(right, left);
1775-
} else {
1776-
TA_EXCEPTION(
1777-
"Tensor<View> scale-inner Hadamard-outer product: a "
1778-
"view result cell cannot be value-assigned a fresh "
1779-
"scaled tensor");
1780-
}
1781-
}
1782-
};
1775+
// N.B. the expression-level scalar prefactor (factor_, != 1 for
1776+
// ScalMult expressions) multiplies the plain operand's element
1777+
auto fallback_op =
1778+
[perm = !this->implicit_permute_inner_ ? inner(this->perm_)
1779+
: Permutation{},
1780+
outer_uses_summa = this->outer_product_uses_summa(),
1781+
factor = this->factor_](result_tile_element_type& result,
1782+
const left_tile_element_type& left,
1783+
const right_tile_element_type& right) {
1784+
if (outer_uses_summa) {
1785+
using TiledArray::axpy_to;
1786+
if constexpr (tot_x_t) {
1787+
if (left.empty()) return; // absent cell: no contribution
1788+
if (perm)
1789+
axpy_to(result, left, right * factor, perm);
1790+
else
1791+
axpy_to(result, left, right * factor);
1792+
} else {
1793+
if (right.empty()) return; // absent cell: no contribution
1794+
if (perm)
1795+
axpy_to(result, right, left * factor, perm);
1796+
else
1797+
axpy_to(result, right, left * factor);
1798+
}
1799+
} else {
1800+
if constexpr (!TiledArray::is_tensor_view_v<
1801+
result_tile_element_type>) {
1802+
using TiledArray::scale;
1803+
if constexpr (tot_x_t)
1804+
result = perm ? scale(left, right * factor, perm)
1805+
: scale(left, right * factor);
1806+
else
1807+
result = perm ? scale(right, left * factor, perm)
1808+
: scale(right, left * factor);
1809+
} else {
1810+
TA_EXCEPTION(
1811+
"Tensor<View> scale-inner Hadamard-outer product: a "
1812+
"view result cell cannot be value-assigned a fresh "
1813+
"scaled tensor");
1814+
}
1815+
}
1816+
};
17831817
if constexpr (arena_eligible_scale) {
17841818
if (this->arena_plan_) {
17851819
if constexpr (tot_x_t)

‎src/TiledArray/expressions/mult_engine.h‎

Lines changed: 33 additions & 61 deletions
Original file line numberDiff line numberDiff line change
@@ -299,48 +299,7 @@ class MultEngine : public ContEngine<MultEngine<Left, Right, Result>> {
299299
// take the no-target init_indices() overload below, which retains the
300300
// bottom-up contraction convention -- general products under
301301
// reductions remain unsupported.
302-
{
303-
auto const avail_l = BinaryEngine_::left_.available_indices();
304-
auto const avail_r = BinaryEngine_::right_.available_indices();
305-
auto demand = [](auto const& avail, auto const& sibling,
306-
auto const& tgt) {
307-
container::svector<std::string> r;
308-
for (auto&& idx : tgt)
309-
if (avail.count(idx)) r.push_back(idx);
310-
for (auto&& idx : avail) {
311-
if (tgt.count(idx)) continue;
312-
if (sibling.count(idx)) r.push_back(idx);
313-
// else: the index is consumed entirely within the child's own
314-
// subtree (contracted deeper down) -- not demanded here. A
315-
// genuinely orphaned index (single occurrence, demanded nowhere =
316-
// implicit trace, unsupported) surfaces as the leaf-level
317-
// target-is-not-a-permutation error.
318-
}
319-
return r;
320-
};
321-
auto bipartite_demand = [&demand](auto const& avail, auto const& sib,
322-
auto const& tgt) {
323-
auto const out = demand(outer(avail), outer(sib), outer(tgt));
324-
auto const in = demand(inner(avail), inner(sib), inner(tgt));
325-
return BipartiteIndexList(IndexList(out.begin(), out.end()),
326-
IndexList(in.begin(), in.end()));
327-
};
328-
// each child dictates the ORDER of its demand (preferred_layout): a
329-
// product child reorders to its canonical (h, eA, eB) layout -- a
330-
// general product cannot host a result permutation, and contraction
331-
// consumers absorb any child layout via the GEMM transpose forms.
332-
// Materialize the demands as named lvalues: some preferred_layout()
333-
// overloads (leaf/binary/unary) return a reference to their argument, so
334-
// binding that to a temporary demand would be needlessly fragile.
335-
const BipartiteIndexList left_demand =
336-
bipartite_demand(avail_l, avail_r, target_indices);
337-
const BipartiteIndexList right_demand =
338-
bipartite_demand(avail_r, avail_l, target_indices);
339-
BinaryEngine_::left_.init_indices(
340-
BinaryEngine_::left_.preferred_layout(left_demand));
341-
BinaryEngine_::right_.init_indices(
342-
BinaryEngine_::right_.preferred_layout(right_demand));
343-
}
302+
BinaryEngine_::init_children_indices(target_indices);
344303

345304
this->product_type_ = compute_product_type(
346305
outer(BinaryEngine_::left_.indices()),
@@ -713,29 +672,34 @@ class ScalMultEngine
713672

714673
/// \param target_indices The target index list for this expression
715674
void init_indices(const BipartiteIndexList& target_indices) {
716-
BinaryEngine_::left_.init_indices();
717-
BinaryEngine_::right_.init_indices();
675+
// deduce the children's index sets top-down (see
676+
// BinaryEngine::init_children_indices), then classify and route exactly
677+
// as MultEngine does (the scalar factor does not affect index roles)
678+
BinaryEngine_::init_children_indices(target_indices);
679+
718680
this->product_type_ = compute_product_type(
719681
outer(BinaryEngine_::left_.indices()),
720682
outer(BinaryEngine_::right_.indices()), outer(target_indices));
683+
this->inner_product_type_ = compute_product_type(
684+
inner(BinaryEngine_::left_.indices()),
685+
inner(BinaryEngine_::right_.indices()), inner(target_indices));
686+
687+
if (this->inner_product_type_ == TensorProduct::General)
688+
TA_EXCEPTION(
689+
"ScalMultEngine: general products (fused + contracted + free "
690+
"indices) between the inner (nested) indices of tensors-of-tensors "
691+
"are not supported");
721692

722693
if (this->product_type() == TensorProduct::Hadamard) {
723694
// since already initialized left and right arg indices assign the target
724695
// indices
725696
BinaryEngine_::perm_indices(target_indices);
726697
} else if (this->product_type() == TensorProduct::General) {
727-
// layout via GeneralPermutationOptimizer (the target determines which
728-
// shared indices are fused vs contracted), then propagate to children
729-
if (!this->implicit_permute()) {
730-
BinaryEngine_::template init_indices_<TensorProduct::General>(
731-
target_indices);
732-
if (BinaryEngine_::left_indices_ != BinaryEngine_::left_.indices())
733-
BinaryEngine_::left_.perm_indices(BinaryEngine_::left_indices_);
734-
if (BinaryEngine_::right_indices_ != BinaryEngine_::right_.indices())
735-
BinaryEngine_::right_.perm_indices(BinaryEngine_::right_indices_);
736-
}
698+
this->perm_indices(target_indices);
737699
} else {
738-
ContEngine_::init_indices(target_indices);
700+
auto children_initialized = true;
701+
ContEngine_::init_indices(children_initialized);
702+
ContEngine_::perm_indices(target_indices);
739703
}
740704
}
741705

@@ -767,12 +731,14 @@ class ScalMultEngine
767731
/// for the result tensor.
768732
/// \param target_indices The target index list for the result tensor
769733
void init_struct(const BipartiteIndexList& target_indices) {
770-
// TODO Phase B (batched Summa): evaluate general products natively
771-
if (this->product_type() == TensorProduct::General)
772-
TA_EXCEPTION(
773-
"ScalMultEngine: evaluation of general products (fused + contracted "
774-
"+ free indices) via the expression layer is not yet implemented; "
775-
"use TiledArray::einsum() instead");
734+
if (this->product_type() == TensorProduct::General) {
735+
// the inner tile op (for tensors-of-tensors) must be initialized
736+
// first; init_struct_general consumes element_nonreturn_op_ and the
737+
// arena plan it builds
738+
this->init_inner_tile_op(inner(target_indices));
739+
ContEngine_::init_struct_general(target_indices);
740+
return;
741+
}
776742

777743
this->init_perm(target_indices);
778744

@@ -794,6 +760,8 @@ class ScalMultEngine
794760
std::shared_ptr<const pmap_interface> pmap) {
795761
if (this->product_type() == TensorProduct::Contraction)
796762
ContEngine_::init_distribution(world, pmap);
763+
else if (this->product_type() == TensorProduct::General)
764+
ContEngine_::init_distribution_general(world, pmap);
797765
else
798766
BinaryEngine_::init_distribution(world, pmap);
799767
}
@@ -804,6 +772,8 @@ class ScalMultEngine
804772
dist_eval_type make_dist_eval() const {
805773
if (this->product_type() == TensorProduct::Contraction)
806774
return ContEngine_::make_dist_eval();
775+
else if (this->product_type() == TensorProduct::General)
776+
return ContEngine_::make_dist_eval_general();
807777
else
808778
return BinaryEngine_::make_dist_eval();
809779
}
@@ -814,6 +784,8 @@ class ScalMultEngine
814784
trange_type make_trange() const {
815785
if (this->product_type() == TensorProduct::Contraction)
816786
return ContEngine_::make_trange();
787+
else if (this->product_type() == TensorProduct::General)
788+
return ContEngine_::make_trange_general();
817789
else
818790
return BinaryEngine_::make_trange();
819791
}

0 commit comments

Comments
 (0)