Skip to content

Make JIT source generation deterministic - #142

Closed
gthyagi wants to merge 1 commit into
underworldcode:developmentfrom
gthyagi:codex/jit-deterministic-source
Closed

gthyagi wants to merge 1 commit into
underworldcode:developmentfrom
gthyagi:codex/jit-deterministic-source

Conversation

@gthyagi

@gthyagi gthyagi commented Apr 25, 2026

Copy link
Copy Markdown
Contributor

Summary

This PR makes UW3 JIT source generation deterministic across MPI ranks by applying stable ordering when collecting symbols, constants, mesh variables, generated includes, and setup metadata.

Motivation

Some parallel runs can generate rank-dependent C source ordering from unordered Python collections. That can trigger JIT hash mismatches even when the mathematical expression is identical on all ranks.

Changes

  • Add stable sorting helpers for JIT source generation inputs.
  • Preserve PETSc field ordering by prioritising field_id when sorting mesh variables.
  • Add regression tests for deterministic ordering.

Validation

  • ./uw python -m pytest tests/test_jit_deterministic_ordering.py tests/test_jit_cache.py -q
  • 9 passed
  • Kramer spherical free-slip case was reproduced locally without requiring PYTHONHASHSEED=0.

Scope

This should not change solver mathematics. It only makes generated JIT source ordering stable across ranks.

Sort unordered symbolic and code-generation collections before they affect emitted JIT C source. This prevents Python hash-randomized set/dict iteration from producing different source text on different MPI ranks.

Preserve PETSc field ordering by sorting mesh variables by field_id before applying stable tie-breakers, then use deterministic ordering for constant extraction, free-symbol processing, UWexpression atom traversal, generated include headers, library directories, and library names.

Add focused regression tests for the stable ordering helper, including field_id precedence and unordered symbol-container ordering.

Validation: py_compile passed for _jitextension.py; tests/test_jit_deterministic_ordering.py passed; tests/test_jit_cache.py passed; Kramer spherical free-slip case2 cellsize 1/8 completed on 8 MPI ranks without PYTHONHASHSEED=0.
@gthyagi
gthyagi requested a review from lmoresi as a code owner April 25, 2026 17:24
@gthyagi gthyagi closed this Apr 25, 2026
@gthyagi
gthyagi deleted the codex/jit-deterministic-source branch April 25, 2026 17:31
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant