Skip to content

Commit 0e04820

Browse files
committed
tensor: fix dropped rows/cols in the ToT x scalar strided scale paths
The GEMM-based ToT x scalar scale paths of Tensor::gemm (and the T x ToT mirror) probe each row (column) for cleanliness; the presence probe stops at the first ABSENT cell, leaving the probed inner size A == -1 when the leading cell is absent. The subsequent 'A <= 0 => empty row, nothing to do' shortcut then dropped the ENTIRE row's contributions even when later cells were present. Rows whose leading contracted cell is absent (common for screened tensor-of-tensors, e.g. PNO-CC CSV intermediates) silently lost their contraction. Fix: when the probe ends with A <= 0, scan the full row (column) for any present cell; only a fully absent row is skipped, anything else takes the per-cell AXPY fallback. Also guard the engine's scale fallback element op against absent cells. This bug predates the general-product work but was masked on the legacy einsum route, whose canonical operand layout fails the NoTranspose gate of the strided path; the expression route's GEMM-canonical layout exposed it. Found with TA_EINSUM_DIFFERENTIAL on c6h14 PNO-CCSD: the opt-in expression-route energy error drops from 1.7 Eh to 6.9e-7 (the small residual is a screening-semantics difference -- the legacy route hard-zeroes sub-threshold result tiles that the engine route keeps -- plus a small systematic difference in phantom-unit denest products, both under review). Adds the CSV-like reproduction test (arena view cells, SparsePolicy, variable inner extents, screened cells, non-leading fused index, interleaved target): expression route and einsum routes agree to 1e-10, deterministically.
1 parent fe6703a commit 0e04820

5 files changed

Lines changed: 459 additions & 3 deletions

File tree

‎src/TiledArray/einsum/tiledarray.h‎

Lines changed: 45 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1096,6 +1096,51 @@ auto einsum(expressions::TsrExpr<ArrayA_> A, expressions::TsrExpr<ArrayB_> B,
10961096
<< (std::string)b + inner.b << " -> " << c_annot
10971097
<< " : diff2 = " << d2 << ", legacy2 = " << ref2
10981098
<< std::endl;
1099+
// per-tile forensics: compare tile norms of the two routes
1100+
auto tile_norm2 = [](auto const &tile) -> double {
1101+
using TileT =
1102+
std::remove_cv_t<std::remove_reference_t<decltype(tile)>>;
1103+
double n2 = 0;
1104+
if constexpr (TiledArray::detail::is_tensor_of_tensor_v<TileT>) {
1105+
for (std::size_t o = 0; o < tile.range().volume(); ++o) {
1106+
auto const &cell = tile.data()[o];
1107+
if (cell.empty()) continue;
1108+
for (std::size_t e = 0; e < cell.range().volume(); ++e)
1109+
n2 += std::abs(cell.data()[e]) * std::abs(cell.data()[e]);
1110+
}
1111+
} else {
1112+
for (std::size_t e = 0; e < tile.range().volume() * tile.nbatch();
1113+
++e)
1114+
n2 += std::abs(tile.data()[e]) * std::abs(tile.data()[e]);
1115+
}
1116+
return n2;
1117+
};
1118+
const auto ntiles = C.array.trange().tiles_range().volume();
1119+
std::size_t n_diff = 0, n_expr_zero = 0, n_legacy_zero = 0,
1120+
n_printed = 0;
1121+
for (std::size_t ord = 0; ord < ntiles; ++ord) {
1122+
const bool ez = expr_route_result->is_zero(ord);
1123+
const bool lz = C.array.is_zero(ord);
1124+
double en2 =
1125+
ez ? 0.0 : tile_norm2(expr_route_result->find(ord).get());
1126+
double ln2 = lz ? 0.0 : tile_norm2(C.array.find(ord).get());
1127+
const double dd = std::abs(en2 - ln2);
1128+
if (dd <= 1e-14 * std::max(std::max(en2, ln2), 1.0)) continue;
1129+
++n_diff;
1130+
if (en2 == 0.0) ++n_expr_zero;
1131+
if (ln2 == 0.0) ++n_legacy_zero;
1132+
if (world.rank() == 0 && n_printed < 8) {
1133+
++n_printed;
1134+
std::cerr << " tile " << ord << " ("
1135+
<< C.array.trange().tiles_range().idx(ord)
1136+
<< "): expr_n2 = " << en2 << ", legacy_n2 = " << ln2
1137+
<< std::endl;
1138+
}
1139+
}
1140+
if (world.rank() == 0)
1141+
std::cerr << " summary: " << n_diff << "/" << ntiles
1142+
<< " tiles differ by norm; expr-zero " << n_expr_zero
1143+
<< ", legacy-zero " << n_legacy_zero << std::endl;
10991144
}
11001145
}
11011146

