tests: cover the measured-leading-dimension guard

blas_suite/measured_ld_fits pins ld_fits()/max_ld() at their boundaries: one
slab needs only the step itself to be representable, more slabs need the
last addressed element (nslab-1)*ld + extent-1 to be, and (under LP64) a
stride past blas_int is rejected outright.

arena_strided_gemm_suite/ce_e_stride_past_blas_int_falls_back drives
arena_strided_gemm_ce_e with two present, uniform-size left k-cells placed
max_ld()+1 elements apart in lazily reserved address space -- the 2-cell run
that motivated the guard. Without the guard the kernel hands the measured
stride to BLAS++ and throws blas::Error; with it the run takes the per-cell
path and the result matches the reference. Skipped under ILP64 or when the
address-space reservation is refused.

Author: kshitij-05

tests/arena_strided_gemm.cpp

@@ -10,10 +10,15 @@
 #include <algorithm>
 #include <cmath>
 #include <complex>
+#include <cstddef>
+#include <cstdint>
 #include <functional>
+#include <limits>
 #include <memory>
 #include <vector>
 
+#include <sys/mman.h>
+
 namespace TA = TiledArray;
 using Inner = TA::ArenaTensor<double, TA::Range>;
 using Outer = TA::Tensor<Inner>;
@@ -376,6 +381,67 @@ 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;
+  if (blas::max_ld() >= std::numeric_limits<blas::integer>::max()) return;
+  const std::size_t M = 1, N = 1, K = 2, P = 3, Q = 4;
+  const std::size_t stride = static_cast<std::size_t>(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;
+  }
+  const auto unmap = [len](void* p) { ::munmap(p, len); };
+  std::unique_ptr<void, decltype(unmap)> mapping(mem, unmap);
+  auto* base = reinterpret_cast<std::byte*>(
+      (reinterpret_cast<std::uintptr_t>(mem) + align - 1) & ~(align - 1));
+  // stride * sizeof(double) is a power of two >= align, so cell 1 is aligned
+  using Cell = Inner::Cell;
+  Cell* c0 = ::new (base) Cell{TA::Range{P}};
+  Cell* c1 = ::new (base + stride * sizeof(double)) Cell{TA::Range{P}};
+  {
+    Inner l0(c0), l1(c1);
+    BOOST_REQUIRE_EQUAL(static_cast<std::size_t>(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<Inner> alloc;
+    Inner* raw = alloc.allocate(K);
+    ::new (raw) Inner(l0);
+    ::new (raw + 1) Inner(l1);
+    std::shared_ptr<Inner[]> ldata(raw, [alloc, K](Inner* p) mutable {
+      for (std::size_t i = 0; i < K; ++i) (p + i)->~Inner();
+      alloc.deallocate(p, K);
+    });
+    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<Outer>(
+        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
+  c1->~Cell();
+  c0->~Cell();
+}
 #ifdef TA_STRIDED_GEMM_COUNT
 BOOST_AUTO_TEST_CASE(ce_e_fires_clean_path) {
   namespace blas = TiledArray::math::blas;

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<integer>(std::numeric_limits<blas_int>::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<integer>::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()