Experimental optional vectorization for SphMap, PolyMap and MatrixMap - #51
Experimental optional vectorization for SphMap, PolyMap and MatrixMap#51embray wants to merge 7 commits into
Conversation
Codecov Report❌ Patch coverage is
Additional details and impacted files@@ Coverage Diff @@
## master #51 +/- ##
==========================================
+ Coverage 61.49% 61.60% +0.11%
==========================================
Files 83 83
Lines 96519 96878 +359
Branches 30595 30705 +110
==========================================
+ Hits 59351 59684 +333
+ Misses 20983 20974 -9
- Partials 16185 16220 +35 ☔ View full report in Codecov by Harness. |
|
This is all well outside my comfort zone, so I've no real comments to make other than it looks very promising. |
|
Very impressive work. I had noodled around with MatrixMap SIMD vectorization but for 2x2 matrices I couldn't get it to make a difference and it wasn't clear we had bigger matrices. This is excellent. |
The MatrixMap improvements are mostly marginal in practice I found. SphMap was the biggest one by far--afterwards callgrind was still showing most of the remaining time spent on MatrixMap so I dove into it anyways, having become a little obsessed with chasing the last mile. Now it looks like |
Mine too honestly, I just tried to see what we could get almost "for free" from the compiler in terms of AVX2 instruction output, though it still required refactoring the kernels of each transform to vectorize efficiently over the input coordinate points. I don't want to make a mess of the code though so if/when you have time I'd value your input on the overall architectural decisions (especially refactoring) and code style. It's not urgent of course. |
|
I am starting to think that for every AST test frameset we have in ast_tester (I know that directory needs to be classified and categorized into subdirectories but I think #37 will be made more complicated if I do that now) we should have a reference forward transform output (maybe of four corners?) so that we can compare outputs from changes like this to what we had before. We could store pixel_x, pixel_y, transform for 100 points in a CSV without blowing up the repo size. At least then you would know if you have substantially changed a mapping unexpectedly. |
|
@timj That would not be a bad idea. For what it's worth, as you can see, this didn't really have a big impact on the existing tests. I don't think this would be a big problem for #37. Even if more tests are added I can either incorporate those or, if we want to just focus on the cmake build that's fine too. |
|
Hmm, as long as we're being persnicketty about satisfying the codecov report, there are several matrix cases that are uncovered by the tests, so I'll try to improve that too. |
|
As far as I'm concerned, this is mostly ready--makes a nice difference already; also improved the tests, boosting coverage a bit. There are two open questions remaining for me:
|
c8feb10 to
b45e4a5
Compare
Baseline x86-64 has no FMA instruction, so compilers round a*b+c twice; arm64 (and Haswell-or-later x86 via -march=x86-64-v3) contract to a fused multiply-add by default. Building one Ubuntu job at v3 makes contraction-driven transform-oracle differences observable as baseline-vs-v3 on identical OS and libm, separating them from the libm differences that dominate ubuntu-vs-macos comparisons. c_flags is a real matrix dimension (base value '') rather than only an include key: an include whose standard keys all match an existing combination would otherwise merge into that job instead of creating a new one. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
- Starlink/ast#37 - Starlink/ast#51 - Starlink/ast#66 - Starlink/ast#67 This includes the experimental SIMD support (enabled by default, and support in AST for libfyaml, allowing us to drop the libyaml requirement). Also updates the minimum supported versions of the ASDF tags supported by AST.
|
I imagine this hasn't been rebased since we added the transform oracle? That would help to solidify whether the SIMD is working properly. I imagine we would also need a way to ensure that we generate oracles from the non-SIMD path? Should we do it as experimental first with an environment variable gate? |
|
I think I have tested it against that, since I did make a private branch with this merged in. But I'll double-check. Needs rebasing anyways. |
This is a squashed commit to clean up several WIP steps taken while experimenting: wip: small win for diagonal matrix case; perhaps more significant win for sphmap but needs more testing temporary (?) workaround for the grid test failure (only fails for me on car3) but even though this rounding gets the tests to pass, visual inspection of the results suggests a slightly lower-quality image in the SIMD case (more jitter in some of the arcs drawn); needs more examination wip: try supporting SIMD options in autotools build Unfortunately automake does not make it easy to set per-object-file flags, so we have to build supported compilation units as their own sub-libraries wip: revert matrixmap changes for now since it wasn't clear if it was getting the results I wanted (compiler was not vectorizing). Needs more careful analysis and should be come back to independently of the sphmap changes wip: performance benchmark script (currently for SphMap but more to be added) wip: first attempt at polynomial vectorization results are mixed -- definite ~4x speedup for smaller N, though at larger N (>~100000) cache pressure seems to dominate, can actually result in slight regression in performance, of course also with the memory tradeoff. Needs more work I think... wip: huge improvement in polynomial performance at large N Instead of processing all points at once we take them in chunks -- an effort is made to determine the available L2 cache capacity; should be perfectly reliable on any x86_64 processor; on other platforms YMMV. Chunk size is determined by this and an additional outer loop processes each chunk in cache maintaining a high throughput and constant memory charge for any N. build: fix autotools build for polymap SIMD optimizations wip: matrixmap enhancements for SIMD wip: cleanup, consolidation of SIMD-supported mappings - add astCPUCacheSize to memory.c for use in different mappings and expanded it to support LD1 and L3 cache sizes for possible future use. - Restructured the supported mappings to add a TransformLoop virtual method, allowing cleanly separating out the original implementations from the SIMD implementations - Also add UseSIMD attribute for each transform to allow swapping out the implementation at runtime (when supported) Some duplication here--if we go further down this route might be worth adding these either to the Mapping base or maybe more likely a "subclass". - Added simple benchmarks for matrixmap (again derived from my GWCS benchmarks) and a more helpful summary table printout (including printing in markdown format for easy pasting into GitHub, etc.) - Cleaned up the failing grid tests -- now that SIMD can be disabled at runtime it's cleaner to allow disabling it on specific tests (only one I'm aware of, grid_car3) where minor numerical differences caused the test to fail. Better to make this explicit than the previous approach of just trying to smooth over the differences with rounding. In fact that particular plot I think (subjectively?) looks worse when produced with UseSIMD=1; maybe worth further investigation. - Consolidated History entries wip: improve benchmarking and minor performance fixes - The benchmarks at small N were not reliable due to clock jitter, etc. so make several calls in batches to smooth out the per-time call; this removes a lot of noise in repeated benchmark runs. - Identified obvious performance degradation in PolyMap at small N due to the additional mallocs; set a minimal N to switch to SIMD version similarly to SphMap - Tightened up the non-diagonial MatrixMap a bit avoiding initial zeroing and vectorizing the scan for AST__BAD in the input.
…MD support plot: disable SIMD in plotting due to negative impacts on curve plotting in some cases. After some investigation found that the tiny 1-4 ULP jitter in the vectorized trig functions result in just enough jitter that it upsets the (dl^2 * (1 - cos^2(theta))) < Crv_limit, resulting in curves that are overly sub-divided in some places, resulting in a "staircase" effect. I don't feel confident to try to fiddle too much with the curve drawing parameters that are otherwise working, so it seems the safest approach to preserve reproducibility is to disable the SIMD computations entirely just for plotting. This allows reverting the earlier changes I made to the tests to make them pass by artificially disabling SIMD just in the test; now plotting is just not broken in general.
This test runs whether or not compiled with SIMD support; mostly it just exercises the UseSIMD attribute for the affected mappings, and makes sure also to expercise the non-SIMD code paths for those mappings, improving overall test coverage.
Found this out actually by the coverage report in the CI. Curiously the `__x86_64__` block was being entered, but not finding anything in cpuid leaf 4. Turns out AMD CPUs use an extended leaf 0x8000001D for this, which otherwise works the same way. The GitHub CI runs don't tell you by default what actual hardware it's running on so I'm just guessing this was the reason, but it seems a reasonable guess, and worth fixing anyways. Thanks to https://github.com/m-j-w/CpuId.jl/blob/master/src/CpuId.jl to helping figure out how to do this. Adds LCOV_EXCL block around the sysconf fallbacks, which I think now should be reasonable on most hardware this runs on...but we'll see...
These additional cases exercise more different matrix shapes and code paths in matrixmap.c that weren't covered. Debatably this should go in a more dedicated test for matrixmap, but just put it for now in the test program I was already working on.
Including it when it's not covered by the tests warps the patch coverage reports.
b45e4a5 to
61e3b6f
Compare
|
@timj Rebased, and confirmed that the oracle tests are passing. Also confirmed via a coverage build that my new code paths were being exercised in the test. I also re-ran my benchmarks and they were consistent with previous results, with some apparent slight improvement in the polynomial evaluations, and a slight regression in the non-diagonal matrix benchmarks. I haven't looked into whether there are any code changes that could account for that. But I was running the benchmark just on my laptop with my system under average load and performance settings so it's not a clean benchmark; just a sanity check. |
3 participants
Background
A couple weeks ago I was doing some benchmarking of AST (wrapped in my libasdf-gwcs library) against the Python GWCS package. I found, to no surprise, that for many cases AST far outperformed Python (on a single thread). However, at about N=100000 points the Python code started to outperform slightly.
I concluded that this was due to NumPy's built-in SIMD support for many of its ufuncs (and potentially BLAS as well, though I don't think that was a major contributor since there are only some small matrix multiplications performed). On my machine N=100000 happened to be where the Python overhead is amortized by the SIMD enhancements.
I found that the major contributors in my case were PolyMap, MatrixMap, and SphMap, with the biggest potential win coming from SphMap. So I set about adding some bare minimum compiler (GCC) generated vectorization in targeted areas to see if we could get any quick wins, starting with SphMap, and indeed it did make a big difference.
I've added a little benchmark program adapted from my libasdf-gwcs benchmark. I then went on to see if I could get any easy wins with MatrixMap and PolyMap. These were a little trickier to get right, especially PolyMap. The key observation in all these cases is that the main outer loop of each
Transformmethod is over the input points. In a way this makes sense for cache efficiency, but ignoring that for the first pass I trudged ahead, and did close the gap with Python up to a point--eventually at even larger N the performance was being dominated by DRAM moves.At least for MatrixMap and PolyMap I added a utility function to return the system's L2 cache size (should work for any current x86-64, YMMV on other platforms), and add a separate outer loop to chunk the points into sets small enough to fit in cache (assuming not huge numbers of polynomial coeffs or matrix elements, other overhead), and this proved a big win--AST at parity with or better than NumPy, with near consistent throughput and memory usage even for large N.
End-to-end impact (libasdf-gwcs vs Python GWCS)
The real motivation, and the most representative measurement, is full Roman WCS evaluation through libasdf-gwcs (which drives AST's
astTran2over a complete WCS pipeline on real Roman L2 calibration files). The plots below show single-thread throughput vs Python GWCS across N, before and after this change.Before:
After, with AST compiled with SIMD features:
The above was just from one run of the benchmark but YMMV. In effect this closes the gap with NumPy, and I had some runs depending on the conditions that outperformed it entirely.
The per-transform microbenchmarks further down isolate where the gains come from; this end-to-end view is what actually matters for a WCS pipeline, and is also where the L2 chunking pays off most (each pipeline stage's output stays cache-resident as the next stage's input, rather than being re-streamed from DRAM).
What's in this PR
Opt-in build flag. A new
AST_ENABLE_SIMDoption (CMake and autotools), OFF by default. When enabled, SIMD-friendly compiler flags (-fopenmp-simd -ffast-math -fno-associative-math -march=x86-64-v3) are applied only tosphmap.c,polymap.candmatrixmap.c, so the floating-point semantics of the rest of the library are unchanged. The flags are only accepted on a GCC/glibc toolchain that supports them; otherwise the option is a no-op.TransformLoopvtable slot. Each of SphMap, PolyMap and MatrixMap gains a per-classTransformLoopslot.Transformis now a thin wrapper (validation, invert handling, point extraction) that dispatches to eitherTransformLoopScalarorTransformLoopSIMD. This keeps the scalar and vectorized inner loops cleanly separated instead of#ifdef-ing through the body ofTransform.Runtime
UseSIMDattribute. A per-instance integer attribute on each of those classes (default 1 when built with SIMD support, 0 otherwise) lets a caller force the scalar path at runtime -- useful for bit-exact reproducibility or debugging without rebuilding. SettingUseSIMD=1on a library built without SIMD reportsAST__ATSER. The attribute is deliberately not serialized (it's just a runtime tuning knob, not part of the mapping description). I debated making this some kind of global setting but opted to make it specific to the supported transforms for now.TransformLoopandUseSIMDthere's a bunch of duplicative code. I didn't want to do any more refactoring for now before waiting to see what you have to say about this, but it seems to merit a base class or even just moved into the baseMappingclass.astCPUCacheSize(level)inmemory.c: returns the L1D/L2/L3 cache size via CPUID (x86-64), with asysconffallback and hard-coded defaults. PolyMap and MatrixMap use the L2 size to pick a chunk size.The vectorized kernels themselves (the substance):
atan2/sqrt(forward) andsin/cos(inverse) under#pragma omp simd, which GCC routes through glibclibmvec; a scalar fixup pass handles poles and bad values.PolyPowers). This could easily be extended to ChebyMap too but I skipped it for now since my usecase doesn't use any Chebychev polynomials, at least for now.VFMADDvia#pragma omp simd, L2 chunking, and a bad-value fixup that runs only when a vectorized scan finds a bad input; diagonal matrix: a branchless compare/blend.perf/bench_simd.c-- a standalone benchmark (details below).Benchmark
perf/bench_simdruns each transform twice in a single process -- scalar (UseSIMD=0) and SIMD (UseSIMD=1) via the runtime attribute -- over an N sweep, and prints a scalar-vs-SIMD summary (plain text, or--markdown).A note on methodology, since it matters for reading the small-N numbers: each measurement times a calibrated batch of repeated
astTranPcalls (the batch grows until it runs for at least 20 ms) and reports the amortized per-call time. At small N a single call is dominated by call overhead and clock resolution rather than throughput, so batching is what makes those rows meaningful -- it removes measurement error, not the per-call cost itself. The input arrays stay hot across a batch, so small/medium N measures compute-bound throughput and large N (which exceeds the last-level cache) measures bandwidth-bound throughput. Reported figures are the median over the reps (-r, default 5).Hardware/build for the numbers below: Intel i7-7820HQ (Kaby Lake, AVX2), L2 = 1 MiB/core, L3 = 8 MiB; GCC 13.3.0 (i.e. just my laptop);
-DCMAKE_BUILD_TYPE=RelWithDebInfo -DAST_ENABLE_SIMD=ON.Headline
sphmap_fwdatan2/sqrtsphmap_invsin/cospoly5_fwdpoly1_fwdmatfull2_fwdmatdiag2_fwdA few things worth calling out:
libmvecvectorizes that directly.Full sweep
scalar vs SIMD, Mpx/s, N = 1 ... 16.7M (median of 9 reps)
Correctness and testing
ctestsuite passes; the only failures are ones that pre-date this branch and are unrelated to it.-ffast-mathlets GCC calllibmvec, whose x86_64 vector math functions are documented to within 4 ULP, so a result can differ from scalarlibmby a few ULP). This is normally irrelevant. The one test I ran that was sensitive to that (grid_car3) runs withUseSIMD=0so it stays bit-exact against its committed reference; the rest of the suite runs with SIMD enabled.-Wall-style warnings and ASan/UBSan (no new warnings, no sanitizer findings) in addition to the normalRelWithDebInfobuild.Caveats and limitations
-march=x86-64-v3flag, andAST_ENABLE_SIMDis off by default, so on any other toolchain it is simply a no-op. Clang can almost certainly do the same (it also supports vector-ABI /libmvecmath), but I haven't tested it with Clang at all yet. glibc also ships alibmvecfor aarch64 these days, but this build isn't set up or tested for that either.astCPUCacheSize's CPUID path is x86-64; elsewhere it falls back tosysconfand then to conservative defaults (L1 32 KiB, L2 512 KiB, L3 8 MiB).