diff --git a/src/TiledArray/math/blas.h b/src/TiledArray/math/blas.h index f5b2afbd00..202e8f5360 100644 --- a/src/TiledArray/math/blas.h +++ b/src/TiledArray/math/blas.h @@ -29,6 +29,7 @@ #include #include +#include // blas_int #include #include #include @@ -36,6 +37,7 @@ #include #include +#include namespace TiledArray::math::blas { @@ -48,6 +50,43 @@ static constexpr auto NoTranspose = Op::NoTrans; static constexpr auto Transpose = Op::Trans; static constexpr auto ConjTranspose = Op::ConjTrans; +/// The largest leading dimension the linked BLAS can accept. `integer` is +/// 64-bit here unconditionally, but BLAS++ narrows every dimension to its own +/// `blas_int` -- 32-bit under an LP64 BLAS -- and throws +/// (`blas::Error: ldb, in function to_blas_int_`) past this value. Callers that +/// MEASURE a leading dimension (as a distance between two addresses, rather +/// than deriving it from an extent) have to check it themselves; see +/// ld_fits(). +inline constexpr integer max_ld() { + return static_cast(std::numeric_limits::max()); +} + +/// The largest row offset `o` for which element o*ld + extent-1 of a +/// row-major matrix at leading dimension `ld` is still indexable by a +/// `blas_int`, i.e. the largest `o` with ld_fits(ld, o+1, extent). Phrased as +/// a division so the bound itself cannot overflow when `blas_int` is 64-bit +/// (ILP64), where max_ld() is INT64_MAX. Callers that grow a run one row at a +/// time should take this bound ONCE, while `ld` and `extent` are fixed, +/// rather than calling ld_fits() per row -- that divides every time round the +/// loop. +inline constexpr integer max_ld_offset(integer ld, integer extent) { + if (ld <= 0 || extent < 1 || extent - 1 > max_ld()) return 0; + return (max_ld() - (extent - 1)) / ld; +} + +/// True if a row-major matrix of `nslab` rows and `extent` columns at leading +/// dimension `ld` addresses only elements a `blas_int` can index. max_ld() +/// alone bounds ONE row step; the GEMM reaches element +/// (nslab-1)*ld + extent-1, and a 32-bit-indexed BLAS must be able to form +/// THAT index too, so a stride comfortably under the cap still overflows it +/// once there are enough rows. +inline constexpr bool ld_fits(integer ld, integer nslab, integer extent) { + if (ld < 0 || extent < 1 || extent - 1 > max_ld() || ld > max_ld()) + return false; + if (nslab <= 1) return true; + return nslab - 1 <= max_ld_offset(ld, extent); +} + /// converts Op to ints in manner useful for bit manipulations /// NoTranspose -> 0, Transpose->1, ConjTranspose->2 inline int64_t to_int(Op op) { diff --git a/src/TiledArray/tensor/arena_einsum.h b/src/TiledArray/tensor/arena_einsum.h index e4194f763f..411654aade 100644 --- a/src/TiledArray/tensor/arena_einsum.h +++ b/src/TiledArray/tensor/arena_einsum.h @@ -279,18 +279,45 @@ inline void phase_stop(std::atomic& acc, /// 2 = nonuniform: present, but inner-cell sizes differ along the run /// 3 = stride : present + uniform, but cells are NOT at a constant /// page-jump-free stride (the strided-GEMM precondition) +/// 4 = ld range : constant stride, but the run spans more elements than the +/// BLAS integer can index (see ld_fits) /// 0 = run looks clean (so the rejection came from the OTHER operand run) inline std::atomic g_fall_runs_ce_ce{0}; inline std::atomic g_fall_res_absent_ce_ce{0}; // 1 inline std::atomic g_fall_res_nonuniform_ce_ce{0}; // 2 inline std::atomic g_fall_res_stride_ce_ce{0}; // 3 +inline std::atomic g_fall_res_ldrange_ce_ce{0}; // 4 inline std::atomic g_fall_op_absent_ce_ce{0}; // 13 inline std::atomic g_fall_op_nonuniform_ce_ce{0}; // 14 inline std::atomic g_fall_op_stride_ce_ce{0}; // 15 inline std::atomic g_fall_op_acrossk_ce_ce{0}; // 16 inline std::atomic g_fall_both_clean_ce_ce{0}; // 17 +inline std::atomic g_fall_op_ldrange_ce_ce{0}; // 18 -/// Classify a strided run: 1=absent, 2=nonuniform size, 3=bad stride, 0=clean. +// --------------------------------------------------------------------------- +// Measured leading dimensions vs. the BLAS index range. +// +// The strided kernels below MEASURE their leading dimensions as pointer +// differences between neighbouring cells, so -- unlike a kernel that derives +// them from extents -- they can hand blaspp an ld it cannot convert to +// `blas_int` (`blas::Error: ldb, in function to_blas_int_`). Two present +// cells that are not arena neighbours -- e.g. individually allocated inner +// tensors the allocator placed gigabytes apart -- satisfy every other +// precondition of a 2-cell run (uniform size, stride >= cell size, trivially +// constant) yet yield exactly such an address delta; and a stride well inside +// `blas_int` still overflows the BLAS index once the run is long enough, +// because the GEMM reaches element (nslab-1)*ld + extent-1. Such a run is not +// strided-GEMM material: the walkers break it up and the per-cell path +// handles it. +// +// math::blas::ld_fits() is the whole-run test. The segment walkers instead +// take math::blas::max_ld_offset() once per segment -- the same bound solved +// for the slab offset -- because ld_fits() divides, and the admission loop +// runs once per cell. +// --------------------------------------------------------------------------- + +/// Classify a strided run: 1=absent, 2=nonuniform size, 3=bad stride, +/// 4=stride past the BLAS index range (see ld_fits), 0=clean. template inline int classify_run(GetCell getcell, std::size_t n) { if (n == 0) return 0; @@ -313,6 +340,8 @@ inline int classify_run(GetCell getcell, std::size_t n) { for (std::size_t i = 0; i < n; ++i) if (getcell(i).data() != base + static_cast(i) * st) return 3; // non-constant stride + // ...and the whole run must stay inside the BLAS index range + if (!math::blas::ld_fits(st, static_cast(n), s0)) return 4; return 0; } @@ -320,7 +349,8 @@ inline int classify_run(GetCell getcell, std::size_t n) { /// is the strided operand run (length `nrun`, per outer-contraction k); getL(k) /// is the per-k single (non-strided) operand cell, expected size P*Q. Returns /// 13=absent/size, 14=nonuniform, 15=bad stride within a k, 16=stride varies -/// across k, 17=clean (gate rejected a run this re-check finds valid). +/// across k, 18=stride past the BLAS index range (see ld_fits), 17=clean (gate +/// rejected a run this re-check finds valid). template inline int classify_operand(GetR getR, GetL getL, std::size_t nrun, std::size_t nK, long P) { @@ -352,6 +382,7 @@ inline int classify_operand(GetR getR, GetL getL, std::size_t nrun, for (std::size_t i = 0; i < nrun; ++i) if (getR(k, i).data() != base + static_cast(i) * sk) return 15; + if (!math::blas::ld_fits(sk, static_cast(nrun), Q)) return 18; if (k == 0) sR = sk; else if (sk != sR) @@ -489,6 +520,7 @@ inline void measure_segments(GetC getC, GetR getR, GetL getL, std::size_t nrun, const auto* rb = r0.data(); std::size_t end = mu + 1; long sC = -1, sR = -1; + long max_off = 0; // set with the strides, below while (end < nrun) { const auto& ce = getC(end); const auto& re = getR(k, end); @@ -503,9 +535,15 @@ inline void measure_segments(GetC getC, GetR getR, GetL getL, std::size_t nrun, sC = dc; sR = dr; if (sC < P || sR < Q) break; + // A segment of off+1 slabs reaches element off*sC+P-1 of C and + // off*sR+Q-1 of R, both of which the BLAS has to be able to index. + // The strides and extents are fixed from here on, so bound off once. + max_off = std::min(math::blas::max_ld_offset(sC, P), + math::blas::max_ld_offset(sR, Q)); } else if (dc != off * sC || dr != off * sR) { break; } + if (off > max_off) break; ++end; } const std::size_t len = end - mu; @@ -528,10 +566,12 @@ inline void record_ce_ce_fallback(int why) { case 1: g_fall_res_absent_ce_ce.fetch_add(1, std::memory_order_relaxed); break; case 2: g_fall_res_nonuniform_ce_ce.fetch_add(1, std::memory_order_relaxed); break; case 3: g_fall_res_stride_ce_ce.fetch_add(1, std::memory_order_relaxed); break; + case 4: g_fall_res_ldrange_ce_ce.fetch_add(1, std::memory_order_relaxed); break; case 13: g_fall_op_absent_ce_ce.fetch_add(1, std::memory_order_relaxed); break; case 14: g_fall_op_nonuniform_ce_ce.fetch_add(1, std::memory_order_relaxed); break; case 15: g_fall_op_stride_ce_ce.fetch_add(1, std::memory_order_relaxed); break; case 16: g_fall_op_acrossk_ce_ce.fetch_add(1, std::memory_order_relaxed); break; + case 18: g_fall_op_ldrange_ce_ce.fetch_add(1, std::memory_order_relaxed); break; default: g_fall_both_clean_ce_ce.fetch_add(1, std::memory_order_relaxed); break; } } @@ -540,7 +580,8 @@ inline void record_ce_ce_fallback(int why) { // ce+e phase timers + fallback diagnosis (mirror of the ce+ce instrumentation). // In ce+e a "run" is one result cell (m,n); the clean check is over the k-slabs // of L (stride ldA) and R (stride ldB). Fallback reasons classify those k-runs: -// L-run: 1 absent / 2 nonuniform / 3 stride; R-run: 11 / 12 / 13; 17 clean. +// L-run: 1 absent / 2 nonuniform / 3 stride / 4 ld range; +// R-run: 11 / 12 / 13 / 14; 17 clean. inline std::atomic g_kernel_ns_ce_e{0}; inline std::atomic g_check_ns_ce_e{0}; inline std::atomic g_fallback_ns_ce_e{0}; @@ -548,9 +589,11 @@ inline std::atomic g_e_fall_runs{0}; inline std::atomic g_e_l_absent{0}; inline std::atomic g_e_l_nonuniform{0}; inline std::atomic g_e_l_stride{0}; +inline std::atomic g_e_l_ldrange{0}; inline std::atomic g_e_r_absent{0}; inline std::atomic g_e_r_nonuniform{0}; inline std::atomic g_e_r_stride{0}; +inline std::atomic g_e_r_ldrange{0}; inline std::atomic g_e_both_clean{0}; inline void record_ce_e_fallback(int why) { @@ -560,9 +603,11 @@ inline void record_ce_e_fallback(int why) { case 1: g_e_l_absent.fetch_add(1, std::memory_order_relaxed); break; case 2: g_e_l_nonuniform.fetch_add(1, std::memory_order_relaxed); break; case 3: g_e_l_stride.fetch_add(1, std::memory_order_relaxed); break; + case 4: g_e_l_ldrange.fetch_add(1, std::memory_order_relaxed); break; case 11: g_e_r_absent.fetch_add(1, std::memory_order_relaxed); break; case 12: g_e_r_nonuniform.fetch_add(1, std::memory_order_relaxed); break; case 13: g_e_r_stride.fetch_add(1, std::memory_order_relaxed); break; + case 14: g_e_r_ldrange.fetch_add(1, std::memory_order_relaxed); break; default: g_e_both_clean.fetch_add(1, std::memory_order_relaxed); break; } } @@ -659,12 +704,16 @@ struct GemmTimingDumper { << " (" << fp(L(g_e_l_nonuniform)) << "%)\n"; std::cerr << "[ce+e-fallback] L-run bad stride : " << L(g_e_l_stride) << " (" << fp(L(g_e_l_stride)) << "%)\n"; + std::cerr << "[ce+e-fallback] L-run ld range : " << L(g_e_l_ldrange) + << " (" << fp(L(g_e_l_ldrange)) << "%)\n"; std::cerr << "[ce+e-fallback] R-run absent : " << L(g_e_r_absent) << " (" << fp(L(g_e_r_absent)) << "%)\n"; std::cerr << "[ce+e-fallback] R-run nonuniform : " << L(g_e_r_nonuniform) << " (" << fp(L(g_e_r_nonuniform)) << "%)\n"; std::cerr << "[ce+e-fallback] R-run bad stride : " << L(g_e_r_stride) << " (" << fp(L(g_e_r_stride)) << "%)\n"; + std::cerr << "[ce+e-fallback] R-run ld range : " << L(g_e_r_ldrange) + << " (" << fp(L(g_e_r_ldrange)) << "%)\n"; std::cerr << "[ce+e-fallback] both runs clean : " << L(g_e_both_clean) << " (" << fp(L(g_e_both_clean)) << "%)\n"; } @@ -683,6 +732,9 @@ struct GemmTimingDumper { std::cerr << "[ce+ce-fallback] result bad stride : " << L(g_fall_res_stride_ce_ce) << " (" << fp(L(g_fall_res_stride_ce_ce)) << "%)\n"; + std::cerr << "[ce+ce-fallback] result ld range : " + << L(g_fall_res_ldrange_ce_ce) << " (" + << fp(L(g_fall_res_ldrange_ce_ce)) << "%)\n"; std::cerr << "[ce+ce-fallback] operand absent/size : " << L(g_fall_op_absent_ce_ce) << " (" << fp(L(g_fall_op_absent_ce_ce)) << "%)\n"; @@ -695,6 +747,9 @@ struct GemmTimingDumper { std::cerr << "[ce+ce-fallback] operand stride X k : " << L(g_fall_op_acrossk_ce_ce) << " (" << fp(L(g_fall_op_acrossk_ce_ce)) << "%)\n"; + std::cerr << "[ce+ce-fallback] operand ld range : " + << L(g_fall_op_ldrange_ce_ce) << " (" + << fp(L(g_fall_op_ldrange_ce_ce)) << "%)\n"; std::cerr << "[ce+ce-fallback] both runs clean (!) : " << L(g_fall_both_clean_ce_ce) << " (" << fp(L(g_fall_both_clean_ce_ce)) << "%)\n"; @@ -1249,6 +1304,9 @@ void arena_strided_gemm_ce_e(ResultOuter& C, const LeftOuter& L, ldA = static_cast(lc[lbase + a_off(m, 1)].data() - l0.data()); ldB = static_cast(rc[rbase + b_off(1, n)].data() - r0.data()); if (ldA < P || ldB < Q) clean = false; + if (!math::blas::ld_fits(ldA, static_cast(K), P) || + !math::blas::ld_fits(ldB, static_cast(K), Q)) + clean = false; for (std::size_t k = 0; clean && k < K; ++k) { if (lc[lbase + a_off(m, k)].data() != l0.data() + static_cast(k) * ldA) @@ -1536,6 +1594,7 @@ void arena_strided_gemm_ce_ce_right(ResultOuter& C, const LeftOuter& L, // reuse a run-wide stale stride). std::size_t end = mu + 1; long sR = -1, sC = -1; + long max_off = 0; // set with the strides, below while (end < Mmu) { const auto& rce = rc[rbase + r_off(k, end)]; const auto& cce = cc[cbase + c_off(m, end)]; @@ -1550,9 +1609,16 @@ void arena_strided_gemm_ce_ce_right(ResultOuter& C, const LeftOuter& L, sR = dR; sC = dC; if (sR < Q || sC < P) break; // page-jump / overlap + // A segment of off+1 slabs reaches element off*sR+Q-1 of A and + // off*sC+P-1 of C, both of which the BLAS has to be able to + // index. The strides and extents are fixed from here on, so + // bound off once rather than dividing per admitted cell. + max_off = std::min(blas::max_ld_offset(sR, Q), + blas::max_ld_offset(sC, P)); } else if (dR != off * sR || dC != off * sC) { break; } + if (off > max_off) break; ++end; } const std::size_t Mseg = end - mu; @@ -1779,6 +1845,7 @@ void arena_strided_gemm_ce_ce_left(ResultOuter& C, const LeftOuter& L, // reuse a run-wide stale stride). std::size_t end = m + 1; long sA = -1, sC = -1; + long max_off = 0; // set with the strides, below while (end < Mo) { const auto& lce = lc[lbase + l_off(end, k)]; const auto& cce = cc[cbase + c_off(end, n)]; @@ -1793,9 +1860,16 @@ void arena_strided_gemm_ce_ce_left(ResultOuter& C, const LeftOuter& L, sA = dA; sC = dC; if (sA < Q || sC < P) break; // page-jump / overlap + // A segment of off+1 slabs reaches element off*sA+Q-1 of A and + // off*sC+P-1 of C, both of which the BLAS has to be able to + // index. The strides and extents are fixed from here on, so + // bound off once rather than dividing per admitted cell. + max_off = std::min(blas::max_ld_offset(sA, Q), + blas::max_ld_offset(sC, P)); } else if (dA != off * sA || dC != off * sC) { break; } + if (off > max_off) break; ++end; } const std::size_t Mseg = end - m; diff --git a/src/TiledArray/tensor/tensor.h b/src/TiledArray/tensor/tensor.h index 3b033a4e4c..e1c8d75278 100644 --- a/src/TiledArray/tensor/tensor.h +++ b/src/TiledArray/tensor/tensor.h @@ -3704,6 +3704,20 @@ class Tensor { if (N > 1) ldc = static_cast(rc0[1].data() - rc0[0].data()); if (ldb < A || ldc < A) clean = false; // sanity + // ...and the strides must be indexable by the BLAS. ldb/ldc + // are MEASURED (a distance between two cell addresses), not + // derived from an extent, so an un-compacted ToT tile whose + // cells the allocator placed far apart can produce one that + // BLAS++ cannot narrow to blas_int -- `blas::Error: ldb, in + // function to_blas_int_` under LP64. The gemm below reads + // element (K-1)*ldb*cw + A*cw-1 of the left slab and writes + // (N-1)*ldc*cw + A*cw-1 of the result slab, so bound the whole + // span, not just one step. + if (!TiledArray::math::blas::ld_fits( + ldb * cw, K, static_cast(A) * cw) || + !TiledArray::math::blas::ld_fits( + ldc * cw, N, static_cast(A) * cw)) + clean = false; const std::ptrdiff_t sb = ldb, sc = ldc; for (integer k = 0; clean && k != K; ++k) if (lc(k).data() != lc(0).data() + k * sb) clean = false; @@ -3903,6 +3917,15 @@ class Tensor { ldc = static_cast(this_data[N + n].data() - this_data[n].data()); if (sbc < A || ldc < A) clean = false; + // measured strides, so bound their BLAS-indexable span too + // (see the ToT x scalar mirror above): the gemm reads element + // (K-1)*sbc*cw + A*cw-1 of the right slab and writes + // (M-1)*ldc*cw + A*cw-1 of the result slab. + if (!TiledArray::math::blas::ld_fits( + sbc * cw, K, static_cast(A) * cw) || + !TiledArray::math::blas::ld_fits( + ldc * cw, M, static_cast(A) * cw)) + clean = false; const std::ptrdiff_t sb = sbc, sc = ldc; for (integer k = 0; clean && k != K; ++k) if (right_data[rcell(k, n)].data() != diff --git a/tests/CMakeLists.txt b/tests/CMakeLists.txt index e0da3296b7..e77f7cf1b0 100644 --- a/tests/CMakeLists.txt +++ b/tests/CMakeLists.txt @@ -97,6 +97,7 @@ set(ta_test_src_files ta_test.cpp dot_inner.cpp general_product.cpp linalg.cpp + math_blas.cpp cp.cpp btas.cpp arena.cpp diff --git a/tests/arena_strided_gemm.cpp b/tests/arena_strided_gemm.cpp index 802e933f09..68e2003c58 100644 --- a/tests/arena_strided_gemm.cpp +++ b/tests/arena_strided_gemm.cpp @@ -10,10 +10,15 @@ #include #include #include +#include +#include #include +#include #include #include +#include + namespace TA = TiledArray; using Inner = TA::ArenaTensor; using Outer = TA::Tensor; @@ -376,6 +381,88 @@ BOOST_AUTO_TEST_CASE(ce_e_multi_external_inner) { #endif } +BOOST_AUTO_TEST_CASE(ce_e_stride_past_blas_int_falls_back) { + // Two present, uniform-size left k-cells at a constant stride that does not + // fit blas_int: individually allocated inner tensors the allocator placed + // gigabytes apart pass every other strided-GEMM precondition of a 2-cell + // run. The kernel must reject the run (per-cell path) rather than hand + // BLAS++ the measured stride (`blas::Error: ldb, in function to_blas_int_` + // under LP64). Only expressible when blas_int is narrower than integer. The + // second cell sits in lazily reserved address space, so only the two cells' + // pages are ever touched; if the reservation is refused the case is skipped. + namespace blas = TiledArray::math::blas; + const std::size_t M = 1, N = 1, K = 2, P = 3, Q = 4; + // Always run, so neither skip below leaves the case without a single + // assertion: an ordinary contiguous 2-cell run is still accepted. + BOOST_CHECK(blas::ld_fits(static_cast(Q), 2, + static_cast(Q))); + if (blas::max_ld() >= std::numeric_limits::max()) { + BOOST_TEST_MESSAGE( + "ILP64: blas_int is as wide as integer, so no stride can outrun it; " + "skipped"); + return; + } + const std::size_t stride = static_cast(blas::max_ld()) + 1; + const std::size_t align = Inner::cell_alignment(); + const std::size_t len = stride * sizeof(double) + Inner::cell_size(P) + align; + void* mem = ::mmap(nullptr, len, PROT_READ | PROT_WRITE, + MAP_PRIVATE | MAP_ANONYMOUS | MAP_NORESERVE, -1, 0); + if (mem == MAP_FAILED) { + BOOST_TEST_MESSAGE("cannot reserve " << len + << " bytes of address space; skipped"); + return; + } + using Cell = Inner::Cell; + // Owns the reservation and the placement-new'd cells together: a failing + // BOOST_REQUIRE below must not skip ~Cell(), whose TA::Range owns a heap + // buffer -- under this repo's ASan debug build that would bury the real + // failure in leak reports. + struct Reservation { + void* mem; + std::size_t len; + Cell* c0 = nullptr; + Cell* c1 = nullptr; + ~Reservation() { + if (c1) c1->~Cell(); + if (c0) c0->~Cell(); + ::munmap(mem, len); + } + } reservation{mem, len}; + auto* base = reinterpret_cast( + (reinterpret_cast(mem) + align - 1) & ~(align - 1)); + // stride * sizeof(double) is a power of two >= align, so cell 1 is aligned + Cell* c0 = reservation.c0 = ::new (base) Cell{TA::Range{P}}; + Cell* c1 = reservation.c1 = + ::new (base + stride * sizeof(double)) Cell{TA::Range{P}}; + { + Inner l0(c0), l1(c1); + BOOST_REQUIRE_EQUAL(static_cast(l1.data() - l0.data()), + stride); + for (std::size_t p = 0; p < P; ++p) { + l0.data()[p] = 1.0 + p; + l1.data()[p] = 2.0 + p; + } + std::allocator alloc; + Inner* raw = alloc.allocate(K); + ::new (raw) Inner(l0); + ::new (raw + 1) Inner(l1); + std::shared_ptr ldata(raw, [alloc, n = K](Inner* p) mutable { + for (std::size_t i = 0; i < n; ++i) (p + i)->~Inner(); + alloc.deallocate(p, n); + }); + Outer L(TA::Range{M, K}, /*nbatch=*/1, std::move(ldata)); + Outer R = make_filled( + TA::Range{N, K}, [&](std::size_t) { return TA::Range{Q}; }, 2.0); + Outer C = TA::detail::arena_outer_init( + TA::Range{M, N}, 1, [&](std::size_t) { return TA::Range{P, Q}; }); + BOOST_REQUIRE_NO_THROW(TA::detail::arena_strided_gemm_ce_e( + C, L, R, M, N, K, blas::NoTranspose, blas::Transpose, 1.0)); + const auto ref = ref_ce_e(L, R, 0, 0, K, P, Q, 1.0); + const double* got = C.data()[0].data(); + for (std::size_t e = 0; e < P * Q; ++e) + BOOST_CHECK_CLOSE(got[e], ref[e], 1e-12); + } // the views die before their cells, which ~Reservation then destroys +} #ifdef TA_STRIDED_GEMM_COUNT BOOST_AUTO_TEST_CASE(ce_e_fires_clean_path) { namespace blas = TiledArray::math::blas; diff --git a/tests/math_blas.cpp b/tests/math_blas.cpp index 0ae418f8ad..76c0462ed4 100644 --- a/tests/math_blas.cpp +++ b/tests/math_blas.cpp @@ -353,4 +353,45 @@ BOOST_AUTO_TEST_CASE_TEMPLATE(complex_gemm_ld, T, floating_point_types) { delete[] c; } +BOOST_AUTO_TEST_CASE(measured_ld_fits) { + // ld_fits() guards the leading dimensions the arena strided kernels + // (arena_einsum.h) and Tensor::gemm's arena scale paths MEASURE as a + // distance between two inner-cell addresses: BLAS++ narrows every dimension + // to blas_int (max_ld()), and a GEMM of nslab rows at leading dimension ld + // reaches element (nslab - 1) * ld + extent - 1, which must be + // representable too. + namespace blas = TiledArray::math::blas; + const integer cap = blas::max_ld(); + BOOST_CHECK_EQUAL(cap, + static_cast(std::numeric_limits::max())); + + // the ordinary contiguous slab always fits + BOOST_CHECK(blas::ld_fits(64, 1, 64)); + BOOST_CHECK(blas::ld_fits(64, 1000, 64)); + // degenerate arguments never do + BOOST_CHECK(!blas::ld_fits(-1, 1, 1)); + BOOST_CHECK(!blas::ld_fits(64, 1, 0)); + // one slab: only the step itself must be representable + BOOST_CHECK(blas::ld_fits(cap, 1, 1)); + BOOST_CHECK(blas::ld_fits(cap, 1, cap)); + // more slabs: the last addressed element must be representable + BOOST_CHECK(blas::ld_fits(cap, 2, 1)); // reaches cap + BOOST_CHECK(!blas::ld_fits(cap, 2, 2)); // reaches cap + 1 + BOOST_CHECK(blas::ld_fits(cap - 1, 2, 2)); // reaches cap + BOOST_CHECK(!blas::ld_fits(cap - 1, 2, 3)); // reaches cap + 1 + const integer nslab = 1000, extent = 7; + const integer ld_max = (cap - (extent - 1)) / (nslab - 1); + BOOST_CHECK(blas::ld_fits(ld_max, nslab, extent)); + BOOST_CHECK(!blas::ld_fits(ld_max + 1, nslab, extent)); + // past the cap outright: only expressible when blas_int is narrower than + // integer (LP64); under ILP64 the cap is integer's own maximum + if (cap < std::numeric_limits::max()) { + BOOST_CHECK(!blas::ld_fits(cap + 1, 1, 1)); + // the motivating case: a 2-cell run whose measured stride is the + // distance between two individually allocated inner tensors + BOOST_CHECK(!blas::ld_fits(cap + 1, 2, 64)); + BOOST_CHECK(!blas::ld_fits(1, 1, cap + 2)); + } +} + BOOST_AUTO_TEST_SUITE_END()