tinyasm: avoid copying the block dofmap in BlockJacobi::solve()

[manishpaulish]

dofsPerBlock is a vector<vector<PetscInt>>, so

auto dofmap = dofsPerBlock[p];

deduces by value and heap-allocates plus copies the block's dofmap on every block, on every application of the preconditioner. solve() runs once per Krylov iteration and block counts are typically in the thousands, so this is on the order of a million small allocations per linear solve.

Binding by const reference instead. No behavioural change, dofmap is only read.

A benchmark driving the same loop, including the dof < 6 branch, with
vertex-star patch sizes taken from a real tetrahedral mesh rather than
invented ones, puts the copy at roughly 14% of solve() for P1 (median
block 15 dofs) and 2 to 3% for P2 (median block 65 dofs). Details and
caveats in the comments below. Note this is a fraction of solve(), not
of an end-to-end solve.

Independent of the timing, the change removes one heap allocation per
block per preconditioner application, for a value the loop only reads.

[pbrubeck]

It'd be good to motivate this PR with an actual benchmark and more context on how and why you came across this bug.

[manishpaulish]

Fair on both counts, and pushing on the benchmark was the right call. It changed the answer.

How I came across it. I was looking for the C++ in Firedrake, since that is my background, and tinyasm is essentially all of it. I read BlockJacobi::solve() rather than profiled it, noticed auto dofmap = dofsPerBlock[p] deduces vector<PetscInt> by value, and went looking for a number afterwards. So it was reading, not a profile, and the number in the description came second.

The number in the description was misleading and I should correct it. I generated block sizes uniformly in [4, 20], which I picked out of the air. Real vertex-star patches are not distributed like that. Rebuilding the benchmark on patches from an actual mesh, a 12^3 cube Kuhn-subdivided into 10368 tets giving 2197 vertex-star patches:

	median block	copy cost
P1	15 dofs	14% of solve()
P2	65 dofs	2 to 3% of solve()

Same loop as solve() including the dof < 6 branch, openblas dgemv, -O3 -march=native, 400 applications, best of 15 interleaved runs, three independent repeats. P1 lands at 13.8, 14.3, 14.5. P2 at 2.0, 2.7, 3.0.

So my original 11 to 13 percent was roughly right for small blocks and overstates the case that matters most by several times. At P2 the dgemv is 65x65 and dominates so completely that the copy nearly disappears. At 50 applications it was inside the noise entirely, bouncing between 3.2, 0.7 and -0.0 percent across runs, and I had to go to 400 to measure it at all.

Two honest caveats. This is still not a Firedrake solve. The matrices are random rather than assembled, and more importantly this is a percentage of solve(), not of an end-to-end solve, so once assembly, Krylov vector operations and MPI are in the picture the share is smaller again. I tried to build Firedrake to give you the real end-to-end figure and did not get PETSc to configure here, so I cannot.

What the change does buy, independent of the timing, is 2197 heap allocations per preconditioner application, one per block, going to zero. That is the part I would still defend: the loop allocates once per block per application for a value it only reads.

Whether 2 to 3 percent of the apply justifies the churn is your call and I will not argue if the answer is no. Happy to add the benchmark under tests/, or to drop it, and happy to correct the description above so the 11 to 13 percent figure is not left standing.

[connorjward]

Assuming that tests pass I see no issue with this and can believe that this affects performance.

[connorjward]

Tests should hopefully pass if you merge in main.

[manishpaulish]

CI finished. Two jobs are red and I do not think either comes from this PR.

Build documentation is red on main too. Run 32027292800 on 5f5a7d1, which is the merge base here, has the same job failing, as does run 32047606659 on an unrelated PR that ran alongside mine.

Test default failed on one test:

tests/firedrake/multigrid/test_adaptive_multigrid.py::test_DG0[firedrake-square-nprocs=4-inject]
AssertionError: assert np.float64(nan) <= 1e-12

That test builds two DG0 spaces, interpolates, calls inject and takes an errornorm. It never constructs a preconditioner, so BlockJacobi::solve() is not reached at all, and the diff here is a const reference binding inside that function. The tinyasm case in the same file is a different test and passed. Test default also passed on main's own run of 5f5a7d1 and on the concurrent PR above, so I read this as the parallel adaptive refinement path rather than anything this touches.

Happy to push an empty commit to re-run, or leave it if you would rather look first. If the nan is worth chasing separately I can open an issue with the log.

[connorjward]

Yeah CI appears to be in a bad place currently distinct from this PR. I will aim to merge this once upstream things are resolved. Thanks!

[pbrubeck]

I think this PR fixes that CI error #5337

[connorjward]

@manishpaulish sorry to keep asking but please merge in main one more time. Hopefully there will be fewer CI errors.

[manishpaulish]

Done, merged main in at 6e2746e. The branch is level with main now and the diff is still the single line in tinyasm/tinyasm.cpp. No apology needed, this was CI rather than you.