Skip to content

Vectorize continuity assembly to remove per-edge scalar NumPy dispatch #3

Description

@emircbngl

Summary

The Python edge loop in assemble_continuity() dominates the current 2D steady solve. Each edge calls the NumPy-backed scalar bernoulli() twice, so a small array operation and its dispatch overhead are repeated tens of thousands of times.

A repository-external vectorized prototype preserved the assembled system exactly and reduced the default end-to-end 2D solve median by 5.87x.

Measured on public main at 5100ab00cb8b2aa5d43682ab0e9b871bee8f7ac1 on an Apple-silicon CPU.

Profile

A profiled default 2D CLI solve showed:

_solve_pn2d_steady                 0.472 s
pn2d.solve_bias                    0.449 s
assemble_continuity (20 calls)     0.374 s
flux.bernoulli (64,640 calls)      0.316 s

The hot path is src/tarhan/numerics/assemble.py::assemble_continuity, particularly the scalar edge loop and the two scalar bernoulli() calls per active edge.

Prototype measurements

Warm medians:

Mesh Current assembly Vectorized prototype Assembly speedup
625 nodes / 1,616 edges 13.184 ms 0.259 ms 50.94x
8,125 nodes / 23,996 edges 198.917 ms 3.437 ms 57.88x

Default 625-node pn2d.solve_bias:

current:     314.095 ms median
prototype:    53.535 ms median
speedup:       5.87x

The prototype bulk-extracted edge endpoints and weights, evaluated bernoulli() on vectors, accumulated the residual with ordered indexed additions, and built the COO arrays in the same interleaved order as the current loop.

Correctness evidence

This is prototype evidence, not yet a formal OptiProof result: the generic OptiProof input generator cannot construct TARHAN's domain-specific Mesh objects without a fixture/workload adapter.

Manual differential coverage used 60 seeded valid cases spanning:

  • electron and hole carriers;
  • random density and potential arrays;
  • zero edge coefficients;
  • source terms;
  • subdomains;
  • Dirichlet rows.

For every case, residual, rows, cols, and vals were array_equal with maximum absolute difference 0.0.

The end-to-end default solve also produced exactly equal psi, phi_n, phi_p, n_hat, p_hat, i_n, i_p, i, psi_step, and current_rel_change.

Required semantic gates

An implementation should preserve:

  1. Exact COO row/column/value ordering and duplicate triplets.
  2. Residual accumulation order and the existing conservation behaviour.
  3. Zero-transmissibility edge skipping.
  4. Subdomain zero-flux behaviour and outside-node identity rows.
  5. Dirichlet row replacement and current validation errors.
  6. Source-volume behaviour.
  7. Existing exception types and messages for invalid inputs.

validation/layer0/numerics/test_assemble.py is the natural regression base. A domain-specific OptiProof workload for Mesh would make the speed claim reproducible in the standard performance gate.

Metadata

Metadata

Assignees

No one assigned

    Labels

    enhancementNew feature or request

    Projects

    No projects

    Milestone

    No milestone

    Relationships

    None yet

    Development

    No branches or pull requests

    Issue actions