‎src/TiledArray/expressions/cont_engine.h‎

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1679,11 +1679,13 @@ class ContEngine : public BinaryEngine<Derived> {
16791679
if (outer_uses_summa) {
16801680
using TiledArray::axpy_to;
16811681
if constexpr (tot_x_t) {
1682+
if (left.empty()) return; // absent cell: no contribution
16821683
if (perm)
16831684
axpy_to(result, left, right, perm);
16841685
else
16851686
axpy_to(result, left, right);
16861687
} else {
1688+
if (right.empty()) return; // absent cell: no contribution
16871689
if (perm)
16881690
axpy_to(result, right, left, perm);
16891691
else

‎src/TiledArray/tensor/tensor.h‎

Lines changed: 23 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -3374,7 +3374,19 @@ class Tensor {
33743374
}
33753375
detail::scale_phase_stop(detail::g_scale[0].check_str_ns,
33763376
_scale_tcs);
3377-
if (A <= 0) continue; // empty row -> nothing to do
3377+
if (A <= 0) {
3378+
// The presence probe stops at the first empty cell, so
3379+
// A <= 0 only means no non-empty cell was seen BEFORE the
3380+
// probe stopped -- NOT that the row is empty. The row
3381+
// contributes nothing only if EVERY left cell is absent;
3382+
// else it must take the per-cell AXPY fallback (whose
3383+
// element op skips absent cells).
3384+
bool row_empty = true;
3385+
for (integer k = 0; row_empty && k != K; ++k)
3386+
if (!lc0[k].empty()) row_empty = false;
3387+
if (row_empty) continue;
3388+
clean = false;
3389+
}
33783390
if (clean) {
33793391
// result[m,n][a] += sum_k left[m,k][a] * right[k,n].
33803392
// Row-major gemm: C2(N x A) += right^T(N x K) * L2(K x A),
@@ -3533,7 +3545,16 @@ class Tensor {
35333545
}
35343546
detail::scale_phase_stop(detail::g_scale[1].check_str_ns,
35353547
_scale_tcs);
3536-
if (A <= 0) continue;
3548+
if (A <= 0) {
3549+
// see the ToT x scalar mirror above: A <= 0 does not imply
3550+
// an empty column when the probe stopped at the first
3551+
// absent cell
3552+
bool col_empty = true;
3553+
for (integer k = 0; col_empty && k != K; ++k)
3554+
if (!right_data[k * N + n].empty()) col_empty = false;
3555+
if (col_empty) continue;
3556+
clean = false;
3557+
}
35373558
if (clean) {
35383559
// C_n(M x A) += left(M x K) * B_n(K x A). Row-major gemm.
35393560
if (detail::scale_gemm_timing_enabled()) {

‎src/TiledArray/tile_op/batched_contract_reduce.h‎

Lines changed: 11 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -193,8 +193,18 @@ class BatchedContractReduce {
193193
typename result_type::range_type(lobounds, upbounds));
194194
} else {
195195
// accumulate through a folded, zero-copy view of the result
196-
auto result_folded = result.reshape(fold_range(result.range()), batch);
196+
const auto full_range = result.range();
197+
auto result_folded = result.reshape(fold_range(full_range), batch);
197198
op_(result_folded, left_folded, right_folded);
199+
// the wrapped op may REBIND the result instead of writing in place:
200+
// the arena grow-to-cover path (a later K-panel touching inner cells
201+
// that earlier panels left null, for contracted-dimension-sparse
202+
// tensor-of-tensor operands) replaces the tile wholesale
203+
// (arena_tot_grow_inplace ends with `result = std::move(grown)`).
204+
// That rebinding lands on the folded view; propagate it to the
205+
// full-range result, else every post-growth contribution is lost.
206+
if (result_folded.data() != result.data())
207+
result = result_folded.reshape(full_range);
198208
}
199209
}
200210

0 commit comments

Comments
 (0)