strided ToT kernels: a measured leading dimension must fit the BLAS integer - #589
Merged
Merged
Conversation
The ce+e / ce+ce strided-GEMM kernels take their leading dimensions from pointer differences between neighbouring inner cells. Two present cells that are not arena neighbours (individually allocated inner tensors placed gigabytes apart) pass every other precondition of a 2-cell run -- uniform size, stride >= cell size, trivially constant -- and hand blaspp an ld it cannot convert under LP64: `blas::Error: ldb, in function to_blas_int_` (uranyl/dyall-v3z PNS-MP1 union path, tpns <= 1e-7, MKL LP64). Such a run is not strided-GEMM material: break the segment (or mark the k-run unclean) when a stride exceeds numeric_limits<blas_int>::max(), and let the per-cell path handle it.
…verywhere Three gaps in the previous commit's guard. Scope. Tensor::gemm's two arena "scale" strided paths measure their leading dimensions the same way the arena kernels do -- lc(1).data() - lc(0).data() for the ToT x plain-scalar row slab, right_data[rcell(1,n)].data() - right_data[rcell(0,n)].data() for the plain-scalar x ToT column slab -- into math::blas::integer, which is 64-bit, so nothing truncates and the value reaches blaspp unchecked. The `ld < A` test rejects only negative/overlapping strides, and with K == 2 (or N == 2, M == 2) the constant-stride verification loop is satisfied by the very pair the stride was measured from. The same un-compacted ToT tile that motivated the arena fix therefore still throws `blas::Error: ldb, in function to_blas_int_` through those two paths. Bound. A leading dimension is not the largest index the GEMM forms: it reaches element (nslab-1)*ld + extent-1. blaspp's to_blas_int_ checks each argument in isolation and never that span, so a stride comfortably under numeric_limits<blas_int>::max() overflows a 32-bit-indexed BLAS once the run is long enough -- wrong results or a segfault inside the BLAS rather than a clean throw. Since the motivating workload demonstrably produces strides above the cap, strides just below it are equally reachable there. Home. The cap is a property of the BLAS binding, not of arena einsum, and tensor.h cannot include arena_einsum.h. Both now share math::blas::max_ld() / ld_fits(ld, nslab, extent) in math/blas.h. ld_fits is phrased as a division so the bound itself cannot overflow under ILP64, where blas_int is 64-bit and max_ld() is INT64_MAX (every check then correctly degenerates to a no-op). Typed math::blas::integer -- what each ld is cast to at the call -- rather than long, which is only accidentally 64-bit here. In the two ce+ce segment walkers the check moves out of the `off == 1` branch so it re-evaluates as the segment grows; a segment can then only come out shorter, never longer, and the remainder resumes at the next cell as before. Finally, classify_run / classify_operand still scored an oversized stride as clean, so a run rejected solely by the new gate was re-diagnosed as 17 and counted in g_e_both_clean / g_fall_both_clean_ce_ce -- the bucket documented as "gate rejected a run this re-check finds valid", i.e. the signal for a gate bug. They gain reason codes 4 and 18 with their own counters and dump lines.
Member
|
@kshitij-05 still a draft? |
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.
kshitij-05
changed the title
kshitij-05
marked this pull request as ready for review
September 18, 2026 20:36
evaleev
approved these changes
Sep 18, 2026
…nit test tests/math_blas.cpp is tracked but was never listed in ta_test_src_files, so nothing in it has ever been compiled -- including the measured_ld_fits case added for this guard. Registering it next to linalg.cpp brings the whole blas_suite (17 cases) into the run; it passes as it stands. ld_fits() divides, and the two ce+ce walkers called it twice per admitted cell in the innermost segment-growth loop, whose other work is a handful of loads and compares. The strides and extents are fixed for a segment once off == 1, so take the bound once there instead: max_ld_offset(ld, extent) is the same inequality solved for the slab offset, and admission then costs a compare. The long arena-contiguous segments this kernel exists to hit -- g_seg_len9p_ce_ce -- were paying two divisions per cell for a check that can never fire on them. ld_fits() is now expressed through max_ld_offset(), so the two cannot drift apart; measure_segments takes the same hoist, since it exists to simulate the walker's segmentation exactly. Also, per review: - spell math::blas::ld_fits at the call sites rather than pulling it into TiledArray::detail with a using-declaration: arena_einsum.h is widely included, and the name is of no interest to anything else in detail. The rationale it carried becomes a section comment. - ce_e_stride_past_blas_int_falls_back returned before any assertion on both its skip paths (ILP64, refused reservation), so on those configurations it passed while checking nothing and only Boost's easily-missed "did not check any assertions" note said so. It now asserts that an ordinary contiguous 2-cell run is still accepted before either branch, and the ILP64 skip explains itself like the reservation one already did. - the reservation and the two placement-new'd cells get one RAII owner: a failing BOOST_REQUIRE skipped ~Cell(), leaking each cell's TA::Range buffer, which under this repo's ASan debug build would bury the real failure in leak reports. - drop two lambda captures of constants (`len`, `K`) that -Wall reports as unnecessary. CI runs TA_WERROR=ON without -Wall so they did not gate it, but cmake/toolchains/travis.cmake does add -Wall, and that combination makes them errors.
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
2 participants
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Problem
The strided ToT GEMM kernels in
tensor/arena_einsum.h(ce+e, ce+ce right/left) and the two arena "scale" paths ofTensor::gemm(ToT x scalar, scalar x ToT) take their leading dimensions from the address distance between neighbouring inner cells. Two present cells that are not arena neighbours -- individually allocated inner tensors the allocator placed gigabytes apart -- pass every other precondition of a 2-cell run (uniform size, stride >= cell size, trivially constant), so the measured stride reaches BLAS++ unchecked and cannot be narrowed toblas_intunder an LP64 BLAS:blas::Error: ldb, in function to_blas_int_(uranyl/dyall-v3z PNS-MP1 union path, MKL LP64). A stride below that cap can still overflow a 32-bit BLAS index once the run is long enough, because the GEMM addresses element(nslab-1)*ld + extent-1.Fix
math/blas.h:max_ld()(the largestblas_int) andld_fits(ld, nslab, extent), which bounds the whole addressed span. Phrased as a division so the bound itself cannot overflow under ILP64, where every check degenerates to a no-op.tensor/arena_einsum.h: the ce+e clean check and the two ce+ce segment walkers reject a run (or stop growing a segment) whose measured stride does not fit, so the per-cell path handles those cells; the walkers re-check as the segment grows. The fallback classifiers gain reason codes 4 / 18 ("ld range") so such runs are no longer reported as "both runs clean".tensor/tensor.h: the same guard on the measuredldb/ldc/sbcof the two arena scale paths.Arena-contiguous cells are unaffected.
Testing
blas_suite/measured_ld_fits:ld_fits/max_ldat their boundaries (one slab, many slabs, degenerate arguments, past-the-cap strides under LP64).arena_strided_gemm_suite/ce_e_stride_past_blas_int_falls_back: drivesarena_strided_gemm_ce_ewith two present, uniform-size left k-cells placedmax_ld()+1elements apart in lazily reserved address space (only two pages are touched). On master the kernel throwsblas::Error; with the guard the run takes the per-cell path and matches the reference. Skipped under ILP64 or when the reservation is refused.TA_ASSERT_THROW):ta_test --run_test=blas_suite,arena_strided_gemm_suite-> no errors; the same binary built against master fails the new case with "unexpected exception thrown by arena_strided_gemm_ce_e".