strided ToT kernels: a measured leading dimension must fit the BLAS integer

src/TiledArray/math/blas.h

@@ -29,13 +29,15 @@
 #include <TiledArray/external/eigen.h>
 #include <TiledArray/type_traits.h>
 
+#include <blas/config.h>  // blas_int
 #include <blas/dot.hh>
 #include <blas/gemm.hh>
 #include <blas/scal.hh>
 #include <blas/util.hh>
 #include <blas/wrappers.hh>
 
 #include <cstdint>
+#include <limits>
 
 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<integer>(std::numeric_limits<blas_int>::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) {

src/TiledArray/tensor/arena_einsum.h

@@ -279,18 +279,45 @@ inline void phase_stop(std::atomic<std::uint64_t>& 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<std::uint64_t> g_fall_runs_ce_ce{0};
 inline std::atomic<std::uint64_t> g_fall_res_absent_ce_ce{0};      // 1
 inline std::atomic<std::uint64_t> g_fall_res_nonuniform_ce_ce{0};  // 2
 inline std::atomic<std::uint64_t> g_fall_res_stride_ce_ce{0};      // 3
+inline std::atomic<std::uint64_t> g_fall_res_ldrange_ce_ce{0};     // 4
 inline std::atomic<std::uint64_t> g_fall_op_absent_ce_ce{0};       // 13
 inline std::atomic<std::uint64_t> g_fall_op_nonuniform_ce_ce{0};   // 14
 inline std::atomic<std::uint64_t> g_fall_op_stride_ce_ce{0};       // 15
 inline std::atomic<std::uint64_t> g_fall_op_acrossk_ce_ce{0};      // 16
 inline std::atomic<std::uint64_t> g_fall_both_clean_ce_ce{0};      // 17
+inline std::atomic<std::uint64_t> 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 <typename GetCell>
 inline int classify_run(GetCell getcell, std::size_t n) {
   if (n == 0) return 0;
@@ -313,14 +340,17 @@ 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<std::ptrdiff_t>(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<long>(n), s0)) return 4;
   return 0;
 }
 
 /// Diagnose the OPERAND side (called when the result run is clean). getR(k,i)
 /// 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 <typename GetR, typename GetL>
 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<std::ptrdiff_t>(i) * sk)
           return 15;
+      if (!math::blas::ld_fits(sk, static_cast<long>(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,17 +580,20 @@ 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<std::uint64_t> g_kernel_ns_ce_e{0};
 inline std::atomic<std::uint64_t> g_check_ns_ce_e{0};
 inline std::atomic<std::uint64_t> g_fallback_ns_ce_e{0};
 inline std::atomic<std::uint64_t> g_e_fall_runs{0};
 inline std::atomic<std::uint64_t> g_e_l_absent{0};
 inline std::atomic<std::uint64_t> g_e_l_nonuniform{0};
 inline std::atomic<std::uint64_t> g_e_l_stride{0};
+inline std::atomic<std::uint64_t> g_e_l_ldrange{0};
 inline std::atomic<std::uint64_t> g_e_r_absent{0};
 inline std::atomic<std::uint64_t> g_e_r_nonuniform{0};
 inline std::atomic<std::uint64_t> g_e_r_stride{0};
+inline std::atomic<std::uint64_t> g_e_r_ldrange{0};
 inline std::atomic<std::uint64_t> 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<long>(lc[lbase + a_off(m, 1)].data() - l0.data());
           ldB = static_cast<long>(rc[rbase + b_off(1, n)].data() - r0.data());
           if (ldA < P || ldB < Q) clean = false;
+          if (!math::blas::ld_fits(ldA, static_cast<long>(K), P) ||
+              !math::blas::ld_fits(ldB, static_cast<long>(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<std::ptrdiff_t>(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;

src/TiledArray/tensor/tensor.h

@@ -3704,6 +3704,20 @@ class Tensor {
                 if (N > 1)
                   ldc = static_cast<integer>(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<integer>(A) * cw) ||
+                    !TiledArray::math::blas::ld_fits(
+                        ldc * cw, N, static_cast<integer>(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<integer>(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<integer>(A) * cw) ||
+                    !TiledArray::math::blas::ld_fits(
+                        ldc * cw, M, static_cast<integer>(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() !=

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